oo more flexible (#366)

tequilahub · Sep 26, 2024 · 9d860ff · 9d860ff
1 parent 2262ce3
commit 9d860ff
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/tequila/quantumchemistry/orbital_optimizer.py b/src/tequila/quantumchemistry/orbital_optimizer.py
@@ -208,11 +208,16 @@ def kernel(self, h1, h2, *args, **kwargs):
             H = molecule.make_hardcore_boson_hamiltonian()
         else:
             H = molecule.make_hamiltonian()
+
+        rdm1 = None
+        rdm2 = None
         if self.vqe_solver is not None:
             vqe_solver_arguments = {}
             if self.vqe_solver_arguments is not None:
                 vqe_solver_arguments = self.vqe_solver_arguments
             result = self.vqe_solver(H=H, circuit=self.circuit, molecule=molecule, **vqe_solver_arguments)
+            if hasattr(self.vqe_solver, "compute_rdms"):
+                rdm1, rdm2 = self.vqe_solver.compute_rdms(U=self.circuit, variables=result.variables, molecule=molecule, use_hcb=restrict_to_hcb)
         elif self.circuit is None:
             raise Exception("Orbital Optimizer: Either provide a callable vqe_solver or a circuit")
         else:
@@ -233,8 +238,9 @@ def kernel(self, h1, h2, *args, **kwargs):
             # static ansatz
             U = self.circuit
 
-        rdm1, rdm2 = molecule.compute_rdms(U=U, variables=result.variables, spin_free=True, get_rdm1=True, get_rdm2=True, use_hcb=restrict_to_hcb)
-        rdm2 = self.reorder(rdm2, 'dirac', 'mulliken')
+        if rdm1 is None or rdm2 is None:
+            rdm1, rdm2 = molecule.compute_rdms(U=U, variables=result.variables, spin_free=True, get_rdm1=True, get_rdm2=True, use_hcb=restrict_to_hcb)
+            rdm2 = self.reorder(rdm2, 'dirac', 'mulliken')
         if not self.silent:
             print("{:20} : {}".format("energy", result.energy))
             if len(self.history) > 0:
@@ -259,3 +265,4 @@ def __str__(self):
             else:
                 result += "{:30} : {}\n".format(k, v)
         return result
+