Multiple time evolution operators (#57)

* Delete __init__.py in tests-directory

* Return self when passing in system

* Allow for list of time-evolution operators

  This lets the user set up multiple (independent) time-evolution
  operators and sum these together. An example might be a time-dependent
  potential and a dipole laser acting on a system.

* Add test of empty operators

* Update documentation

* Add documentation to set_time_evolution_operator

quantum_systems/system.py

@@ -1,5 +1,6 @@
 import abc
 import copy
+import typing
 
 
 class QuantumSystem(metaclass=abc.ABCMeta):
@@ -24,7 +25,9 @@ def __init__(self, n, basis_set):
         self.np = self._basis_set.np
         self.set_system_size(n, self._basis_set.l)
 
-        self._time_evolution_operator = None
+        self._time_evolution_operator = []
+        self._add_h_0 = True
+        self._add_u_0 = True
 
     def set_system_size(self, n, l):
         """Function setting the system size. Note that ``l`` should
@@ -133,37 +136,78 @@ def particle_charge(self):
         """Getter returning the electronic charge of the particles."""
         return self._basis_set.particle_charge
 
-    def set_time_evolution_operator(self, time_evolution_operator):
-        # TODO: Support list of time-evolution operators
-        self._time_evolution_operator = time_evolution_operator
-        if not self._time_evolution_operator is None:
-            self._time_evolution_operator.set_system(self)
+    def set_time_evolution_operator(
+        self, time_evolution_operator, add_h_0=True, add_u_0=True
+    ):
+        """Function adding time-dependent terms to the Hamiltonian.
+
+        Parameters
+        ----------
+        time_evolution_operator : TimeEvolutionOperator, list
+            Either a list or a single instance of ``TimeEvolutionOperator``.
+        add_h_0 : bool
+            When ``True`` includes the time-independent part of the one-body
+            Hamiltonian when calling ``QuantumSystem.h_t``. Default is
+            ``True``.
+        add_u_0 : bool
+            When ``True`` includes the time-independent part of the two-body
+            Hamiltonian when calling ``QuantumSystem.u_t``. Default is
+            ``True``.
+
+        See Also
+        --------
+        TimeEvolutionOperator
+        """
+
+        if not isinstance(time_evolution_operator, typing.Iterable):
+            time_evolution_operator = [time_evolution_operator]
+
+        self._add_h_0 = add_h_0
+        self._add_u_0 = add_u_0
+
+        self._time_evolution_operator = [
+            op.set_system(self) for op in time_evolution_operator
+        ]
 
     @property
     def has_one_body_time_evolution_operator(self):
-        if self._time_evolution_operator is None:
-            return False
-
-        return self._time_evolution_operator.is_one_body_operator
+        return any(
+            op.is_one_body_operator for op in self._time_evolution_operator
+        )
 
     @property
     def has_two_body_time_evolution_operator(self):
-        if self._time_evolution_operator is None:
-            return False
-
-        return self._time_evolution_operator.is_two_body_operator
+        return any(
+            op.is_two_body_operator for op in self._time_evolution_operator
+        )
 
     def h_t(self, current_time):
-        if self._time_evolution_operator is None:
-            return self._basis_set.h
+        h_0 = (
+            self._basis_set.h
+            if self._add_h_0
+            else self.np.zeros_like(self._basis_set.h)
+        )
+
+        if not self.has_one_body_time_evolution_operator:
+            return h_0
 
-        return self._time_evolution_operator.h_t(current_time)
+        return h_0 + sum(
+            op.h_t(current_time) for op in self._time_evolution_operator
+        )
 
     def u_t(self, current_time):
-        if self._time_evolution_operator is None:
-            return self._basis_set.u
+        u_0 = (
+            self._basis_set.u
+            if self._add_u_0
+            else self.np.zeros_like(self._basis_set.u)
+        )
 
-        return self._time_evolution_operator.u_t(current_time)
+        if not self.has_two_body_time_evolution_operator:
+            return u_0
+
+        return u_0 + sum(
+            op.u_t(current_time) for op in self._time_evolution_operator
+        )
 
     def transform_one_body_elements(self, h, C, C_tilde=None):
         return self._basis_set.transform_one_body_elements(

quantum_systems/time_evolution_operators/operator.py

@@ -34,11 +34,19 @@ def set_system(self, system):
         Parameters
         ----------
         system : QuantumSystem
-            A QuantumSystem instance to apply the time-evolution operator to.
+            An instance of ``QuantumSystem`` to apply the time-evolution
+            operator to.
+
+        Returns
+        -------
+        self
+            The ``TimeEvolutionOperator``-instance.
         """
 
         self._system = system
 
+        return self
+
     def h_t(self, current_time):
         """Function computing the one-body part of the Hamiltonian for a
         specified time.
@@ -55,7 +63,7 @@ def h_t(self, current_time):
             time-point.
         """
 
-        return self._system.h
+        return 0
 
     def u_t(self, current_time):
         """Function computing the two-body part of the Hamiltonian for a
@@ -73,7 +81,7 @@ def u_t(self, current_time):
             time-point.
         """
 
-        return self._system.u
+        return 0
 
 
 class DipoleFieldInteraction(TimeEvolutionOperator):
@@ -132,9 +140,7 @@ def h_t(self, current_time):
             tmp = self._polarization
             self._polarization = lambda t: tmp
 
-        return self._system.h - self._electric_field_strength(
-            current_time
-        ) * np.tensordot(
+        return -self._electric_field_strength(current_time) * np.tensordot(
             self._polarization(current_time),
             self._system.dipole_moment,
             axes=(0, 0),
@@ -174,4 +180,4 @@ def h_t(self, current_time):
             tmp = self._weight
             self._weight = lambda t: tmp
 
-        return self._system.h + self._weight(current_time) * self._operator
+        return self._weight(current_time) * self._operator

tests/test_time_evolution_operators.py

@@ -0,0 +1,186 @@
+import numpy as np
+
+from quantum_systems import (
+    RandomBasisSet,
+    SpatialOrbitalSystem,
+    GeneralOrbitalSystem,
+)
+from quantum_systems.time_evolution_operators import (
+    DipoleFieldInteraction,
+    CustomOneBodyOperator,
+    AdiabaticSwitching,
+)
+
+
+def test_no_operators():
+    n = 4
+    l = 10
+    dim = 3
+
+    spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
+    gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
+
+    assert not spas.has_one_body_time_evolution_operator
+    assert not gos.has_one_body_time_evolution_operator
+    assert not spas.has_two_body_time_evolution_operator
+    assert not gos.has_two_body_time_evolution_operator
+
+    np.testing.assert_allclose(spas.h_t(10), spas.h)
+    np.testing.assert_allclose(spas.u_t(10), spas.u)
+    np.testing.assert_allclose(gos.h_t(10), gos.h)
+    np.testing.assert_allclose(gos.u_t(10), gos.u)
+
+    spas.set_time_evolution_operator(
+        [],
+        add_h_0=False,
+        add_u_0=False,
+    )
+    gos.set_time_evolution_operator([], add_h_0=False, add_u_0=False)
+
+    assert not spas.has_one_body_time_evolution_operator
+    assert not gos.has_one_body_time_evolution_operator
+    assert not spas.has_two_body_time_evolution_operator
+    assert not gos.has_two_body_time_evolution_operator
+
+    np.testing.assert_allclose(spas.h_t(0), np.zeros_like(spas.h))
+    np.testing.assert_allclose(spas.u_t(0), np.zeros_like(spas.u))
+    np.testing.assert_allclose(gos.h_t(0), np.zeros_like(gos.h))
+    np.testing.assert_allclose(gos.u_t(0), np.zeros_like(gos.u))
+
+
+def test_single_time_evolution_operator():
+    n = 4
+    l = 10
+    dim = 3
+
+    spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
+    gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
+
+    assert not spas.has_one_body_time_evolution_operator
+    assert not gos.has_one_body_time_evolution_operator
+    assert not spas.has_two_body_time_evolution_operator
+    assert not gos.has_two_body_time_evolution_operator
+
+    np.testing.assert_allclose(spas.h_t(10), spas.h)
+    np.testing.assert_allclose(spas.u_t(10), spas.u)
+    np.testing.assert_allclose(gos.h_t(10), gos.h)
+    np.testing.assert_allclose(gos.u_t(10), gos.u)
+
+    spas.set_time_evolution_operator(
+        CustomOneBodyOperator(2, spas.h), add_h_0=False
+    )
+    gos.set_time_evolution_operator(
+        CustomOneBodyOperator(3, gos.h), add_u_0=False
+    )
+
+    assert spas.has_one_body_time_evolution_operator
+    assert gos.has_one_body_time_evolution_operator
+    assert not spas.has_two_body_time_evolution_operator
+    assert not gos.has_two_body_time_evolution_operator
+
+    np.testing.assert_allclose(
+        spas.h_t(0),
+        spas.h * 2,
+    )
+    np.testing.assert_allclose(spas.u_t(0), spas.u)
+
+    np.testing.assert_allclose(
+        gos.h_t(0),
+        gos.h + gos.h * 3,
+    )
+    np.testing.assert_allclose(gos.u_t(0), np.zeros_like(gos.u))
+
+
+def test_single_dipole_time_evolution_operator():
+    n = 4
+    l = 10
+    dim = 3
+
+    omega = 0.25
+
+    spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
+    gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
+
+    field = lambda t: np.sin(omega * 2)
+    polarization = np.zeros(dim)
+    polarization[0] = 1
+
+    spas.set_time_evolution_operator(
+        DipoleFieldInteraction(
+            field,
+            polarization,
+        )
+    )
+    gos.set_time_evolution_operator(
+        DipoleFieldInteraction(
+            field,
+            polarization,
+        )
+    )
+
+    assert spas.has_one_body_time_evolution_operator
+    assert gos.has_one_body_time_evolution_operator
+    assert not spas.has_two_body_time_evolution_operator
+    assert not gos.has_two_body_time_evolution_operator
+
+    for t in [0, 0.1, 0.5, 1.3]:
+        np.testing.assert_allclose(
+            spas.h_t(t),
+            spas.h - field(t) * spas.dipole_moment[0],
+        )
+        np.testing.assert_allclose(
+            gos.h_t(t),
+            gos.h - field(t) * gos.dipole_moment[0],
+        )
+
+        np.testing.assert_allclose(spas.u_t(t), spas.u)
+        np.testing.assert_allclose(gos.u_t(t), gos.u)
+
+
+def test_multiple_time_evolution_operators():
+    n = 4
+    l = 10
+    dim = 3
+
+    spas = SpatialOrbitalSystem(n, RandomBasisSet(l, dim))
+    gos = GeneralOrbitalSystem(n, RandomBasisSet(l, dim))
+
+    assert not spas.has_one_body_time_evolution_operator
+    assert not gos.has_one_body_time_evolution_operator
+    assert not spas.has_two_body_time_evolution_operator
+    assert not gos.has_two_body_time_evolution_operator
+
+    spas.set_time_evolution_operator(
+        [
+            CustomOneBodyOperator(2, spas.h),
+            CustomOneBodyOperator(3, spas.s),
+            AdiabaticSwitching(2),
+        ],
+        add_u_0=False,
+    )
+
+    gos.set_time_evolution_operator(
+        (
+            CustomOneBodyOperator(1, gos.h),
+            CustomOneBodyOperator(3, gos.s),
+            CustomOneBodyOperator(-2, gos.position[0]),
+        ),
+        add_h_0=False,
+    )
+
+    assert spas.has_one_body_time_evolution_operator
+    assert gos.has_one_body_time_evolution_operator
+    assert spas.has_two_body_time_evolution_operator
+    assert not gos.has_two_body_time_evolution_operator
+
+    np.testing.assert_allclose(
+        spas.h_t(0),
+        spas.h + spas.h * 2 + spas.s * 3,
+    )
+    np.testing.assert_allclose(spas.u_t(0), 2 * spas.u)
+
+    np.testing.assert_allclose(
+        gos.h_t(0),
+        gos.h + gos.s * 3 - gos.position[0] * 2,
+    )
+    np.testing.assert_allclose(gos.u_t(0), gos.u)