adding analog noise

pyqtorch/hamiltonians/evolution.py

@@ -16,6 +16,7 @@
 from pyqtorch.embed import ConcretizedCallable, Embedding
 from pyqtorch.primitives import Primitive
 from pyqtorch.quantum_operation import QuantumOperation
+from pyqtorch.time_dependent.mesolve import mesolve
 from pyqtorch.time_dependent.sesolve import sesolve
 from pyqtorch.utils import (
     ATOL,
@@ -140,6 +141,11 @@ class HamiltonianEvolution(Sequence):
         operations: List of operations.
         cache_length: LRU cache cache_length evolution operators for given set
                     of parameter values.
+        duration: Total duration for evolving when using a solver.
+        steps: Number of steps to use when using solver.
+        solver: Time-dependent Lindblad master equation solver.
+        noise_operators: List of tensors or Kraus oeprators adding analog noise
+            when solving with a Shrodinger equation solver.
     """
 
     def __init__(
@@ -151,6 +157,7 @@ def __init__(
         duration: Tensor | str | float | None = None,
         steps: int = 100,
         solver=SolverType.DP5_SE,
+        noise_operators: list[Tensor] = list(),
     ):
         """Initializes the HamiltonianEvolution.
         Depending on the generator argument, set the type and set the right generator getter.
@@ -161,7 +168,13 @@ def __init__(
             qubit_support: The qubits the operator acts on. If generator is a quantum
                 operation or sequence of operations,
                 it will be inferred from the generator.
-            generator_parametric: Whether the generator is parametric or not.
+            cache_length: LRU cache cache_length evolution operators for given set
+                    of parameter values.
+            duration: Total duration for evolving when using a solver.
+            steps: Number of steps to use when using solver.
+            solver: Time-dependent Lindblad master equation solver.
+            noise_operators: List of tensors or Kraus oeprators adding analog noise
+                when solving with a Shrodinger equation solver.
         """
 
         self.solver_type = solver
@@ -243,6 +256,8 @@ def __init__(
         self._cache_hamiltonian_evo: dict[str, Tensor] = dict()
         self.cache_length = cache_length
 
+        self.noise_operators: list[Tensor] = noise_operators
+
     @property
     def generator(self) -> ModuleList:
         """Returns the operations making the generator.
@@ -411,17 +426,28 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 .squeeze(2)
             )
 
-        sol = sesolve(
-            Ht,
-            torch.flatten(state, start_dim=0, end_dim=-2),
-            t_grid,
-            self.solver_type,
-        )
+        if len(self.noise_operators) == 0:
+            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]
+            # Retrieve the last state of shape (2**n_qubits, batch_size)
+            state = sol.states[-1]
 
-        return state.reshape([2] * n_qubits + [batch_size])
+            return state.reshape([2] * n_qubits + [batch_size])
+        else:
+            sol = mesolve(
+                Ht,
+                torch.flatten(state, start_dim=0, end_dim=-2),
+                self.noise_operators,
+                t_grid,
+                self.solver_type,
+            )
+            state = sol.states[-1]
+            return state.reshape([2] * n_qubits * 2 + [batch_size])
 
     def forward(
         self,

pyqtorch/noise/gates.py

@@ -9,7 +9,12 @@
 from pyqtorch.apply import apply_operator_dm
 from pyqtorch.embed import Embedding
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, XMAT, YMAT, ZMAT
-from pyqtorch.utils import DensityMatrix, density_mat, qubit_support_as_tuple
+from pyqtorch.utils import (
+    DensityMatrix,
+    density_mat,
+    promote_operator,
+    qubit_support_as_tuple,
+)
 
 
 class Noise(torch.nn.Module):
@@ -35,11 +40,12 @@ def extra_repr(self) -> str:
     def kraus_operators(self) -> list[Tensor]:
         return [getattr(self, f"kraus_{i}") for i in range(len(self._buffers))]
 
-    def tensor(
-        self,
-    ) -> list[Tensor]:
+    def tensor(self, n_qubit_support: int | None = None) -> list[Tensor]:
         # Since PyQ expects tensor.Size = [2**n_qubits, 2**n_qubits,batch_size].
-        return [kraus_op.unsqueeze(2) for kraus_op in self.kraus_operators]
+        t_ops = [kraus_op.unsqueeze(2) for kraus_op in self.kraus_operators]
+        if n_qubit_support is None:
+            return t_ops
+        return [promote_operator(t, self.target, n_qubit_support) for t in t_ops]
 
     def forward(
         self,

tests/test_analog.py

@@ -18,10 +18,12 @@
     XMAT,
     ZMAT,
 )
+from pyqtorch.noise import Depolarizing
 from pyqtorch.utils import (
     ATOL,
     RTOL,
     SolverType,
+    density_mat,
     is_normalized,
     operator_kron,
     overlap,
@@ -312,6 +314,7 @@ def test_timedependent(
 
     # simulate with time-dependent solver
     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]
@@ -336,6 +339,51 @@ def test_timedependent(
     assert torch.allclose(psi_solver, psi_hamevo, rtol=RTOL, atol=ATOL)
 
 
+@pytest.mark.parametrize("duration", [torch.rand(1), "duration"])
+def test_timedependent_with_noise(
+    tparam: str,
+    param_y: float,
+    duration: float,
+    n_steps: int,
+    torch_hamiltonian: Callable,
+    hamevo_generator: Sequence,
+    sin: tuple,
+    sq: tuple,
+) -> None:
+
+    psi_start = density_mat(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, dur_val[0], n_steps)
+
+    list_ops = Depolarizing(0, error_probability=0.1).tensor(2)
+    solver = SolverType.DP5_ME
+    psi_solver = pyq.mesolve(
+        torch_hamiltonian, psi_start.reshape(-1, 1), list_ops, t_points, 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=tparam,
+        duration=duration,
+        steps=n_steps,
+        solver=solver,
+        noise_operators=list_ops,
+    )
+    values = {"y": param_y, "duration": dur_val}
+    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: