[Feature] Allow time-dependent duration to be passed at runtime

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 name = "pyqtorch"
 description = "An efficient, large-scale emulator designed for quantum machine learning, seamlessly integrated with a PyTorch backend. Please refer to https://pyqtorch.readthedocs.io/en/latest/ for setup and usage info, along with the full documentation."
 readme = "README.md"
-version = "1.5.1"
+version = "1.5.2"
 requires-python = ">=3.8,<3.13"
 license = { text = "Apache 2.0" }
 keywords = ["quantum"]

pyqtorch/hamiltonians/evolution.py

@@ -119,6 +119,8 @@ def evolve(hamiltonian: Operator, time_evolution: Tensor) -> Operator:
     return torch.transpose(evol_operator, 0, -1)
 
 
+# FIXME: Review HamiltonianEvolution inheriting from Sequence.
+# Probably makes more sense to inherit from Parametric.
 class HamiltonianEvolution(Sequence):
     """
     The HamiltonianEvolution corresponds to :math:`t`, returns :math:`exp(-i H, t)` where
@@ -146,7 +148,7 @@ def __init__(
         time: Tensor | str | ConcretizedCallable,
         qubit_support: Tuple[int, ...] | None = None,
         cache_length: int = 1,
-        duration: float | Tensor = 1.0,
+        duration: Tensor | str | float | None = None,
         steps: int = 100,
         solver=SolverType.DP5_SE,
     ):
@@ -165,17 +167,21 @@ def __init__(
         self.solver_type = solver
         self.steps = steps
         self.duration = duration
-        self.is_time_dependent = None
+
+        if isinstance(duration, (str, float, Tensor)) or duration is None:
+            self.duration = duration
+        else:
+            raise ValueError(
+                "Optional argument `duration` should be passed as str, float or Tensor."
+            )
 
         if isinstance(time, (str, Tensor, ConcretizedCallable)):
             self.time = time
         else:
             raise ValueError(
-                "time should be passed as str, Tensor or ConcretizedCallable."
+                "Argument `time` should be passed as str, Tensor or ConcretizedCallable."
             )
 
-        self.has_time_param = self._has_time_param(generator)
-
         if isinstance(generator, Tensor):
             if qubit_support is None:
                 raise ValueError(
@@ -246,21 +252,25 @@ def generator(self) -> ModuleList:
         """
         return self.operations
 
+    @cached_property
+    def is_time_dependent(self) -> bool:
+        return self._is_time_dependent(self.generator)
+
     def flatten(self) -> ModuleList:
         return ModuleList([self])
 
     @property
     def param_name(self) -> Tensor | str:
         return self.time
 
-    def _has_time_param(self, generator: TGenerator) -> bool:
+    def _is_time_dependent(self, generator: TGenerator) -> bool:
         from pyqtorch.primitives import Parametric
 
         res = False
         if isinstance(self.time, Tensor):
             return res
         else:
-            if isinstance(generator, (Sequence, QuantumOperation)):
+            if isinstance(generator, (Sequence, QuantumOperation, ModuleList)):
                 for m in generator.modules():
                     if isinstance(m, (Scale, Parametric)):
                         if self.time in getattr(m.param_name, "independent_args", []):
@@ -368,8 +378,10 @@ def _forward_time(
         values = values or dict()
         n_qubits = len(state.shape) - 1
         batch_size = state.shape[-1]
-        t_grid = torch.linspace(0, float(self.duration), self.steps)
-
+        duration = (
+            values[self.duration] if isinstance(self.duration, str) else self.duration
+        )
+        t_grid = torch.linspace(0, float(duration), self.steps)
         if embedding is not None:
             values.update({embedding.tparam_name: torch.tensor(0.0)})  # type: ignore [dict-item]
             embedded_params = embedding(values)
@@ -390,7 +402,13 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 values[self.time] = torch.as_tensor(t)
                 reembedded_time_values = values
             return (
-                self.generator[0].tensor(reembedded_time_values, embedding).squeeze(2)
+                self.generator[0]
+                .tensor(
+                    reembedded_time_values,
+                    embedding,
+                    full_support=tuple(range(n_qubits)),
+                )
+                .squeeze(2)
             )
 
         sol = sesolve(
@@ -423,7 +441,7 @@ def forward(
             The transformed state.
         """
         values = values or dict()
-        if self.has_time_param or (
+        if self.is_time_dependent or (
             embedding is not None and getattr(embedding, "tparam_name", None)
         ):
             return self._forward_time(state, values, embedding)  # type: ignore [arg-type]

tests/test_analog.py

@@ -292,6 +292,7 @@ def test_hamevo_endianness_cnot() -> None:
     assert torch.allclose(wf_cnot, wf_hamevo, rtol=RTOL, atol=ATOL)
 
 
+@pytest.mark.parametrize("duration", [torch.rand(1), "duration"])
 @pytest.mark.parametrize("ode_solver", [SolverType.DP5_SE, SolverType.KRYLOV_SE])
 def test_timedependent(
     tparam: str,
@@ -307,8 +308,10 @@ def test_timedependent(
 
     psi_start = random_state(2)
 
+    dur_val = duration if isinstance(duration, torch.Tensor) else torch.rand(1)
+
     # simulate with time-dependent solver
-    t_points = torch.linspace(0, duration, n_steps)
+    t_points = torch.linspace(0, dur_val[0], n_steps)
     psi_solver = pyq.sesolve(
         torch_hamiltonian, psi_start.reshape(-1, 1), t_points, ode_solver
     ).states[-1]
@@ -325,7 +328,7 @@ def test_timedependent(
         steps=n_steps,
         solver=ode_solver,
     )
-    values = {"y": param_y}
+    values = {"y": param_y, "duration": dur_val}
     psi_hamevo = hamiltonian_evolution(
         state=psi_start, values=values, embedding=embedding
     ).reshape(-1, 1)
@@ -420,4 +423,4 @@ def apply_hamevo_and_compare_expected(psi, values):
 )
 def test_hamevo_is_time_dependent_generator(generator, time_param, result) -> None:
     hamevo = HamiltonianEvolution(generator, time_param)
-    assert hamevo.has_time_param == result
+    assert hamevo.is_time_dependent == result