Merge pull request #47 from BoothGroup/unrestricted_single_prec_df_fix

Fixes for single precision unrestricted density fitting

ebcc/brueckner.py

@@ -109,7 +109,7 @@ def get_rotation_matrix(self, u_tot=None, diis=None, t1=None):
         if t1 is None:
             t1 = self.cc.t1
         if u_tot is None:
-            u_tot = np.eye(self.cc.space.ncorr)
+            u_tot = np.eye(self.cc.space.ncorr, dtype=types[float])
 
         t1_block = np.zeros((self.cc.space.ncorr, self.cc.space.ncorr), dtype=types[float])
         t1_block[: self.cc.space.ncocc, self.cc.space.ncocc :] = -t1
@@ -122,6 +122,7 @@ def get_rotation_matrix(self, u_tot=None, diis=None, t1=None):
             u_tot[:, 0] *= -1
 
         a = scipy.linalg.logm(u_tot)
+        a = a.real.astype(types[float])
         if diis is not None:
             a = diis.update(a, xerr=t1)
 
@@ -258,7 +259,7 @@ def update_coefficients(self, u_tot, mo_coeff, mo_coeff_ref):
 
         u = util.einsum(
             "pq,pi,pj->ij",
-            self.mf.get_ovlp(),
+            self.mf.get_ovlp().astype(types[float]),
             self.mo_to_correlated(self.mf.mo_coeff),
             mo_coeff_new_corr,
         )
@@ -288,8 +289,8 @@ def kernel(self):
         diis.space = self.options.diis_space
 
         # Initialise coefficients:
-        mo_coeff_new = np.array(self.cc.mo_coeff, copy=True)
-        mo_coeff_ref = np.array(self.cc.mo_coeff, copy=True)
+        mo_coeff_new = np.array(self.cc.mo_coeff, copy=True, dtype=types[float])
+        mo_coeff_ref = np.array(self.cc.mo_coeff, copy=True, dtype=types[float])
         mo_coeff_ref = self.mo_to_correlated(mo_coeff_ref)
         u_tot = None
 
@@ -366,8 +367,8 @@ def get_rotation_matrix(self, u_tot=None, diis=None, t1=None):
             t1 = self.cc.t1
         if u_tot is None:
             u_tot = util.Namespace(
-                aa=np.eye(self.cc.space[0].ncorr),
-                bb=np.eye(self.cc.space[1].ncorr),
+                aa=np.eye(self.cc.space[0].ncorr, dtype=types[float]),
+                bb=np.eye(self.cc.space[1].ncorr, dtype=types[float]),
             )
 
         t1_block = util.Namespace(
@@ -398,6 +399,7 @@ def get_rotation_matrix(self, u_tot=None, diis=None, t1=None):
             ],
             axis=0,
         )
+        a = a.real.astype(types[float])
         if diis is not None:
             xerr = np.concatenate([t1.aa.ravel(), t1.bb.ravel()])
             a = diis.update(a, xerr=xerr)
@@ -469,7 +471,7 @@ def update_coefficients(self, u_tot, mo_coeff, mo_coeff_ref):
         )
         mo_coeff_new = self.mo_update_correlated(mo_coeff, mo_coeff_new_corr)
         mo_coeff_mf_corr = self.mo_to_correlated(self.mf.mo_coeff)
-        ovlp = self.mf.get_ovlp()
+        ovlp = self.mf.get_ovlp().astype(types[float])
 
         u = util.Namespace(
             aa=util.einsum("pq,pi,pj->ij", ovlp, mo_coeff_mf_corr[0], mo_coeff_new_corr[0]),

ebcc/cderis.py

@@ -3,7 +3,7 @@
 import numpy as np
 from pyscf import ao2mo
 
-from ebcc import util
+from ebcc import precision, util
 
 
 class CDERIs(util.Namespace):
@@ -84,17 +84,32 @@ def __getattr__(self, key):
                 coeffs = []
                 for i, k in enumerate(key):
                     coeffs.append(self.mo_coeff[i][:, self.slices[i][k]])
-                ijslice = (
-                    0,
-                    coeffs[0].shape[-1],
-                    coeffs[0].shape[-1],
-                    coeffs[0].shape[-1] + coeffs[1].shape[-1],
-                )
-                coeffs = np.concatenate(coeffs, axis=1)
-                block = ao2mo._ao2mo.nr_e2(
-                    self.mf.with_df._cderi, coeffs, ijslice, aosym="s2", mosym="s1"
-                )
-                block = block.reshape(-1, ijslice[1] - ijslice[0], ijslice[3] - ijslice[2])
+                if precision.types[float] == np.float64:
+                    ijslice = (
+                        0,
+                        coeffs[0].shape[-1],
+                        coeffs[0].shape[-1],
+                        coeffs[0].shape[-1] + coeffs[1].shape[-1],
+                    )
+                    coeffs = np.concatenate(coeffs, axis=1)
+                    block = ao2mo._ao2mo.nr_e2(
+                        self.mf.with_df._cderi, coeffs, ijslice, aosym="s2", mosym="s1"
+                    )
+                    block = block.reshape(-1, ijslice[1] - ijslice[0], ijslice[3] - ijslice[2])
+                else:
+                    shape = (
+                        self.mf.with_df.get_naoaux(),
+                        self.mf.with_df.mol.nao_nr(),
+                        self.mf.with_df.mol.nao_nr(),
+                    )
+                    block = util.decompress_axes(
+                        "Qpp",
+                        self.mf.with_df._cderi.astype(precision.types[float]),
+                        include_diagonal=True,
+                        symmetry="+++",
+                        shape=shape,
+                    )
+                    block = util.einsum("Qpq,pi,qj->Qij", block, coeffs[0].conj(), coeffs[1])
                 self.__dict__[key] = block
             return self.__dict__[key]
         else:

ebcc/rebcc.py

@@ -773,7 +773,7 @@ def energy(self, eris=None, amplitudes=None):
             amplitudes=amplitudes,
         )
 
-        return func(**kwargs)
+        return types[float](func(**kwargs).real)
 
     def energy_perturbative(self, eris=None, amplitudes=None, lambdas=None):
         """
@@ -801,7 +801,7 @@ def energy_perturbative(self, eris=None, amplitudes=None, lambdas=None):
             lambdas=lambdas,
         )
 
-        return func(**kwargs)
+        return types[float](func(**kwargs).real)
 
     def update_amps(self, eris=None, amplitudes=None):
         """

ebcc/util.py

@@ -640,8 +640,8 @@ def decompress_axes(
     symmetry : str, optional
         Symmetry of the output array, with a `"+"` indicating symmetry and
         `"-"` indicating antisymmetry for each dimension in the
-        decompressed array. If `None`, defaults to fully symmetric (i.e.
-        all characters are `"+"`). Default value is `None`.
+        decompressed array. If `None`, defaults to fully antisymmetric
+        (i.e. all characters are `"-"`). Default value is `None`.
     out : numpy.ndarray, optional
         Output array. If `None`, a new array is created, and `shape` must
         be passed. Default value is `None`.