Save global ind of each sample in NS_samples file and in traj file

.github/workflows/pytests.yml

@@ -14,10 +14,10 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ 3.9 ]
+        python-version: [ 3.11 ]
       max-parallel: 5
     env:
-      coverage-on-version: 3.9
+      coverage-on-version: NONE
       use-mpi: True
 
     steps:

pymatnext/__version__.py

@@ -1 +1 @@
-__version__ = "0.1.2"
+__version__ = "0.1.3"

pymatnext/cli/ns_analyse.py

@@ -137,13 +137,14 @@ def colname(colname_str):
                     f = line.split()
                     ## should we ignore certain kinds of malformed lines?
                     ## try:
-                    it = int(f[0])
-                    vals_line = [float(v) for v in f[1:]]
+                    NS_iter = int(f[0])
+                    # NS_global_ind = int(f[1])
+                    vals_line = [float(v) for v in f[2:]]
 
                     if l_i == 0 and len(vals_line) != 1 + len(header["extras"]):
-                        raise ValueError(f"Expecting 1 + {len(header['extras'])} extra fields, but got {len(vals_line)} on first line, refusing to continue")
+                        raise ValueError(f"Expecting 1 (NS quantity) + {len(header['extras'])} extra fields, but got {len(vals_line)} on first line, refusing to continue")
 
-                    iters.append(it)
+                    iters.append(NS_iter)
                     Es.append(vals_line[0])
                     vals.append(vals_line[1:])
                     ## except:

pymatnext/cli/sample.py

@@ -247,10 +247,11 @@ def sample(args, MPI, NS_comm, walker_comm):
             break
 
         # max info should already be set to: ns.rank_of_max, ns.local_ind_of_max, ns.max_val, ns.max_quants
+        global_ind_of_max = ns.global_ind(ns.rank_of_max, ns.local_ind_of_max)
 
         # write quantities for max config which will be culled below
         if NS_comm.rank == 0 and sample_interval > 0 and loop_iter % sample_interval == 0:
-            ns_file.write(f"{loop_iter} {ns.max_val:.10f} " + " ".join([f"{quant:.10f}" for quant in ns.max_quants]) + "\n")
+            ns_file.write(f"{loop_iter} {global_ind_of_max} {ns.max_val:.10f} " + " ".join([f"{quant:.10f}" for quant in ns.max_quants]) + "\n")
             ns_file.flush()
 
         # tune step sizes at some iteration interval
@@ -261,7 +262,6 @@ def sample(args, MPI, NS_comm, walker_comm):
                               adjust_factor=params_step_size_tune["adjust_factor"])
 
         # pick random config as source for clone.
-        global_ind_of_max = ns.global_ind(ns.rank_of_max, ns.local_ind_of_max)
         global_ind_of_clone_source = (global_ind_of_max + 1 + ns.rng_global.integers(0, ns.n_configs_global - 1)) % ns.n_configs_global
         rank_of_clone_source, local_ind_of_clone_source = ns.local_ind(global_ind_of_clone_source)
 
@@ -282,7 +282,7 @@ def sample(args, MPI, NS_comm, walker_comm):
                     ns.extra_config.recv(ns.rank_of_max, ns.comm, MPI)
                     max_config_write = ns.extra_config
 
-                max_config_write.write(traj_file, extra_info={"NS_iter": loop_iter})
+                max_config_write.write(traj_file, extra_info={"NS_iter": loop_iter, "global_ind": global_ind_of_max })
                 traj_file.flush()
 
             elif NS_comm.rank == ns.rank_of_max:

pymatnext/ns_configs/ase_atoms/atoms_contig_store.py

@@ -110,10 +110,10 @@ def _count_info_arrays_storage(self):
         for data_dict in [self.arrays, self.info]:
             for item_name in data_dict:
                 item = data_dict[item_name]
-                # create correct shape view into preallocated space
                 if not isinstance(item, np.ndarray):
                     continue
 
+                # determine correct size for preallocated space
                 if item.dtype == np.int64:
                     n_ints += item.size
                 elif item.dtype == np.float64:
@@ -129,10 +129,10 @@ def _set_info_arrays_storage_as_views(self, int_offset, float_offset, copy):
         for data_dict in [self.arrays, self.info]:
             for item_name in data_dict:
                 item = data_dict[item_name]
-                # create correct shape view into preallocated space
                 if not isinstance(item, np.ndarray):
                     continue
 
+                # create correct shape view into preallocated space
                 prealloc_view = None
                 if item.dtype == np.int64:
                     prealloc_view = self.contig_storage_int[int_offset:int_offset + item.size].reshape(item.shape)

tests/test_sample.py

@@ -97,7 +97,7 @@ def do_Morse_ASE_restart(tmp_path, monkeypatch, using_mpi):
         (tmp_path / (NS_samples_file + '.new')).rename(tmp_path / NS_samples_file)
 
     # run again, should restart from existing files
-    do_Morse_ASE(tmp_path, monkeypatch, using_mpi=using_mpi)
+    do_Morse_ASE(tmp_path, monkeypatch, using_mpi=using_mpi, restart=True)
 
     # check that NS_samples file wasn't overwritten by looking for "restart" key in header
     with open(tmp_path / NS_samples_file) as fin:
@@ -106,7 +106,7 @@ def do_Morse_ASE_restart(tmp_path, monkeypatch, using_mpi):
     assert header["restart"]
 
 
-def do_Morse_ASE(tmp_path, monkeypatch, using_mpi, max_iter=None):
+def do_Morse_ASE(tmp_path, monkeypatch, using_mpi, restart=False, max_iter=None):
     if using_mpi:
         from mpi4py import MPI
         if MPI.COMM_WORLD.size != 2:
@@ -160,9 +160,17 @@ def do_Morse_ASE(tmp_path, monkeypatch, using_mpi, max_iter=None):
 
     # from test run 6/20/2025, when cell_shear_per_rt3_atom default was increased
     if using_mpi:
-        samples_fields_ref = np.asarray([1.09000000e+02, 8.02719236e+00, 1.28609799e+04, 1.60000000e+01])
+        samples_fields_ref = np.asarray([1.09000000e+02, -1, 8.02719236e+00, 1.28609799e+04, 1.60000000e+01])
+        if restart:
+            samples_fields_ref[1] = 85
+        else:
+            samples_fields_ref[1] = 85
     else:
-        samples_fields_ref = np.asarray([1.09000000e+02, 8.06422068e+00, 1.29203058e+04, 1.60000000e+01])
+        samples_fields_ref = np.asarray([1.09000000e+02, -1, 8.06422068e+00, 1.29203058e+04, 1.60000000e+01])
+        if restart:
+            samples_fields_ref[1] = 136
+        else:
+            samples_fields_ref[1] = 135
 
     with open(tmp_path / 'Morse_ASE.test.NS_samples') as fin:
         for l in fin:
@@ -243,9 +251,9 @@ def do_EAM_LAMMPS(tmp_path, monkeypatch, using_mpi, max_iter=None):
 
     # from test run 6/20/2025 when cell shear default was increased
     if using_mpi:
-        fields_ref = np.asarray([2.99000000e+02, -3.90328931e+02,  1.99365520e+04,  1.60000000e+01, 1.87500000e-01,  8.12500000e-01])
+        fields_ref = np.asarray([2.99000000e+02, 55, -3.90328931e+02,  1.99365520e+04,  1.60000000e+01, 1.87500000e-01,  8.12500000e-01])
     else:
-        fields_ref = np.asarray([2.99000000e+02, -3.90994404e+02,  1.45640509e+04,  1.60000000e+01, 1.87500000e-01,  8.12500000e-01])
+        fields_ref = np.asarray([2.99000000e+02, 50, -3.90994404e+02,  1.45640509e+04,  1.60000000e+01, 1.87500000e-01,  8.12500000e-01])
 
     with open(tmp_path / 'EAM_LAMMPS.test.NS_samples') as fin:
         for l in fin:
@@ -358,8 +366,8 @@ def do_sGC(tmp_path, monkeypatch, using_mpi):
                     lfirst = l
             llast = l
 
-        f_13_first, f_29_first = [float(f) for f in lfirst.strip().split()[4:6]]
-        f_13_last, f_29_last = [float(f) for f in llast.strip().split()[4:6]]
+        f_13_first, f_29_first = [float(f) for f in lfirst.strip().split()[5:7]]
+        f_13_last, f_29_last = [float(f) for f in llast.strip().split()[5:7]]
 
         assert f_29_first / f_13_first == 1.0
         assert f_29_last / f_13_last > 4