Add qchem time evolution operation (#455)

* Adds function TimeEvolution to dynamics

* Adds normal-to-local operations to sample functions

* Adds unittest for TimeEvolution function

* Modifies docstring of TimeEvolution function

* Runs black

* Modifies the format of TimeEvolution function

* Adds test to check operation order

* Update strawberryfields/apps/qchem/dynamics.py

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

* Update strawberryfields/apps/qchem/dynamics.py

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

* Update strawberryfields/apps/qchem/dynamics.py

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

* Update strawberryfields/apps/qchem/dynamics.py

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

* Update strawberryfields/apps/qchem/dynamics.py

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

* Update strawberryfields/apps/qchem/dynamics.py

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

* Update strawberryfields/apps/qchem/dynamics.py

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

* Changes the order of time and frequency

* Changes 'sf operation' to 'operation'

* Adds comments from code review

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

strawberryfields/apps/qchem/dynamics.py

@@ -64,8 +64,8 @@
 
 This module contains functions for implementing this algorithm.
 
-- The function :func:`~.evolution` returns a custom ``sf`` operation that contains the required
-  unitary and rotation operations explained in steps 2-4 of the algorithm.
+- The function :func:`~.TimeEvolution` is an operation that contains the required
+  rotation operations explained in step 3 of the algorithm.
 
 - The function :func:`~.sample_fock` generates samples for simulating vibrational quantum dynamics
   in molecules with a Fock input state.
@@ -92,51 +92,41 @@
 from strawberryfields.utils import operation
 
 
-def evolution(modes: int):
-    r"""Generates a custom ``sf`` operation for performing the transformation
-    :math:`U(t) = U_l e^{-i\hat{H}t/\hbar} U_l^\dagger` on a given state.
+def TimeEvolution(w: np.ndarray, t: float):
+    r"""An operation for performing the transformation
+    :math:`e^{-i\hat{H}t/\hbar}` on a given state where :math:`\hat{H} = \sum_i \hbar \omega_i a_i^\dagger a_i`
+    defines a Hamiltonian of independent quantum harmonic oscillators
 
-    The custom operation returned by this function can be used as part of a Strawberry Fields
-    :class:`~.Program` just like any other operation from the :mod:`~.ops` module. Its arguments
-    are:
-
-    - t (float): time in femtoseconds
-    - Ul (array): normal-to-local transformation matrix :math:`U_l`
-    - w (array): normal mode frequencies :math:`\omega` in units of :math:`\mbox{cm}^{-1}` that
-      compose the Hamiltonian :math:`\hat{H} = \sum_i \hbar \omega_i a_i^\dagger a_i`
+    This operation can be used as part of a Strawberry Fields :class:`~.Program` just like any
+    other operation from the :mod:`~.ops` module.
 
     **Example usage:**
 
     >>> modes = 2
-    >>> transform =  evolution(modes)
     >>> p = sf.Program(modes)
     >>> with p.context as q:
     >>>     sf.ops.Fock(1) | q[0]
-    >>>     sf.ops.Fock(2) | q[1]
-    >>>     transform(t, Ul, w) | q
+    >>>     sf.ops.Interferometer(Ul.T) | q
+    >>>     TimeEvolution(w, t) | q
+    >>>     sf.ops.Interferometer(Ul) | q
 
     Args:
-        modes (int): number of modes
-
-    Returns:
-        an ``sf`` operation for enacting the dynamics transformation
-    Return type:
-        op
+        w (array): normal mode frequencies :math:`\omega` in units of :math:`\mbox{cm}^{-1}` that
+            compose the Hamiltonian :math:`\hat{H} = \sum_i \hbar \omega_i a_i^\dagger a_i`
+        t (float): time in femtoseconds
     """
     # pylint: disable=expression-not-assigned
-    @operation(modes)
-    def op(t, Ul, w, q):
+    n_modes = len(w)
 
-        theta = -w * 100.0 * c * 1.0e-15 * t * (2.0 * pi)
+    @operation(n_modes)
+    def op(q):
 
-        sf.ops.Interferometer(Ul.T) | q
+        theta = -w * 100.0 * c * 1.0e-15 * t * (2.0 * pi)
 
-        for i in range(modes):
+        for i in range(n_modes):
             sf.ops.Rgate(theta[i]) | q[i]
 
-        sf.ops.Interferometer(Ul) | q
-
-    return op
+    return op()
 
 
 def sample_fock(
@@ -190,10 +180,10 @@ def sample_fock(
         raise ValueError("Input state must not contain negative values")
 
     if max(input_state) >= cutoff:
-        raise ValueError("Number of photons in each input state mode must be smaller than cutoff")
+        raise ValueError("Number of photons in each input mode must be smaller than cutoff")
 
     modes = len(Ul)
-    op = evolution(modes)
+
     s = []
 
     eng = sf.Engine("fock", backend_options={"cutoff_dim": cutoff})
@@ -206,7 +196,11 @@ def sample_fock(
         for i in range(modes):
             sf.ops.Fock(input_state[i]) | q[i]
 
-        op(t, Ul, w) | q
+        sf.ops.Interferometer(Ul.T) | q
+
+        TimeEvolution(w, t) | q
+
+        sf.ops.Interferometer(Ul) | q
 
         if loss:
             for _q in q:
@@ -268,9 +262,9 @@ def sample_tmsv(
 
     This function generates samples from a GBS device with two-mode squeezed vacuum input states.
     Given :math:`N` squeezing parameters and an :math:`N`-dimensional normal-to-local transformation
-    matrix, a GBS device with :math:`2N` modes is simulated. The evolution operator acts on only
-    the first :math:`N` modes in the device. Samples are generated by measuring the number of photons
-    in each of the :math:`2N` modes.
+    matrix, a GBS device with :math:`2N` modes is simulated. The :func:`~.TimeEvolution` operator
+    acts only on the first :math:`N` modes in the device. Samples are generated by measuring the
+    number of photons in each of the :math:`2N` modes.
 
     **Example usage:**
 
@@ -307,7 +301,6 @@ def sample_tmsv(
         )
 
     N = len(Ul)
-    op = evolution(N)
 
     eng = sf.LocalEngine(backend="gaussian")
     prog = sf.Program(2 * N)
@@ -318,7 +311,11 @@ def sample_tmsv(
         for i in range(N):
             sf.ops.S2gate(r[i][0], r[i][1]) | (q[i], q[i + N])
 
-        op(t, Ul, w) | q[:N]
+        sf.ops.Interferometer(Ul.T) | q[:N]
+
+        TimeEvolution(w, t) | q[:N]
+
+        sf.ops.Interferometer(Ul) | q[:N]
 
         if loss:
             for _q in q:
@@ -380,7 +377,6 @@ def sample_coherent(
         )
 
     modes = len(Ul)
-    op = evolution(modes)
 
     eng = sf.LocalEngine(backend="gaussian")
 
@@ -392,7 +388,11 @@ def sample_coherent(
         for i in range(modes):
             sf.ops.Dgate(alpha[i][0], alpha[i][1]) | q[i]
 
-        op(t, Ul, w) | q
+        sf.ops.Interferometer(Ul.T) | q
+
+        TimeEvolution(w, t) | q
+
+        sf.ops.Interferometer(Ul) | q
 
         if loss:
             for _q in q:

tests/apps/qchem/test_dynamics.py

@@ -114,26 +114,47 @@
 
 @pytest.mark.parametrize("time, unitary, frequency, prob", [(t1, U1, w1, p1), (t2, U2, w2, p2)])
 def test_evolution(time, unitary, frequency, prob):
-    """Test if the function ``strawberryfields.apps.qchem.dynamics.evolution`` gives the correct
+    """Test if the function ``strawberryfields.apps.qchem.dynamics.TimeEvolution`` gives the correct
     probabilities of all possible Fock basis states when used in a circuit"""
 
     modes = len(unitary)
-    op = dynamics.evolution(modes)
 
     eng = sf.Engine("fock", backend_options={"cutoff_dim": 4})
-    gbs = sf.Program(modes)
+    prog = sf.Program(modes)
 
-    with gbs.context as q:
+    with prog.context as q:
 
         sf.ops.Fock(2) | q[0]
 
-        op(time, unitary, frequency) | q
+        sf.ops.Interferometer(unitary.T) | q
 
-        p = eng.run(gbs).state.all_fock_probs()
+        dynamics.TimeEvolution(frequency, time) | q
+
+        sf.ops.Interferometer(unitary) | q
+
+        p = eng.run(prog).state.all_fock_probs()
 
     assert np.allclose(p, prob)
 
 
+@pytest.mark.parametrize("time, frequency", [(t2, w2)])
+def test_evolution_order(time, frequency):
+    """Test if function ``strawberryfields.apps.qchem.dynamics.TimeEvolution``correctly applies the
+    operations"""
+
+    modes = len(frequency)
+
+    prog = sf.Program(modes)
+
+    with prog.context as q:
+
+        dynamics.TimeEvolution(frequency, time) | q
+
+    assert isinstance(prog.circuit[0].op, sf.ops.Rgate)
+    assert isinstance(prog.circuit[1].op, sf.ops.Rgate)
+    assert isinstance(prog.circuit[2].op, sf.ops.Rgate)
+
+
 @pytest.mark.parametrize("d", [d1, d2])
 class TestSampleFock:
     """Tests for the function ``strawberryfields.apps.qchem.dynamics.sample_fock``"""
@@ -174,7 +195,7 @@ def test_invalid_input_photon(self, d):
         mode is not smaller than cutoff."""
         with pytest.raises(
             ValueError,
-            match="Number of photons in each input state mode must be smaller than cutoff",
+            match="Number of photons in each input mode must be smaller than cutoff",
         ):
             in_state, t, U, w, ns, cf = d
             dynamics.sample_fock([i * 2 for i in in_state], t, U, w, ns, 3)