DMET + frozen orbitals for each fragment

tangelo/problem_decomposition/dmet/dmet_problem_decomposition.py

@@ -26,6 +26,7 @@
 from tangelo.problem_decomposition.dmet.fragment import SecondQuantizedDMETFragment
 from tangelo.algorithms import FCISolver, CCSDSolver, VQESolver
 from tangelo.toolboxes.post_processing.mc_weeny_rdm_purification import mcweeny_purify_2rdm
+from tangelo.toolboxes.molecular_computation.rdms import pad_rdms_with_frozen_orbitals
 
 
 class Localization(Enum):
@@ -56,6 +57,9 @@ class DMETProblemDecomposition(ProblemDecomposition):
             system.
         fragment_solvers (list or str): List of solvers for each fragment. If
             only a string is detected, this solver is used for all fragments.
+        fragment_frozen_orbitals (list): List of frozen orbitals. List of
+            integers freezes the n first orbitals in the fragment. Nested list
+            of integers freezes orbitals based on indices.
         optimizer (function handle): A function defining the classical optimizer
             and its behavior.
         initial_chemical_potential (float) : Initial value for the chemical
@@ -79,6 +83,7 @@ def __init__(self, opt_dict):
                            "electron_localization": Localization.meta_lowdin,
                            "fragment_atoms": list(),
                            "fragment_solvers": "ccsd",
+                           "fragment_frozen_orbitals": list(),
                            "optimizer": self._default_optimizer,
                            "initial_chemical_potential": 0.0,
                            "solvers_options": list(),
@@ -146,6 +151,11 @@ def __init__(self, opt_dict):
         if self.molecule.natm != sum(self.fragment_atoms):
             raise RuntimeError("The number of fragment sites is not equal to the number of atoms in the molecule")
 
+        if not len(self.fragment_frozen_orbitals):
+            self.fragment_frozen_orbitals = [None] * len(self.fragment_atoms)
+        elif len(self.fragment_frozen_orbitals) != len(self.fragment_atoms):
+            raise RuntimeError("The number of elements in the frozen orbitals list does not match the number of fragments.")
+
         # Check that the number of solvers matches the number of fragments.
         # If a single string is detected, it is converted to a list.
         # If a list is detected, it must have the same length than the fragment_atoms.
@@ -419,7 +429,9 @@ def _oneshot_loop(self, chemical_potential, save_results=False, resample=False,
             # We create a dummy SecondQuantizedMolecule with a DMETFragment class.
             # It has the same important attributes and methods to be used with
             # functions of this package.
-            dummy_mol = SecondQuantizedDMETFragment(mol_frag, mf_fragment, fock, fock_frag_copy, t_list, one_ele, two_ele, False)
+            dummy_mol = SecondQuantizedDMETFragment(mol_frag, mf_fragment, fock,
+                fock_frag_copy, t_list, one_ele, two_ele, False,
+                self.fragment_frozen_orbitals[i])
 
             if self.verbose:
                 print("\t\tFragment Number : # ", i + 1)
@@ -467,9 +479,9 @@ def _oneshot_loop(self, chemical_potential, save_results=False, resample=False,
                     self.rdm_measurements[i] = self.solver_fragment_dict[i].rdm_freq_dict
 
             # Compute the fragment energy and sum up the number of electrons
-            fragment_energy, _, one_rdm = self._compute_energy_restricted(mf_fragment, onerdm, twordm,
-                fock_frag_copy, t_list, one_ele, two_ele, fock)
-            n_electron_frag = np.trace(one_rdm[: t_list[0], : t_list[0]])
+            onerdm_padded, twordm_padded = pad_rdms_with_frozen_orbitals(dummy_mol, onerdm, twordm)
+            fragment_energy, onerdm = self._compute_energy_restricted(dummy_mol, onerdm_padded, twordm_padded)
+            n_electron_frag = np.trace(onerdm[: t_list[0], : t_list[0]])
 
             number_of_electron += n_electron_frag
 
@@ -505,7 +517,10 @@ def get_resources(self):
 
             # Unpacking the information for the selected fragment.
             mf_fragment, fock_frag_copy, mol_frag, t_list, one_ele, two_ele, fock = info_fragment
-            dummy_mol = SecondQuantizedDMETFragment(mol_frag, mf_fragment, fock, fock_frag_copy, t_list, one_ele, two_ele, False)
+
+            dummy_mol = SecondQuantizedDMETFragment(mol_frag, mf_fragment, fock,
+                fock_frag_copy, t_list, one_ele, two_ele, False,
+                self.fragment_frozen_orbitals[i])
 
             # Buiding SCF fragments and quantum circuit. Resources are then
             # estimated. For classical sovlers, this functionality is not
@@ -526,109 +541,94 @@ def get_resources(self):
 
         return resources_fragments
 
-    def _compute_energy_restricted(self, mf_frag, onerdm, twordm, fock_frag_copy, t_list, oneint, twoint, fock):
+    def _compute_energy_restricted(self, dmet_fragment, onerdm, twordm):
         """Calculate the fragment energy.
 
         Args:
-            mf_frag (pyscf.scf.RHF): The mean field of the fragment.
+            dmet_fragment (SecondQuantizedDMETFragment): Self-explanatory.
             onerdm (numpy.array): one-particle reduced density matrix (float).
             twordm (numpy.array): two-particle reduced density matrix (float).
-            fock_frag_copy (numpy.array): Fock matrix with the chemical
-                potential subtracted (float).
-            t_list (list): List of number of fragment and bath orbitals (int).
-            oneint (numpy.array): One-electron integrals of fragment (float).
-            twoint (numpy.array): Two-electron integrals of fragment (float).
-            fock (numpy.array): Fock matrix of fragment (float).
 
         Returns:
             float: Fragment energy (fragment_energy).
-            float: Total energy for fragment using RDMs (total_energy_rdm).
-            numpy.array: One-particle RDM for a fragment (one_rdm, float).
+            numpy.array: One-particle RDM for a fragment (onerdm, float).
         """
 
-        # Execute CCSD calculation
+        fock = dmet_fragment.fock
+        t_list = dmet_fragment.t_list
+        oneint = dmet_fragment.one_ele
+        twoint = dmet_fragment.two_ele
+        mo_coeff = dmet_fragment.mean_field.mo_coeff
+
         norb = t_list[0]
 
         # Calculate the one- and two- RDM for DMET energy calculation (Transform to AO basis)
-        one_rdm = mf_frag.mo_coeff @ onerdm @ mf_frag.mo_coeff.T
+        onerdm = mo_coeff @ onerdm @ mo_coeff.T
 
-        twordm = np.einsum("pi,ijkl->pjkl", mf_frag.mo_coeff, twordm)
-        twordm = np.einsum("qj,pjkl->pqkl", mf_frag.mo_coeff, twordm)
-        twordm = np.einsum("rk,pqkl->pqrl", mf_frag.mo_coeff, twordm)
-        twordm = np.einsum("sl,pqrl->pqrs", mf_frag.mo_coeff, twordm)
-
-        # Calculate the total energy based on RDMs
-        total_energy_rdm = np.einsum("ij,ij->", fock_frag_copy, one_rdm) \
-            + 0.5 * np.einsum("ijkl,ijkl->", twoint, twordm)
+        twordm = np.einsum("pi,ijkl->pjkl", mo_coeff, twordm)
+        twordm = np.einsum("qj,pjkl->pqkl", mo_coeff, twordm)
+        twordm = np.einsum("rk,pqkl->pqrl", mo_coeff, twordm)
+        twordm = np.einsum("sl,pqrl->pqrs", mo_coeff, twordm)
 
         # Calculate fragment expectation value
-        fragment_energy_one_rdm = np.einsum("ij,ij->", one_rdm[: norb, :], fock[: norb, :] + oneint[: norb, :]) \
-            + np.einsum("ij,ij->", one_rdm[:, : norb], fock[:, : norb] + oneint[:, : norb])
-        fragment_energy_one_rdm *= 0.25
+        fragment_energy_onerdm = np.einsum("ij,ij->", onerdm[: norb, :], fock[: norb, :] + oneint[: norb, :]) \
+            + np.einsum("ij,ij->", onerdm[:, : norb], fock[:, : norb] + oneint[:, : norb])
+        fragment_energy_onerdm *= 0.25
 
-        fragment_energy_two_rdm = np.einsum("ijkl,ijkl->", twordm[: norb, :, :, :], twoint[: norb, :, :, :]) \
+        fragment_energy_twordm = np.einsum("ijkl,ijkl->", twordm[: norb, :, :, :], twoint[: norb, :, :, :]) \
             + np.einsum("ijkl,ijkl->", twordm[:, : norb, :, :], twoint[:, : norb, :, :]) \
             + np.einsum("ijkl,ijkl->", twordm[:, :, : norb, :], twoint[:, :, : norb, :]) \
             + np.einsum("ijkl,ijkl->", twordm[:, :, :, : norb], twoint[:, :, :, : norb])
-        fragment_energy_two_rdm *= 0.125
+        fragment_energy_twordm *= 0.125
 
-        fragment_energy = fragment_energy_one_rdm + fragment_energy_two_rdm
+        fragment_energy = fragment_energy_onerdm + fragment_energy_twordm
 
-        return fragment_energy, total_energy_rdm, one_rdm
+        return fragment_energy, onerdm
 
-    def _compute_energy_unrestricted(self, mf_frag, onerdm, twordm, fock_frag_copy, t_list, oneint, twoint, fock):
-        """Calculate the fragment energy.
+    def _compute_energy_unrestricted(self, dmet_fragment, onerdm, twordm):
+        """Calculate the fragment energy (unrestricted mean-field).
 
         Args:
-            mf_frag (pyscf.scf.UHF): The mean field of the fragment.
-            onerdm (numpy.array): one-particle reduced density matrix (float64).
-            twordm (numpy.array): two-particle reduced density matrix (float64).
-            fock_frag_copy (numpy.array): Fock matrix with the chemical potential subtracted (float64).
-            t_list (list): List of number of fragment and bath orbitals (int).
-            oneint (numpy.array): One-electron integrals of fragment (float64).
-            twoint (numpy.array): Two-electron integrals of fragment (float64).
-            fock (numpy.array): Fock matrix of fragment (float64).
+            dmet_fragment (SecondQuantizedDMETFragment): Self-explanatory.
+            onerdm (numpy.array): one-particle reduced density matrix (float).
+            twordm (numpy.array): two-particle reduced density matrix (float).
 
         Returns:
             float64: Fragment energy (fragment_energy).
-            float64: Total energy for fragment using RDMs (total_energy_rdm).
             numpy.array: One-particle (alpha) RDM for a fragment (onerdm_a, float64).
             numpy.array: One-particle (beta) RDM for a fragment (onerdm_b, float64).
         """
 
-        # Execute CCSD calculation
+        fock = dmet_fragment.fock
+        t_list = dmet_fragment.t_list
+        oneint = dmet_fragment.one_ele
+        twoint = dmet_fragment.two_ele
+        mo_coeff = dmet_fragment.mean_field.mo_coeff
+
         norb = t_list[0]
 
         # Calculate the one- and two- RDM for DMET energy calculation (Transform to AO basis)
-        one_rdm_a, one_rdm_b = onerdm
-
-        onerdm_a = mf_frag.mo_coeff[0] @ one_rdm_a @ mf_frag.mo_coeff[0].T
-        onerdm_b = mf_frag.mo_coeff[1] @ one_rdm_b @ mf_frag.mo_coeff[1].T
+        onerdm_a, onerdm_b = onerdm
 
-        two_rdm_aa, two_rdm_ab, two_rdm_bb = twordm
+        onerdm_a = mo_coeff[0] @ onerdm_a @ mo_coeff[0].T
+        onerdm_b = mo_coeff[1] @ onerdm_b @ mo_coeff[1].T
 
-        twordm_aa = np.einsum("pi,ijkl->pjkl", mf_frag.mo_coeff[0], two_rdm_aa)
-        twordm_aa = np.einsum("qj,pjkl->pqkl", mf_frag.mo_coeff[0], twordm_aa)
-        twordm_aa = np.einsum("rk,pqkl->pqrl", mf_frag.mo_coeff[0], twordm_aa)
-        twordm_aa = np.einsum("sl,pqrl->pqrs", mf_frag.mo_coeff[0], twordm_aa)
+        twordm_aa, twordm_ab, twordm_bb = twordm
 
-        twordm_ab = np.einsum("pi,ijkl->pjkl", mf_frag.mo_coeff[0], two_rdm_ab)
-        twordm_ab = np.einsum("qj,pjkl->pqkl", mf_frag.mo_coeff[0], twordm_ab)
-        twordm_ab = np.einsum("rk,pqkl->pqrl", mf_frag.mo_coeff[1], twordm_ab)
-        twordm_ab = np.einsum("sl,pqrl->pqrs", mf_frag.mo_coeff[1], twordm_ab)
+        twordm_aa = np.einsum("pi,ijkl->pjkl", mo_coeff[0], twordm_aa)
+        twordm_aa = np.einsum("qj,pjkl->pqkl", mo_coeff[0], twordm_aa)
+        twordm_aa = np.einsum("rk,pqkl->pqrl", mo_coeff[0], twordm_aa)
+        twordm_aa = np.einsum("sl,pqrl->pqrs", mo_coeff[0], twordm_aa)
 
-        twordm_bb = np.einsum("pi,ijkl->pjkl", mf_frag.mo_coeff[1], two_rdm_bb)
-        twordm_bb = np.einsum("qj,pjkl->pqkl", mf_frag.mo_coeff[1], twordm_bb)
-        twordm_bb = np.einsum("rk,pqkl->pqrl", mf_frag.mo_coeff[1], twordm_bb)
-        twordm_bb = np.einsum("sl,pqrl->pqrs", mf_frag.mo_coeff[1], twordm_bb)
+        twordm_ab = np.einsum("pi,ijkl->pjkl", mo_coeff[0], twordm_ab)
+        twordm_ab = np.einsum("qj,pjkl->pqkl", mo_coeff[0], twordm_ab)
+        twordm_ab = np.einsum("rk,pqkl->pqrl", mo_coeff[1], twordm_ab)
+        twordm_ab = np.einsum("sl,pqrl->pqrs", mo_coeff[1], twordm_ab)
 
-        # Calculate the total energy based on RDMs
-        total_energy_rdm = np.einsum("ij,ij->", fock_frag_copy, onerdm_a) \
-            + np.einsum("ij,ij->", fock_frag_copy, onerdm_b) \
-            + 0.5 * np.einsum("ijkl,ijkl->", twoint, twordm_aa) \
-            + 0.5 * np.einsum("ijkl,ijkl->", twoint, twordm_ab) \
-            + 0.5 * np.einsum("ijkl,klij->", twoint, twordm_ab) \
-            + 0.5 * np.einsum("ijkl,ijkl->", twoint, twordm_bb)
+        twordm_bb = np.einsum("pi,ijkl->pjkl", mo_coeff[1], twordm_bb)
+        twordm_bb = np.einsum("qj,pjkl->pqkl", mo_coeff[1], twordm_bb)
+        twordm_bb = np.einsum("rk,pqkl->pqrl", mo_coeff[1], twordm_bb)
+        twordm_bb = np.einsum("sl,pqrl->pqrs", mo_coeff[1], twordm_bb)
 
         # Calculate fragment expectation value
         fragment_energy_one = np.einsum("ij,ij->", onerdm_a[: norb, :], fock[: norb, :] + oneint[: norb, :]) \
@@ -657,7 +657,7 @@ def _compute_energy_unrestricted(self, mf_frag, onerdm, twordm, fock_frag_copy,
 
         fragment_energy = fragment_energy_one + fragment_energy_two
 
-        return fragment_energy, total_energy_rdm, onerdm_a, onerdm_b
+        return fragment_energy, onerdm_a, onerdm_b
 
     def _default_optimizer(self, func, var_params):
         """Function used as a default optimizer for DMET when user does not

tangelo/problem_decomposition/dmet/fragment.py

@@ -20,13 +20,14 @@
 
 import openfermion
 from openfermion.chem.molecular_data import spinorb_from_spatial
-import openfermion.ops.representations as reps
 from openfermion.utils import down_index, up_index
+import openfermion.ops.representations as reps
 import numpy as np
 import pyscf
 from pyscf import ao2mo
 
 from tangelo.toolboxes.qubit_mappings.mapping_transform import get_fermion_operator
+from tangelo.toolboxes.molecular_computation.frozen_orbitals import convert_frozen_orbitals
 
 
 @dataclass
@@ -48,27 +49,70 @@ class SecondQuantizedDMETFragment:
 
     uhf: bool
 
+    frozen_orbitals: list
+
     n_active_electrons: int = field(init=False)
     n_active_sos: int = field(init=False)
     q: int = field(init=False)
     spin: int = field(init=False)
 
     basis: str = field(init=False)
     n_active_mos: int = field(init=False)
-    frozen_mos: None = field(init=False)
 
     def __post_init__(self):
-        self.n_active_electrons = self.molecule.nelectron
-        self.n_active_ab_electrons = self.mean_field.nelec if self.uhf else self.n_active_electrons
+
         self.q = self.molecule.charge
         self.spin = self.molecule.spin
         self.active_spin = self.spin
 
         self.basis = self.molecule.basis
-        self.n_active_mos = len(self.mean_field.mo_energy) if not self.uhf else (len(self.mean_field.mo_energy[0]), len(self.mean_field.mo_energy[1]))
-        self.n_active_sos = 2*self.n_active_mos if not self.uhf else max(2*self.n_active_mos[0], 2*self.n_active_mos[1])
 
-        self.frozen_mos = None
+        self.n_mos = len(self.mean_field.mo_energy)
+        self.mo_occ = self.mean_field.mo_occ
+
+        list_of_active_frozen = convert_frozen_orbitals(self, self.frozen_orbitals)
+        self.active_occupied = list_of_active_frozen[0]
+        self.frozen_occupied = list_of_active_frozen[1]
+        self.active_virtual = list_of_active_frozen[2]
+        self.frozen_virtual = list_of_active_frozen[3]
+
+        if self.uhf:
+            self.active_mos = [self.active_occupied[i]+self.active_virtual[i] for i in range(2)]
+            self.n_active_mos = [len(self.active_mos[0]), len(self.active_mos[1])]
+            self.n_active_sos = max(2*self.n_active_mos[0], 2*self.n_active_mos[1])
+            self.n_active_ab_electrons = (int(sum([self.mo_occ[0][i] for i in self.active_occupied[0]])),
+                int(sum([self.mo_occ[1][i] for i in self.active_occupied[1]])))
+        else:
+            self.active_mos = self.active_occupied + self.active_virtual
+            self.n_active_mos = len(self.active_mos)
+            self.n_active_sos = 2*self.n_active_mos
+
+            n_active_electrons = int(sum([self.mo_occ[i] for i in self.active_occupied]))
+            n_alpha = n_active_electrons//2 + self.spin//2 + (n_active_electrons % 2)
+            n_beta = n_active_electrons//2 - self.spin//2
+            self.n_active_ab_electrons = (n_alpha, n_beta)
+        self.n_active_electrons = sum(self.n_active_ab_electrons)
+
+    @property
+    def frozen_mos(self):
+        """This property returns MOs indexes for the frozen orbitals. It was
+        written to take into account if one of the two possibilities (occ or
+        virt) is None. In fact, list + None, None + list or None + None return
+        an error. An empty list cannot be sent because PySCF mp2 returns
+        "IndexError: list index out of range".
+
+        Returns:
+            list: MOs indexes frozen (occupied + virtual).
+        """
+        if self.frozen_occupied and self.frozen_virtual:
+            return (self.frozen_occupied + self.frozen_virtual if not self.uhf else
+                    [self.frozen_occupied[0] + self.frozen_virtual[0], self.frozen_occupied[1] + self.frozen_virtual[1]])
+        elif self.frozen_occupied:
+            return self.frozen_occupied
+        elif self.frozen_virtual:
+            return self.frozen_virtual
+        else:
+            return None
 
     @property
     def fermionic_hamiltonian(self):
@@ -102,6 +146,12 @@ def _fermionic_hamiltonian_restricted(self):
         # made before performing the truncation for frozen orbitals.
         two_electron_integrals = two_electron_integrals.transpose(0, 2, 3, 1)
 
+        core_offset, one_electron_integrals, two_electron_integrals = reps.get_active_space_integrals(
+            one_electron_integrals, two_electron_integrals, self.frozen_occupied, self.active_mos)
+
+        # Adding frozen electron contribution to core constant.
+        core_constant += core_offset
+
         one_body_coefficients, two_body_coefficients = spinorb_from_spatial(one_electron_integrals, two_electron_integrals)
         fragment_hamiltonian = reps.InteractionOperator(core_constant, one_body_coefficients, 0.5 * two_body_coefficients)
 

tangelo/problem_decomposition/tests/dmet/test_dmet.py

@@ -245,6 +245,37 @@ def test_retrieving_quantum_data(self):
         self.assertAlmostEqual(np.trace(rho), n_elec * (n_elec - 1), delta=1e-3,
                                msg="Trace of two_rdm does not match n_elec * (n_elec-1)")
 
+    def test_dmet_frozen_orbitals(self):
+        """Tests the DMET energy for an H10 ring in 3-21G with frozen orbitals."""
+
+        opt_dmet = {"molecule": mol_H10_321g,
+                    "fragment_atoms": [1]*10,
+                    "fragment_solvers": "fci",
+                    # Make every fragment a 2 level problem in this basis.
+                    "fragment_frozen_orbitals": [[0, 3, 4, 5]]*10,
+                    "verbose": False
+                    }
+
+        solver = DMETProblemDecomposition(opt_dmet)
+        solver.build()
+        energy = solver.simulate()
+        self.assertAlmostEqual(energy, -4.41503, places=4)
+
+    def test_dmet_wrong_number_frozen_orbitals(self):
+        """Tests if the program raises the error when the number of frozen
+        orbital elements is not equal to the number of fragment.
+        """
+
+        opt_dmet = {"molecule": mol_H10_321g,
+                    "fragment_atoms": [1]*10,
+                    "fragment_solvers": "fci",
+                    "fragment_frozen_orbitals": [[0, 3, 4, 5]]*9,
+                    "verbose": False
+                    }
+
+        with self.assertRaises(RuntimeError):
+            DMETProblemDecomposition(opt_dmet)
+
 
 if __name__ == "__main__":
     unittest.main()