add test density_matrix

pasqal-io · Nov 11, 2024 · 9680d56 · 9680d56
1 parent 5062c1e
commit 9680d56
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
@@ -29,6 +29,8 @@
     finitediff,
     is_diag,
     is_parametric,
+    DensityMatrix,
+    density_mat,
 )
 
 BATCH_DIM = 2
@@ -439,6 +441,8 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
 
             return state.reshape([2] * n_qubits + [batch_size])
         else:
+            if not isinstance(state, DensityMatrix):
+                state = density_mat(state)
             sol = mesolve(
                 Ht,
                 torch.flatten(state, start_dim=0, end_dim=-2),