Codex/gpu perf setup hamiltonian

pyfrsoc/compute_so_mat.py

@@ -157,15 +157,16 @@ def setup_pps(self, pp_files):
         """
         Setup q. Only Gamma point at this moment
         """
+        g_pp = self.ks_up.g_loc
+        npw_pp = self.ks_up.npw
+
         collect_pps = dict()
         for i in range(len(pp_files)):
             pp = pseudo_potential(pp_files[i], verbose=(self.rank == 0))
             pp.compute_dvan()
 
-            pp.interp_beta(
-                self.ks_up.ngw, self.ks_up.g, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01
-            )
-            pp.beta_ang(self.ks_up.ngw, self.ks_up.g, self.q)
+            pp.interp_beta(npw_pp, g_pp, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01)
+            pp.beta_ang(npw_pp, g_pp, self.q)
 
             pp.dvan_so = np.asarray(pp.dvan_so)
             pp.vkb1 = np.asarray(pp.vkb1)

pyfrsoc/dso_sos.py

@@ -146,15 +146,16 @@ def setup_pps(self, pp_files):
         """
         Setup q. Only Gamma point
         """
+        g_pp = self.ks_up.g_loc
+        npw_pp = self.ks_up.npw
+
         collect_pps = dict()
         for i in range(len(pp_files)):
             pp = pseudo_potential(pp_files[i])
             pp.compute_dvan()
 
-            pp.interp_beta(
-                self.ks_up.ngw, self.ks_up.g, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01
-            )
-            pp.beta_ang(self.ks_up.ngw, self.ks_up.g, self.q)
+            pp.interp_beta(npw_pp, g_pp, self.q, self.ks_up.omega, self.ks_up.tpiba, dbeta=False, dq=0.01)
+            pp.beta_ang(npw_pp, g_pp, self.q)
 
             pp.dvan_so = np.asarray(pp.dvan_so)
             pp.vkb1 = np.asarray(pp.vkb1)

pyfrsoc/dso_stern.py

@@ -251,7 +251,6 @@ def compute_h_soc_psi(self):
         gc.collect()
         if lGPU:
             synchronize()
-            np.get_default_memory_pool().free_all_blocks()
         self.comm.Barrier()
 
         return [h_soc_psi_0, h_soc_psi_1]
@@ -316,7 +315,6 @@ def solve_stern(self, stern, h_soc_psi, threshold):
                     self.comm.Barrier()
                     if lGPU:
                         synchronize()
-                        np.get_default_memory_pool().free_all_blocks()
                     gc.collect()
                     self.comm.Barrier()
 

pyfrsoc/dso_stern_hybrid.py

@@ -303,7 +303,6 @@ def compute_h_soc_psi(self):
         gc.collect()
         if lGPU:
             synchronize()
-            np.get_default_memory_pool().free_all_blocks()
         self.comm.Barrier()
 
         return [h_soc_psi_0, h_soc_psi_1]
@@ -398,7 +397,6 @@ def solve_stern(self, stern, h_soc_psi, threshold):
                 self.comm.Barrier()
                 if lGPU:
                     synchronize()
-                    np.get_default_memory_pool().free_all_blocks()
                 gc.collect()
                 self.comm.Barrier()
 

pyfrsoc/hamiltonian/hamiltonian.py

@@ -1,4 +1,5 @@
 import gc
+import warnings
 
 from mpi4py import MPI
 
@@ -22,25 +23,39 @@ def load_cube(ks, fname):
     headl = noa + 6
     num_elements = fft_grid[0] * fft_grid[1] * fft_grid[2]
 
+    data = None
     with open(fname, "r") as f:
         # Skip the header lines
         for _ in range(headl):
             next(f)
 
-        # Preallocate an array for the data
-        data = np2.zeros(num_elements)
-
-        # Index for filling the array
-        index = 0
-
-        # Read and process the file line by line
-        for line in f:
-            line_data = np2.fromiter(map(float, line.split()), dtype=np2.float64)
-            data[index : index + len(line_data)] = line_data
-            index += len(line_data)
-
-        # Reshape to match the fft_grid
-        potential = data.reshape(fft_grid)
+        # Fast path: C parser over the remaining flat numeric payload.
+        try:
+            with warnings.catch_warnings():
+                # Numpy emits DeprecationWarning for unmatched text and will raise
+                # ValueError in future versions. Treat it as parse failure now.
+                warnings.simplefilter("error", DeprecationWarning)
+                data = np2.fromfile(f, dtype=np2.float64, sep=" ")
+        except (TypeError, ValueError, DeprecationWarning):
+            data = None
+
+    if data is None or data.size != num_elements:
+        with open(fname, "r") as f:
+            for _ in range(headl):
+                next(f)
+
+            data = np2.zeros(num_elements)
+            index = 0
+            for line in f:
+                line_data = np2.fromiter(map(float, line.split()), dtype=np2.float64)
+                data[index : index + len(line_data)] = line_data
+                index += len(line_data)
+
+            if index != num_elements:
+                raise ValueError(f"Unexpected cube data size in {fname}: got {index}, expected {num_elements}")
+
+    # Reshape to match the fft_grid
+    potential = data.reshape(fft_grid)
 
     potential = np.asarray(potential)
 
@@ -121,14 +136,18 @@ def setup_pps(self, pp_files):
         """
         Setup q. Only Gamma point
         """
+        g_pp = self.ks.g_loc
+        npw_pp = self.ks.npw
+
         collect_pps = dict()
         for i in range(len(pp_files)):
             # average the pseudo potential
             pp = pseudo_potential(pp_files[i], average_pp=True, verbose=(self.rank == 0))
             pp.compute_dvan()
 
-            pp.interp_beta(self.ks.ngw, self.ks.g, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01)
-            pp.beta_ang(self.ks.ngw, self.ks.g, self.q)
+            # Build beta projectors only on local PW shards to avoid redundant per-rank work.
+            pp.interp_beta(npw_pp, g_pp, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01)
+            pp.beta_ang(npw_pp, g_pp, self.q)
 
             if pp.spin_orbit:
                 pp.dvan_so = np.asarray(pp.dvan_so)
@@ -1047,6 +1066,5 @@ def compute_h_psi(self, psi):
         self.comm.Barrier()
         if lGPU:
             synchronize()
-            np.get_default_memory_pool().free_all_blocks()
         gc.collect()
         self.comm.Barrier()

pyfrsoc/hamiltonian/soc_hamiltonian.py

@@ -58,13 +58,16 @@ def setup_pps(self, pp_files):
         """
         Setup q. Only Gamma point at this moment
         """
+        g_pp = self.ks.g_loc
+        npw_pp = self.ks.npw
+
         collect_pps = dict()
         for i in range(len(pp_files)):
             pp = pseudo_potential(pp_files[i])
             pp.compute_dvan()
 
-            pp.interp_beta(self.ks.ngw, self.ks.g, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01)
-            pp.beta_ang(self.ks.ngw, self.ks.g, self.q)
+            pp.interp_beta(npw_pp, g_pp, self.q, self.ks.omega, self.ks.tpiba, dbeta=False, dq=0.01)
+            pp.beta_ang(npw_pp, g_pp, self.q)
 
             if pp.spin_orbit:
                 pp.dvan_so = np.asarray(pp.dvan_so)

pyfrsoc/sternheimer/sternheimer.py

@@ -290,7 +290,6 @@ def cg_solver(self, pert_evc, threshold=1e-16, max_iter=100, debug=False):
         self.comm.Barrier()
         if lGPU:
             synchronize()
-            np.get_default_memory_pool().free_all_blocks()
         gc.collect()
         self.comm.Barrier()
 

pyfrsoc/sternheimer/sternheimer_hybrid.py

@@ -1000,7 +1000,6 @@ def cg_solver(self, pert_evc, threshold=1e-16, max_iter=100, debug=False):
         self.comm.Barrier()
         if lGPU:
             synchronize()
-            np.get_default_memory_pool().free_all_blocks()
         gc.collect()
         self.comm.Barrier()
 

pyfrsoc/utils/compute_vnl.py

@@ -95,8 +95,14 @@ def solid_vkb_loc(atoms, pps, q, g, mill, eigts, fft_grid, pw_index, npwx, neg=F
     Compute projectors for the solid, including the structure factor
     To be used with parallelization
     """
-    g_loc = g[pw_index]
-    mill_loc = mill[pw_index]
+    if g.shape[0] == len(pw_index):
+        g_loc = g
+    else:
+        g_loc = g[pw_index]
+    if mill.shape[0] == len(pw_index):
+        mill_loc = mill
+    else:
+        mill_loc = mill[pw_index]
 
     nr1 = fft_grid[0]
     nr2 = fft_grid[1]
@@ -121,11 +127,15 @@ def solid_vkb_loc(atoms, pps, q, g, mill, eigts, fft_grid, pw_index, npwx, neg=F
         pp = pps[element]
 
         vkb1 = pp.vkb1[: pp.number_of_proj_ang, :]
+        if vkb1.shape[1] == npw:
+            vkb1_loc = vkb1
+        else:
+            vkb1_loc = vkb1[:, pw_index]
 
         if neg:
-            vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = np.conj(vkb1[:, pw_index]) * sk[na, :] * phases[na]
+            vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = np.conj(vkb1_loc) * sk[na, :] * phases[na]
         else:
-            vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = vkb1[:, pw_index] * sk[na, :] * phases[na]
+            vkb[ikb : ikb + pp.number_of_proj_ang, :npw] = vkb1_loc * sk[na, :] * phases[na]
 
         ikb += pp.number_of_proj_ang