From 612cea8869e50b1ecd89e8d76b48956fc66dce35 Mon Sep 17 00:00:00 2001 From: Noam Bernstein Date: Wed, 18 Jun 2025 10:04:44 -0400 Subject: [PATCH] Fix bug in step size tuning having to do with order of attempt/success counts Switch to keeping track of step_size parameters and attempt/success statistics in dicts, rather than arrays, to reduce chance of inconsistencies in the order, which are what caused https://github.com/libAtoms/pymatnext/issues/24 --- pymatnext/ns.py | 47 +++++++++---------- pymatnext/ns_configs/ase_atoms/__init__.py | 32 ++++++++----- .../ns_configs/ase_atoms/walks_ase_calc.py | 12 +++++ 3 files changed, 56 insertions(+), 35 deletions(-) diff --git a/pymatnext/ns.py b/pymatnext/ns.py index 6c55ccf..af89d75 100644 --- a/pymatnext/ns.py +++ b/pymatnext/ns.py @@ -324,14 +324,11 @@ def step_size_tune(self, n_configs=1, min_accept_rate=0.25, max_accept_rate=0.5, adjust_factor: float, default 1.25 factor to adjust step size by """ - max_step_size = list(self.local_configs[0].max_step_size.values()) - step_size = [v / m for v, m in zip(self.local_configs[0].step_size.values(), max_step_size)] - step_size_names = list(self.local_configs[0].step_size.keys()) + max_step_size = self.local_configs[0].max_step_size + step_size = {k: self.local_configs[0].step_size[k] / max_step_size[k] for k in max_step_size} - n_params = len(step_size) - - last_too_small = [False] * n_params - last_too_big = [False] * n_params + last_too_small = {k: False for k in max_step_size} + last_too_big = {k: False for k in max_step_size} # save data from local_configs[0], which will be used for all pilot walks local_configs_0_data = self.local_configs[0].backup() @@ -339,7 +336,7 @@ def step_size_tune(self, n_configs=1, min_accept_rate=0.25, max_accept_rate=0.5, first_iter = True while True: - accept_freq = np.zeros((n_params, 2)) + accept_freq = {k: np.zeros(2, dtype=int) for k in max_step_size} for ns_config_i in range(n_configs): ns_config = self.local_configs[ns_config_i] if ns_config_i == 0: @@ -349,13 +346,16 @@ def step_size_tune(self, n_configs=1, min_accept_rate=0.25, max_accept_rate=0.5, 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) + for k in accept_freq: + accept_freq[k] += accept_freq_contribution[k] - accept_freq = self.comm.allreduce(accept_freq, self.MPI.SUM) + # order of dict must be same among MPI tasks, but this should really be a safe thing to assume + accept_freq_values = self.comm.allreduce(np.asarray(list(accept_freq.values())), self.MPI.SUM) + accept_freq = {k: v for k, v in zip(accept_freq.keys(), accept_freq_values)} 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", size, "max", max_size, "freq", freq) + for param_name, max_val in max_step_size.items(): + print(f"step_size_tune initial {param_name} size {self.local_configs[0].step_size[param_name]} max {max_val} freq {accept_freq[param_name]}") first_iter = False # It looks like the following should always give the same values, hence exit @@ -366,20 +366,19 @@ def step_size_tune(self, n_configs=1, min_accept_rate=0.25, max_accept_rate=0.5, # https://github.com/libAtoms/pymatnext/issues/20), a deadlock may occur. # Only fix is to make sure this doesn't happen (https://github.com/libAtoms/pymatnext/pull/23) done = [] - for param_i in range(n_params): - if accept_freq[param_i][0] > 0: - accept_rate_i = accept_freq[param_i, 1] / accept_freq[param_i, 0] - # only adjust if some steps were attempted - step_size[param_i], done_i, last_too_small[param_i], last_too_big[param_i] = self._tune_from_accept_rate( - step_size[param_i], last_too_small[param_i], last_too_big[param_i], accept_rate_i, - min_accept_rate, max_accept_rate, adjust_factor) + for param_name in max_step_size: + if accept_freq[param_name][0] > 0: + accept_rate = accept_freq[param_name][1] / accept_freq[param_name][0] + step_size[param_name], done_i, last_too_small[param_name], last_too_big[param_name] = self._tune_from_accept_rate( + step_size[param_name], last_too_small[param_name], last_too_big[param_name], accept_rate, + min_accept_rate, max_accept_rate, adjust_factor) else: done_i = True done.append(done_i) # set actual step sizes by rescaling by maximum - new_step_size = {k: v * m for k, v, m in zip(step_size_names, step_size, max_step_size)} + new_step_size = {k: step_size[k] * max_step_size[k] for k in max_step_size} for ns_config in self.local_configs: ns_config.step_size = new_step_size # make sure that config used as buffer also has correct step_size @@ -392,12 +391,12 @@ def step_size_tune(self, n_configs=1, min_accept_rate=0.25, max_accept_rate=0.5, if all(done): break - if any(np.asarray(step_size) < 1.0e-12): - raise RuntimeError(f"Stepsize got too small with automatic tuning {step_size} {step_size_names}") + if any(np.asarray(list(step_size.values())) < 1.0e-12): + raise RuntimeError(f"Stepsize got too small with automatic tuning {step_size}") if self.comm.rank == 0: - for name, size in zip(step_size_names, self.local_configs[0].step_size.values()): - print("Final step_size_tune", name, size) + for param_name, max_val in max_step_size.items(): + print(f"step_size_tune final {param_name} size {self.local_configs[0].step_size[param_name]}") # restore to original config self.local_configs[0].restore(local_configs_0_data) diff --git a/pymatnext/ns_configs/ase_atoms/__init__.py b/pymatnext/ns_configs/ase_atoms/__init__.py index ac17410..a8b6c71 100644 --- a/pymatnext/ns_configs/ase_atoms/__init__.py +++ b/pymatnext/ns_configs/ase_atoms/__init__.py @@ -44,14 +44,27 @@ class NSConfig_ASE_Atoms(): allocate_only: bool, default False do not actually initialize content, just allocate [NOTE: maybe refactor to separate allocation and initialization?] + + Class Attributes + ----------------- + filename_suffix: str filename suffix (including leading ".") for ase.io.write that preserves + all info in structure + n_quantities: int number of (float) quantities that need to be exchanged when communicating + config. Here, NS "energy", cell volume, number of atoms, and composition (if more than + 1 species) + _step_size_params: list(str) step size parameters, fixing order when step-size tuning quantities are + returned in a list + _max_E_hist: collections.deque history of max_E for effective temperature calculation + _walk_moves: list(str) types of walk moves + _Zs: list(int) atomic numbers in system """ filename_suffix = ".extxyz" n_quantities = -1 + _step_size_params = ["pos_gmc_each_atom", "cell_volume_per_atom", "cell_shear", "cell_stretch"] _max_E_hist = collections.deque(maxlen=1000) _walk_moves = ["gmc", "cell", "type"] - _step_size_params = ["pos_gmc_each_atom", "cell_volume_per_atom", "cell_shear", "cell_stretch"] _Zs = [] @@ -240,7 +253,7 @@ 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) + self.n_att_acc = {k: np.zeros(2, dtype=int) for k in NSConfig_ASE_Atoms._step_size_params} def end_calculator(self): @@ -416,6 +429,7 @@ def _prep_walk(self, params, vol_per_atom=None): # max step sizes self.max_step_size = params["max_step_size"].copy() + assert set(list(self.max_step_size.keys())) == set(self._step_size_params) # max step size for position GMC and cell volume defaults are scaled to volume per atom if self.max_step_size["pos_gmc_each_atom"] < 0.0: self.max_step_size["pos_gmc_each_atom"] = (vol_per_atom ** (1.0/3.0)) / 10.0 @@ -424,12 +438,10 @@ def _prep_walk(self, params, vol_per_atom=None): # actual step sizes self.step_size = params["step_size"].copy() + assert set(list(self.step_size.keys())) == set(self._step_size_params) # default to half the max for each type self.step_size = {k: (v if v >= 0.0 else self.max_step_size[k] / 2.0) for k, v in self.step_size.items()} - assert set(list(self.max_step_size.keys())) == set(self._step_size_params) - assert set(list(self.step_size.keys())) == set(self._step_size_params) - # store function pointers for moves self.walk_func = {} if self.calc_type == "ASE": @@ -689,10 +701,8 @@ def walk(self, Emax, walk_len, rng): Returns ------- - np.ndarray(n_move_types, 2): array containing number of attempted moves and - number of successful moves (2nd index 0 and 1) for each move type, with value - of 1st index from position of that step size type in - NSConfig_ASE_Atoms._step_size_params + dict(str param_name: ndarray([int n_attempts, int n_success])) dict containing number of attempted moves + and number of successful moves for each move type """ # if we fixed the number of steps for every move type, and only varied proportions, # we could do all the rng in a single call @@ -701,10 +711,10 @@ def walk(self, Emax, walk_len, rng): while walk_len_so_far < walk_len: move = rng.choice(NSConfig_ASE_Atoms._walk_moves, p=self.walk_prob) + # returns list of tuples with move param attempt/success statistics n_att_acc_walk = self.walk_func[move](self, Emax, rng) - # can this be done with a single numpy call somehow? for param, n_att, n_acc in n_att_acc_walk: - self.n_att_acc[NSConfig_ASE_Atoms._step_size_params.index(param)] += (n_att, n_acc) + self.n_att_acc[param] += (n_att, n_acc) walk_len_so_far += self.walk_traj_len[move] diff --git a/pymatnext/ns_configs/ase_atoms/walks_ase_calc.py b/pymatnext/ns_configs/ase_atoms/walks_ase_calc.py index beed512..8fd54e8 100644 --- a/pymatnext/ns_configs/ase_atoms/walks_ase_calc.py +++ b/pymatnext/ns_configs/ase_atoms/walks_ase_calc.py @@ -18,6 +18,10 @@ def walk_pos_gmc(ns_atoms, Emax, rng): maximum shifted energy rng: numpy.Generator random number generator + + Returns + ------- + [("pos_gmc_each_atom", int n_attempt, int n_success)] info on move params and attempts/successes """ atoms = ns_atoms.atoms # new random velocities @@ -119,6 +123,10 @@ def walk_cell(ns_atoms, Emax, rng): maximum shifted energy rng: numpy.Generator random number generator + + Returns + ------- + [("cell_volume_per_atom", int n_attempt, int n_success), ("cell_shear", ...), ("cell_stretch", ...)] info on move params and attempts/successes """ atoms = ns_atoms.atoms N_atoms = len(atoms) @@ -204,6 +212,10 @@ def walk_type(ns_atoms, Emax, rng): maximum shifted energy rng: numpy.Generator random number generator + + Returns + ------- + [] empty list (nominally cell move param attempt/success) """ atoms = ns_atoms.atoms N_atoms = len(atoms)