Merge branch 'dev' into lassi_clib

examples/laspdft/c2h4n4_si_laspdft.py

@@ -30,8 +30,8 @@
 lsi.kernel()
 
 # LASSI-PDFT
-mc = mcpdft.LASSI(lsi, 'tPBE', (3, 3), ((2,1),(1,2)), states=[0, 1])
-mc.kernel() 
+mc = mcpdft.LASSI(lsi, 'tPBE', states=[0, 1])
+mc.kernel()
 
 # CASCI-PDFT in las orbitals
 from pyscf import mcpdft

my_pyscf/mcpdft/__init__.py

@@ -107,8 +107,10 @@ def LASSCFPDFT(mc_or_mf_or_mol, ot, ncas_sub, nelecas_sub,  ncore=None, spin_sub
     return _laspdftEnergy(LASSCF,  mc_or_mf_or_mol, ot, ncas_sub, nelecas_sub, ncore=ncore,
                           spin_sub=spin_sub, frozen=frozen, **kwargs)
 
-def LASSIPDFT(mc_or_mf_or_mol, ot, ncas_sub, nelecas_sub, ncore=None, spin_sub=None, frozen=None,
-        states=None, **kwargs):
+def LASSIPDFT(mc_or_mf_or_mol, ot, ncas_sub=None, nelecas_sub=None, ncore=None, spin_sub=None,
+              frozen=None, states=None, **kwargs):
+    if ncas_sub is None: ncas_sub = getattr (mc_or_mf_or_mol, 'ncas_sub', None)
+    if nelecas_sub is None: nelecas_sub = getattr (mc_or_mf_or_mol, 'nelecas_sub', None)
     from mrh.my_pyscf.mcscf.lasscf_o0 import LASSCF
     return _lassipdftEnergy(LASSCF,  mc_or_mf_or_mol, ot, ncas_sub, nelecas_sub, DoLASSI=True, ncore=ncore,
                           spin_sub=spin_sub, frozen=frozen, states=states, **kwargs)

my_pyscf/mcpdft/laspdft.py

@@ -133,24 +133,36 @@ def compute_pdft_energy_(self, mo_coeff=None, ci=None, ot=None, otxc=None,
         if _mc_class.DoLASSI:
 
             '''
-            Current RDM function for LASSI is generating the rdm1 and 2 for all the states.
-            The cost of this function is similar to LASSI diagonalization step. Therefore,
+            The cost of the RDM build is similar to LASSI diagonalization step. Therefore,
             calling it 2n time for n-states becomes prohibitively expensive. One alternative 
             can be just call it once and store all the generated casdm1 and casdm2 and later on
             just call a reader function which will read the rdms from this temp file.
-            I have to make sure to delete or close this tempfile after the calculation, I 
-            will do that later.
             '''
             def _store_rdms(self):
-                rdm1s, rdm2s = lassi.roots_make_rdm12s(self, self.ci, self.si)
+                # MRH: I made it loop over blocks of states to handle the O(N^5) memory cost
+                # If there's enough memory it'll still do them all at once
+                log = lib.logger.new_logger (self, self.verbose)
+                mem_per_state = (2*(self.ncas**2) + 4*(self.ncas**4)) / 1e6
+                current_mem = lib.current_memory ()[0]
+                if current_mem > self.max_memory:
+                    log.warn ("Current memory usage (%d MB) exceeds maximum memory (%d MB)",
+                              current_mem, self.max_memory)
+                    nblk = 1
+                else:
+                    nblk = int ((self.max_memory - current_mem) / mem_per_state)
                 rdmstmpfile = self.rdmstmpfile
                 with h5py.File(rdmstmpfile, 'w') as f:
-                    for i in range(len(self.e_states)):
-                        rdm1s_dname = f'rdm1s_{i}'
-                        f.create_dataset(rdm1s_dname, data=rdm1s[i])
-                        rdm2s_dname = f'rdm2s_{i}'
-                        f.create_dataset(rdm2s_dname, data=rdm2s[i])
-
+                    for i in range (0, len (self.e_states), nblk):
+                        j = min (i+nblk, len (self.e_states))
+                        rdm1s, rdm2s = lassi.root_make_rdm12s(self, self.ci, self.si,
+                                                              state=list(range(i,j)))
+                        for k in range (i, j):
+                            rdm1s_dname = f'rdm1s_{k}'
+                            f.create_dataset(rdm1s_dname, data=rdm1s[k])
+                            rdm2s_dname = f'rdm2s_{k}'
+                            f.create_dataset(rdm2s_dname, data=rdm2s[k])
+                        rdm1s = rdm2s = None     
+
             # # This code doesn't seem efficent, have to calculate the casdm1 and casdm2 in different functions.
             # def make_one_casdm1s(self, ci=None, state=0, **kwargs):
                 # with lib.temporary_env (self, verbose=2):