add memmap to avoid storing 2e-integ in DRAM

PyDuck.py

@@ -6,6 +6,8 @@
 from pyscf import gto
 import numpy as np
 import subprocess
+import time
+
 
 #Find the value of the environnement variable QUACK_ROOT. If not present we use the current repository
 if "QUACK_ROOT" not in os.environ:
@@ -23,6 +25,7 @@
 parser.add_argument('-c', '--charge', type=int, default=0, help='Total charge of the molecule. Specify negative charges with "m" instead of the minus sign, for example m1 instead of -1. Default is 0')
 parser.add_argument('--cartesian', default=False, action='store_true', help='Add this option if you want to use cartesian basis functions.')
 parser.add_argument('--print_2e', default=False, action='store_true', help='Add this option if you want to print 2e-integrals.')
+parser.add_argument('--mmap_2e', default=False, action='store_true', help='If True, avoid using DRAM when generating 2e-integrals.')
 parser.add_argument('-fc', '--frozen_core', type=bool, default=False, help='Freeze core MOs. Default is false')
 parser.add_argument('-m', '--multiplicity', type=int, default=1, help='Spin multiplicity. Default is 1 therefore singlet')
 parser.add_argument('--working_dir', type=str, default=QuAcK_dir, help='Set a working directory to run the calculation.')
@@ -38,6 +41,7 @@
 xyz=args.xyz + '.xyz'
 cartesian=args.cartesian
 print_2e=args.print_2e
+mmap_2e=args.mmap_2e
 working_dir=args.working_dir
 
 #Read molecule
@@ -129,33 +133,47 @@ def write_matrix_to_file(matrix,size,file,cutoff=1e-15):
 subprocess.call(['rm', '-f', working_dir + '/int/z.dat'])
 write_matrix_to_file(z,norb,working_dir+'/int/z.dat')
 
-eri_ao = mol.intor('int2e')
-
-def write_tensor_to_file(tensor,size,file,cutoff=1e-15):
-    f = open(file, 'w')
+def write_tensor_to_file(tensor,size,file_name,cutoff=1e-15):
+    f = open(file_name, 'w')
     for i in range(size):
         for j in range(i,size):
             for k in range(i,size):
                 for l in range(j,size):
                     if abs(tensor[i][k][j][l]) > cutoff:
-                        #f.write(str(i+1)+' '+str(j+1)+' '+str(k+1)+' '+str(l+1)+' '+"{:.16E}".format(tensor[i][k][j][l]))
                         f.write(str(i+1)+' '+str(j+1)+' '+str(k+1)+' '+str(l+1)+' '+"{:.16E}".format(tensor[i][k][j][l]))
                         f.write('\n')
     f.close()
 
-# Write two-electron integrals
+# Write two-electron integrals to HD
+ti_2e = time.time()
 if print_2e:
     # (formatted)
-    subprocess.call(['rm', '-f', working_dir + '/int/ERI.dat'])
-    write_tensor_to_file(eri_ao, norb, working_dir + '/int/ERI.dat')
+    output_file_path = working_dir + '/int/ERI.dat'
+    subprocess.call(['rm', '-f', output_file_path])
+    eri_ao = mol.intor('int2e')
+    write_tensor_to_file(eri_ao, norb, output_file_path)
 else:
     # (binary)
-    subprocess.call(['rm', '-f', working_dir + '/int/ERI.bin'])
-    # chem -> phys notation
-    eri_ao = eri_ao.transpose(0, 2, 1, 3)
-    f = open(working_dir + '/int/ERI.bin', 'w')
-    eri_ao.tofile(working_dir + '/int/ERI.bin')
-    f.close()
+    output_file_path = working_dir + '/int/ERI.bin'
+    subprocess.call(['rm', '-f', output_file_path])
+    if(mmap_2e):
+        # avoid using DRAM
+        eri_shape = (norb, norb, norb, norb)
+        eri_mmap = np.memmap(output_file_path, dtype='float64', mode='w+', shape=eri_shape)
+        mol.intor('int2e', out=eri_mmap)
+        for i in range(norb):
+            transposed_chunk = eri_mmap[i, :, :, :].transpose(1, 0, 2)
+            eri_mmap[i, :, :, :] = transposed_chunk
+        eri_mmap.flush()
+        del eri_mmap
+    else:
+        eri_ao = mol.intor('int2e').transpose(0, 2, 1, 3) # chem -> phys
+        f = open(output_file_path, 'w')
+        eri_ao.tofile(output_file_path)
+        f.close()
+te_2e = time.time()
+print("Wall time for writing 2e-integrals (physicist notation) to disk: {:.3f} seconds".format(te_2e - ti_2e))
+
 
 
 #Execute the QuAcK fortran program