Skip to content

Fix bug in step size tuning having to do with order of attempt/successs counts #25

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bernstei merged 1 commit into from
Jun 20, 2025
Merged

Fix bug in step size tuning having to do with order of attempt/successs counts #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 23 additions & 24 deletions pymatnext/ns.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -324,22 +324,19 @@ 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()

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:
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down
32 changes: 21 additions & 11 deletions pymatnext/ns_configs/ase_atoms/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []


Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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":
Expand Down Expand Up @@ -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
Expand All @@ -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]

Expand Down
12 changes: 12 additions & 0 deletions pymatnext/ns_configs/ase_atoms/walks_ase_calc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down