From 9680d56b7b9ba13f822535c2c1760fe68203023f Mon Sep 17 00:00:00 2001
From: Charles MOUSSA <charles.moussa@nl-psq-031.local>
Date: Mon, 11 Nov 2024 13:23:22 +0100
Subject: [PATCH] add test density_matrix

---
 pyqtorch/hamiltonians/evolution.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
index a8985b03..0c94f8ab 100644
--- 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),