Permit nfrag>32 LASSI-PDFT (#46, #131). Add test.

my_pyscf/lassi/op_o1/rdm.py

@@ -119,20 +119,18 @@ def sprint_profile (self):
         return profile
 
     def get_single_rootspace_sivec (self, iroot):
-        '''A single-rootspace slice of the SI vectors, reshaped to expose the lroots.
+        '''A single-rootspace slice of the SI vectors.
 
         Args:
             iroot: integer
                 Rootspace index
 
         Returns:
-            sivec: col-major ndarray of shape (lroots[0,iroot], lroots[1,iroot], ...,
-                                               nroots_si)
+            sivec: col-major ndarray of shape (np.prod (lroots[:,iroot], nroots_si)
                 SI vectors
         '''
         i, j = self.offs_lroots[iroot]
-        vecshape = list (self.lroots[:,iroot]) + [self.nroots_si,]
-        return self.si[i:j,:].reshape (vecshape, order='F')
+        return self.si[i:j,:]
 
     _lowertri = True
 
@@ -151,8 +149,18 @@ def get_frag_transposed_sivec (self, iroot, *inv):
                 SI vectors with the faster dimension iterating over states of fragments not in
                 inv and the slower dimension iterating over states of fragments in inv 
         '''
-        axesorder = [i for i in range (self.nfrags) if not (i in inv)] + list (inv) + [self.nfrags,]
-        sivec = self.get_single_rootspace_sivec (iroot).transpose (*axesorder)
+        sivec = self.get_single_rootspace_sivec (iroot)
+        rdim = self.nroots_si
+        lroots = self.lroots[:,iroot].copy ()
+        for i in list (inv)[::-1]:
+            lrl = np.prod (lroots[:i])
+            lrc = lroots[i]
+            lrr = np.prod (lroots[i+1:])
+            sivec = sivec.reshape ((lrl,lrc,lrr,rdim), order='F')
+            sivec = sivec.transpose (0,2,1,3)
+            sivec = sivec.reshape ((lrl*lrr,lrc*rdim), order='F')
+            rdim = rdim * lrc
+            lroots[i] = 1
         nprods = np.prod (self.lroots[:,iroot])
         nrows = np.prod (self.lroots[inv,iroot])
         ncols = nprods // nrows

tests/lassi/test_33frag.py

@@ -0,0 +1,41 @@
+import copy
+import unittest
+import numpy as np
+from scipy import linalg
+from pyscf import lib, gto, scf
+from pyscf.lo import orth
+from mrh.my_pyscf.mcscf import lasscf_o0
+from mrh.my_pyscf.lassi import LASSIrq, LASSIS, lassi
+from mrh.my_pyscf import mcpdft
+
+# The purpose of this test is to make sure the LASSI kernel doesn't try to build an array with
+# nfrags dimensions at any point in the calculation.
+
+def setUpModule():
+    global las
+    atom = '\n'.join (['H {} 0 0'.format (i) for i in range (33)])
+    mol = gto.M (atom=atom, basis='6-31g', symmetry=False, spin=33, verbose=0, output='/dev/null')
+    mf = scf.RHF (mol)
+    nlas = [2,]*33
+    spin_sub = [2,]*33
+    nelelas = [[1,0],]*33
+    las = lasscf_o0.LASSCF (mf, nlas, nelelas, spin_sub=spin_sub)
+    las.max_cycle_macro = 1
+    las.mo_coeff = orth.orth_ao (mol, 'meta_lowdin', s=mf.get_ovlp ())
+    las.lasci ()
+
+def tearDownModule():
+    global las
+    del las
+
+class KnownValues(unittest.TestCase):
+
+    def test_lassi01 (self):
+        lsi = LASSIrq (las, 0, 1).run ()
+        rdm1s, rdm2s = lassi.root_make_rdm12s(lsi, lsi.ci, lsi.si, state=0)
+
+if __name__ == "__main__":
+    print("Full Tests for LASSI 33-fragment calculation")
+    unittest.main()
+
+