Add qchem vibronic transition operation (#451)

* Adds function VibronicTransition to vibronic module

* Adds vibronic operation to sample function

* Adds description to docstring

* Adds unit test for VibronicTransition

* Adds corrections to the docstring

* Adds comments from code review

* Adds comments from code review

* Update strawberryfields/apps/qchem/vibronic.py

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* Update strawberryfields/apps/qchem/vibronic.py

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* Update strawberryfields/apps/qchem/vibronic.py

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* Adds test to check operation order

* Removes return type from docstring

* Changes 'sf operation' to 'operation'

* Adds description to docstring

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>
Co-authored-by: Josh Izaac <josh146@gmail.com>

strawberryfields/apps/qchem/vibronic.py

@@ -49,6 +49,9 @@
 function computes two-mode squeezing parameters :math:`t`, from the molecule's temperature, which
 are required by the GBS algorithm to compute vibronic spectra for molecules at finite temperature.
 The :func:`sample` function then takes the computed GBS parameters and generates samples.
+The :func:`~.VibronicTransition` function is an operation that can be used within the
+conventional :class:`~.Program` interface for application of the Doktorov operator, allowing for
+greater freedom on the choice of input state and any subsequent quantum operations.
 
 Energies from samples
 ---------------------
@@ -59,11 +62,11 @@
 import warnings
 from typing import Tuple, Union
 
-from scipy.constants import c, h, k
-
 import numpy as np
+from scipy.constants import c, h, k
 
 import strawberryfields as sf
+from strawberryfields.utils import operation
 
 
 def gbs_params(
@@ -138,6 +141,46 @@ def energies(samples: list, w: np.ndarray, wp: np.ndarray) -> Union[list, float]
     return [np.dot(s[: len(s) // 2], wp) - np.dot(s[len(s) // 2 :], w) for s in samples]
 
 
+def VibronicTransition(U1: np.ndarray, r: np.ndarray, U2: np.ndarray, alpha: np.ndarray):
+    r"""An operation for applying the Doktorov operator
+    :math:`\hat{U}_{\text{Dok}} = \hat{D}({\alpha}) \hat{R}(U_2) \hat{S}({r}) \hat{R}(U_1)` on a
+    given state.
+
+    The Doktorov operator describes the transformation between the initial and the final vibronic
+    states of a molecule when it undergoes a vibronic transition.
+
+    **Example usage:**
+
+    >>> modes = 2
+    >>> p = sf.Program(modes)
+    >>> with p.context as q:
+    >>>     VibronicTransition(U1, r, U2, alpha) | q
+
+    Args:
+        U1 (array): unitary matrix for the first interferometer
+        r (array): squeezing parameters
+        U2 (array): unitary matrix for the second interferometer
+        alpha (array): displacement parameters
+    """
+    # pylint: disable=expression-not-assigned
+    n_modes = len(U1)
+
+    @operation(n_modes)
+    def op(q):
+
+        sf.ops.Interferometer(U1) | q
+
+        for i in range(n_modes):
+            sf.ops.Sgate(r[i]) | q[i]
+
+        sf.ops.Interferometer(U2) | q
+
+        for i in range(n_modes):
+            sf.ops.Dgate(np.abs(alpha[i]), np.angle(alpha[i])) | q[i]
+
+    return op()
+
+
 def sample(
     t: np.ndarray,
     U1: np.ndarray,
@@ -207,15 +250,7 @@ def sample(
             for i in range(n_modes):
                 sf.ops.S2gate(t[i]) | (q[i], q[i + n_modes])
 
-        sf.ops.Interferometer(U1) | q[:n_modes]
-
-        for i in range(n_modes):
-            sf.ops.Sgate(r[i]) | q[i]
-
-        sf.ops.Interferometer(U2) | q[:n_modes]
-
-        for i in range(n_modes):
-            sf.ops.Dgate(np.abs(alpha[i]), np.angle(alpha[i])) | q[i]
+        VibronicTransition(U1, r, U2, alpha) | q[:n_modes]
 
         if loss:
             for _q in q:

tests/apps/qchem/test_vibronic.py

@@ -16,13 +16,13 @@
 """
 # pylint: disable=protected-access,no-self-use,unused-argument,redefined-outer-name
 from unittest import mock
+
 import numpy as np
 import pytest
-
 from scipy.constants import c, h, k
 
-from strawberryfields.apps.qchem import vibronic
 import strawberryfields as sf
+from strawberryfields.apps.qchem import vibronic
 
 pytestmark = pytest.mark.apps
 
@@ -232,3 +232,65 @@ def test_sample_integration(p, integration_sample_number):
     assert dims == (integration_sample_number, len(alpha) * 2)
     assert samples.dtype == "int"
     assert (samples >= 0).all()
+
+
+class TestOperation:
+    """Tests for the function ``strawberryfields.apps.qchem.vibronic.VibronicTransition``"""
+
+    U1 = np.array([[-0.17905859, 0.98383841], [0.98383841, 0.17905859]])
+
+    r = np.array([0.00387885, -0.20415111])
+
+    U2 = np.array([[-0.33417234, 0.94251199], [0.94251199, 0.33417234]])
+
+    alpha = np.array([0.38558533, 2.99728422])
+
+    prob = np.array(
+        [
+            [1.49817854e-04, 1.27328834e-03, 5.43505231e-03, 1.55353201e-02],
+            [2.26762512e-06, 2.68454082e-05, 1.52080070e-04, 5.56423985e-04],
+            [2.80701727e-06, 2.52203270e-05, 1.14702166e-04, 3.51760396e-04],
+            [1.14518212e-07, 1.37859685e-06, 7.96673530e-06, 2.98259165e-05],
+        ]
+    )
+
+    p1 = [U1, r, U2, alpha, prob]
+
+    @pytest.mark.parametrize("params", [p1])
+    def test_op_prob(self, params):
+        """Test if the function ``strawberryfields.apps.qchem.vibronic.VibronicTransition`` gives
+        the correct probabilities of all possible Fock basis states when used in a circuit"""
+
+        U1, r, U2, alpha, prob = params
+
+        eng = sf.LocalEngine(backend="gaussian")
+
+        gbs = sf.Program(len(U1))
+
+        with gbs.context as q:
+
+            vibronic.VibronicTransition(U1, r, U2, alpha) | q
+
+            p = eng.run(gbs).state.all_fock_probs(cutoff=4)
+
+        assert np.allclose(p, prob)
+
+    @pytest.mark.parametrize("params", [p1])
+    def test_op_order(self, params):
+        """Test if function ``strawberryfields.apps.qchem.vibronic.VibronicTransition``correctly
+        applies the operations"""
+
+        U1, r, U2, alpha, prob = params
+
+        gbs = sf.Program(len(U1))
+
+        with gbs.context as q:
+
+            vibronic.VibronicTransition(U1, r, U2, alpha) | q
+
+        assert isinstance(gbs.circuit[0].op, sf.ops.Interferometer)
+        assert isinstance(gbs.circuit[1].op, sf.ops.Sgate)
+        assert isinstance(gbs.circuit[2].op, sf.ops.Sgate)
+        assert isinstance(gbs.circuit[3].op, sf.ops.Interferometer)
+        assert isinstance(gbs.circuit[4].op, sf.ops.Dgate)
+        assert isinstance(gbs.circuit[5].op, sf.ops.Dgate)