From a397e53cf5f5752fbad321b82dc42d57bb26c366 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 17 Aug 2026 17:29:26 +0200 Subject: [PATCH 1/3] Fix twelve findings from the 2026-08-17 review Replica exchange: - Add the de Broglie cross-term 3*sum_s(dN_s)*ln(Lambda_i/Lambda_j) to the grand-canonical swap criterion in both BatchedReplicaExchange and the MPI ReplicaExchange (lambda_dbs now rides on get_state). Mu ladders were already exact (the term cancels at shared T); temperature and joint ladders with fluctuating N were biased. - Reject non-finite swap exponents with a warning instead of silently inverting the NaN behavior (the old min(1, exp(nan)) accepted every swap). - Warn per rung on zero second-half swap acceptance (midpoint tally snapshot): the whole-run check could not see a ladder dead at one rung. - Report swap attempts, not doubled per-slot tallies, in the degeneracy warning, and emit it as a single log record. - Refuse factories that share atoms, move_selector or cells across replicas under batching. State and cells: - set_state restores step and exchange tallies when present (lossless restart round-trip); ReplicaExchange strips them on swaps instead. - is_point_exchangeable defaults on BaseCell, and the molecule moves fall back to is_point_inside for duck-typed cells (pre-1.4.0 contract). Alchemi calculators: - Raise when head= is passed with a pre-loaded MACEWrapper (it was silently ignored, evaluating multihead models on the pretrain head). - Drop the impossible workaround from the energy_only guard message. - Refuse forces-stripped models in AlchemiFCalculator and in run_md. Docs: - Ladder-spacing reference: include the mu_max endpoint (fencepost), sort descending isotherms, clip noise-negative dN/dmu. --- docs/replica_exchange_ladder_spacing.rst | 17 ++- mcpy/calculators/_alchemi_common.py | 9 ++ mcpy/calculators/alchemi_calculator.py | 15 +- mcpy/calculators/alchemi_f_calculator.py | 10 ++ mcpy/cell/base_cell.py | 16 +++ mcpy/cell/cell.py | 14 -- mcpy/ensembles/batched_replica_exchange.py | 101 ++++++++++++- mcpy/ensembles/canonical_ensemble.py | 10 +- mcpy/ensembles/grand_canonical_ensemble.py | 35 +++-- mcpy/ensembles/replica_exchange.py | 22 ++- mcpy/moves/molecule_displacement_move.py | 5 +- mcpy/moves/molecule_utils.py | 17 ++- tests/test_audit_regressions.py | 41 ++++-- tests/test_batched_re_acceptance.py | 157 ++++++++++++++++++++- tests/test_molecule_moves.py | 13 ++ 15 files changed, 421 insertions(+), 61 deletions(-) diff --git a/docs/replica_exchange_ladder_spacing.rst b/docs/replica_exchange_ladder_spacing.rst index 54ea9d4..d4ee24e 100644 --- a/docs/replica_exchange_ladder_spacing.rst +++ b/docs/replica_exchange_ladder_spacing.rst @@ -79,7 +79,9 @@ and neighbouring rungs exchange when their distributions overlap, i.e. when \Delta\mu(\mu) = \frac{1}{\sqrt{\beta\,dN/d\mu}} , \qquad - n_\text{rungs} = \int \sqrt{\beta\,\frac{dN}{d\mu}}\; d\mu . + n_\text{gaps} = \int \sqrt{\beta\,\frac{dN}{d\mu}}\; d\mu , + \qquad + n_\text{rungs} = n_\text{gaps} + 1 . Use :math:`dN/d\mu` from a measured isotherm, **not** the per-rung :math:`\sigma_N` observed in an unconverged run: a stuck rung reports a @@ -118,14 +120,23 @@ is the disease being diagnosed. """ beta = 1.0 / (8.617333e-5 * temperature) mu, n_ads = np.asarray(mu, float), np.asarray(n_ads, float) + order = np.argsort(mu) # desorption scans arrive descending + mu, n_ads = mu[order], n_ads[order] slope = np.diff(n_ads) / np.diff(mu) # dN/dmu per interval midpoints = 0.5 * (mu[1:] + mu[:-1]) fine = np.linspace(mu_min, mu_max, 2001) - density = np.sqrt(beta * np.interp(fine, midpoints, slope)) + # Clip noise-induced negative dN/dmu to zero: a NaN here would + # propagate into ``cumulative`` and crash np.interp cryptically. + density = np.sqrt(beta * np.maximum(np.interp(fine, midpoints, slope), + 0.0)) cumulative = np.concatenate([[0.0], np.cumsum( 0.5 * (density[1:] + density[:-1]) * np.diff(fine))]) - return np.interp(np.arange(0.0, cumulative[-1], 1.0), cumulative, fine) + # linspace, not arange: the ladder must reach mu_max, so the last gap + # shrinks rather than the top rung being dropped. + targets = np.linspace(0.0, cumulative[-1], + int(np.ceil(cumulative[-1])) + 1) + return np.interp(targets, cumulative, fine) Calibration diff --git a/mcpy/calculators/_alchemi_common.py b/mcpy/calculators/_alchemi_common.py index b79f0d6..0002d1f 100644 --- a/mcpy/calculators/_alchemi_common.py +++ b/mcpy/calculators/_alchemi_common.py @@ -48,6 +48,15 @@ def _load_model( head: Union[str, int, None] = None, ) -> MACEWrapper: if isinstance(checkpoint, MACEWrapper): + if head is not None: + raise ValueError( + 'head= cannot be applied to a pre-loaded MACEWrapper: the ' + 'head is baked in at wrap time, so it would be silently ' + 'ignored and a multihead model would evaluate on head 0 ' + '(usually the pretrain head). Load from the .model path ' + 'with head=, or wrap the raw model in _HeadMACEWrapper ' + 'yourself before sharing it.' + ) return checkpoint if compile_model: # torch >= 2.12.1 donates compiled-backward buffers that the FIRE diff --git a/mcpy/calculators/alchemi_calculator.py b/mcpy/calculators/alchemi_calculator.py index e9cebd2..79bc357 100644 --- a/mcpy/calculators/alchemi_calculator.py +++ b/mcpy/calculators/alchemi_calculator.py @@ -56,7 +56,7 @@ def __init__( ) -> None: if energy_only and isinstance(checkpoint, MACEWrapper): # energy_only mutates the wrapper's own model_config (below), and a - # pre-loaded wrapper is shared by every calculator built from it — + # pre-loaded wrapper is shared by every calculator built from it -- # including an AlchemiFCalculator whose FIRE relaxation needs the # forces this would switch off. There is no way to scope the change # to one calculator, so refuse instead of breaking the other one @@ -65,8 +65,8 @@ def __init__( "energy_only=True cannot be combined with a pre-loaded " "MACEWrapper: dropping 'forces' from its active outputs would " "also disable forces for every other calculator sharing that " - "wrapper. Pass the checkpoint path so this calculator loads " - "its own model, or set energy_only on all of the sharers." + "wrapper. Pass the checkpoint path instead, so this " + "calculator loads its own model." ) self.device = device self.dtype = dtype @@ -163,6 +163,15 @@ def run_md( model is loaded. ``friction`` is in 1/fs, ``dt`` in fs, ``temperature`` in K. ``FixAtoms`` constraints are honored. """ + if self.energy_only: + # The integrator would otherwise die steps deep inside nvalchemi + # ("NVTLangevin requires forces...") without naming the cause. + raise ValueError( + 'run_md needs forces, but this calculator was built with ' + 'energy_only=True (forces are stripped from the model ' + 'outputs). Build a separate calculator without energy_only ' + 'for MD.' + ) _run_langevin_md( self.model, self._nl_config, atoms, temperature=temperature, friction=friction, dt=dt, steps=steps, diff --git a/mcpy/calculators/alchemi_f_calculator.py b/mcpy/calculators/alchemi_f_calculator.py index 2a9856d..3029805 100644 --- a/mcpy/calculators/alchemi_f_calculator.py +++ b/mcpy/calculators/alchemi_f_calculator.py @@ -97,6 +97,16 @@ def __init__( self._optimizer_cls = _ALCHEMI_OPTIMIZERS[optimizer] self.optimizer_name = optimizer self.model = _load_model(checkpoint, device, dtype, enable_cueq, compile_model, head) + if 'forces' not in self.model.model_config.active_outputs: + # A model reused from an energy_only AlchemiCalculator (its + # ``.model`` attribute) has forces stripped; FIRE would only fail + # steps deep inside nvalchemi without naming the cause. + raise ValueError( + "the provided model has 'forces' disabled in its active " + 'outputs (an energy_only AlchemiCalculator strips them), but ' + 'FIRE relaxation needs forces. Load a fresh wrapper for this ' + 'calculator instead of sharing an energy-only one.' + ) self._nl_config = self.model.model_config.neighbor_config def get_potential_energy(self, atoms: Atoms) -> float: diff --git a/mcpy/cell/base_cell.py b/mcpy/cell/base_cell.py index 9c82ee2..a30afe6 100644 --- a/mcpy/cell/base_cell.py +++ b/mcpy/cell/base_cell.py @@ -39,3 +39,19 @@ def get_volume(self): :return: Volume of the cell. """ raise NotImplementedError("This method should be implemented in subclasses.") + + def is_point_exchangeable(self, point): + """Whether a molecule whose center of mass sits at ``point`` may be + exchanged with the reservoir (counted, deleted, displaced). + + The point counterpart of ``get_atoms_specie_inside_cell``, and the + predicate the molecule moves use. It is separate from + ``is_point_inside`` because a region may deliberately accept + molecules it would never *propose* -- :class:`CustomCell` drops the z + upper bound so a desorbed molecule stays deletable instead of + accumulating forever. Defined here on the ABC so cells written + against the original single-predicate contract (only + ``is_point_inside``) keep working; cells whose two regions coincide + (the box, the sphere, the dome) inherit it unchanged. + """ + return self.is_point_inside(point) diff --git a/mcpy/cell/cell.py b/mcpy/cell/cell.py index 581a785..e0c45a4 100644 --- a/mcpy/cell/cell.py +++ b/mcpy/cell/cell.py @@ -95,17 +95,3 @@ def is_point_inside(self, point): any cell type (the region cells implement a real test). """ return True - - def is_point_exchangeable(self, point): - """Whether a molecule whose center of mass sits at ``point`` may be - exchanged with the reservoir (counted, deleted, displaced). - - The point counterpart of :meth:`get_atoms_specie_inside_cell`, and the - predicate the molecule moves use. It is separate from - :meth:`is_point_inside` because a region may deliberately accept - molecules it would never *propose* — :class:`CustomCell` drops the z - upper bound so a desorbed molecule stays deletable instead of - accumulating forever. Cells whose two regions coincide (the box, the - sphere, the dome) inherit this delegation unchanged. - """ - return self.is_point_inside(point) diff --git a/mcpy/ensembles/batched_replica_exchange.py b/mcpy/ensembles/batched_replica_exchange.py index b5e4ae0..28a1ca4 100644 --- a/mcpy/ensembles/batched_replica_exchange.py +++ b/mcpy/ensembles/batched_replica_exchange.py @@ -75,6 +75,26 @@ def __init__( gcmc_factory(mu=mu, rank=i) for i, mu in enumerate(self.mus) ] + # Batching shares one process across replicas, so nothing isolates + # them but the factory contract above. A factory written for the MPI + # class (module-level atoms/moves, safe under one-process-per-rank) + # silently corrupts every replica here; refuse it up front. + shared = [ + attr for attr in ('atoms', 'move_selector') + if len({id(getattr(r, attr, None)) for r in self.replicas}) + < self.n_replicas + ] + cell_ids = [id(c) for r in self.replicas for c in getattr(r, 'cells', [])] + if len(set(cell_ids)) < len(cell_ids): + shared.append('cells') + if shared: + raise ValueError( + f"replicas share {', '.join(shared)}; gcmc_factory must build " + 'a fresh ensemble with its own atoms, cells and move_selector ' + 'per replica (see the module docstring). Shared objects mix ' + 'the replicas\' volumes, counters and RNG streams silently.' + ) + if not hasattr(calculator, 'get_potential_energies'): raise TypeError( "calculator must implement get_potential_energies(atoms_list). " @@ -92,6 +112,11 @@ def __init__( self.exchange_attempts = [0] * self.n_replicas self.exchange_successes = [0] * self.n_replicas + # Midpoint snapshot of the tallies, for the second-half per-rung + # degeneracy check (early free swaps inflate the cumulative rate + # forever -- docs/replica_exchange_ladder_spacing.rst). + self._mid_attempts = None + self._mid_successes = None self._outfile = outfile self._outfile_handle = None @@ -120,6 +145,9 @@ def run(self) -> None: try: for step in range(self.gcmc_steps): + if step == self.gcmc_steps // 2: + self._mid_attempts = list(self.exchange_attempts) + self._mid_successes = list(self.exchange_successes) self._batched_gcmc_step() if step > 0 and step % self.exchange_interval == 0: @@ -306,6 +334,12 @@ def _accept_swap(self, i: int, j: int) -> bool: identically zero and would accept every swap, destroying the ladder. With shared β it reduces to β Σ_s (μ_i,s - μ_j,s)(N_Y,s - N_X,s), the rule ``ReplicaExchange._exchange_prob_mu`` implements for MPI. + + At unequal temperatures the grand-canonical weight also carries the + de Broglie factors Λ_s^(-3 N_s), and Λ depends on T, so swapping + configurations with different particle counts shifts the joint weight + by a further 3 Σ_s (N_X,s - N_Y,s) ln(Λ_i,s / Λ_j,s). The term + vanishes on a μ ladder (shared T) and in LJ units (Λ = 1). """ ri, rj = self.replicas[i], self.replicas[j] beta_i, beta_j = ri.units.beta, rj.units.beta @@ -315,6 +349,17 @@ def _accept_swap(self, i: int, j: int) -> bool: phi_ij = rj._minimum_score(ri.atoms, ri.E_old) # config X, μ of j delta = (beta_i * phi_ii + beta_j * phi_jj - beta_i * phi_ji - beta_j * phi_ij) + if beta_i != beta_j: + delta += self._de_broglie_cross_term(ri, rj) + if not np.isfinite(delta): + # A NaN would silently reject every draw from here on (the old + # min(1, exp(delta)) form silently *accepted*); a dead rung must + # not be quieter than a diverged relaxation. + self.logger.warning( + 'swap %d<->%d: non-finite acceptance exponent (delta=%r); ' + 'rejecting. A NaN/inf replica energy usually means a ' + 'diverged relaxation.', i, j, delta) + return False # exp() overflows to inf for a large positive delta; short-circuit. p = 1.0 if delta >= 0.0 else float(np.exp(delta)) self.logger.debug( @@ -324,6 +369,23 @@ def _accept_swap(self, i: int, j: int) -> bool: ) return self.rng.get_uniform() < p + @staticmethod + def _de_broglie_cross_term(ri, rj) -> float: + """3 Σ_s (N_X,s - N_Y,s) ln(Λ_i,s / Λ_j,s) for configs X (slot i) and + Y (slot j). Species are counted exactly as the replicas' own + acceptance counts them (molecules for molecular μ keys, atoms + otherwise), via ``GrandCanonicalEnsemble._species_count``. A replica + without μ (e.g. canonical) exchanges nothing with a reservoir and + contributes no Λ factors.""" + term = 0.0 + for specie in getattr(ri, '_mu', None) or {}: + dn = (ri._species_count(ri.atoms, specie) + - rj._species_count(rj.atoms, specie)) + if dn: + term += 3.0 * dn * np.log( + ri.units.lambda_dbs[specie] / rj.units.lambda_dbs[specie]) + return term + @staticmethod def _grand_potential(r) -> float: """Φ = E - Σ_s μ_s N_s for a grand-canonical replica; bare E without μ. @@ -400,8 +462,17 @@ def _warn_if_ladder_degenerate(self) -> None: adsorbate that has not begun to adsorb) legitimately accepts every early swap, and a warning that fires during equilibration would be noise. The trade is that this is a post-mortem. + + A ladder can also die at individual rungs while the others stay + alive, which the whole-run tally cannot show (those early free swaps + inflate it forever); ``_warn_dead_rungs`` covers that with the + second-half per-rung test from + docs/replica_exchange_ladder_spacing.rst. """ - attempts, successes = sum(self.exchange_attempts), sum(self.exchange_successes) + # Both slots of a pair are tallied per attempt, so the slot sums are + # exactly twice the number of swap attempts; report attempts. + attempts = sum(self.exchange_attempts) // 2 + successes = sum(self.exchange_successes) // 2 if attempts < _SWAP_DEGENERACY_MIN_ATTEMPTS: return if successes == attempts: @@ -412,12 +483,32 @@ def _warn_if_ladder_degenerate(self) -> None: reason = ('no swap was ever accepted -- rungs are too far apart ' 'to exchange, so each replica sampled in isolation') else: + self._warn_dead_rungs() return self.logger.warning( - 'replica-exchange ladder looks degenerate: %d/%d swaps accepted. ' - 'Widen or tighten the %s spacing.', - successes, attempts, 'mu' if self.mus is not None else 'temperature') - self.logger.warning(' %s', reason) + 'replica-exchange ladder looks degenerate: %d/%d swap attempts ' + 'accepted. Widen or tighten the %s spacing: %s', + successes, attempts, + 'mu' if self.mus is not None else 'temperature', reason) + + def _warn_dead_rungs(self) -> None: + """Per-rung second-half check: a rung whose swaps all fail after the + midpoint is dead even when the cumulative tally looks mixed (early + equilibration swaps accept freely and inflate it permanently).""" + if self._mid_attempts is None: + return + for k, (att0, suc0) in enumerate( + zip(self._mid_attempts, self._mid_successes)): + att = self.exchange_attempts[k] - att0 + suc = self.exchange_successes[k] - suc0 + # Half the whole-run floor: a rung sees roughly half the rounds. + if att >= _SWAP_DEGENERACY_MIN_ATTEMPTS // 2 and suc == 0: + self.logger.warning( + 'replica-exchange rung %d accepted no swaps in the second ' + 'half of the run (0/%d): the ladder is dead there; ' + 'respace the %s ladder around that rung ' + '(docs/replica_exchange_ladder_spacing.rst).', + k, att, 'mu' if self.mus is not None else 'temperature') def _log_summary(self) -> None: """Final consolidated summary: per-replica move acceptances.""" diff --git a/mcpy/ensembles/canonical_ensemble.py b/mcpy/ensembles/canonical_ensemble.py index 4465bac..5cd3a2b 100644 --- a/mcpy/ensembles/canonical_ensemble.py +++ b/mcpy/ensembles/canonical_ensemble.py @@ -95,8 +95,14 @@ def set_state(self, state): # config swap, without relying on the pickled .info surviving MPI. self.atoms.info.setdefault("key_value_pairs", {}) self.atoms.info["key_value_pairs"]["potential_energy"] = state["energy"] - # Step count and exchange statistics stay with the slot, not the - # config -- see GrandCanonicalEnsemble.set_state. + # Restored only when present, so the restart round-trip is lossless; + # ReplicaExchange strips these slot-bound keys before a swap -- see + # GrandCanonicalEnsemble.set_state. + self._step = state.get("step", self._step) + self.exchange_attempts = state.get( + "exchange_attempts", self.exchange_attempts) + self.exchange_successes = state.get( + "exchange_successes", self.exchange_successes) def _acceptance_condition(self, potential_diff: float) -> bool: if potential_diff <= 0: diff --git a/mcpy/ensembles/grand_canonical_ensemble.py b/mcpy/ensembles/grand_canonical_ensemble.py index fac1b6d..024b5f9 100644 --- a/mcpy/ensembles/grand_canonical_ensemble.py +++ b/mcpy/ensembles/grand_canonical_ensemble.py @@ -89,6 +89,7 @@ def get_state(self) -> Dict[str, any]: "mu": self._mu, "temperature": self._temperature, "beta": self.units.beta, + "lambda_dbs": self.units.lambda_dbs, "step": self._step, "exchange_attempts": self.exchange_attempts, "exchange_successes": self.exchange_successes, @@ -98,11 +99,16 @@ def set_state(self, state: Dict[str, any]) -> None: self.atoms = state["atoms"] self.E_old = state["energy"] self.n_atoms = state["n_atoms"] - # Only the configuration travels. Step count and exchange statistics - # describe the *slot*, not the config: ``ReplicaExchange`` passes a - # full ``get_state()`` dict here on every accepted swap, so restoring - # them would trade the two ranks' swap tallies back and forth and make - # the per-rank "Accepted Exchange (%)" column meaningless. + # Step count and exchange statistics are restored only when the dict + # carries them, so the get_state()/set_state() restart round-trip is + # lossless. ``ReplicaExchange`` strips them before applying a swap: + # they describe the *slot*, not the config, and trading them between + # ranks made the per-rank "Accepted Exchange (%)" column meaningless. + self._step = state.get("step", self._step) + self.exchange_attempts = state.get( + "exchange_attempts", self.exchange_attempts) + self.exchange_successes = state.get( + "exchange_successes", self.exchange_successes) # The configuration changed under the cells: refresh their free # volumes now, or the next insertion/deletion acceptance would use # the previous configuration's volume (GCMC only recalculates on @@ -169,20 +175,23 @@ def write_outfile(self, step: int = None, energy: float = None) -> None: except (OSError, AttributeError): self.logger.exception("Error writing to file %s", self._outfile) + def _species_count(self, atoms: Atoms, specie: str) -> int: + """Particle count for one μ species: molecules for molecular keys, + atoms otherwise. The one counting convention shared by + ``_minimum_score`` and the replica-exchange swap criterion.""" + if specie in self.units.molecules: + template = self.units.molecules[specie] + return len(find_molecules( + atoms, sorted(template.get_chemical_symbols()))) + return atoms.get_chemical_symbols().count(specie) + def _minimum_score(self, atoms: Atoms, energy: float) -> float: """Grand potential Ω = E − Σ μ_i N_i. Comparing raw E across moves that change N is not meaningful in the grand canonical ensemble. Molecular species count molecules, atomic species count atoms.""" score = energy - symbols = atoms.get_chemical_symbols() for specie, mu in self._mu.items(): - if specie in self.units.molecules: - template = self.units.molecules[specie] - n = len(find_molecules( - atoms, sorted(template.get_chemical_symbols()))) - else: - n = symbols.count(specie) - score -= mu * n + score -= mu * self._species_count(atoms, specie) return score def _acceptance_condition(self, diff --git a/mcpy/ensembles/replica_exchange.py b/mcpy/ensembles/replica_exchange.py index 6856dbe..5412ef5 100644 --- a/mcpy/ensembles/replica_exchange.py +++ b/mcpy/ensembles/replica_exchange.py @@ -118,6 +118,19 @@ def _exchange_prob_T(self, state1, state2): energy2 -= mu2[specie] * state2['atoms'].symbols.count(specie) delta = (beta2 - beta1) * (energy2 - energy1) + + # Grand-canonical weights also carry Lambda_s^(-3 N_s) per exchangeable + # species, and Lambda depends on T, so configs with different particle + # counts shift the joint weight by a de Broglie cross-term (zero in LJ + # units, where Lambda = 1). lambda_dbs rides on each replica's state; + # a units-less ensemble (CanonicalEnsemble) has no mu and skips this. + l1, l2 = state1.get('lambda_dbs'), state2.get('lambda_dbs') + if mu1 and mu2 and l1 and l2: + for specie in mu1: + dn = (state1['atoms'].symbols.count(specie) + - state2['atoms'].symbols.count(specie)) + if dn: + delta += 3.0 * dn * np.log(l1[specie] / l2[specie]) exchange_prob = min(1.0, np.exp(delta)) self.logger.debug( f"beta1: {beta1:.3f}, beta2: {beta2:.3f}, E1: {energy1:.3f}, " @@ -199,7 +212,14 @@ def do_exchange(self): if u < exchange_prob: self.logger.debug(f"Accepted exchange with rank {partner_rank}") - self.gcmc.set_state(partner_state) + # Only the configuration travels on a swap. step and the exchange + # tallies describe the slot, not the config; set_state restores + # them when present (the restart round-trip), so strip them here. + config_state = { + k: v for k, v in partner_state.items() + if k not in ('step', 'exchange_attempts', 'exchange_successes') + } + self.gcmc.set_state(config_state) return True else: self.logger.debug(f"Rejected exchange with rank {partner_rank}") diff --git a/mcpy/moves/molecule_displacement_move.py b/mcpy/moves/molecule_displacement_move.py index 500e7d0..f5625e2 100644 --- a/mcpy/moves/molecule_displacement_move.py +++ b/mcpy/moves/molecule_displacement_move.py @@ -3,7 +3,8 @@ from ase.geometry import find_mic, wrap_positions from .base_move import BaseMove -from .molecule_utils import find_molecules, molecule_com, random_rotation_matrix +from .molecule_utils import (exchangeable_predicate, find_molecules, + molecule_com, random_rotation_matrix) from ..cell import Cell @@ -68,7 +69,7 @@ def do_trial_move(self, atoms) -> Atoms: new_com = com + shift if np.any(np.asarray(atoms.pbc)): new_com = wrap_positions([new_com], atoms.cell, pbc=atoms.pbc)[0] - if not self.cell.is_point_exchangeable(new_com): + if not exchangeable_predicate(self.cell)(new_com): return False, 0, self.name if self.max_angle is None: rotation = random_rotation_matrix(self.rng) diff --git a/mcpy/moves/molecule_utils.py b/mcpy/moves/molecule_utils.py index 2f2a176..c925f05 100644 --- a/mcpy/moves/molecule_utils.py +++ b/mcpy/moves/molecule_utils.py @@ -37,6 +37,17 @@ def molecule_com(atoms, members): return com +def exchangeable_predicate(cell): + """The cell's exchangeable-region point predicate. + + Falls back to ``is_point_inside`` for user cells written against the + original single-predicate duck-type contract (pre-1.4.0), which + ``BaseCell`` subclasses inherit as a default but plain duck-typed cells + do not. + """ + return getattr(cell, 'is_point_exchangeable', None) or cell.is_point_inside + + def find_molecules(atoms, template_symbols, cell=None): """Index groups of whole molecules matching ``template_symbols``. @@ -45,11 +56,12 @@ def find_molecules(atoms, template_symbols, cell=None): the cell counts as exchangeable (``cell.is_point_exchangeable``) are returned; ``None`` skips the spatial filter (used by the grand-potential bookkeeping). The exchangeable region is not always the proposal region: - see :meth:`mcpy.cell.Cell.is_point_exchangeable`. + see :meth:`mcpy.cell.BaseCell.is_point_exchangeable`. """ ids = atoms.arrays.get('molecule_id') if ids is None: return [] + predicate = None if cell is None else exchangeable_predicate(cell) symbols = np.asarray(atoms.get_chemical_symbols()) groups = [] for mid in np.unique(ids): @@ -58,8 +70,7 @@ def find_molecules(atoms, template_symbols, cell=None): members = np.where(ids == mid)[0] if sorted(symbols[members]) != list(template_symbols): continue - if cell is not None and not cell.is_point_exchangeable( - molecule_com(atoms, members)): + if predicate is not None and not predicate(molecule_com(atoms, members)): continue groups.append(members) return groups diff --git a/tests/test_audit_regressions.py b/tests/test_audit_regressions.py index ac2bfa5..abf670d 100644 --- a/tests/test_audit_regressions.py +++ b/tests/test_audit_regressions.py @@ -455,11 +455,15 @@ def get_potential_energies(self, atoms_list): np.testing.assert_allclose(replica.atoms.positions, before) -# Replica exchange swaps configurations, not slot bookkeeping (bug: get_state -# emitted step/exchange counters and set_state restored them, so accepted -# swaps traded the two ranks' tallies) +# Replica exchange swaps configurations, not slot bookkeeping (bug: accepted +# swaps traded the two ranks' tallies). ReplicaExchange strips the slot-bound +# keys before applying a swap; a full get_state() dict (the restart +# round-trip) restores everything. -def test_gcmc_set_state_keeps_local_step_and_exchange_counters(): +_SLOT_KEYS = ('step', 'exchange_attempts', 'exchange_successes') + + +def test_gcmc_set_state_swap_payload_keeps_local_counters(): g = _gcmc(_h2(), [NullCell()], MoveSelector([1], [_DirtyBailingMove()]), {'H': 0.0}) g._step, g.exchange_attempts, g.exchange_successes = 40, 8, 3 @@ -467,21 +471,42 @@ def test_gcmc_set_state_keeps_local_step_and_exchange_counters(): partner = _gcmc(_h2(), [NullCell()], MoveSelector([1], [_DirtyBailingMove()]), {'H': 0.0}) partner._step, partner.exchange_attempts, partner.exchange_successes = 40, 8, 7 - g.set_state(partner.get_state()) + swap_payload = {k: v for k, v in partner.get_state().items() + if k not in _SLOT_KEYS} # what ReplicaExchange applies + g.set_state(swap_payload) assert (g._step, g.exchange_attempts, g.exchange_successes) == (40, 8, 3) assert g.atoms is partner.atoms # the configuration still travels -def test_canonical_set_state_keeps_local_step_and_exchange_counters(tmp_path): +def test_gcmc_set_state_full_dict_round_trips_counters(): + """Restart contract: get_state() -> set_state() must be lossless, or an + appended outfile renumbers its rows from 0 after a restart.""" + g = _gcmc(_h2(), [NullCell()], MoveSelector([1], [_DirtyBailingMove()]), + {'H': 0.0}) + partner = _gcmc(_h2(), [NullCell()], + MoveSelector([1], [_DirtyBailingMove()]), {'H': 0.0}) + partner._step, partner.exchange_attempts, partner.exchange_successes = 40, 8, 7 + g.set_state(partner.get_state()) + assert (g._step, g.exchange_attempts, g.exchange_successes) == (40, 8, 7) + + +def test_canonical_set_state_swap_payload_keeps_local_counters(tmp_path): mc = _canonical(tmp_path) mc._step, mc.exchange_attempts, mc.exchange_successes = 12, 4, 1 - mc.set_state({'atoms': _h2(), 'energy': -1.0, 'step': 99, - 'exchange_attempts': 40, 'exchange_successes': 39}) + mc.set_state({'atoms': _h2(), 'energy': -1.0}) assert (mc._step, mc.exchange_attempts, mc.exchange_successes) == (12, 4, 1) assert mc._current_energy == -1.0 +def test_canonical_set_state_full_dict_round_trips_counters(tmp_path): + mc = _canonical(tmp_path) + mc._step, mc.exchange_attempts, mc.exchange_successes = 12, 4, 1 + mc.set_state({'atoms': _h2(), 'energy': -1.0, 'step': 99, + 'exchange_attempts': 40, 'exchange_successes': 39}) + assert (mc._step, mc.exchange_attempts, mc.exchange_successes) == (99, 40, 39) + + # Per-interval acceptance counters are cleared even with the outfile disabled # (bug: write_outfile returned before reset_counters, so interval_ratios() # silently degraded into total_ratios()) diff --git a/tests/test_batched_re_acceptance.py b/tests/test_batched_re_acceptance.py index 479b93d..04d6e1e 100644 --- a/tests/test_batched_re_acceptance.py +++ b/tests/test_batched_re_acceptance.py @@ -6,6 +6,7 @@ with lightweight replica stubs. """ import logging +from collections import defaultdict import numpy as np import pytest @@ -14,9 +15,13 @@ class _FakeUnits: - def __init__(self, beta): + def __init__(self, beta, lambda_dbs=None): self.beta = beta self.molecules = {} + # Lambda = 1 unless a test opts in: keeps the de Broglie cross-term + # zero so the Lambda-free expectations below stay exact. + self.lambda_dbs = (lambda_dbs if lambda_dbs is not None + else defaultdict(lambda: 1.0)) class _FakeAtoms: @@ -28,8 +33,8 @@ def get_chemical_symbols(self): class _FakeReplica: - def __init__(self, beta, energy, mu, symbols): - self.units = _FakeUnits(beta) + def __init__(self, beta, energy, mu, symbols, lambda_dbs=None): + self.units = _FakeUnits(beta, lambda_dbs) self.E_old = energy self._mu = mu self.atoms = _FakeAtoms(symbols) @@ -39,11 +44,14 @@ def _minimum_score(self, atoms, energy): (these fakes carry no molecular species) so _grand_potential's delegation to it can be exercised without a real ensemble.""" score = energy - symbols = atoms.get_chemical_symbols() for specie, mu in (self._mu or {}).items(): - score -= mu * symbols.count(specie) + score -= mu * self._species_count(atoms, specie) return score + def _species_count(self, atoms, specie): + """Mirrors GrandCanonicalEnsemble._species_count (atomic only).""" + return atoms.get_chemical_symbols().count(specie) + class _FakeRng: def __init__(self, u): @@ -193,6 +201,8 @@ def test_temperature_ladder_swap_unchanged_by_the_mu_fix(): def test_temperature_ladder_uphill_swap_matches_legacy_probability(): + """With Lambda = 1 (the fakes' default, and LJ units in production) the + de Broglie cross-term vanishes and the Lambda-free form is exact.""" ri = _FakeReplica(beta=1.0, energy=10.0, mu={'Ag': 1.0}, symbols=['Ag'] * 18) rj = _FakeReplica(beta=0.5, energy=20.0, mu={'Ag': 1.0}, symbols=['Ag'] * 5) re = _bare_re() @@ -205,6 +215,56 @@ def test_temperature_ladder_uphill_swap_matches_legacy_probability(): assert not re._accept_swap(0, 1) +def test_temperature_ladder_swap_includes_de_broglie_term(): + """Grand-canonical T-ladder swaps carry the Lambda_s^(-3N_s) factors of + the stationary weight: with Lambda_i/Lambda_j = e^0.1 and N differing by + 13, the exponent gains 3 * 13 * 0.1 = 3.9 over the Lambda-free form.""" + ri = _FakeReplica(beta=1.0, energy=10.0, mu={'Ag': 1.0}, + symbols=['Ag'] * 18, + lambda_dbs={'Ag': float(np.exp(0.1))}) + rj = _FakeReplica(beta=0.5, energy=20.0, mu={'Ag': 1.0}, + symbols=['Ag'] * 5, lambda_dbs={'Ag': 1.0}) + re = _bare_re() + re.replicas = [ri, rj] + # Lambda-free delta = -11.5 (test above); cross-term adds +3.9. + p = float(np.exp(-11.5 + 3.9)) + re.rng = _FakeRng(p * 0.99) + assert re._accept_swap(0, 1) + re.rng = _FakeRng(p * 1.01) + assert not re._accept_swap(0, 1) + + +def test_de_broglie_term_is_symmetric_in_slot_order(): + ri = _FakeReplica(beta=1.0, energy=10.0, mu={'Ag': 1.0}, + symbols=['Ag'] * 18, + lambda_dbs={'Ag': float(np.exp(0.1))}) + rj = _FakeReplica(beta=0.5, energy=20.0, mu={'Ag': 1.0}, + symbols=['Ag'] * 5, lambda_dbs={'Ag': 1.0}) + re = _bare_re() + re.replicas = [ri, rj] + p = float(np.exp(-11.5 + 3.9)) + for u in (p * 0.99, p * 1.01): + re.rng = _FakeRng(u) + assert re._accept_swap(0, 1) == re._accept_swap(1, 0) + + +def test_nan_energy_rejects_the_swap_and_warns(caplog): + """A diverged relaxation (NaN energy) must sever the swap loudly: the + old min(1, exp(nan)) form silently *accepted* every such swap, and the + plain comparison form would silently reject.""" + re = _bare_re() + re.replicas = [ + _FakeReplica(beta=1.0, energy=float('nan'), mu={'Ag': 1.0}, + symbols=['Ag'] * 5), + _FakeReplica(beta=0.5, energy=20.0, mu={'Ag': 1.0}, + symbols=['Ag'] * 5), + ] + re.rng = _FakeRng(1e-12) + with caplog.at_level(logging.WARNING, logger='test'): + assert re._accept_swap(0, 1) is False + assert 'non-finite' in ' '.join(r.getMessage() for r in caplog.records) + + def test_consolidated_status_line(caplog): """One console line covers all replicas; per-replica detail stays in files.""" import logging @@ -240,6 +300,8 @@ def _tally_re(attempts, successes, mus=None): pt.logger = logging.getLogger('mcpy.ensembles.batched_replica_exchange') pt.exchange_attempts = attempts pt.exchange_successes = successes + pt._mid_attempts = None + pt._mid_successes = None pt.mus = mus return pt @@ -272,8 +334,12 @@ def test_degenerate_ladder_warning_names_the_ladder_kind(caplog): with caplog.at_level(logging.WARNING, logger='mcpy.ensembles.batched_replica_exchange'): pt._warn_if_ladder_degenerate() - text = ' '.join(r.getMessage() for r in caplog.records) - assert '80/80' in text and 'mu spacing' in text + # One record (log aggregators must not strand the reason line), counting + # swap attempts, not the doubled per-slot tallies (both slots of a pair + # are incremented per attempt: 40 attempts, not 80). + assert len(caplog.records) == 1 + text = caplog.records[0].getMessage() + assert '40/40' in text and 'mu spacing' in text caplog.clear() pt = _tally_re([40, 40], [40, 40], mus=None) @@ -282,3 +348,80 @@ def test_degenerate_ladder_warning_names_the_ladder_kind(caplog): pt._warn_if_ladder_degenerate() assert 'temperature spacing' in ' '.join( r.getMessage() for r in caplog.records) + + +def test_dead_rung_warned_from_second_half_tally(caplog): + """The documented real failure: the cumulative tally looks mixed (early + free swaps inflate it forever), but one rung accepted nothing after the + midpoint. The whole-run check cannot see it; the per-rung second-half + check must.""" + pt = _tally_re([40, 80, 80, 80, 40], [20, 30, 30, 20, 0], + mus=[{'CO': -1.0}] * 5) + pt._mid_attempts = [20, 40, 40, 40, 20] + pt._mid_successes = [10, 20, 22, 12, 0] + with caplog.at_level(logging.WARNING, + logger='mcpy.ensembles.batched_replica_exchange'): + pt._warn_if_ladder_degenerate() + text = ' '.join(r.getMessage() for r in caplog.records) + assert 'rung 4' in text + assert 'rung 0' not in text # rung 0 kept accepting in the second half + + +def test_no_dead_rung_warning_without_midpoint_snapshot(caplog): + """A run too short to snapshot (or a healthy mixed tally) stays quiet.""" + pt = _tally_re([40, 40], [12, 12]) + with caplog.at_level(logging.WARNING, + logger='mcpy.ensembles.batched_replica_exchange'): + pt._warn_if_ladder_degenerate() + assert not caplog.records + + +# -------------------------------------------------------------------------- +# Replica isolation. A factory written for the MPI class (module-level shared +# atoms/moves, safe under one-process-per-rank) silently corrupts every +# replica under batching; the constructor must refuse it. +# -------------------------------------------------------------------------- + +class _StubBatchCalc: + def get_potential_energies(self, atoms_list): + return [0.0] * len(atoms_list) + + +def test_shared_replica_state_is_refused(): + from types import SimpleNamespace + + shared_atoms = object() + + def factory(T, rank): + return SimpleNamespace(atoms=shared_atoms, move_selector=object(), + cells=[]) + + with pytest.raises(ValueError, match='share'): + BatchedReplicaExchange(factory, _StubBatchCalc(), + temperatures=[300.0, 400.0]) + + +def test_shared_cells_are_refused(): + from types import SimpleNamespace + + shared_cell = object() + + def factory(T, rank): + return SimpleNamespace(atoms=object(), move_selector=object(), + cells=[shared_cell]) + + with pytest.raises(ValueError, match='cells'): + BatchedReplicaExchange(factory, _StubBatchCalc(), + temperatures=[300.0, 400.0]) + + +def test_distinct_replica_state_is_accepted(): + from types import SimpleNamespace + + def factory(T, rank): + return SimpleNamespace(atoms=object(), move_selector=object(), + cells=[object()]) + + re = BatchedReplicaExchange(factory, _StubBatchCalc(), + temperatures=[300.0, 400.0]) + assert re.n_replicas == 2 diff --git a/tests/test_molecule_moves.py b/tests/test_molecule_moves.py index 633a0c5..1218111 100644 --- a/tests/test_molecule_moves.py +++ b/tests/test_molecule_moves.py @@ -93,6 +93,19 @@ def is_point_exchangeable(self, point): assert find_molecules(atoms, sorted(['O', 'H', 'H']), _NowhereCell()) == [] +def test_duck_typed_cell_with_only_is_point_inside_still_works(): + """Pre-1.4.0 contract: a user cell exposing only ``is_point_inside`` + (no BaseCell inheritance) must keep working with the molecule moves via + the fallback predicate.""" + class _LegacyDuckCell: + def is_point_inside(self, point): + return False + + atoms = _water_box() + assert find_molecules(atoms, sorted(['O', 'H', 'H']), + _LegacyDuckCell()) == [] + + def test_random_rotation_matrix_is_rotation(): rng = RandomNumberGenerator(seed=7) r1 = random_rotation_matrix(rng) From 918ac5463004be77101e1ce23e3d252fc0fb3558 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Tue, 18 Aug 2026 15:47:54 +0200 Subject: [PATCH 2/3] Fix the 2026-08-18 nvalchemi audit findings and harden the relax loop - eval() the wrapper on local checkpoint loads: the nn.Module train-mode default made MACE run the force autograd with create_graph=True and retain_graph=True on every FIRE and Langevin step (values unaffected) - drop the stale manual requires_grad in AlchemiCalculator and run energy_only forwards under no_grad, restoring the graph-free forward (peak 242 vs 342 MB per forward on Ag225) - zero FixAtoms forces after the bootstrap compute so step 0 sees the same frozen-force state as every hook-managed step - remove the donated_buffer=False workaround: the crash it papered over was a train-mode symptom; verified with compile+cuEq on the alias and local checkpoint paths across varying-N GCMC steps - add NaNDetectorHook to the FIRE and Langevin loops and warn when a relaxation returns at the step cap (unconverged-energy bias trap) - correct the FIRE dt docstring (nvalchemi defaults dt_max to 10*dt, unlike ASE's 1 fs cap) and document why AlchemiBrownianMove needs no position wrapping Verified: verify_compact_parity PASS (max compact-vs-serial delta 0.0002 eV, FixAtoms rows bit-identical), 198 tests pass, flake8 clean. Post-fix numbers on the RTX 5090: 3.7 s/GCMC step at 3.7k atoms (4.9 GiB allocated), single-relax ceiling 32k atoms. --- mcpy/calculators/_alchemi_common.py | 34 +++++++++++++++++------- mcpy/calculators/alchemi_calculator.py | 17 ++++++++---- mcpy/calculators/alchemi_f_calculator.py | 34 ++++++++++++++++++++---- mcpy/moves/alchemi_brownian_move.py | 6 +++++ 4 files changed, 71 insertions(+), 20 deletions(-) diff --git a/mcpy/calculators/_alchemi_common.py b/mcpy/calculators/_alchemi_common.py index 0002d1f..b11f7c3 100644 --- a/mcpy/calculators/_alchemi_common.py +++ b/mcpy/calculators/_alchemi_common.py @@ -16,7 +16,7 @@ from nvalchemi.hooks import DynamicsContext from nvalchemi.dynamics import NVTLangevin, initialize_velocities from nvalchemi.dynamics.base import DynamicsStage -from nvalchemi.dynamics.hooks import FreezeAtomsHook +from nvalchemi.dynamics.hooks import FreezeAtomsHook, NaNDetectorHook class _HeadMACEWrapper(MACEWrapper): @@ -58,12 +58,11 @@ def _load_model( 'yourself before sharing it.' ) return checkpoint - if compile_model: - # torch >= 2.12.1 donates compiled-backward buffers that the FIRE - # force backward still needs; GCMC runs crash within a few steps - # ("compiled with non-empty donated buffers"). Zero measured cost. - from torch._functorch import config as _functorch_config - _functorch_config.donated_buffer = False + # The donated_buffer=False workaround (b333c8e) is gone: the crash it + # papered over came from the wrapper running in train mode + # (create_graph=True made the FIRE force backward need the buffers the + # compiled backward donates); with wrapper.eval() below, donation is safe + # (verified compile+cuEq, alias and local paths, varying-N GCMC). # Local .model file: load directly. MACEWrapper.from_checkpoint treats the # string as a download alias, so it cannot open local paths. ``head`` # selects a multihead model's head by name or index. @@ -96,7 +95,14 @@ def _load_model( for param in raw.parameters(): param.requires_grad = False raw = torch.compile(raw) - return MACEWrapper(raw) if idx is None else _HeadMACEWrapper(raw, idx) + wrapper = MACEWrapper(raw) if idx is None else _HeadMACEWrapper(raw, idx) + # MACEWrapper.forward passes training=self.training down to MACE, where + # it becomes create_graph/retain_graph on the force autograd. nn.Module + # defaults to train mode, and from_checkpoint's wrapper.eval() (alias + # path below) never runs here — without this every force pass builds a + # retained second-order graph. + wrapper.eval() + return wrapper if head is not None: raise ValueError( "head= is only supported when loading a local .model path; " @@ -214,7 +220,8 @@ def _run_langevin_md( random_seed=seed, ) - freeze_hooks = _freeze_hook_for(batch, _fixed_indices(atoms)) + fixed = _fixed_indices(atoms) + freeze_hooks = _freeze_hook_for(batch, fixed) nl_hook = NeighborListHook(nl_config, max_neighbors=max_neighbors) opt = NVTLangevin( model=model, @@ -223,13 +230,20 @@ def _run_langevin_md( friction=friction, random_seed=seed, n_steps=steps, - hooks=freeze_hooks, + # NaN forces/energy would otherwise propagate silently into the + # proposal the ensemble scores; the hook raises at the offending step. + hooks=(freeze_hooks or []) + [NaNDetectorHook()], ) opt.register_hook(nl_hook, stage=DynamicsStage.BEFORE_COMPUTE) # Bootstrap: build NL and compute initial forces before the MD loop. _build_nl(batch, nl_hook) opt.compute(batch) + if fixed: + # FreezeAtomsHook zeros frozen forces only at AFTER_POST_UPDATE; the + # bootstrap compute() runs outside the hook loop, so step 0 would + # otherwise kick and displace the frozen rows once. + batch.forces[fixed] = 0.0 opt.run(batch) _write_back_positions(atoms, batch) diff --git a/mcpy/calculators/alchemi_calculator.py b/mcpy/calculators/alchemi_calculator.py index 79bc357..3150f4c 100644 --- a/mcpy/calculators/alchemi_calculator.py +++ b/mcpy/calculators/alchemi_calculator.py @@ -82,6 +82,16 @@ def __init__( # fp32 rounding. See docs/superpowers/specs/2026-06-09-...-design.md. self.model.model_config.active_outputs.discard('forces') + def _forward(self, batch): + """Model forward. The wrapper manages positions grad itself when forces + are active; with energy_only nothing needs grad, and no_grad keeps the + forward graph-free even when model parameters still require grad + (compile_model=False local loads).""" + if self.energy_only: + with torch.no_grad(): + return self.model(batch) + return self.model(batch) + def get_potential_energy(self, atoms: Atoms) -> float: """ Single forward pass — no geometry relaxation. @@ -97,11 +107,9 @@ def get_potential_energy(self, atoms: Atoms) -> float: Potential energy in eV. """ batch = _make_batch(atoms, self.device, self.dtype) - # positions must have grad enabled — MACEWrapper always computes forces via autograd - batch.positions.requires_grad_(True) nl_hook = NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) _build_nl(batch, nl_hook) - out = self.model(batch) + out = self._forward(batch) return float(out['energy'].sum().item()) def get_potential_energies( @@ -140,10 +148,9 @@ def get_potential_energies( for start, stop in chunk_ranges(len(atoms_list), cs): chunk = atoms_list[start:stop] batch = _make_multi_batch(chunk, self.device, self.dtype) - batch.positions.requires_grad_(True) nl_hook = NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) _build_nl(batch, nl_hook) - model_out = self.model(batch) + model_out = self._forward(batch) out.append(_per_graph_energies(model_out['energy'], len(chunk))) return np.concatenate(out) diff --git a/mcpy/calculators/alchemi_f_calculator.py b/mcpy/calculators/alchemi_f_calculator.py index 3029805..ff011bf 100644 --- a/mcpy/calculators/alchemi_f_calculator.py +++ b/mcpy/calculators/alchemi_f_calculator.py @@ -11,6 +11,7 @@ from nvalchemi.hooks.neighbor_list import NeighborListHook from nvalchemi.dynamics import FIRE as AlchemiFIRE, FIRE2 as AlchemiFIRE2, ConvergenceHook from nvalchemi.dynamics.base import DynamicsStage +from nvalchemi.dynamics.hooks import NaNDetectorHook from ..utils.chunking import chunk_ranges from ._alchemi_common import ( @@ -56,8 +57,11 @@ class AlchemiFCalculator: compile_model : bool torch.compile. Best to pre-warm in __init__ if reusing. dt : float - FIRE initial timestep (default 1.0). Matches ASE FIRE's dtmax; - benchmarks show dt=1.0 converges in ~half the steps vs dt=0.1. + FIRE initial timestep in fs (default 1.0, the value of ASE FIRE's + dtmax; benchmarks show dt=1.0 converges in ~half the steps vs 0.1). + Note nvalchemi grows the adaptive timestep up to dt_max = 10*dt, + unlike ASE's 1 fs cap; per-step displacement stays clamped by + maxstep=0.2 A either way. optimizer : str 'fire' (default, classic FIRE) or 'fire2' (Guénolé et al variant — typically converges in fewer steps). @@ -131,7 +135,8 @@ def get_potential_energy(self, atoms: Atoms) -> float: batch.forces = torch.zeros_like(batch.positions) batch.energy = torch.zeros(1, 1, device=self.device, dtype=self.dtype) - freeze_hooks = _freeze_hook_for(batch, _fixed_indices(atoms)) + fixed = _fixed_indices(atoms) + freeze_hooks = _freeze_hook_for(batch, fixed) nl_hook = NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) opt = self._optimizer_cls( @@ -139,17 +144,30 @@ def get_potential_energy(self, atoms: Atoms) -> float: dt=self.dt, convergence_hook=ConvergenceHook.from_fmax(self.fmax), n_steps=self.steps, - hooks=freeze_hooks, + # NaN forces/energy would otherwise flow silently into the GCMC + # acceptance decision; the hook raises at the offending step. + hooks=(freeze_hooks or []) + [NaNDetectorHook()], ) opt.register_hook(nl_hook, stage=DynamicsStage.BEFORE_COMPUTE) # Bootstrap: build NL and compute initial forces before the FIRE loop _build_nl(batch, nl_hook) opt.compute(batch) + if fixed: + # FreezeAtomsHook zeros frozen forces only at AFTER_POST_UPDATE; + # the bootstrap compute() runs outside the hook loop, so step 0 + # would otherwise displace the frozen rows once. + batch.forces[fixed] = 0.0 opt.run(batch) self.last_relax_steps = int(opt.step_count) self.total_relax_steps += self.last_relax_steps + if self.last_relax_steps >= self.steps: + # Known bias trap: a too-tight cap returns truncated relaxations + # whose energies skew the GCMC acceptance. + logger.warning('FIRE reached the %d-step cap (fmax=%.3g eV/A); ' + 'returned energy may be unconverged', + self.steps, self.fmax) logger.debug("FIRE relaxation: %d/%d steps (fmax=%.3g eV/A)", self.last_relax_steps, self.steps, self.fmax) _write_back_positions(atoms, batch) @@ -177,12 +195,14 @@ def _relax_batch(self, atoms_list: List[Atoms]) -> tuple: dt=self.dt, convergence_hook=ConvergenceHook.from_fmax(self.fmax), n_steps=self.steps, - hooks=freeze_hooks, + hooks=(freeze_hooks or []) + [NaNDetectorHook()], ) opt.register_hook(nl_hook, stage=DynamicsStage.BEFORE_COMPUTE) _build_nl(batch, nl_hook) opt.compute(batch) + if fixed: + batch.forces[fixed] = 0.0 # see get_potential_energy bootstrap return self._run_compacted(opt, batch, atoms_list) @@ -224,6 +244,10 @@ def _run_compacted(self, opt, batch, atoms_list: List[Atoms]) -> tuple: # or the next hook context builds an out-of-bounds mask. opt._last_converged = None if alive.size: # step cap hit: harvest the stragglers as-is + logger.warning('%d/%d graphs reached the %d-step FIRE cap ' + '(fmax=%.3g eV/A); their energies may be ' + 'unconverged', alive.size, len(atoms_list), + self.steps, self.fmax) self._harvest(batch, np.arange(alive.size), alive, energies, atoms_list) finally: diff --git a/mcpy/moves/alchemi_brownian_move.py b/mcpy/moves/alchemi_brownian_move.py index 03f5a6b..2272001 100644 --- a/mcpy/moves/alchemi_brownian_move.py +++ b/mcpy/moves/alchemi_brownian_move.py @@ -15,6 +15,12 @@ class AlchemiBrownianMove(BaseMove): This move is pure Python; the nvalchemi dependency enters only through ``calculator.run_md`` at call time. + Positions are not wrapped back into the box after the MD trajectory. + This is safe: the periodic ``Cell`` selects deletion candidates by + species only (no position test), MACE handles PBC through neighbor-list + shifts, and region cells (DomeCell, CustomCell) treat out-of-region + drift exactly as they do for any displacement move. + Parameters ---------- calculator : AlchemiCalculator | AlchemiFCalculator From d268d7aae57ab6261165e8a671c11964cfe723d9 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Tue, 18 Aug 2026 16:24:45 +0200 Subject: [PATCH 3/3] Record the metric-capture mechanism research for issue #27 --- docs/research/metric-capture.md | 102 ++++++++++++++++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 docs/research/metric-capture.md diff --git a/docs/research/metric-capture.md b/docs/research/metric-capture.md new file mode 100644 index 0000000..4edbb49 --- /dev/null +++ b/docs/research/metric-capture.md @@ -0,0 +1,102 @@ +# Metric-capture mechanisms for the batched-RE regression benchmark + +Research ticket: [#27](https://github.com/farrisric/mcpy/issues/27), child of wayfinder map [#23](https://github.com/farrisric/mcpy/issues/23). +Date: 2026-08-18. +Line numbers refer to commit `918ac54` (branch point of `research/metric-capture`). + +## Question + +How should the harness capture peak allocated GPU memory, a GPU-utilization figure, per-step wall time split into relaxation vs overhead, and per-replica relax iteration counts, without perturbing the timings it measures? + +## Recommendations (one mechanism per metric) + +### (a) Peak allocated GPU memory: `torch.cuda.reset_peak_memory_stats()` after warmup + `torch.cuda.max_memory_allocated()` at the end + +This is the established pattern in `benchmark/internal/re_batched_scaling.py`. +The peak counter is reset before the run at line 173 (after `torch.cuda.empty_cache()` at line 172) and reset again inside the timed step wrapper at line 167 once warmup completes, so first-step allocation and compile transients are excluded. +The single read happens once per config at line 189 (`peak_torch = torch.cuda.max_memory_allocated() / 1024 ** 2`). +`remeasure_supported_np.py:54-61` uses the same reset-then-read pattern per relaxation. +Use allocated, never reserved: `gcmc_memory_growth.py:5-14` documents that reserved tracks the caching-allocator pool, which never shrinks and grows with fragmentation, and the repo's history shows earlier reserved-metric numbers were misleading. +Reading a counter costs nothing on the timed path, so this mechanism is perturbation-free by construction. +Nothing is missing; the harness copies the `re_batched_scaling.py` pattern verbatim. + +### (b) GPU utilization: the threaded nvidia-smi sampler (`GpuSampler`), windowed to post-warmup + +`re_batched_scaling.py:64-93` already defines a `GpuSampler` thread that polls `nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader,nounits` via subprocess every 0.15 s and stores `(perf_counter, mem_MB, util_pct)` tuples. +The post-warmup window is selected at line 192 (`sampler.window(warm_done_t[0])`) and mean/max utilization are computed at lines 194-196. +The sampler runs in its own thread and subprocess, never touches the CUDA stream, and its CPU cost is negligible, so it does not perturb the timings. +It also yields whole-device memory as a free cross-check of the allocated figure (the docstring at lines 9-10 notes smi includes the CUDA context and resident model). +No pynvml or `torch.cuda.utilization` is used anywhere in the repo; introducing one would add a dependency for no gain. +Nothing is missing; lift `GpuSampler` into the harness and report the post-warmup mean (with max as a secondary column if wanted). + +### (c) Per-step wall time split into relaxation vs overhead: synced `perf_counter` brackets at two nesting levels + +The house timing pattern is `time.perf_counter()` before, `torch.cuda.synchronize()` immediately before the stop read, and warmup steps excluded from the aggregate. +It appears in `re_batched_scaling.py:161-168` (step wrapper), `benchmark_alchemi_vs_mace.py:69-82` (`time_call`), and `remeasure_supported_np.py:54-60`. +For the total step time, wrap `BatchedReplicaExchange._batched_gcmc_step` exactly as `re_batched_scaling.py:159-170` already does; the method (defined at `mcpy/ensembles/batched_replica_exchange.py:208-233`) has no timing of its own. +For the relaxation share, wrap `calculator.get_potential_energies`, which is the single funnel every batched relaxation passes through: the contract is enforced at `batched_replica_exchange.py:98-102` and the call sites are `_rebatch_initial_energies` (line 202) and `_batched_single_move` (line 278). +Overhead is then total step time minus the summed relax time inside that step; no third timer is needed. +The added `torch.cuda.synchronize()` at the end of the relax bracket is effectively free because the calculator already forces a device-to-host copy to return Python floats; the sync only moves the wait to a named point. +Do not rely on the library's own `_last_step_seconds` (`mcpy/ensembles/grand_canonical_ensemble.py:356-359`): it is unsynced, only reaches the console logger, and stays 0.0 under batched RE because `_batched_gcmc_step` bypasses `GrandCanonicalEnsemble._run` entirely (`batched_replica_exchange.py:228-233`). +Missing piece: only the two harness-side wrappers themselves; no library change is required, and keeping the timing in the harness keeps the production path unperturbed. + +### (d) Per-replica relax iteration counts: record graph retirement steps in `_run_compacted` + +The existing metric surface is the `last_relax_steps` / `total_relax_steps` counter pair on `AlchemiFCalculator` (`mcpy/calculators/alchemi_f_calculator.py:95-96`), set from `opt.step_count` on the single-graph path (lines 163-164) and aggregated across chunks at lines 315-316. +`remeasure_supported_np.py:64,128` and `benchmark_alchemi_vs_mace.py:148-159` already consume these counters. +They are not per-replica in batched mode: `last_relax_steps` becomes the per-chunk max and `total_relax_steps` a sum, because nvalchemi's `BaseDynamics.step_count` is batch-global (`nvalchemi/dynamics/base.py:1455,1950`). +The per-graph signal does exist transiently: `_run_compacted` steps FIRE manually and receives the converged-graph indices each iteration (`batch, conv = opt.step(batch)` at `alchemi_f_calculator.py:231`, `n_steps += 1` at 232), then retires those graphs at lines 240-242 without recording when each one converged. +`BaseDynamics.step` returning the converged-index tensor (`nvalchemi/dynamics/base.py:1865`) is the only per-graph iteration signal nvalchemi exposes on the paths mcpy uses, so no upstream plumbing is needed. +Missing piece: a small addition to `AlchemiFCalculator` that appends `(original graph row -> n_steps)` at the existing retirement point into a `last_relax_steps_per_graph` list (plus the cap-hit case for graphs that never converge), which the harness reads after each step and maps back to replicas via the chunk layout. +The bookkeeping is a few integer appends at retirement points where CPU-side work (`_harvest`, lines 257-271) already happens, so it does not perturb the timed path. +If windowed reporting is wanted, follow the house counter style of dual interval-plus-cumulative tallies with a reset method, as in `MoveSelector` (`mcpy/moves/move_selector.py:66-71,138-165`). + +## Inventory of existing measurement code + +### `benchmark/internal/re_batched_scaling.py` + +Peak allocated memory via `reset_peak_memory_stats` (lines 167, 173) and `max_memory_allocated` (line 189), with resident model-plus-context memory captured separately via `memory_allocated()` after a warm forward (line 236) and attached to every row (line 245). +Whole-device memory and utilization via the threaded `GpuSampler` nvidia-smi poller (lines 64-93), windowed post-warmup (lines 192-196). +Step timing via a monkey-patched `_batched_gcmc_step` wrapper with `synchronize()` before the stop read (lines 159-168) and warmup exclusion (line 197, default warmup 2 at line 383). +OOM configs are caught and recorded as `status='oom'` (lines 180-183). +No relax-iteration counting and no alloc-conf handling. + +### `benchmark/internal/gcmc_memory_growth.py` + +Per-step instantaneous triple read after `synchronize()` (lines 84-91): `memory_allocated`, `memory_reserved`, and a one-shot nvidia-smi helper (lines 41-48). +Its docstring (lines 5-14) is the canonical allocated-vs-reserved pitfall text and documents running under `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` set in the environment. +End-of-run `empty_cache()` probe distinguishes fragmentation from a leak (lines 99-106). +No timing, no utilization, no iteration counting. + +### `benchmark/internal/remeasure_supported_np.py` + +Per-relax `synchronize(); empty_cache(); reset_peak_memory_stats()` then reads of both `max_memory_allocated` and `max_memory_reserved` (lines 50-65, repeated for the batched point at lines 96-104). +Relax timing is median-of-3 `perf_counter` with pre-stop `synchronize()` (lines 57-60, 69) and explicit uncounted compile warmups (lines 87, 91). +GCMC throughput is aggregate: one `perf_counter` pair around `gcmc.run(GCMC_STEPS)` divided by step count (lines 123-127), which cannot split relax from overhead. +Reads `calc.last_relax_steps` per relax (line 64) and `calc.total_relax_steps` after the run (line 128). + +### `benchmark/internal/benchmark_alchemi_vs_mace.py` + +Generic `time_call(fn, n_warmup, n_repeats)` helper: untimed warmups, then `sync(); t0; fn(); sync(); append` (lines 69-82), aggregated as a mean (lines 116-117, 153-154). +Resets `total_relax_steps = 0` per system and averages over warmup-plus-repeats calls (lines 148-149, 157-159); note the warmup relaxations are included in the step average even though excluded from timing. +No memory or utilization capture. + +### `mcpy/ensembles/` + +`BaseEnsemble` initializes `_last_step_seconds` (`base_ensemble.py:51`); `GrandCanonicalEnsemble._run` fills it unsynced and logs it to the console only, on outfile-interval steps (`grand_canonical_ensemble.py:356-365`). +`BatchedReplicaExchange` has no timing, no memory, and no iteration instrumentation; `_batched_gcmc_step` (lines 208-233) is bare and calls each replica's `write_outfile` directly, bypassing `_run`. +`write_outfile` rows carry step, atom counts, energy, and move-acceptance ratios only (`grand_canonical_ensemble.py:143-176`; `batched_replica_exchange.py:525-543`). +The exchange counters with midpoint snapshots (`batched_replica_exchange.py:113-119,148-150`) are a ready-made template for windowed metrics. +No ensemble reads `last_relax_steps` or `total_relax_steps`. + +### `mcpy/calculators/` + +`AlchemiFCalculator` holds the only relax-iteration counters (details in recommendation d). +`MACE_F_Calculator` mirrors the counter pair from ASE `opt.nsteps` (`mace_f_calculator.py:27-28,54-55`). +`BaseCalculator` discards its optimizer and exposes no counters (`base_calculator.py:14-22`). +`AlchemiCalculator` is energy-only with no relaxation or instrumentation; `energy_only=True` drops forces so no autograd graph is built (`alchemi_calculator.py:78-83`). +`_alchemi_common.py` imports `DynamicsContext` and the freeze/NaN hooks (lines 16, 19) but records no metrics; a custom `AFTER_STEP` or `ON_CONVERGE` hook receiving `DynamicsContext` (`nvalchemi/hooks/_context.py:60`, carrying `step_count` and `converged_mask`) is a viable alternative capture point for (d), but the `_run_compacted` retirement point is simpler because the loop already owns `n_steps` and the compaction bookkeeping. + +### Orphan data note + +`benchmark/internal/relax_step_spread.csv` and `relax_step_spread_smoke.csv` exist but no script in the repo references them; the script that produced per-graph relax-step spread data appears to have been deleted.