Merge pull request #53 from BoothGroup/fno

Frozen natural orbitals

FEATURES.md

@@ -2,8 +2,9 @@
 - Lambda equation solver
 - Equation-of-motion solver
 - Density matrices
-- Brueckner orbital calculations
 - Frozen and active space constraints
+- Brueckner orbital calculations
+- Frozen natural orbital calculations
 - Single- and mixed-precision calculations
 
 The following table summarises the available methods and routines for the ansatz currently treated by code generation, in the three spin cases:

ebcc/space.py

@@ -1,6 +1,10 @@
 """Space definition."""
 
+from pyscf.mp import MP2
+
 from ebcc import numpy as np
+from ebcc import util
+from ebcc.precision import types
 
 
 class Space:
@@ -295,3 +299,134 @@ def naocc(self):
     def navir(self):
         """Get the number of virtual active orbitals."""
         return np.sum(self.active_virtual)
+
+
+def construct_default_space(mf):
+    """
+    Construct a default space.
+
+    Parameters
+    ----------
+    mf : pyscf.scf.hf.SCF
+        PySCF mean-field object.
+
+    Returns
+    -------
+    mo_coeff : np.ndarray
+        The molecular orbital coefficients.
+    mo_occ : np.ndarray
+        The molecular orbital occupation numbers.
+    space : Space
+        The default space.
+    """
+
+    def _construct(mo_occ):
+        # Build the default space
+        frozen = np.zeros_like(mo_occ, dtype=bool)
+        active = np.zeros_like(mo_occ, dtype=bool)
+        space = Space(
+            occupied=mo_occ > 0,
+            frozen=frozen,
+            active=active,
+        )
+        return space
+
+    # Construct the default space
+    if np.ndim(mf.mo_occ) == 2:
+        space_a = _construct(mf.mo_occ[0])
+        space_b = _construct(mf.mo_occ[1])
+        space = (space_a, space_b)
+    else:
+        space = _construct(mf.mo_occ)
+
+    return mf.mo_coeff, mf.mo_occ, space
+
+
+def construct_fno_space(mf, occ_tol=1e-5, occ_frac=None, amplitudes=None):
+    """
+    Construct a frozen natural orbital space.
+
+    Parameters
+    ----------
+    mf : pyscf.scf.hf.SCF
+        PySCF mean-field object.
+    occ_tol : float, optional
+        Threshold in the natural orbital occupation numbers. Default
+        value is `1e-5`.
+    occ_frac : float, optional
+        Fraction of the natural orbital occupation numbers to be
+        retained. Overrides `occ_tol` if both are specified. Default
+        value is `None`.
+    amplitudes : Namespace, optional
+        Cluster amplitudes. If provided, use these amplitudes when
+        calculating the MP2 1RDM. Default value is `None`.
+
+    Returns
+    -------
+    no_coeff : np.ndarray
+        The natural orbital coefficients.
+    no_occ : np.ndarray
+        The natural orbital occupation numbers.
+    no_space : Space
+        The frozen natural orbital space.
+    """
+
+    # Get the MP2 1RDM
+    solver = MP2(mf)
+    if amplitudes is None:
+        solver.kernel()
+        dm1 = solver.make_rdm1()
+    else:
+        if isinstance(amplitudes.t2, util.Namespace):
+            t2 = (amplitudes.t2.aaaa, amplitudes.t2.abab, amplitudes.t2.bbbb)
+            dm1 = solver.make_rdm1(t2=t2)
+        else:
+            dm1 = solver.make_rdm1(t2=amplitudes.t2)
+
+    def _construct(dm1, mo_energy, mo_coeff, mo_occ):
+        # Get the number of occupied orbitals
+        nocc = np.sum(mo_occ > 0)
+
+        # Calculate the natural orbitals
+        n, c = np.linalg.eigh(dm1[nocc:, nocc:])
+        n, c = n[::-1], c[:, ::-1]
+
+        # Truncate the natural orbitals
+        if occ_frac is None:
+            active_vir = n > occ_tol
+        else:
+            active_vir = np.cumsum(n / np.sum(n)) <= occ_frac
+        num_active_vir = np.sum(active_vir)
+
+        # Canonicalise the natural orbitals
+        fock_vv = np.diag(mo_energy[nocc:]).astype(types[float])
+        fock_vv = util.einsum("ab,au,bv->uv", fock_vv, c, c)
+        _, c_can = np.linalg.eigh(fock_vv[active_vir][:, active_vir])
+
+        # Transform the natural orbitals
+        no_coeff_avir = np.linalg.multi_dot((mo_coeff[:, nocc:], c[:, :num_active_vir], c_can))
+        no_coeff_fvir = np.dot(mo_coeff[:, nocc:], c[:, num_active_vir:])
+        no_coeff_occ = mo_coeff[:, :nocc]
+        no_coeff = np.hstack((no_coeff_occ, no_coeff_avir, no_coeff_fvir)).astype(types[float])
+
+        # Build the natural orbital space
+        frozen = np.zeros_like(mo_occ, dtype=bool)
+        frozen[nocc + num_active_vir :] = True
+        no_space = Space(
+            occupied=mo_occ > 0,
+            frozen=frozen,
+            active=np.zeros_like(mo_occ, dtype=bool),
+        )
+
+        return no_coeff, no_space
+
+    # Construct the natural orbitals
+    if np.ndim(mf.mo_occ) == 2:
+        no_coeff_a, no_space_a = _construct(dm1[0], mf.mo_energy[0], mf.mo_coeff[0], mf.mo_occ[0])
+        no_coeff_b, no_space_b = _construct(dm1[1], mf.mo_energy[1], mf.mo_coeff[1], mf.mo_occ[1])
+        no_coeff = (no_coeff_a, no_coeff_b)
+        no_space = (no_space_a, no_space_b)
+    else:
+        no_coeff, no_space = _construct(dm1, mf.mo_energy, mf.mo_coeff, mf.mo_occ)
+
+    return no_coeff, mf.mo_occ, no_space

examples/13-fno_ccsd.py

@@ -0,0 +1,26 @@
+"""
+Example of a simple FNO-CCSD calculation.
+"""
+
+import numpy as np
+from pyscf import gto, scf
+
+from ebcc import EBCC
+from ebcc.space import construct_fno_space
+
+# Define the molecule using PySCF
+mol = gto.Mole()
+mol.atom = "H 0 0 0; F 0 0 1.1"
+mol.basis = "cc-pvdz"
+mol.build()
+
+# Run a Hartree-Fock calculation using PySCF
+mf = scf.RHF(mol)
+mf.kernel()
+
+# Construct the FNOs
+no_coeff, no_occ, no_space = construct_fno_space(mf, occ_tol=1e-3)
+
+# Run a FNO-CCSD calculation using EBCC
+ccsd = EBCC(mf, mo_coeff=no_coeff, mo_occ=no_occ, space=no_space)
+ccsd.kernel()

tests/test_RCCSD.py

@@ -13,6 +13,7 @@
 from pyscf import cc, gto, lib, scf
 
 from ebcc import REBCC, NullLogger, Space
+from ebcc.space import construct_fno_space
 
 
 @pytest.mark.reference
@@ -234,6 +235,61 @@ def test_rdm2(self):
     #    self.assertAlmostEqual(e1[0], e2[0], 6)
 
 
+@pytest.mark.reference
+class FNORCCSD_PySCF_Tests(RCCSD_PySCF_Tests):
+    """Test FNO-RCCSD against the PySCF values.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        mol = gto.Mole()
+        mol.atom = "O 0.0 0.0 0.11779; H 0.0 0.755453 -0.471161; H 0.0 -0.755453 -0.471161"
+        #mol.atom = "Li 0 0 0; H 0 0 1.4"
+        mol.basis = "cc-pvdz"
+        mol.verbose = 0
+        mol.build()
+
+        mf = scf.RHF(mol)
+        mf.conv_tol = 1e-12
+        mf.kernel()
+
+        ccsd_ref = cc.FNOCCSD(mf, thresh=1e-3)
+        ccsd_ref.conv_tol = 1e-10
+        ccsd_ref.conv_tol_normt = 1e-14
+        ccsd_ref.max_cycle = 200
+        ccsd_ref.kernel()
+        ccsd_ref.solve_lambda()
+
+        no_coeff, no_occ, no_space = construct_fno_space(mf, occ_tol=1e-3)
+        # Use the PySCF coefficients in case the phases are different
+        no_coeff = ccsd_ref.mo_coeff
+
+        ccsd = REBCC(
+                mf,
+                mo_coeff=no_coeff,
+                mo_occ=no_occ,
+                space=no_space,
+                ansatz="CCSD",
+                log=NullLogger(),
+        )
+        ccsd.options.e_tol = 1e-10
+        eris = ccsd.get_eris()
+        ccsd.kernel(eris=eris)
+        ccsd.solve_lambda(eris=eris)
+
+        cls.mf, cls.ccsd_ref, cls.ccsd, cls.eris = mf, ccsd_ref, ccsd, eris
+
+    def test_rdm1(self):
+        a = self.ccsd_ref.make_rdm1(with_frozen=False)
+        b = self.ccsd.make_rdm1_f(eris=self.eris)
+        np.testing.assert_almost_equal(a, b, 6, verbose=True)
+
+    def test_rdm2(self):
+        a = self.ccsd_ref.make_rdm2(with_frozen=False)
+        b = self.ccsd.make_rdm2_f(eris=self.eris)
+        np.testing.assert_almost_equal(a, b, 6, verbose=True)
+
+
 @pytest.mark.reference
 class RCCSD_PySCF_Frozen_Tests(unittest.TestCase):
     """Test RCCSD against the PySCF values with frozen orbitals.

tests/test_UCCSD.py

@@ -290,6 +290,63 @@ def test_eom_ee(self):
         self.assertAlmostEqual(e1[0], e2[0], 5)
 
 
+# Disabled until PySCF fix bug  # TODO
+#@pytest.mark.reference
+#class FNOUCCSD_PySCF_Tests(UCCSD_PySCF_Tests):
+#    """Test FNO-UCCSD against the PySCF values.
+#    """
+#
+#    @classmethod
+#    def setUpClass(cls):
+#        mol = gto.Mole()
+#        mol.atom = "O 0 0 0; O 0 0 1"
+#        mol.basis = "6-31g"
+#        mol.spin = 2
+#        mol.verbose = 0
+#        mol.build()
+#
+#        mf = scf.RHF(mol)
+#        mf.conv_tol = 1e-12
+#        mf.kernel()
+#        mf = mf.to_uhf()
+#
+#        ccsd_ref = cc.FNOCCSD(mf)
+#        ccsd_ref.conv_tol = 1e-12
+#        ccsd_ref.conv_tol_normt = 1e-12
+#        ccsd_ref.kernel()
+#        ccsd_ref.solve_lambda()
+#
+#        no_coeff, no_occ, no_space = construct_fno_space(mf, occ_tol=1e-3)
+#
+#        ccsd = UEBCC(
+#                mf,
+#                mo_coeff=no_coeff,
+#                mo_occ=no_occ,
+#                space=no_space,
+#                ansatz="CCSD",
+#                log=NullLogger(),
+#        )
+#        ccsd.options.e_tol = 1e-10
+#        eris = ccsd.get_eris()
+#        ccsd.kernel(eris=eris)
+#        ccsd.solve_lambda(eris=eris)
+#
+#        cls.mf, cls.ccsd_ref, cls.ccsd, cls.eris = mf, ccsd_ref, ccsd, eris
+#
+#    def test_rdm1(self):
+#        a = self.ccsd_ref.make_rdm1(with_frozen=False)
+#        b = self.ccsd.make_rdm1_f(eris=self.eris)
+#        np.testing.assert_almost_equal(a[0], b.aa, 6)
+#        np.testing.assert_almost_equal(a[1], b.bb, 6)
+#
+#    def test_rdm2(self):
+#        a = self.ccsd_ref.make_rdm2(with_frozen=False)
+#        b = self.ccsd.make_rdm2_f(eris=self.eris)
+#        np.testing.assert_almost_equal(a[0], b.aaaa, 6)
+#        np.testing.assert_almost_equal(a[1], b.aabb, 6)
+#        np.testing.assert_almost_equal(a[2], b.bbbb, 6)
+
+
 @pytest.mark.reference
 class UCCSD_PySCF_Frozen_Tests(unittest.TestCase):
     """Test UCCSD against the PySCF values with frozen orbitals.