More efficient handling of overlap matrices

salted/cp2k/density_fitting.py

@@ -1,16 +1,43 @@
-import numpy as np
-import time
-import sys
-import os
 import glob
+import os
 import os.path as osp
+import sys
+import time
 
-from salted.constants import bohr2angs
-from salted.sys_utils import ParseConfig, read_system, get_atom_idx, check_MPI_tasks_count, detect_mpi, distribute_jobs
-from salted.cp2k.utils import gto_rec, gto_rec_prim, gto_rec_g0, gto_rec_ewald, build_gvec, get_basis_set_info_numba, read_local_pseudo, setup_pyscf_species, setup_pyscf_ewald
-from salted.cp2k.utils import build_contraction_matrix, get_w_prim, get_rho_n, compute_hartree_energy, build_gcutoff, pair_cutoffs, overlap_identity, overlap_coulomb_rec, overlap_coulomb_real
+import numpy as np
 from mpi4py import MPI
 
+from salted.constants import bohr2angs
+from salted.cp2k.utils import (
+    build_contraction_matrix,
+    build_gcutoff,
+    build_gvec,
+    compute_hartree_energy,
+    get_basis_set_info_numba,
+    get_rho_n,
+    get_w_prim,
+    gto_rec,
+    gto_rec_ewald,
+    gto_rec_g0,
+    gto_rec_prim,
+    overlap_coulomb_real,
+    overlap_coulomb_rec,
+    overlap_identity,
+    pair_cutoffs,
+    read_local_pseudo,
+    setup_pyscf_ewald,
+    setup_pyscf_species,
+)
+from salted.sys_utils import (
+    ParseConfig,
+    check_MPI_tasks_count,
+    detect_mpi,
+    distribute_jobs,
+    get_atom_idx,
+    read_system,
+    save_overlap,
+)
+
 inp = ParseConfig().parse_input()
 df_metric = inp.qm.dfmetric
 
@@ -102,7 +129,7 @@
     dr = np.array([dx, dy, dz])
 
     if rho_KS.size != npoints:
-        print(f"ERROR: inconsistent number of grid points!")
+        print("ERROR: inconsistent number of grid points!")
         sys.exit(0)
 
     rho_KS = rho_KS.reshape((nx, ny, nz))
@@ -179,7 +206,7 @@
 
     # Save data
     np.save(os.path.join(inp.salted.saltedpath, "coefficients", f"coefficients_conf{iconf}.npy"), c)
-    np.save(os.path.join(inp.salted.saltedpath, "overlaps", f"overlap_conf{iconf}.npy"), S)
+    save_overlap(os.path.join(inp.salted.saltedpath, "overlaps", f"overlap_conf{iconf}.npy"), S) # lower triangle only
 
     # Hartree energy calculation
     time_c = time.time()

salted/hessian_matrix.py

@@ -1,7 +1,6 @@
 import os
 import os.path as osp
 import random
-import sys
 import time
 
 import numpy as np
@@ -15,6 +14,7 @@
     distribute_jobs,
     format_index_ranges,
     get_atom_idx,
+    load_overlap,
     read_system,
 )
 
@@ -66,7 +66,7 @@ def build():
         saltedpath, rdir, f"training_set_N{inp.gpr.Ntrain}.txt"
     ), trainrangetot, fmt='%i')
     ntrain = round(inp.gpr.trainfrac*inp.gpr.Ntrain)
-    trainrange = trainrangetot[:ntrain]
+    trainrange = sorted(trainrangetot[:ntrain])
 
     """
     Calculate regression matrices in parallel or serial mode.
@@ -132,7 +132,7 @@ def _compute_sparse_operations(psivec, ref_projs, over, sparse_algorithm):
         N_df, K_rkhs = psivec.shape
         avec_contrib = psivec.T @ ref_projs               # O(nnz) vector multiply
         engine       = get_hessian_engine(N_df, K_rkhs)
-        bmat_contrib = engine.compute(over, psivec).copy()
+        bmat_contrib = engine.compute(over, psivec)
         return avec_contrib, bmat_contrib, "numba"
 
     # Dense fallback (original behavior)
@@ -159,11 +159,11 @@ def matrices(trainrange,ntrain,av_coefs,rank):
 
     if inp.salted.saltedtype=="density-response":
         p = sparse.load_npz(osp.join(
-            saltedpath, fdir, f"M{Menv}_zeta{zeta}", f"psi-nm_conf0_x.npz"
+            saltedpath, fdir, f"M{Menv}_zeta{zeta}", "psi-nm_conf0_x.npz"
         ))
     else:
         p = sparse.load_npz(osp.join(
-            saltedpath, fdir, f"M{Menv}_zeta{zeta}", f"psi-nm_conf0.npz"
+            saltedpath, fdir, f"M{Menv}_zeta{zeta}", "psi-nm_conf0.npz"
         ))
 
     species, lmax, nmax, llmax, nnmax, ndata, atomic_symbols, atomic_coords, natoms, natmax = read_system()
@@ -178,22 +178,28 @@ def matrices(trainrange,ntrain,av_coefs,rank):
 
     Avec = np.zeros(totsize)
     Bmat = np.zeros((totsize,totsize))
+    total_start_time = time.time()
+    total_io_time, total_compute_time = 0.0, 0.0
     for iconf in trainrange:
 
         start_time = time.time()
+        io_time, compute_time = 0.0, 0.0
 
         if inp.salted.saltedtype=="density":
 
             # load reference QM data
+            t0 = time.time()
             ref_coefs = np.load(osp.join(
                 saltedpath, "coefficients", f"coefficients_conf{iconf}.npy"
             ))
-            over = np.load(osp.join(
+            over = load_overlap(osp.join(
                 saltedpath, "overlaps", f"overlap_conf{iconf}.npy"
             ))
             psivec = sparse.load_npz(osp.join(
                 saltedpath, fdir, f"M{Menv}_zeta{zeta}", f"psi-nm_conf{iconf}.npz"
             ))
+            t1 = time.time()
+            io_time += t1 - t0
 
             if inp.system.average:
 
@@ -224,22 +230,28 @@ def matrices(trainrange,ntrain,av_coefs,rank):
                 sparse_algorithm = algorithm_used
             Avec += avec_contrib
             Bmat += bmat_contrib
+            compute_time += time.time() - t1
 
 
         elif inp.salted.saltedtype=="density-response":
 
-            over = np.load(osp.join(
+            t0 = time.time()
+            over = load_overlap(osp.join(
                 saltedpath, "overlaps", f"overlap_conf{iconf}.npy"
             ))
+            io_time += time.time() - t0
 
             for icart in ["x","y","z"]:
 
+                t0 = time.time()
                 ref_coefs = np.load(osp.join(
                     saltedpath, "coefficients", f"{icart}/coefficients_conf{iconf}.npy"
                 ))
                 psivec = sparse.load_npz(osp.join(
                     saltedpath, fdir, f"M{Menv}_zeta{zeta}", f"psi-nm_conf{iconf}_{icart}.npz"
                 ))
+                t1 = time.time()
+                io_time += t1 - t0
 
                 ref_projs = np.dot(over,ref_coefs)
 
@@ -249,8 +261,15 @@ def matrices(trainrange,ntrain,av_coefs,rank):
                 )
                 Avec += avec_contrib
                 Bmat += bmat_contrib
+                compute_time += time.time() - t1
 
-        if inp.salted.verbose: print(f"conf {iconf}, time = {(time.time() - start_time):.2f} s", flush=True)
+        del over, psivec
+
+        total_io_time += io_time
+        total_compute_time += compute_time
+        if inp.salted.verbose: print(f"conf {iconf}, time = {(time.time() - start_time):.2f} s (I/O = {io_time:.2f} s, compute = {compute_time:.2f} s)", flush=True)
+
+    if inp.salted.verbose: print(f"Task {rank}, total time for {len(trainrange)} structures = {(time.time() - total_start_time):.2f} s (I/O = {total_io_time:.2f} s, compute = {total_compute_time:.2f} s)", flush=True)
 
     Avec /= float(ntrain)
     Bmat /= float(ntrain)

salted/minimize_loss.py

@@ -1,6 +1,7 @@
 import os
 import os.path as osp
 import random
+import sys
 import time
 
 import numpy as np
@@ -14,6 +15,7 @@
     distribute_jobs,
     format_index_ranges,
     get_atom_idx,
+    load_overlap,
     read_system,
 )
 
@@ -60,9 +62,8 @@ def build():
             )
 
     dirpath = os.path.join(saltedpath, rdir, f"M{Menv}_zeta{zeta}")
-    if rank == 0:
-        if not os.path.exists(dirpath):
-            os.makedirs(dirpath, exist_ok=True)
+    if rank == 0 and not os.path.exists(dirpath):
+        os.makedirs(dirpath, exist_ok=True)
     if parallel:
         comm.Barrier()
 
@@ -74,7 +75,7 @@ def build():
                 f"than are present in the input data {ndata=}."
             )
         else:
-            exit()
+            sys.exit()
     dataset = list(range(ndata))
     if inp.gpr.trainsel == "sequential":
         trainrangetot = dataset[:Ntrain]
@@ -101,7 +102,7 @@ def build():
             print(f"Task {rank} handles the following structures: {format_index_ranges(trainrange,True)}", flush=True)
     else:
         trainrange = trainrangetot[:ntraintot]
-    ntrain = int(len(trainrange))
+    ntrain = len(trainrange)
 
     def loss_func(weights, ovlp_list, psi_list):
         """Given the weight-vector of the RKHS, compute the gradient of the electron-density loss function."""
@@ -300,7 +301,7 @@ def curv_func(cg_dire, ovlp_list, psi_list):
 
         totsize = psi_list[0].shape[1]
 
-        Ad = np.zeros((totsize))
+        Ad = np.zeros(totsize)
 
         if saltedtype=="density":
 
@@ -331,7 +332,7 @@ def curv_func(cg_dire, ovlp_list, psi_list):
     psi_list = []
     for iconf in trainrange:
         ovlp_list.append(
-            np.load(osp.join(saltedpath, "overlaps", f"overlap_conf{iconf}.npy"))
+            load_overlap(osp.join(saltedpath, "overlaps", f"overlap_conf{iconf}.npy"))
         )
         # load feature vector as a scipy sparse object
         if saltedtype=="density":
@@ -439,7 +440,7 @@ def curv_func(cg_dire, ovlp_list, psi_list):
             loss = loss_func(w, ovlp_list, psi_list)
             if loss>loss_old:
                 if rank == 0:
-                    print(f"WARNING: loss function increased, search direction reset as the steepest descent.")
+                    print("WARNING: loss function increased, search direction reset as the steepest descent.")
                 r = -grad_func(w, ovlp_list, psi_list)
                 if rank == 0:
                     print(f"step {i+1}, gradient norm: {np.linalg.norm(r):.3e}, loss: {loss:.3e}", flush=True)

salted/solve_regression.py

@@ -1,6 +1,4 @@
-import os
 import os.path as osp
-import sys
 import time
 
 import numpy as np
@@ -34,13 +32,13 @@ def build():
 
     start = time.time()
 
-    w = np.linalg.solve(Bmat+np.eye(totsize)*regul,Avec)
+    Bmat[np.diag_indices_from(Bmat)] += regul
+    w = np.linalg.solve(Bmat,Avec)
 
-    print(f"regression time: {((time.time()-start)/60):.3f} minutes",flush=True)
+    print(f"regression time: {time.time() - start} seconds",flush=True)
 
     np.save(osp.join(saltedpath, rdir, f"M{Menv}_zeta{zeta}", f"weights_N{ntrain}_reg{int(np.log10(regul))}.npy"), w)
 
-    return
 
 if __name__ == "__main__":
-    build()
+    build()

salted/sys_utils.py

@@ -538,6 +538,56 @@ def get_feats_projs_response(species, lmax):
     return Vmat, Mspe, power_env_sparse, power_env_sparse_antisymm
 
 
+def save_overlap(path: str, S: np.ndarray):
+    #Save a symmetric overlap matrix storing only its lower triangle.
+
+    assert S.ndim == 2 and S.shape[0] == S.shape[1] # must be a square matrix
+    S = np.ascontiguousarray(S, dtype=np.float64)
+    N = S.shape[0] # number of rows/columns
+    header = {
+        "descr": np.lib.format.dtype_to_descr(np.dtype(np.float64)), # dtype descriptor
+        "fortran_order": False, # not Fortran order
+        "shape": (N * (N + 1) // 2,), # shape of the packed lower triangle
+    }
+    with open(path, "wb") as f:
+        np.lib.format.write_array_header_1_0(f, header) # write header
+        f.writelines(S[i, : i + 1] for i in range(N)) # write lower triangle rows sequentially
+
+
+def load_overlap(path: str, tile: int = 128) -> np.ndarray:
+    # Load an overlap matrix as a full dense (N,N) array.
+
+    with open(path, "rb") as f:
+        version = np.lib.format.read_magic(f) # read magic number to determine format version
+        if version == (1, 0):
+            shape, fortran_order, dtype = np.lib.format.read_array_header_1_0(f) # read header for version 1.0
+        else:
+            shape, fortran_order, dtype = np.lib.format.read_array_header_2_0(f) # read header for version 2.0
+        if len(shape) == 2:
+            return np.load(path)  # full matrix (older files, other codes)
+        L = shape[0]
+        N = (int(np.sqrt(8 * L + 1)) - 1) // 2 # compute N from L = N*(N+1)/2
+        if len(shape) != 1 or dtype != np.float64 or N * (N + 1) // 2 != L:
+            raise ValueError(f"{path}: shape {shape} and dtype {dtype} do not match a packed float64 overlap matrix")
+
+        # read the lower triangle sequentially, each row straight into S
+        S = np.empty((N, N))
+        for i in range(N):
+            if f.readinto(S[i, : i + 1]) != (i + 1) * 8:
+                raise ValueError(f"{path}: file is truncated")
+
+    # mirror lower triangle into the upper one, in small tiles that fit in cache
+    for i0 in range(0, N, tile):
+        i1 = min(i0 + tile, N) # compute the end index of the tile
+        diag = S[i0:i1, i0:i1] # extract the diagonal tile
+        diag[...] = np.tril(diag) + np.tril(diag, -1).T # mirror the diagonal tile
+        for j0 in range(i1, N, tile): # for each tile in the upper triangle
+            j1 = min(j0 + tile, N) # compute the end index of the tile
+            S[i0:i1, j0:j1] = S[j0:j1, i0:i1].T # mirror the off-diagonal tile
+
+    return S
+
+
 class AttrDict:
     """Access dict keys as attributes