Lots of changes, most of them trivial, but a few major bugfixes

Add applied pressure to LAMMPS fix ns/cell_mc, add corresponding pytest

Fix serious sign error in lammps fix ns/type (detected when debugging pressure in cell_mc). Hopefully correct derivation documented in comment in ns_configs/ase_atoms/walks_lammps.py. NOTE: should probably add a test comparing ASE and LAMMPS dynamics for fixed mu, to ensure they behave the same

Fix serious bug in step-size tuning that failed to reset walk counters before each pilot walk.

Fix bug in initializing rngs when restarting from snapshot.  Note that restart still isn't perfect because snapshots are in extxyz which has insufficient precision.  This affects some of the restart tests, so loosen tolerance.

HACK: handle calculators that create Calculator.results["energy"] when "free_energy" is passed to Calculator.calculate in properties

Switch from ASE LJ to somewhat faster ASE Morse for pytests (update expected values)

Fix bug in composition vs. full_composition pytest

Fix searching for potential module in sampling pytests

Author: bernstei

examples/EAM_LAMMPS/params.toml

@@ -32,7 +32,7 @@
 
         [configs.calculator.args]
 
-            cmds = ["pair_style eam/alloy", "pair_coeff * * /home/cluster2/bernstei/src/work/pymatnext/examples/EAM_LAMMPS/AlCu_Zhou04.eam.alloy Al Cu"]
+            cmds = ["pair_style eam/alloy", "pair_coeff * * AlCu_Zhou04.eam.alloy Al Cu"]
 
             types.13 = 1
             types.29 = 2

examples/EAM_LAMMPS/params_sGC.toml

@@ -32,7 +32,7 @@
 
         [configs.calculator.args]
 
-            cmds = ["pair_style eam/alloy", "pair_coeff * * /home/cluster2/bernstei/src/work/pymatnext/examples/EAM_LAMMPS/AlCu_Zhou04.eam.alloy Al Cu"]
+            cmds = ["pair_style eam/alloy", "pair_coeff * * AlCu_Zhou04.eam.alloy Al Cu"]
 
             types.13 = 1
             types.29 = 2

lammps_patches/new/src/NS/fix_ns_cellmc.cpp

@@ -371,9 +371,9 @@ void FixNSCellMC::final_integrate()
 
   // if potential energy - d(P V) is above Emax then reject move
   // need to include previous steps' cumulative dPV contributions, as well as current ones
-  if (ecurrent - cumulative_dPV - dPV >= Emax) {
+  if (ecurrent + cumulative_dPV + dPV >= Emax) {
 #ifdef DEBUG
-    std::cout << "REJECT E == " << ecurrent << " - " << cumulative_dPV + dPV << " >= Emax == " << Emax << std::endl;
+    std::cout << "REJECT E == " << ecurrent << " + " << cumulative_dPV + dPV << " >= Emax == " << Emax << std::endl;
 #endif
     // reject move, so don't touch cumulative_dPV, since type change that led to current dPV was reverted
 
@@ -421,7 +421,7 @@ new_cell[1][2] = 0.0;
 new_cell[2][0] = domain->xz;
 new_cell[2][1] = domain->yz;
 new_cell[2][2] = domain->boxhi[2] - domain->boxlo[2];
-    std::cout << "ACCEPT E == " << ecurrent " - " << cumulative_dPV << " < Emax == " << Emax << " min_aspect " << min_aspect_ratio_val(new_cell) << std::endl;
+    std::cout << "ACCEPT E == " << ecurrent " + " << cumulative_dPV << " < Emax == " << Emax << " min_aspect " << min_aspect_ratio_val(new_cell) << std::endl;
     // std::cout << "final cell " << new_cell[0][0] << " " << new_cell[0][1] << " " << new_cell[0][2] << std::endl;
     // std::cout << "           " << new_cell[1][0] << " " << new_cell[1][1] << " " << new_cell[1][2] << std::endl;
     // std::cout << "           " << new_cell[2][0] << " " << new_cell[2][1] << " " << new_cell[2][2] << std::endl;

lammps_patches/new/src/NS/fix_ns_type.cpp

@@ -263,9 +263,9 @@ void FixNSType::final_integrate()
 
   // if potential energy + d(mu N) is above Emax then reject move
   // need to include previous steps' cumulative dmuN contributions, as well as current ones
-  if (ecurrent + cumulative_dmuN + dmuN >= Emax) {
+  if (ecurrent - cumulative_dmuN - dmuN >= Emax) {
 #ifdef DEBUG
-    std::cout << "REJECT E == " << ecurrent << " + " << cumulative_dmuN + dmuN << " >= Emax == " << Emax << " " << std::endl;
+    std::cout << "REJECT E == " << ecurrent << " - " << cumulative_dmuN + dmuN << " >= Emax == " << Emax << " " << std::endl;
 #endif
     // reject move, so don't touch cumulative_dmuN, since type change that led to current dmuN was reverted
 

pymatnext/cli/sample.py

@@ -75,7 +75,7 @@ def parse_args(args_list=None):
     parser.add_argument("--random_seed", "-s", type=int, help="random seed overriding params.global.random_seed file")
     parser.add_argument("--output_filename_postfix", "-p", help="string to append to params.global.output_filename_prefix", default="")
     parser.add_argument("--max_iter", "-i", type=int, help="max number of NS iterations, overriding params")
-    parser.add_argument("--restart_diff_nproc", "-d", action="store_true", help="allow restarts to use a different number of processors that previous partial run")
+    parser.add_argument("--restart_diff_nproc", "-d", action="store_true", help="allow restarts to use a different number of processors than previous partial run")
     parser.add_argument("input", help="input json/yaml file")
     args = parser.parse_args(args_list)
 

pymatnext/extras/lammps.patch

@@ -1,6 +1,6 @@
 diff -N -u -r orig/src/NS/fix_ns_cellmc.cpp new/src/NS/fix_ns_cellmc.cpp
 --- orig/src/NS/fix_ns_cellmc.cpp	1969-12-31 19:00:00.000000000 -0500
-+++ new/src/NS/fix_ns_cellmc.cpp	2023-04-03 08:16:30.856268936 -0400
++++ new/src/NS/fix_ns_cellmc.cpp	2023-04-06 12:03:21.386180157 -0400
 @@ -0,0 +1,474 @@
 +/* ----------------------------------------------------------------------
 +   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
@@ -375,9 +375,9 @@ diff -N -u -r orig/src/NS/fix_ns_cellmc.cpp new/src/NS/fix_ns_cellmc.cpp
 +
 +  // if potential energy - d(P V) is above Emax then reject move
 +  // need to include previous steps' cumulative dPV contributions, as well as current ones
-+  if (ecurrent - cumulative_dPV - dPV >= Emax) {
++  if (ecurrent + cumulative_dPV + dPV >= Emax) {
 +#ifdef DEBUG
-+    std::cout << "REJECT E == " << ecurrent << " - " << cumulative_dPV + dPV << " >= Emax == " << Emax << std::endl;
++    std::cout << "REJECT E == " << ecurrent << " + " << cumulative_dPV + dPV << " >= Emax == " << Emax << std::endl;
 +#endif
 +    // reject move, so don't touch cumulative_dPV, since type change that led to current dPV was reverted
 +
@@ -425,7 +425,7 @@ diff -N -u -r orig/src/NS/fix_ns_cellmc.cpp new/src/NS/fix_ns_cellmc.cpp
 +new_cell[2][0] = domain->xz;
 +new_cell[2][1] = domain->yz;
 +new_cell[2][2] = domain->boxhi[2] - domain->boxlo[2];
-+    std::cout << "ACCEPT E == " << ecurrent " - " << cumulative_dPV << " < Emax == " << Emax << " min_aspect " << min_aspect_ratio_val(new_cell) << std::endl;
++    std::cout << "ACCEPT E == " << ecurrent " + " << cumulative_dPV << " < Emax == " << Emax << " min_aspect " << min_aspect_ratio_val(new_cell) << std::endl;
 +    // std::cout << "final cell " << new_cell[0][0] << " " << new_cell[0][1] << " " << new_cell[0][2] << std::endl;
 +    // std::cout << "           " << new_cell[1][0] << " " << new_cell[1][1] << " " << new_cell[1][2] << std::endl;
 +    // std::cout << "           " << new_cell[2][0] << " " << new_cell[2][1] << " " << new_cell[2][2] << std::endl;
@@ -478,7 +478,7 @@ diff -N -u -r orig/src/NS/fix_ns_cellmc.cpp new/src/NS/fix_ns_cellmc.cpp
 +}
 diff -N -u -r orig/src/NS/fix_ns_cellmc.h new/src/NS/fix_ns_cellmc.h
 --- orig/src/NS/fix_ns_cellmc.h	1969-12-31 19:00:00.000000000 -0500
-+++ new/src/NS/fix_ns_cellmc.h	2023-04-01 17:15:39.278299788 -0400
++++ new/src/NS/fix_ns_cellmc.h	2023-04-04 21:47:03.744996335 -0400
 @@ -0,0 +1,74 @@
 +/* -*- c++ -*- ----------------------------------------------------------
 +   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
@@ -826,7 +826,7 @@ diff -N -u -r orig/src/NS/fix_ns_gmc.h new/src/NS/fix_ns_gmc.h
 +*/
 diff -N -u -r orig/src/NS/fix_ns_type.cpp new/src/NS/fix_ns_type.cpp
 --- orig/src/NS/fix_ns_type.cpp	1969-12-31 19:00:00.000000000 -0500
-+++ new/src/NS/fix_ns_type.cpp	2023-04-01 17:14:14.838079400 -0400
++++ new/src/NS/fix_ns_type.cpp	2023-04-06 12:03:37.722019077 -0400
 @@ -0,0 +1,341 @@
 +/* ----------------------------------------------------------------------
 +   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
@@ -1093,9 +1093,9 @@ diff -N -u -r orig/src/NS/fix_ns_type.cpp new/src/NS/fix_ns_type.cpp
 +
 +  // if potential energy + d(mu N) is above Emax then reject move
 +  // need to include previous steps' cumulative dmuN contributions, as well as current ones
-+  if (ecurrent + cumulative_dmuN + dmuN >= Emax) {
++  if (ecurrent - cumulative_dmuN - dmuN >= Emax) {
 +#ifdef DEBUG
-+    std::cout << "REJECT E == " << ecurrent << " + " << cumulative_dmuN + dmuN << " >= Emax == " << Emax << " " << std::endl;
++    std::cout << "REJECT E == " << ecurrent << " - " << cumulative_dmuN + dmuN << " >= Emax == " << Emax << " " << std::endl;
 +#endif
 +    // reject move, so don't touch cumulative_dmuN, since type change that led to current dmuN was reverted
 +

pymatnext/ns.py

@@ -339,14 +339,15 @@ def step_size_tune(self, n_configs=1, min_accept_rate=0.25, max_accept_rate=0.5,
                 else:
                     self.local_configs[0].copy_contents(ns_config)
 
+                self.local_configs[0].reset_walk_counters()
                 accept_freq_contribution = self.local_configs[0].walk(self.max_val, self.local_walk_length, self.rng_local)
                 accept_freq += np.asarray(accept_freq_contribution)
 
             accept_freq = self.comm.allreduce(accept_freq, self.MPI.SUM)
 
             if first_iter and self.comm.rank == 0:
                 for name, size, max_size, freq in zip(step_size_names, self.local_configs[0].step_size.values(), max_step_size, accept_freq):
-                    print("step_size_tune initial", name, size, max_size, freq)
+                    print("step_size_tune initial", name, "size", size, "max", max_size, "freq", freq)
                 first_iter = False
 
             done = []
@@ -374,7 +375,7 @@ def step_size_tune(self, n_configs=1, min_accept_rate=0.25, max_accept_rate=0.5,
                 break
 
             if any(np.asarray(step_size) < 1.0e-12):
-                raise RuntimeError(f"Stepsize got too small with automatic tuning {step_size}")
+                raise RuntimeError(f"Stepsize got too small with automatic tuning {step_size} {step_size_names}")
 
         if self.comm.rank == 0:
             for name, size in zip(step_size_names, self.local_configs[0].step_size.values()):
@@ -529,15 +530,15 @@ def read_rngs(self, bit_generator_states, different_nlocal=False):
             bit_generator_states = {"global": None, "locals": []}
 
         # bcast globals
-        self.rng_global.bit_generator.state = self.comm.bcast(bit_generator_states["global"], root=0)
+        new_state = self.comm.bcast(bit_generator_states["global"], root=0)
+        self.rng_global.bit_generator.state = new_state
 
         # scatter or generate locals
         n_locals = self.comm.bcast(len(bit_generator_states["locals"]), root = 0)
-        # print("BOB", n_locals, self.comm.size)
-        if n_locals > self.comm.size:
+        if n_locals == self.comm.size:
             self.rng_local.bit_generator.state = self.comm.scatter(bit_generator_states["locals"][:self.comm.size], root=0)
         else:
-            # need more local generators than we have,generate all from global rng
+            # number of read-in generators doesn't match number needed, start over by generating all from global rng
             self.rng_global, self.rng_local = new_rngs(self.comm.rank, None, self.rng_global)
 
 

pymatnext/ns_configs/ase_atoms/__init__.py

@@ -72,7 +72,7 @@ def _parse_composition(composition):
         """
         # parse composition and set up symbols and cls._Zs
         if isinstance(composition, str):
-            composition_p = re.split(r"([A-Z][a-z]?[0-9]*)", re.sub("\s+", "", composition))
+            composition_p = re.split(r"([A-Z][a-z]?[0-9]*)", re.sub(r"\s+", "", composition))
             if any([s != "" for s in composition_p[0::2]]):
                 raise ValueError(f"Unknown characters in composition {composition}")
             composition_p = composition_p[1::2]
@@ -125,7 +125,8 @@ def initialize(cls, params):
     def __init__(self, params, compression=np.inf, source="random", rng=None, kB=ase.units.kB, allocate_only=False):
         check_fill_defaults(params, param_defaults_ase_atoms, label="configs")
 
-        assert len(self._Zs) > 0
+        if len(self._Zs) == 0:
+            NSConfig_ASE_Atoms.initialize(params)
 
         # Save params for later construction of the calculator
         # Also save calc_type, since it's needed for some other setup tasks,
@@ -231,8 +232,16 @@ def __init__(self, params, compression=np.inf, source="random", rng=None, kB=ase
         self.atoms.prev_positions = np.zeros(self.atoms.positions.shape)
         self.atoms.prev_cell = np.zeros(self.atoms.cell.array.shape)
 
+        self.reset_walk_counters()
+
+
+    def reset_walk_counters(self):
+        """Reset attempted and successful step counters
+        """
+
         self.n_att_acc = np.zeros((len(NSConfig_ASE_Atoms._step_size_params), 2), dtype=np.int64)
 
+
     def init_calculator(self):
         """Initialize calculator.  Not part of constructor since some calculators (e.g. LAMMPS)
         cannot be pickled for mpi4py communication.
@@ -281,8 +290,8 @@ def init_calculator(self):
                 Z_types[Z] = species_type
 
             # make numpy arrays for (presumably?) faster lookup
-            self.type_of_Z = np.zeros(max(Z_types.keys()) + 1, dtype=np.int)
-            self.Z_of_type = np.zeros(max(Z_types.values()) + 1, dtype=np.int)
+            self.type_of_Z = np.zeros(max(Z_types.keys()) + 1, dtype=int)
+            self.Z_of_type = np.zeros(max(Z_types.values()) + 1, dtype=int)
             for Z, Z_type in Z_types.items():
                 self.type_of_Z[Z] = Z_type
                 self.Z_of_type[Z_type] = Z
@@ -431,7 +440,7 @@ def initial_calc_and_store(self):
         """
         if self.calc_type == "ASE":
             self.calc.calculate(self.atoms, properties=["free_energy", "forces"])
-            self.atoms.info["NS_energy"][...] = self.calc.results["free_energy"]
+            self.atoms.info["NS_energy"][...] = self.calc.results.get("free_energy", self.calc.results.get("energy"))
             self.atoms.arrays["NS_forces"][...] = self.calc.results["forces"]
         elif self.calc_type == "LAMMPS":
             self.calc.reset_box([0.0, 0.0, 0.0], np.diag(self.atoms.cell), self.atoms.cell[0, 1], self.atoms.cell[1, 2], self.atoms.cell[0, 2])
@@ -573,7 +582,6 @@ def header_dict(self):
                         "flat_V_prior": self.move_params["cell"]["flat_V_prior"],
                         "extras": [ "volume", "natoms"] + ([f"x_{Z}" for Z in self._Zs] if len(self._Zs) > 1 else []) }
 
-        print("BOB NSConfigs_ASE_Atoms returning header_dict", header_dict)
         return header_dict
 
 

pymatnext/ns_configs/ase_atoms/walks_ase_calc.py

@@ -32,7 +32,7 @@ def walk_pos_gmc(ns_atoms, Emax, rng):
         # step and evaluate new energy, forces
         atoms.positions += atoms.arrays["NS_velocities"]
         atoms.calc.calculate(atoms, properties=["free_energy", "forces"])
-        E = atoms.calc.results["free_energy"]
+        E = atoms.calc.results.get("free_energy", atoms.calc.results.get("energy"))
         F = atoms.calc.results["forces"]
 
         if E >= Emax: # reflect or fail
@@ -93,7 +93,7 @@ def _eval_and_accept_or_signal_revert(atoms, Emax, delta_PV=0.0, delta_muN=0.0):
     revert: bool, True if move is rejected and must be reverted
     """
     atoms.calc.calculate(atoms, properties=["free_energy", "forces"])
-    E = atoms.calc.results["free_energy"]
+    E = atoms.calc.results.get("free_energy", atoms.calc.results.get("energy"))
     E_shift = atoms.info["NS_energy_shift"] + delta_PV - delta_muN
     if E + E_shift < Emax:
         # accept can happen here, because it's always the same action

pymatnext/ns_configs/ase_atoms/walks_lammps.py

@@ -4,6 +4,24 @@
 
 import lammps
 
+# avoiding every lammps fix from having to do the full energy shift is useful
+# we want
+#   E + P V - mu N < Emax
+# therefore, we modify Emax passed to lammps defining Emax^lammps, with V0 and N0 
+#        defined as the initial volume and species numbers
+#   E < Emax - P V0 + mu N0 = Emax^lammps
+# this is sufficient for position moves, where V and N don't change
+#
+# for volume changes, we want
+#   E + P (V0 + dV) - mu N0 < Emax
+# where dV is the cumulative difference of the current steps' volume and the step where
+#        the fix trajectory started
+#   E < Emax - P (V0 + dV) + mu N0
+#   E < Emax - P V0 + mu N0 - P dV
+#   E < Emax^lammps - P dV
+# which is equivalent to
+#   E + P dV < Emax^lammps
+
 def set_lammps_from_atoms(ns_atoms):
     """set lammps internal configuration data (cell, types, positions, velocities) from
     ns_atoms object
@@ -184,7 +202,7 @@ def walk_cell(ns_atoms, Emax, rng):
     # NOTE: would it be faster to construct more of this command ahead of time?
     # would either still have to do step_sizes here, or recreate command
     # every time step sizes are changed
-    fix_cmd = (f"fix NS all ns/cellmc {lammps_seed} {Emax} {ns_atoms.pressure} {move_params_cell['min_aspect_ratio']} "
+    fix_cmd = (f"fix NS all ns/cellmc {lammps_seed} {Emax} {move_params_cell['min_aspect_ratio']} {ns_atoms.pressure} "
                f"{submove_probs['volume']} {step_size_volume} "
                f"{submove_probs['stretch']} {step_size_stretch} "
                f"{submove_probs['shear']} {step_size_shear} "