[Feature] Time-dependence Port

pyqtorch/hamiltonians/evolution.py

@@ -15,15 +15,18 @@
 from pyqtorch.embed import Embedding
 from pyqtorch.primitives import Primitive
 from pyqtorch.quantum_operation import QuantumOperation
+from pyqtorch.time_dependent.sesolve import sesolve
 from pyqtorch.utils import (
     ATOL,
     Operator,
+    SolverType,
     State,
     StrEnum,
     _round_operator,
     expand_operator,
     finitediff,
     is_diag,
+    is_parametric,
 )
 
 BATCH_DIM = 2
@@ -141,8 +144,9 @@ def __init__(
         generator: TGenerator,
         time: Tensor | str,
         qubit_support: Tuple[int, ...] | None = None,
-        generator_parametric: bool = False,
         cache_length: int = 1,
+        steps: int = 100,
+        solver=SolverType.DP5_SE,
     ):
         """Initializes the HamiltonianEvolution.
         Depending on the generator argument, set the type and set the right generator getter.
@@ -154,6 +158,9 @@ def __init__(
             generator_parametric: Whether the generator is parametric or not.
         """
 
+        self.solver_type = solver
+        self.steps = steps
+
         if isinstance(generator, Tensor):
             if qubit_support is None:
                 raise ValueError(
@@ -178,7 +185,7 @@ def __init__(
                     "Taking support from generator and ignoring qubit_support input."
                 )
             qubit_support = generator.qubit_support
-            if generator_parametric:
+            if is_parametric(generator):
                 generator = [generator]
                 self.generator_type = GeneratorType.PARAMETRIC_OPERATION
             else:
@@ -315,6 +322,55 @@ def spectral_gap(self) -> Tensor:
         spectral_gap = torch.unique(torch.abs(torch.tril(diffs)))
         return spectral_gap[spectral_gap.nonzero()]
 
+    def _forward(
+        self,
+        state: Tensor,
+        values: dict[str, Tensor] | ParameterDict = dict(),
+        embedding: Embedding | None = None,
+    ) -> State:
+        evolved_op = self.tensor(values, embedding)
+        return apply_operator(
+            state=state, operator=evolved_op, qubit_support=self.qubit_support
+        )
+
+    def _forward_time(
+        self,
+        state: Tensor,
+        values: dict[str, Tensor] | ParameterDict = dict(),
+        embedding: Embedding = Embedding(),
+    ) -> State:
+        n_qubits = len(state.shape) - 1
+        batch_size = state.shape[-1]
+        t_grid = torch.linspace(0, float(self.time), self.steps)
+
+        values.update({embedding.tparam_name: torch.tensor(0.0)})  # type: ignore [dict-item]
+        embedded_params = embedding(values)
+
+        def Ht(t: torch.Tensor) -> torch.Tensor:
+            """Accepts a value 't' for time and returns
+            a (2**n_qubits, 2**n_qubits) Hamiltonian evaluated at time 't'.
+            """
+            # We use the origial embedded params and return a new dict
+            # where we reembedded all parameters depending on time with value 't'
+            reembedded_time_values = embedding.reembed_tparam(
+                embedded_params, torch.as_tensor(t)
+            )
+            return (
+                self.generator[0].tensor(reembedded_time_values, embedding).squeeze(2)
+            )
+
+        sol = sesolve(
+            Ht,
+            torch.flatten(state, start_dim=0, end_dim=-2),
+            t_grid,
+            self.solver_type,
+        )
+
+        # Retrieve the last state of shape (2**n_qubits, batch_size)
+        state = sol.states[-1]
+
+        return state.reshape([2] * n_qubits + [batch_size])
+
     def forward(
         self,
         state: Tensor,
@@ -327,17 +383,16 @@ def forward(
         Arguments:
             state: Input state.
             values: Values of parameters.
+            embedding: Embedding of parameters.
 
         Returns:
             The transformed state.
         """
+        if embedding is not None and getattr(embedding, "tparam_name", None):
+            return self._forward_time(state, values, embedding)
 
-        evolved_op = self.tensor(values, embedding)
-        return apply_operator(
-            state=state,
-            operator=evolved_op,
-            qubit_support=self.qubit_support,
-        )
+        else:
+            return self._forward(state, values, embedding)
 
     def tensor(
         self,

pyqtorch/utils.py

@@ -15,11 +15,13 @@
 from numpy import arange, argsort, array, delete, log2
 from numpy import ndarray as NDArray
 from torch import Tensor, moveaxis
+from typing_extensions import TypeAlias
 
+import pyqtorch as pyq
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, DEFAULT_REAL_DTYPE, IMAT
 
-State = Tensor
-Operator = Tensor
+State: TypeAlias = Tensor
+Operator: TypeAlias = Tensor
 
 ATOL = 1e-06
 ATOL_embedding = 1e-03
@@ -741,3 +743,17 @@ class SolverType(StrEnum):
 
     KRYLOV_SE = "krylov_se"
     """Uses Krylov Schrodinger equation solver"""
+
+
+def is_parametric(operation: pyq.Sequence) -> bool:
+    from pyqtorch.primitives import Parametric
+
+    params = []
+    for m in operation.modules():
+        if isinstance(m, (pyq.Scale, Parametric)):
+            params.append(m.param_name)
+
+    res = False
+    if any(isinstance(p, str) for p in params):
+        res = True
+    return res

tests/conftest.py

@@ -9,6 +9,8 @@
 from pytest import FixtureRequest
 from torch import Tensor
 
+from pyqtorch.composite import Add, Scale, Sequence
+from pyqtorch.embed import ConcretizedCallable
 from pyqtorch.noise import (
     AmplitudeDamping,
     BitFlip,
@@ -483,3 +485,38 @@ def hamiltonian_t(t: float, args: Any) -> qutip.Qobj:
         )
 
     return hamiltonian_t
+
+
+@pytest.fixture
+def tparam() -> str:
+    return "t"
+
+
+@pytest.fixture
+def sin(tparam: str) -> tuple[str, ConcretizedCallable]:
+    sin_t, sin_fn = "sin_x", ConcretizedCallable("sin", [tparam])
+    return sin_t, sin_fn
+
+
+@pytest.fixture
+def sq(tparam: str) -> tuple[str, ConcretizedCallable]:
+    t_sq, t_sq_fn = "t_sq", ConcretizedCallable("mul", [tparam, tparam])
+    return t_sq, t_sq_fn
+
+
+@pytest.fixture
+def hamevo_generator(
+    omega: float, param_x: float, param_y: float, sin: tuple, sq: tuple
+) -> Sequence:
+    sin_t, _ = sin
+    t_sq, _ = sq
+    generator = Scale(
+        Add(
+            [
+                Scale(Scale(X(0), sin_t), "y"),
+                Scale(Scale(Y(1), t_sq), param_x),
+            ]
+        ),
+        omega,
+    )
+    return generator

tests/test_analog.py

@@ -8,6 +8,7 @@
 from helpers import calc_mat_vec_wavefunction, random_pauli_hamiltonian
 
 import pyqtorch as pyq
+from pyqtorch.composite import Sequence
 from pyqtorch.hamiltonians import GeneratorType
 from pyqtorch.matrices import (
     DEFAULT_MATRIX_DTYPE,
@@ -19,6 +20,7 @@
 from pyqtorch.utils import (
     ATOL,
     RTOL,
+    SolverType,
     is_normalized,
     operator_kron,
     overlap,
@@ -289,6 +291,46 @@ def test_hamevo_endianness_cnot() -> None:
     assert torch.allclose(wf_cnot, wf_hamevo, rtol=RTOL, atol=ATOL)
 
 
+@pytest.mark.parametrize("ode_solver", [SolverType.DP5_SE, SolverType.KRYLOV_SE])
+def test_timedependent(
+    tparam: str,
+    param_y: float,
+    duration: float,
+    n_steps: int,
+    torch_hamiltonian: Callable,
+    hamevo_generator: Sequence,
+    sin: tuple,
+    sq: tuple,
+    ode_solver: SolverType,
+) -> None:
+
+    psi_start = random_state(2)
+
+    # simulate with time-dependent solver
+    t_points = torch.linspace(0, duration, n_steps)
+    psi_solver = pyq.sesolve(
+        torch_hamiltonian, psi_start.reshape(-1, 1), t_points, ode_solver
+    ).states[-1]
+
+    # simulate with HamiltonianEvolution
+    embedding = pyq.Embedding(
+        tparam_name=tparam,
+        var_to_call={sin[0]: sin[1], sq[0]: sq[1]},
+    )
+    hamiltonian_evolution = pyq.HamiltonianEvolution(
+        generator=hamevo_generator,
+        time=torch.as_tensor(duration),
+        steps=n_steps,
+        solver=ode_solver,
+    )
+    values = {"y": param_y}
+    psi_hamevo = hamiltonian_evolution(
+        state=psi_start, values=values, embedding=embedding
+    ).reshape(-1, 1)
+
+    assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)
+
+
 @pytest.mark.parametrize("n_qubits", [2, 4, 6])
 @pytest.mark.parametrize("batch_size", [1, 2])
 def test_hamevo_parametric_gen(n_qubits: int, batch_size: int) -> None:
@@ -302,9 +344,7 @@ def test_hamevo_parametric_gen(n_qubits: int, batch_size: int) -> None:
 
     param_list.append("t")
 
-    hamevo = pyq.HamiltonianEvolution(
-        generator, tparam, generator_parametric=True, cache_length=2
-    )
+    hamevo = pyq.HamiltonianEvolution(generator, tparam, cache_length=2)
 
     assert hamevo.generator_type == GeneratorType.PARAMETRIC_OPERATION
     assert len(hamevo._cache_hamiltonian_evo) == 0

tests/test_solvers.py

@@ -25,6 +25,7 @@ def test_sesolve(
     qutip_hamiltonian: Callable,
     ode_solver: SolverType,
 ) -> None:
+
     psi0_qutip = qutip.basis(4, 0)
 
     # simulate with torch-based solver

tests/test_tensor.py

@@ -243,7 +243,7 @@ def test_hevo_pauli_tensor(
     assert torch.allclose(psi_star, psi_expected, rtol=RTOL, atol=ATOL)
     # Test the hamiltonian evolution
     tparam = "t"
-    operator = HamiltonianEvolution(generator, tparam, generator_parametric=make_param)
+    operator = HamiltonianEvolution(generator, tparam)
     if make_param:
         assert operator.generator_type == GeneratorType.PARAMETRIC_OPERATION
     else: