From cf70f6500913472c460575af03b01bee54874c05 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Tue, 18 Aug 2026 16:24:45 +0200 Subject: [PATCH 01/20] 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. From 4121970298edd9eeedc1e6f192644bc06fb891a1 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:10:35 +0200 Subject: [PATCH 02/20] Restore FixAtoms rows in float64 after an Alchemi relaxation The batch position array is float32, so assigning the float64 frozen coordinates into it truncated them: the "restored" FixAtoms rows came back about 5e-7 A off their pinned positions at slab-sized coordinates, in x and y, growing with the coordinate magnitude. It did not accumulate across steps (each restore re-truncates an already-truncated value) and it flipped no acceptance decision in the Ag(111) fixture, but every Alchemi-relaxed frozen substrate was slightly off its constraint. Widen to float64 before restoring, in both write-back paths (the single structure and the batched harvest). Frozen drift goes from 5.4e-7 A to 2.0e-15 A. Found by tests/test_ag111_gcmc_golden.py on its first run, which is also the first test ever to execute this package. Also folds the duplicated inline FixAtoms loop into _fixed_indices. --- mcpy/calculators/_alchemi_common.py | 13 ++++++------- mcpy/calculators/alchemi_f_calculator.py | 3 ++- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/mcpy/calculators/_alchemi_common.py b/mcpy/calculators/_alchemi_common.py index b11f7c3..938486e 100644 --- a/mcpy/calculators/_alchemi_common.py +++ b/mcpy/calculators/_alchemi_common.py @@ -148,14 +148,13 @@ def _build_nl(batch: Batch, nl_hook: NeighborListHook) -> None: def _write_back_positions(atoms: Atoms, batch: Batch) -> None: """Copy relaxed positions from batch back to atoms, skipping FixAtoms-constrained indices.""" - relaxed = batch.positions.detach().cpu().numpy() - fixed: set[int] = set() - for c in atoms.constraints: - if isinstance(c, FixAtoms): - fixed.update(c.index.tolist()) + # float64: the batch is float32, and assigning the float64 frozen rows + # into a float32 array would truncate them, leaving the "restored" atoms + # about 5e-7 A off their pinned positions at slab-sized coordinates. + relaxed = batch.positions.detach().cpu().numpy().astype(np.float64) + fixed = _fixed_indices(atoms) if fixed: - relaxed = relaxed.copy() - relaxed[list(fixed)] = atoms.positions[list(fixed)] + relaxed[fixed] = atoms.positions[fixed] atoms.positions = relaxed diff --git a/mcpy/calculators/alchemi_f_calculator.py b/mcpy/calculators/alchemi_f_calculator.py index ff011bf..2900f92 100644 --- a/mcpy/calculators/alchemi_f_calculator.py +++ b/mcpy/calculators/alchemi_f_calculator.py @@ -259,7 +259,8 @@ def _harvest(batch, rows, alive, energies, atoms_list) -> None: """Write energy + relaxed positions of ``rows`` (batch row indices) back to their originating Atoms objects, restoring FixAtoms rows.""" e = _per_graph_energies(batch.energy, int(batch.num_graphs)) - pos = batch.positions.detach().cpu().numpy() + # float64 before restoring frozen rows: see _write_back_positions. + pos = batch.positions.detach().cpu().numpy().astype(np.float64) batch_idx = batch.batch_idx.detach().cpu().numpy() for row in rows: g = int(alive[row]) From 4b8786e7a981e37074338b9adef691ea8399d5c4 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:10:55 +0200 Subject: [PATCH 03/20] Add the Ag(111)/O golden test for the GPU stack Fourteen end-to-end assertions on a 4x4x3 Ag(111) slab with the bottom layer fixed, run through AlchemiFCalculator with the MACE small-density-agnesi checkpoint. Two fixtures: single-replica GCMC, and a three-rung mu ladder through BatchedReplicaExchange. This is the first test to execute mcpy/calculators at all (the package sat at 0% coverage), and the first to exercise the Alchemi FIRE path, FreezeAtomsHook, the batched swap rule and the batched forward pass. Everything is seeded, so the accept/reject sequence is exact rather than statistical: same seeds reproduce every discrete decision, with positions agreeing to ~4e-6 A and energies to ~1e-4 eV (float32 GPU sums are not bitwise reproducible, the decisions are). Parameters are the calibrated set from the old examples/gcmc_custom_cell.py: a=4.165, the O-Ag distance 2.068 used both as the exclusion radius and as the region depth below the top layer, height 7, min_insert 0.5, and the -0.176 bulk Ag correction. delta_mu_O = -0.5 is the balance point where both insertion and deletion accept; past about -0.7 nothing adsorbs and the test would silently stop covering insertion, which test_both_move_types_accept guards. The ladder test asserts coverage is monotonic in mu (O = 8, 5, 2 for delta_mu = -0.45, -0.50, -0.55). That is the assertion a degenerate ladder fails while every other number still looks plausible: the all-accept mode averages the rungs together and flattens exactly this gradient. Opt-in, not CI: needs a GPU, nvalchemi and a local checkpoint, and takes about six minutes. Gated on MCPY_GPU_TESTS rather than an import probe, because importing a misconfigured torch raises SIGBUS, which no try/except can catch and which would take the whole collection down. --- tests/test_ag111_gcmc_golden.py | 340 ++++++++++++++++++++++++++++++++ 1 file changed, 340 insertions(+) create mode 100644 tests/test_ag111_gcmc_golden.py diff --git a/tests/test_ag111_gcmc_golden.py b/tests/test_ag111_gcmc_golden.py new file mode 100644 index 0000000..5a6e482 --- /dev/null +++ b/tests/test_ag111_gcmc_golden.py @@ -0,0 +1,340 @@ +"""Golden end-to-end test: O GCMC on Ag(111) with the production GPU stack. + +Opt-in, not CI: needs a GPU, nvalchemi, and a local MACE checkpoint. Run with + + MCPY_GPU_TESTS=1 /path/to/envs/alchemi/bin/python -m pytest \ + tests/test_ag111_gcmc_golden.py -q + +The gate is an environment variable rather than an import probe on purpose: +importing torch in an env where it is misconfigured raises SIGBUS, which no +``try/except`` can catch, so a probe would take the whole collection down. + +What this covers that the torch-free suite cannot: the Alchemi FIRE relaxation +with ``FreezeAtomsHook``, the real O/Ag(111) energetics, and the fact that the +whole GCMC loop (CustomCell free volume with periodic images, ``min_insert`` +overlap rejection, de Broglie acceptance, arrays+constraints rollback, +trajectory writing) reproduces exactly when seeded. + +Parameters come from ``examples/gcmc_custom_cell.py``, the calibrated set for +this checkpoint: a=4.165, the O-Ag exclusion radius 2.068 used both as the +radius and as the region depth below the top layer, height 7, min_insert 0.5, +and the -0.176 bulk Ag correction. +""" +import os +import pathlib + +import numpy as np +import pytest +from ase.build import fcc111 +from ase.constraints import FixAtoms +from ase.geometry import find_mic +from ase.io import read + +from mcpy.cell import CustomCell +from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble +from mcpy.moves import DeletionMove, InsertionMove +from mcpy.moves.move_selector import MoveSelector + +AG_LATTICE = 4.165 +R_O_AG = 2.068 # O-Ag exclusion radius, and the region depth below z_top +CELL_HEIGHT = 7.0 +MIN_INSERT = 0.5 +SEED = 777 +STEPS = 60 + +CHECKPOINT = pathlib.Path(os.environ.get( + 'MCPY_TEST_CHECKPOINT', + pathlib.Path.home() / '.cache/mace/macesmalldensityagnesistressmodel')) + +# Derived from this checkpoint by ``derive_mu_bulk``/``derive_mu_gas``; stable +# across runs, so they are pinned rather than re-derived on every invocation +# (that would add eight relaxations per run). A different checkpoint changes +# these, and the goldens below must then be regenerated together with them. +MU_AG = -2.9929 # derive_mu_bulk('Ag', a=4.165) - 0.176 +MU_O = -4.9093 # derive_mu_gas('O2') = E(O2)/2 +DELTA_MU_O = -0.5 # balance point: both insertion and deletion accept + +# Golden values: seed 777, 60 steps, 4x4x3 slab, bottom layer fixed. +# Energies are float32 GPU sums, so they repeat to ~1e-4 eV rather than +# bitwise; every discrete decision repeats exactly. Regenerate only on a +# deliberate physics change, never by pasting whatever a new run produced. +GOLDEN_N = [48, 53, 53, 51, 52, 52, 53] +GOLDEN_ACCEPTED = [5, 10] +GOLDEN_ENERGY = -151.923111 + +pytestmark = pytest.mark.skipif( + not os.environ.get('MCPY_GPU_TESTS') or not CHECKPOINT.exists(), + reason='set MCPY_GPU_TESTS=1 and provide a MACE checkpoint (GPU required)') + + +def _build(seed_cell): + atoms = fcc111('Ag', a=AG_LATTICE, size=(4, 4, 3), periodic=True, vacuum=10.0) + # Bottom layer only: the two upper layers stay free to relax and to open up + # for subsurface O. + atoms.set_constraint(FixAtoms(indices=[a.index for a in atoms if a.tag == 3])) + z_top = float(atoms.positions[atoms.get_tags() == 1, 2].max()) + cell = CustomCell(atoms, custom_height=CELL_HEIGHT, bottom_z=z_top - R_O_AG, + species_radii={'Ag': R_O_AG, 'O': 0.0}, seed=seed_cell) + return atoms, cell, z_top + + +def _run(tmp_path, steps=STEPS, write_interval=10): + from mcpy.calculators import AlchemiFCalculator + + tmp_path.mkdir(parents=True, exist_ok=True) + seed_del, seed_ins, seed_cell, seed_sel, seed_ens = ( + int(s) for s in np.random.SeedSequence(SEED).generate_state(5, dtype=np.uint32) + ) + atoms, cell, _ = _build(seed_cell) + + move_selector = MoveSelector( + [1, 1], + [DeletionMove(cell, species=['O'], seed=seed_del), + InsertionMove(cell, species=['O'], min_insert=MIN_INSERT, seed=seed_ins)], + seed=seed_sel, + ) + calculator = AlchemiFCalculator( + checkpoint=str(CHECKPOINT), steps=300, fmax=0.05, + device='cuda', compile_model=False, optimizer='fire2', + ) + gcmc = GrandCanonicalEnsemble( + atoms=atoms, + cells=[cell], + calculator=calculator, + mu={'Ag': MU_AG, 'O': MU_O + DELTA_MU_O}, + units_type='metal', + species=['O'], + temperature=500.0, + move_selector=move_selector, + random_seed=seed_ens, + traj_file=str(tmp_path / 'traj.xyz'), + outfile=str(tmp_path / 'out.out'), + trajectory_write_interval=write_interval, + outfile_write_interval=write_interval, + ) + gcmc.run(steps) + return gcmc + + +@pytest.fixture(scope='module') +def run(tmp_path_factory): + """One run shared by every assertion below: it costs minutes, not seconds.""" + return _run(tmp_path_factory.mktemp('gcmc')) + + +def test_golden_trajectory(run): + """Every accept/reject decision is pinned by the seeds.""" + counts = [len(f) for f in read(run._traj_file, ':')] + assert counts[:len(GOLDEN_N)] == GOLDEN_N + assert run.move_selector.move_acceptance_total == GOLDEN_ACCEPTED + assert run.E_old == pytest.approx(GOLDEN_ENERGY, abs=1e-3) + + +def test_both_move_types_accept(run): + """A fixture where one branch never fires stops testing that branch.""" + accepted_del, accepted_ins = run.move_selector.move_acceptance_total + assert accepted_del > 0, 'no deletion accepted: mu drifted off the balance point' + assert accepted_ins > 0, 'no insertion accepted: mu drifted off the balance point' + + +def test_fixed_layer_never_moves(run): + """FixAtoms must hold across accepted moves, rejected moves and relaxation. + + Compared under the minimum-image convention, not on raw coordinates: + accepted moves wrap the cell, so an atom crossing x=0 shows a box-sized + coordinate difference while never having moved. + """ + atoms = run.atoms + fixed = atoms.constraints[0].index + pristine, _, _ = _build(0) + drift = find_mic(atoms.positions[fixed] - pristine.positions[fixed], + atoms.cell, pbc=atoms.pbc)[0] + assert np.linalg.norm(drift, axis=1).max() < 1e-9 + + +def test_only_oxygen_is_exchanged(run): + """Ag is not in the move species list, so its count must never change.""" + symbols = np.asarray(run.atoms.get_chemical_symbols()) + assert (symbols == 'Ag').sum() == 48 + assert len(run.atoms) == 48 + (symbols == 'O').sum() + + +def test_oxygen_stays_deletable(run): + """An O outside the exchangeable region can never be removed again.""" + atoms = run.atoms + o_idx = np.where(np.asarray(atoms.get_chemical_symbols()) == 'O')[0] + assert len(o_idx) > 0 + assert all(run.cells[0].is_point_exchangeable(atoms.positions[i]) for i in o_idx) + + +def test_oxygen_chemisorbs(run): + """Real energetics: O binds at a sane distance instead of collapsing into + the slab or recombining into O2.""" + atoms = run.atoms + symbols = np.asarray(atoms.get_chemical_symbols()) + o_idx = np.where(symbols == 'O')[0] + d = atoms.get_all_distances(mic=True) + assert d[np.ix_(o_idx, np.where(symbols == 'Ag')[0])].min() > 1.8 + if len(o_idx) > 1: + assert min(d[i][j] for i in o_idx for j in o_idx if i < j) > 2.0 + + +def test_free_volume_is_non_degenerate(run): + """A collapsed free volume silently breaks the deletion prefactor + N*Lambda^3/V (inf or nan, see BaseCell._clamp_free_volume).""" + cell = run.cells[0] + assert 0.0 < cell.get_volume() < cell.cell_volume + + +# ------------------------------------------------- batched replica exchange + +# The single-GPU flagship: every replica's trial energy in one batched forward +# pass, and the general swap rule that works on both ladders. A mu ladder +# (shared temperature) is used because that is the mode the swap rule's +# cross-terms exist for -- the shortened (beta_j - beta_i)(Phi_j - Phi_i) form +# is identically zero there and would accept every swap. +# +# Spacing is 0.05 eV, deliberately tight: uniform 0.2 eV steps kill the ladder +# (docs/replica_exchange_ladder_spacing.rst) and the test would then only prove +# that nothing swaps. +RE_DELTA_MU_LADDER = [-0.45, -0.50, -0.55] +RE_STEPS = 30 +RE_EXCHANGE_INTERVAL = 5 + +# Golden values: seed 777, ladder above, 30 steps. +RE_GOLDEN_N = [56, 53, 50] # final atom count per replica +RE_GOLDEN_O = [8, 5, 2] # O count per replica, monotonic in mu +RE_GOLDEN_ENERGY = [-167.4708, -151.8001, -136.2122] + + +def _replica_factory(tmp_path, calculator): + """``gcmc_factory(mu=..., rank=...)`` for BatchedReplicaExchange. + + Every replica gets its own atoms, cells, move_selector and seed stream: + sharing any of them silently mixes the replicas' volumes and counters, and + the class rejects a factory that does (see its module docstring). + """ + def factory(mu, rank): + seeds = np.random.SeedSequence(SEED + rank).generate_state(5, dtype=np.uint32) + seed_del, seed_ins, seed_cell, seed_sel, seed_ens = (int(s) for s in seeds) + atoms, cell, _ = _build(seed_cell) + move_selector = MoveSelector( + [1, 1], + [DeletionMove(cell, species=['O'], seed=seed_del), + InsertionMove(cell, species=['O'], min_insert=MIN_INSERT, seed=seed_ins)], + seed=seed_sel, + ) + return GrandCanonicalEnsemble( + atoms=atoms, + cells=[cell], + calculator=calculator, + mu=mu, + units_type='metal', + species=['O'], + temperature=500.0, + move_selector=move_selector, + random_seed=seed_ens, + traj_file=str(tmp_path / f'replica_{rank}.xyz'), + outfile=str(tmp_path / f'replica_{rank}.out'), + trajectory_write_interval=10, + outfile_write_interval=10, + ) + return factory + + +@pytest.fixture(scope='module') +def batched_re(tmp_path_factory): + """One batched replica-exchange run, shared by the assertions below.""" + from mcpy.calculators import AlchemiFCalculator + from mcpy.ensembles.batched_replica_exchange import BatchedReplicaExchange + + tmp_path = tmp_path_factory.mktemp('batched_re') + calculator = AlchemiFCalculator( + checkpoint=str(CHECKPOINT), steps=300, fmax=0.05, + device='cuda', compile_model=False, optimizer='fire2', + ) + re = BatchedReplicaExchange( + gcmc_factory=_replica_factory(tmp_path, calculator), + calculator=calculator, + mus=[{'Ag': MU_AG, 'O': MU_O + d} for d in RE_DELTA_MU_LADDER], + gcmc_steps=RE_STEPS, + exchange_interval=RE_EXCHANGE_INTERVAL, + outfile=str(tmp_path / 're.log'), + write_out_interval=10, + seed=SEED, + global_minimum_file=str(tmp_path / 'global_minimum.xyz'), + ) + re.run() + return re, tmp_path + + +def test_batched_re_golden(batched_re): + """Per-replica outcome is pinned by the seeds. + + A swap moves a whole configuration between slots, so a single flipped + borderline acceptance changes which config a slot holds -- the atom counts + would break first, well before the energy tolerance mattered. + """ + re, _ = batched_re + assert [len(r.atoms) for r in re.replicas] == RE_GOLDEN_N + for r, expected in zip(re.replicas, RE_GOLDEN_ENERGY): + assert r.E_old == pytest.approx(expected, abs=1e-2) + + +def test_batched_re_actually_swaps(batched_re): + """A ladder that never swaps is not replica exchange, it is N lone runs.""" + re, _ = batched_re + assert sum(re.exchange_attempts) > 0, 'no swap attempted: check exchange_interval' + assert sum(re.exchange_successes) > 0, ( + 'no swap accepted: the mu ladder spacing is too wide ' + '(docs/replica_exchange_ladder_spacing.rst)') + + +def test_batched_re_replicas_stay_independent(batched_re): + """Shared atoms, cells or move_selector would silently mix the replicas.""" + re, _ = batched_re + assert len({id(r.atoms) for r in re.replicas}) == len(re.replicas) + assert len({id(r.move_selector) for r in re.replicas}) == len(re.replicas) + cell_ids = [id(c) for r in re.replicas for c in r.cells] + assert len(set(cell_ids)) == len(cell_ids) + + +def test_batched_re_each_replica_keeps_its_mu(batched_re): + """mu is pinned to the slot: a swap moves the configuration, not the mu.""" + re, _ = batched_re + assert [r._mu['O'] for r in re.replicas] == [MU_O + d for d in RE_DELTA_MU_LADDER] + + +def test_batched_re_conserves_silver(batched_re): + """Only O is exchanged, in every replica, swaps included.""" + re, _ = batched_re + for r in re.replicas: + symbols = np.asarray(r.atoms.get_chemical_symbols()) + assert (symbols == 'Ag').sum() == 48 + + +def test_batched_re_writes_global_minimum(batched_re): + """The lowest grand potential across all replicas is written once.""" + re, tmp_path = batched_re + best = tmp_path / 'global_minimum.xyz' + assert best.exists() + frames = read(str(best), ':') + assert len(frames) == 1 + assert min(r._best_score for r in re.replicas if r._best_atoms is not None) < float('inf') + + +def test_batched_re_coverage_follows_mu(batched_re): + """Thermodynamic sanity: a higher mu_O must hold more O. + + This is the assertion a broken ladder fails while everything else still + looks plausible -- the all-accept degeneracy averages the rungs together + and flattens exactly this gradient (BatchedReplicaExchange warns about it, + but only as a post-mortem). + """ + re, _ = batched_re + o_counts = [sum(1 for s in r.atoms.get_chemical_symbols() if s == 'O') + for r in re.replicas] + assert o_counts == RE_GOLDEN_O + assert o_counts == sorted(o_counts, reverse=True), ( + f'coverage not monotonic in mu_O: {o_counts} for ' + f'delta_mu {RE_DELTA_MU_LADDER}') From b3dd96eb1bd96ca9717267c2832bf52460bc8acb Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:10:55 +0200 Subject: [PATCH 04/20] Record the repo-wide over-engineering audit Forty-three findings with paths and line numbers: dead code, duplication, hand-rolled stdlib, and config nobody sets. About -2,700 lines available, no dependency removable. Also lists what looked bloated but earns its lines, so a later pass does not "simplify" a guard that encodes a real past bug. Two hot spots worth naming here: the serial and batched GCMC accept/reject loops are the same sixty-five lines kept in step by hand, and so are the single structure and batched Alchemi relaxation paths (which have already drifted). The suggested order puts the dead-code deletions first and the shared-helper extractions last, behind the golden test. Findings are listed only; nothing is applied. --- docs/research/ponytail-audit.md | 146 ++++++++++++++++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 docs/research/ponytail-audit.md diff --git a/docs/research/ponytail-audit.md b/docs/research/ponytail-audit.md new file mode 100644 index 0000000..ed75d3b --- /dev/null +++ b/docs/research/ponytail-audit.md @@ -0,0 +1,146 @@ +# Over-engineering audit (ponytail-audit) + +Date: 2026-08-24. +Line numbers refer to commit `cf70f65` (tip of `research/metric-capture`). +Scope: tracked files only (`mcpy/`, `docs/`, `examples/`, `tests/`, `benchmark/`, `scripts/`). +Untracked working dirs (`scratch/`, `runs/`, `private_thesis/`, `benchmark/internal/`) were not audited. + +This is a complexity audit, not a code review. +Correctness, security and performance were out of scope, with two exceptions noted at the end because they were found on the way. +Nothing here has been applied. + +Reviewable HTML version of the same findings: `.lavish/ponytail-audit.html` (untracked). + +Net total: about **-2,735 lines across 43 findings**, **-0 dependencies** (`ase` and `scipy` both earn their place). + +Tags follow the ponytail vocabulary. +`delete` = dead code or speculative feature, replacement is nothing. +`stdlib` = hand-rolled thing the standard library ships. +`native` = code doing what the platform or an existing dependency already does. +`yagni` = abstraction with one implementation, config nobody sets, layer with one caller. +`shrink` = same logic, fewer lines. + +## Top three + +1. `delete` `mcpy/moves/go_moves.py`, whole file, **-187**. + Three of its classes (`PermutationMove`, `ShakeMove`, `BrownianMove`) are stale copy-based twins of the exported ones. + The other four (`BallMove`, `ShellMove`, `BondMove`, `HighEnergyAtomsMove`) have no caller, no export and no test. + Two doc mentions (`docs/moves.rst:109-113`, `docs/reference/moves.rst:151-152`) go with them. +2. `native` `docs/reference/*.rst`, **-688**. + Five hand-typed API pages restate docstrings that already exist in the code. + `sphinx.ext.autodoc` is already enabled in `docs/conf.py:23` and never used. + Replace the bodies with `automodule` plus a short intro each. +3. `shrink` `docs/examples/*.rst`, **-400**. + Nine pages paste the `examples/*.py` source into code blocks by hand, so every example now lives twice and drifts. + Replace with `.. literalinclude:: ../../examples/x.py`. + +## Main findings + +| # | Tag | What to cut | Replacement | Lines | Path | +|---|-----|-------------|-------------|-------|------| +| 1 | delete | Whole module: 3 stale duplicates of exported moves + 4 never-referenced moves | Nothing | -187 | `mcpy/moves/go_moves.py` | +| 2 | native | Hand-written API reference duplicating every docstring | `sphinx.ext.autodoc` | -688 | `docs/reference/*.rst` | +| 3 | shrink | Example source pasted into 9 doc pages by hand | `literalinclude` | -400 | `docs/examples/*.rst` | +| 4 | delete | Second batched-RE example differing only by the `chunk_size` kwarg | A `--chunk-size` flag on `re_gcmc_batched.py` | -169 | `examples/re_gcmc_batched_chunked.py` | +| 5 | yagni | Two special-case swap rules (`_exchange_prob_T`, `_exchange_prob_mu`) for one decision | `BatchedReplicaExchange._accept_swap` already covers both ladders, and carries the de Broglie term the T-rule only half-implements | -95 | `mcpy/ensembles/replica_exchange.py:101-168` | +| 6 | shrink | Free-volume MC sampler written three times (sample, per-radius cKDTree, free fraction) | One `Cell._free_fraction(points, atoms)`; each cell keeps only its region mask | -75 | `cell/spherical_cell.py`, `cell/dome_cell.py`, `cell/custom_cell.py` | +| 7 | delete | Second MACE wrapper hand-building `AtomicData`/batch dicts; zero callers in the repo | Nothing; `MACE_F_Calculator` already uses `mace.calculators.MACECalculator` | -58 | `mcpy/calculators/mace_calculator.py` | +| 8 | shrink | RE header block written twice, once through `logger.info` and once to the outfile | Build the block once as a string, log it and write it | -40 | `replica_exchange.py:292-330,356-395` | +| 9 | shrink | `get_atoms_specie_inside_cell`: 4 copies of the same symbol-mask and str/list normalisation | One base method calling a per-cell `_inside_mask(positions)` | -40 | `mcpy/cell/*.py` | +| 10 | native | Hand-rolled extended-XYZ writer (`write_xyz`) | `ase.io.write(handle, atoms, format='extxyz')` with the extras in `atoms.info`; keep the hand version only if the per-frame cost the comment claims is re-measured | -40 | `ensembles/base_ensemble.py:263-305` | +| 11 | delete | Unused helpers `find_surface_indices`, `sphere_volume`, `total_volume`, `total_volume_with_overlap` (analytic overlap volume, superseded by MC sampling) | Nothing; `overlap_volume` survives only through one audit test, delete both | -35 | `mcpy/utils/utils.py:20-78` | +| 12 | yagni | `BaseCalculator` and `MACE_F_Calculator` are the same relax-then-energy wrapper; the second also accepts a prebuilt calculator | One wrapper with `optimizer=` and the step counters; keep `MACE_F_Calculator` as an alias | -30 | `calculators/base_calculator.py`, `calculators/mace_f_calculator.py` | +| 13 | delete | Dead knobs and methods: `BaseEnsemble(units_type=, user_tag=)`, `RandomNumberGenerator(warm_up=)`, `MoveSelector.get_operator`, `.calculate_volumes`, `.get_acceptance_ration` plus its test | Nothing. `units_type` appears only in the signature and is never read; `user_tag` is stored and never read | -30 | `base_ensemble.py:21,23`, `move_selector.py:116-160` | +| 14 | native | `DisplacementMove(constraints=[indices])` plus its own index cache, re-implementing a frozen-atom list the file's own docstring warns is fragile under insertion and deletion | `ase.constraints.FixAtoms`, which the GCMC loop already snapshots and restores | -20 | `moves/displacement_move.py:26-52` | +| 15 | delete | Duplicate `get_species()` (identical to `Cell`'s) and a repeated `self.species_radii = species_radii` that `super().__init__` already did | Inherit it | -18 | `spherical_cell.py`, `dome_cell.py`, `custom_cell.py` | +| 16 | native | Physical constants and unit factors hardcoded (Planck, Boltzmann, amu to kg, `sqrt(e)*1e10`) | `ase.units.kB`, `_hplanck`, `_amu`, `_e`; ase is already a hard dependency and `CanonicalEnsemble` already imports `ase.units.kB` | -12 | `utils/set_unit_constant.py:63-84` | +| 17 | native | Hand-built quaternion to rotation matrix | `scipy.spatial.transform.Rotation.random().as_matrix()`; scipy is already used for `cKDTree`. Pass a seeded generator to keep reproducibility | -12 | `moves/molecule_utils.py:80-88` | +| 18 | stdlib | Weighted pick built from `np.cumsum` + `np.searchsorted` + an index clamp | `random.choices(move_list, cum_weights=..., k=1)`; also drops numpy from the module | -10 | `moves/move_selector.py:60-78` | +| 19 | yagni | `exchangeable_predicate`'s `getattr(..., 'is_point_exchangeable', None) or ...` fallback | `BaseCell` now defines `is_point_exchangeable`, so plain attribute access is enough. Only pure duck-typed non-`BaseCell` cells need the fallback; decide whether that contract is still promised | -8 | `moves/molecule_utils.py:39-48` | +| 20 | yagni | `NullCell`: a whole cell class whose three methods return `0`, existing only because `BaseMove` insists on a cell | Make `BaseMove.cell` optional and have `get_volume` return `0.0`. Note `get_random_point()` returns the int `0`, which is not a point; nothing calls it | -35 | `cell/null_cell.py`, `moves/base_move.py` | + +## Ensembles (1,827 lines across 5 files, the biggest subpackage) + +One theme dominates. +The MC accept/reject loop and the replica-exchange bookkeeping each exist twice, once serial and once batched or MPI, kept in step by hand. +The repo already carries paired regression tests (`test_gcmc_rolls_back_a_move_that_mutates_then_reports_failure` and its `test_batched_re_` twin) that exist purely because the copies can drift. +That is the smell. +Findings 5 and 8 above belong to this group and are not counted again here. + +| # | Tag | What to cut | Replacement | Lines | Path | +|---|-----|-------------|-------------|-------|------| +| E1 | shrink | `_batched_single_move` re-implements `do_gcmc_step` line for line: the arrays and constraints snapshot, the `is False or is None` sentinel check, the "returned a different Atoms object" guard, the de Broglie count fallback, and the whole accept branch (`wrap`, `n_atoms`, `E_old`, counter, volumes, `_record_minimum`), comments included | Split the serial step into `_propose(atoms)` and `_commit_or_rollback(E_new, meta)` on `GrandCanonicalEnsemble`; both loops then call the same two methods and cannot drift | -65 | `grand_canonical_ensemble.py:243-311` vs `batched_replica_exchange.py:235-304` | +| E2 | shrink | The seven-column replica row format `"{:<5} {:<10} {:<25} {:<15.6f} ..."` is typed out four times (header plus rows, in each RE class), and `summarize_states` builds a dict only to reformat the same values twice | One `_RE_ROW` format constant and one row-builder shared by both classes | -35 | `replica_exchange.py:311-330,384-395,404-418`, `batched_replica_exchange.py:525-543` | +| E3 | yagni | `CanonicalEnsemble` keeps the current energy in three places at once: `self._current_energy`, `atoms.info['key_value_pairs']['potential_energy']` (an ASE-GA convention nothing else here reads), and `self.lowest_energy` beside the base class's `_best_energy`/`_best_score` | Keep `_current_energy` and the base `_best_*`; drop the `key_value_pairs` round-trip and `lowest_energy` | -24 | `canonical_ensemble.py:88-96,112-125,140-152` | +| E4 | shrink | `_write_global_minimum` written twice; only the "find the best replica" line differs (MPI gather vs local min), the open/write/except tail is identical | One `write_global_minimum(path, atoms, energy, score, rank, logger)` helper next to `write_xyz` | -18 | `replica_exchange.py:274-300`, `batched_replica_exchange.py:176-195` | +| E5 | shrink | `_write_initial_row` re-types `write_outfile`'s format string just to emit `N/A` placeholders | Call `write_outfile` with a placeholder ratio list | -14 | `grand_canonical_ensemble.py:328-341` | +| E6 | yagni | `consolidate_logging=True` reaches for a logger by hardcoded module-path string, mutates its level for the duration of `run()`, and restores it in a `finally`: a global side effect to tidy console output | Have the replicas log their lifecycle at DEBUG and let the application choose the level, the way `mcpy/utils/logging.py` already promises | -14 | `batched_replica_exchange.py:120-171` | +| E7 | delete | `BaseEnsemble.__del__`: a third cleanup path after `finalize_run` and `__exit__`, which swallows every exception | Nothing; the context manager and `run()`'s `finally` already close the handles | -8 | `base_ensemble.py:175-181` | +| E8 | shrink | Parameters optional in the signature but mandatory in fact: `CanonicalEnsemble(optimizer=None, move_selector=None)`, either of which reaches a `TypeError` inside `relax()`/`do_mutation()`. Same for the one-caller helper `_exchange_pairs` | Make them required positional arguments; inline the helper | -8 | `canonical_ensemble.py:30,33`, `batched_replica_exchange.py:306-309` | + +Kept on purpose, do not touch: the arrays and constraints rollback snapshot, the `is not atoms` identity guard, the two-tier interval and total acceptance counters, the swap-payload key stripping in `set_state`, and both ladder-degeneracy warnings. +Each encodes a bug that already happened once. + +## Alchemi calculators (780 lines across 3 files) + +The three-file split is right. +What is wrong is that the single-structure and batched paths were written twice and have already drifted apart, and that the FIRE and MD bootstrap ritual appears three times. + +| # | Tag | What to cut | Replacement | Lines | Path | +|---|-----|-------------|-------------|-------|------| +| A1 | shrink | Single-structure relaxation re-implements the batched one: same pre-allocation, same freeze hooks, same NL bootstrap, same optimizer build | `get_potential_energy = _relax_batch([atoms])[0]`. The two have already drifted: the single path calls `opt.run` (whole-batch), the batch path steps manually and compacts. One of them is the good one; keep only it | -45 | `alchemi_f_calculator.py:116-205` | +| A2 | shrink | Bootstrap ritual written three times: pre-allocate `forces`/`energy`, `_freeze_hook_for`, `NeighborListHook`, `_build_nl`, `opt.compute`, then zero frozen forces, with the same three-line comment each time | One `_bootstrap(opt, batch, fixed, nl_hook)` in `_alchemi_common` | -25 | `alchemi_f_calculator.py:132-162,178-205`, `_alchemi_common.py:212-246` | +| A3 | delete | `run_md` exists twice, byte-identical body and docstring, in both calculators | One method on a small shared mixin; both classes already import `_run_langevin_md` | -22 | `alchemi_calculator.py:157-187`, `alchemi_f_calculator.py:323-344` | +| A4 | yagni | `_ALCHEMI_OPTIMIZERS` registry and `optimizer=` validation for a second optimizer (`fire2`) that no example, benchmark or test selects | Import `FIRE` directly; re-add the switch when something asks for FIRE2. Note `scratch/gcmc_ag111_alchemi.py` does default to `fire2`, so confirm before cutting | -8 | `alchemi_f_calculator.py:32,96-101` | +| A5 | shrink | `_write_back_positions` re-derives the frozen indices with its own inline `FixAtoms` loop; `_fixed_indices` is defined ten lines below it | Call `_fixed_indices(atoms)` | -8 | `_alchemi_common.py:149-168` | +| A6 | delete | `_make_batch` is exactly `_make_multi_batch([atoms])` | Keep the multi form only | -5 | `_alchemi_common.py:120-127` | +| A7 | shrink | `_per_graph_energies`: six lines of comment describing a scatter-reduced fallback, above a branch that only raises | Keep the raise, drop the fiction | -6 | `_alchemi_common.py:130-142` | + +Not cuts, deliberately kept: `_HeadMACEWrapper`, the `wrapper.eval()` line, the cuEq import guard, the `energy_only`/`forces` pre-flight errors and the FIRE step-cap warning. +Every one of them has a documented failure it prevents (wrong head, retained second-order graph, silent truncation bias). + +## Tests (3,505 lines, 18 files, no conftest) + +| # | Tag | What to cut | Replacement | Lines | Path | +|---|-----|-------------|-------------|-------|------| +| T1 | shrink | There is no `tests/conftest.py`, so the same throwaway doubles are re-declared file by file: `StubCalc`, `_StubCalc`, `_UphillCalc`, `_HugeUphillCalc`, `StubGasCalculator`, `_PairedCalculator`, plus eight stub cell classes | One `conftest.py` with a `flat_calc` / `uphill_calc` / `stub_cell` fixture trio | -80 | `tests/` (5 files) | +| T2 | delete | Tests that exist only to hold dead code alive: `test_overlap_volume_matches_analytic`, `test_overlap_volume_limits`, the two `nullcell` tests, `test_moveselector_legacy_typo_alias_warns` | Nothing; they fall out for free with findings 11, 13 and 20 | -60 | `tests/test_audit_regressions.py` | +| T3 | yagni | Tests asserting on internals rather than behaviour: `test_no_sklearn_import` (guards a dependency the project never had), `test_max_attempts_constant_exists`, `test_no_rejection_loop` | The behavioural tests beside them already cover it (`test_impossible_min_insert_does_not_loop_forever`, `test_distribution_is_uniform`) | -20 | `tests/test_optimizations.py:93,127,174` | + +The suite is otherwise in good shape. +`test_audit_regressions.py` reads like a grab bag but every case names a real past bug, and the acceptance, molecule and batched-RE files carry the physics that example scripts cannot check. +No coverage worth deleting was found. + +## Phase diagram, examples, benchmarks + +| # | Tag | What to cut | Replacement | Lines | Path | +|---|-----|-------------|-------------|-------|------| +| P1 | delete | The superseded phase-diagram entry point `analyze_phase_diagram_results` (190 lines) plus the two helpers only it uses (`free_en`, and `from_mu_to_press` which has zero callers anywhere) | Nothing; `docs/phase_diagrams.rst:104` already says "Prefer `plot_phase_diagram` for new work". Caveat: `notebooks/phase_diagram.ipynb` still calls it, so port that one cell first | -205 | `mcpy/utils/phase_diagram.py:32-231` | +| P2 | shrink | A four-flag `argparse` CLI inside the library module, exposing 4 of about 20 knobs, with no `console_scripts` entry in `pyproject.toml`, so it is reachable only as `python -m` by someone who already read the source | Drop it; the notebook and `docs/make_figures.py` call the function directly | -25 | `phase_diagram.py:541-563` | +| X1 | shrink | Every example calls `configure_logging()` at import time, between its imports, which forces about five `# noqa: E402` suppressions per file across 13 files | Call it on the first line of `main()`. `configure` is idempotent and library modules only ever `getLogger`, so nothing needs it at import time. Imports return to the top and every noqa disappears | -65 | `examples/*.py` (13 files) | +| B1 | shrink | The two tracked parity benchmarks copy-paste their scaffolding: `block_stats` is byte-identical, `parse_thermo` differs only by a dropped docstring, `run_lammps` by one env tweak, `compare` by a field width, and `ZeroCalc` is in both | One `benchmark/_parity_common.py` | -90 | `benchmark/lammps_gcmc_parity.py`, `benchmark/mace_gcmc_parity.py` | +| X2 | shrink | Examples import deep module paths (`mcpy.moves.move_selector`, `mcpy.ensembles.grand_canonical_ensemble`) even though both names are exported from the package `__init__` | Import from `mcpy.moves` / `mcpy.ensembles`; the examples then demonstrate the public API | -0 | `examples/*.py` | + +## Looked bloated, is not + +- The Alchemi and MACE calculator split itself. Energy-only vs relaxing, and CPU-ASE vs GPU-resident, are four genuinely different call paths. Do not merge the classes; the waste is inside them (A1 to A7). +- `chunk_ranges`, `_clamp_free_volume`, `molecule_com`. Small, single-purpose, each with a named failure it prevents (a `0.0` volume turning the deletion prefactor into `inf`; a molecule split across a periodic face). +- The optional-backend `try/except ImportError` in `calculators/__init__.py`. Looks like defensive noise; it is what keeps the torch-free CI matrix working. +- Narrative `docs/*.rst` vs `docs/reference/*`. Genuinely different registers (concepts vs signatures), not duplication. Only the reference half should become autodoc. + +## Out of scope, found anyway + +Neither is a cut, so neither is counted above. + +- `scipy` is imported by three cell modules (`cKDTree`) but is not in `pyproject.toml` dependencies. It currently arrives only as an ase or mace transitive. +- `analyze_phase_diagram_results` carries about 20 Ag/O-specific magic defaults in its signature (`e_host=-2.82894684`, `idx_ref=2400`, `z_threshold=13.4`, `output_plot_path="lines_phases_mace.png"`). Any other system silently gets wrong physics rather than an error. If P1 is not applied, make those parameters required. + +## Suggested order + +Lowest risk first, and each step is independently shippable. + +1. Findings 1, 7, 11, 13, 15, T2: pure dead code and dead parameters. No behaviour change, about -350 lines. +2. Findings 2, 3, X1, B1, P2: docs and script hygiene. No library change, about -1,270 lines. +3. E1, then E2, E4, E5: the ensemble duplication, guarded by the existing paired regression tests. +4. A1 to A7: the Alchemi internals, guarded by `benchmark/internal/verify_compact_parity.py`. +5. Findings 5, 6, 9, 10, 16, 17, 18, 20: the shared-helper extractions and the native or stdlib swaps. These change behaviour at the margins (RNG streams, float formatting), so they need the golden regression test in place first. From 171e2c665eb392461a005400e53bcc012a380b28 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:10:55 +0200 Subject: [PATCH 05/20] Keep only the Ag(111) and batched-RE examples Promotes the calibrated Ag(111)/O script out of the ignored scratch/ directory into examples/gcmc_ag111.py, so the setup behind the golden test is tracked. It carries the calibrated geometry, derives its reference chemical potentials from the running potential instead of hardcoding them, and seeds all five RNGs (cells, both moves, the selector and the ensemble) so a run with --seed is actually repeatable; previously --seed reached only the two move RNGs. Every other example moves to a local directory pending a decision on what to ship. Note that eleven of them are still referenced by docs/ or README.md, so those references now point at files that are not in the tree. --- examples/Al2O3.poscar | 38 ---- examples/gcmc.py | 143 ------------- examples/gcmc_ag111.py | 199 +++++++++++++++++ examples/gcmc_big_alchemi.py | 131 ------------ examples/gcmc_bimet_nano.py | 86 -------- examples/gcmc_custom_cell.py | 90 -------- examples/gcmc_molecule.py | 81 ------- examples/gcmc_molecule_mace.py | 109 ---------- examples/gcmc_nano.py | 78 ------- examples/gcmc_nano_alchemi.py | 143 ------------- examples/gcmc_nano_supported.py | 91 -------- examples/plot_supported_np.py | 109 ---------- examples/re_gcmc.py | 112 ---------- examples/re_gcmc_batched_chunked.py | 169 --------------- examples/re_gcmc_co_cupd_batched.py | 317 ---------------------------- examples/re_gcmc_nano_alchemi.py | 149 ------------- 16 files changed, 199 insertions(+), 1846 deletions(-) delete mode 100644 examples/Al2O3.poscar delete mode 100644 examples/gcmc.py create mode 100644 examples/gcmc_ag111.py delete mode 100644 examples/gcmc_big_alchemi.py delete mode 100644 examples/gcmc_bimet_nano.py delete mode 100644 examples/gcmc_custom_cell.py delete mode 100644 examples/gcmc_molecule.py delete mode 100644 examples/gcmc_molecule_mace.py delete mode 100644 examples/gcmc_nano.py delete mode 100644 examples/gcmc_nano_alchemi.py delete mode 100644 examples/gcmc_nano_supported.py delete mode 100644 examples/plot_supported_np.py delete mode 100644 examples/re_gcmc.py delete mode 100644 examples/re_gcmc_batched_chunked.py delete mode 100644 examples/re_gcmc_co_cupd_batched.py delete mode 100644 examples/re_gcmc_nano_alchemi.py diff --git a/examples/Al2O3.poscar b/examples/Al2O3.poscar deleted file mode 100644 index 8073055..0000000 --- a/examples/Al2O3.poscar +++ /dev/null @@ -1,38 +0,0 @@ -Al12 O18 -1.0 - 2.4025136104259079 -4.1612756391334127 0.0000000000000000 - 2.4025136104259079 4.1612756391334127 0.0000000000000000 - 0.0000000000000000 0.0000000000000000 13.1162532478223355 -Al O -12 18 -direct - 0.3333333333333330 0.6666666666666661 0.8145706666666651 Al3+ - 0.6666666666666661 0.3333333333333330 0.6854293333333340 Al3+ - 0.0000000000000000 0.0000000000000000 0.6479039999999990 Al3+ - 0.3333333333333330 0.6666666666666661 0.5187626666666670 Al3+ - 0.0000000000000000 0.0000000000000000 0.1479039999999990 Al3+ - 0.3333333333333330 0.6666666666666661 0.0187626666666670 Al3+ - 0.6666666666666661 0.3333333333333330 0.9812373333333321 Al3+ - 0.0000000000000000 0.0000000000000000 0.8520960000000001 Al3+ - 0.6666666666666661 0.3333333333333330 0.4812373333333320 Al3+ - 0.0000000000000000 0.0000000000000000 0.3520959999999990 Al3+ - 0.3333333333333330 0.6666666666666661 0.3145706666666650 Al3+ - 0.6666666666666661 0.3333333333333330 0.1854293333333340 Al3+ - 0.3605211666666670 0.3333333333333330 0.5833333333333330 O2- - 0.6938544999999990 0.6938544999999990 0.7500000000000000 O2- - 0.9728121666666660 0.6394788333333320 0.5833333333333330 O2- - 0.6666666666666661 0.0271878333333340 0.5833333333333330 O2- - 0.0000000000000000 0.3061454999999990 0.7500000000000000 O2- - 0.3061454999999990 0.0000000000000000 0.7500000000000000 O2- - 0.0271878333333330 0.6666666666666661 0.9166666666666661 O2- - 0.3605211666666670 0.0271878333333330 0.0833333333333330 O2- - 0.6394788333333320 0.9728121666666660 0.9166666666666661 O2- - 0.3333333333333330 0.3605211666666670 0.9166666666666661 O2- - 0.6666666666666661 0.6394788333333320 0.0833333333333330 O2- - 0.9728121666666650 0.3333333333333330 0.0833333333333330 O2- - 0.6938544999999990 0.0000000000000000 0.2500000000000000 O2- - 0.0271878333333330 0.3605211666666670 0.4166666666666660 O2- - 0.3061454999999990 0.3061454999999990 0.2500000000000000 O2- - 0.0000000000000000 0.6938544999999990 0.2500000000000000 O2- - 0.3333333333333330 0.9728121666666650 0.4166666666666660 O2- - 0.6394788333333320 0.6666666666666661 0.4166666666666660 O2- diff --git a/examples/gcmc.py b/examples/gcmc.py deleted file mode 100644 index 65bace2..0000000 --- a/examples/gcmc.py +++ /dev/null @@ -1,143 +0,0 @@ -"""GCMC of metal + gas adsorption on a generic fcc surface with two CustomCell regions. - -Supports fcc(111), fcc(100), fcc(110), fcc(211) terminations. -""" -import argparse -import os - -import numpy as np -from ase.build import fcc100, fcc110, fcc111, fcc211 -from ase.constraints import FixAtoms - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.calculators import MACE_F_Calculator # noqa: E402 -from mcpy.cell import CustomCell as Cell # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--model-path', required=True, help='Path to the MACE model file') - p.add_argument('--delta-mu-gas', type=float, required=True, help='Shift applied to mu_gas (eV)') - p.add_argument('--device', default='cuda', help='Torch device for MACE') - p.add_argument('--metal', default='Au', help='Metal species symbol') - p.add_argument('--gas', default='O', help='Gas species symbol') - p.add_argument('--lattice-param', type=float, default=4.1755, help='fcc lattice constant (Å)') - p.add_argument('--surface-type', default='111', choices=['111', '100', '110', '211']) - p.add_argument('--surface-size', type=int, nargs=3, default=[4, 4, 3], - help='Slab size; fcc(211) needs the first dim divisible by 3, e.g. 6 4 3') - p.add_argument('--vacuum', type=float, default=8.0, help='Vacuum spacing (Å)') - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--mu-metal', type=float, default=-3.26999, help='Bulk metal DFT energy (eV)') - p.add_argument('--delta-mu-metal', type=float, default=-0.16, help='Shift to mu_metal (eV)') - p.add_argument('--mu-gas', type=float, default=-4.91, help='Reference mu_gas (eV)') - p.add_argument('--moves-per-step', type=int, default=25, - help='Insertion+deletion moves per species per GCMC step') - p.add_argument('--gcmc-steps', type=int, default=200, help='Number of GCMC steps') - p.add_argument('--write-interval', type=int, default=1, help='Outfile/traj write interval') - p.add_argument('--min-insert', type=float, default=0.5, - help='Min insertion distance to existing atoms (Å)') - p.add_argument('--cell-height', type=float, default=5.0, help='Insertion cell height (Å)') - p.add_argument('--cell-bottom-z', type=float, default=10.411, help='Cell bottom z (Å)') - p.add_argument('--metal-radius', type=float, default=2.75) - p.add_argument('--gas-radius', type=float, default=2.11) - p.add_argument('--rel-max-steps', type=int, default=40, help='LBFGS relaxation steps') - p.add_argument('--rel-fmax', type=float, default=0.1, help='LBFGS force convergence (eV/Å)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def build_slab(args): - size = tuple(args.surface_size) - if args.surface_type == '111': - atoms = fcc111(args.metal, a=args.lattice_param, size=size, - periodic=True, vacuum=args.vacuum) - bottom = [a.index for a in atoms if a.tag == size[-1]] - elif args.surface_type == '100': - atoms = fcc100(args.metal, a=args.lattice_param, size=size, - periodic=True, vacuum=args.vacuum) - bottom = [a.index for a in atoms if a.tag == size[-1]] - elif args.surface_type == '110': - # 110 termination: fix last two layers (more open surface) - atoms = fcc110(args.metal, a=args.lattice_param, size=size, - periodic=True, vacuum=args.vacuum) - bottom = [a.index for a in atoms if a.tag in (size[-1], size[-1] - 1)] - elif args.surface_type == '211': - # ASE fcc211 needs size[0] divisible by 3 and assigns no layer tags: - # fix the lowest-z layer's worth of atoms instead. - atoms = fcc211(args.metal, a=args.lattice_param, size=size, - vacuum=args.vacuum) - layer = size[0] * size[1] - bottom = list(np.argsort(atoms.positions[:, 2])[:layer]) - else: - raise ValueError(f'Invalid surface type: {args.surface_type}') - atoms.set_constraint(FixAtoms(indices=bottom)) - return atoms - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seeds = [int(s) for s in ss.generate_state(4, dtype=np.uint32)] - - atoms = build_slab(args) - - metal_radii = {args.metal: args.metal_radius, args.gas: 0.0} - gas_radii = {args.metal: args.gas_radius, args.gas: 0.0} - - cell_metal = Cell(atoms, custom_height=args.cell_height, bottom_z=args.cell_bottom_z, - species_radii=metal_radii) - cell_gas = Cell(atoms, custom_height=args.cell_height, bottom_z=args.cell_bottom_z, - species_radii=gas_radii) - - calculator = MACE_F_Calculator( - model_paths=args.model_path, - steps=args.rel_max_steps, - fmax=args.rel_fmax, - cueq=False, - device=args.device, - ) - - n = args.moves_per_step - move_selector = MoveSelector( - [n, n, n, n], - [DeletionMove(cell_metal, species=[args.metal], seed=seeds[0]), - DeletionMove(cell_gas, species=[args.gas], seed=seeds[1]), - InsertionMove(cell_metal, species=[args.metal], min_insert=args.min_insert, seed=seeds[2]), - InsertionMove(cell_gas, species=[args.gas], min_insert=args.min_insert, seed=seeds[3])], - ) - - mus = { - args.metal: args.mu_metal + args.delta_mu_metal, - args.gas: args.mu_gas + args.delta_mu_gas, - } - - tag = f'{atoms.get_chemical_formula()}_{args.gas}_dmu_{args.delta_mu_gas}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[cell_metal, cell_gas], - calculator=calculator, - mu=mus, - units_type='metal', - species=[args.metal, args.gas], - temperature=args.T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_relax_{tag}.out'), - trajectory_write_interval=args.write_interval, - outfile_write_interval=args.write_interval, - traj_file=os.path.join(args.outdir, f'gcmc_relax_{tag}.xyz'), - ) - - gcmc.run(args.gcmc_steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_ag111.py b/examples/gcmc_ag111.py new file mode 100644 index 0000000..4aa2dd0 --- /dev/null +++ b/examples/gcmc_ag111.py @@ -0,0 +1,199 @@ +"""GCMC of O adsorption on an Ag(111) surface using the nvalchemi GPU stack. + +The geometry constants below (lattice constant, O-Ag exclusion radius, region +height, the -0.176 bulk-Ag correction) are the calibrated set for the MACE +small-density-agnesi checkpoint; pass that checkpoint with ``--checkpoint`` to +reproduce ``tests/test_ag111_gcmc_golden.py``. Reference chemical potentials +are derived from whatever potential is passed, so another checkpoint stays +self-consistent, but its balance point for ``--delta-mu-O`` will differ. + +Builds a 16×16×4 Ag(111) slab (1024 atoms) by default, fixes the bottom two +layers, and runs GCMC inserting/deleting O in a 7 Å region that starts one +O-Ag distance *below* the top layer, so subsurface sites are reachable. + +Requirements: + pip install 'nvalchemi-toolkit[mace]' + +Typical usage: + # FIRE relaxation per step (recommended, GPU-native so fast enough): + python gcmc_ag111.py --checkpoint medium-mpa-0 --device cuda + + # Energy-only (fastest, no geometry relaxation): + python gcmc_ag111.py --checkpoint medium-mpa-0 --no-relax + + # Adjust chemical potential sweep: + python gcmc_ag111.py --checkpoint medium-mpa-0 --delta-mu-O -1.0 +""" +import argparse +import os + +import numpy as np +from ase.build import fcc111 +from ase.constraints import FixAtoms + +import logging + +from mcpy.utils.logging import configure as configure_logging + +configure_logging() +logger = logging.getLogger('mcpy') + +from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 +from mcpy.moves.move_selector import MoveSelector # noqa: E402 +from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 +from mcpy.calculators import AlchemiCalculator, AlchemiFCalculator # noqa: E402 +from mcpy.cell import CustomCell # noqa: E402 +from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 + +# Calibrated set for O on Ag(111) with the small-density-agnesi MACE model, +# taken from examples/gcmc_custom_cell.py. Fixed on purpose: this is fixture +# geometry, not a per-checkpoint number. Only mu is derived at run time. +AG_LATTICE = 4.165 # Ag fcc lattice constant (Angstrom) +# Relaxed O-Ag pair distance: the free-volume exclusion radius, and also how +# far below the top layer the region starts so O can go subsurface (the +# oxidation path). Anchoring at the top layer makes the region pure vacuum. +R_O_AG = 2.068 +SPECIES_RADII = {'Ag': R_O_AG, 'O': 0.0} +# Bulk Ag reference correction (examples/gcmc_custom_cell.py). +MU_AG_CORRECTION = -0.176 + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument('--checkpoint', default='medium-mpa-0', + help='Named checkpoint (e.g. medium-mpa-0) or path to a .pt file') + p.add_argument('--no-relax', action='store_true', + help='Use AlchemiCalculator (energy-only); default uses AlchemiFCalculator') + p.add_argument('--device', default='cuda', help='Torch device (cuda or cpu)') + p.add_argument('--no-cueq', action='store_true', help='Disable cuEquivariance kernel fusion') + p.add_argument('--no-compile', action='store_true', help='Disable torch.compile') + p.add_argument('--fmax', type=float, default=0.05, help='FIRE force threshold (eV/Å)') + p.add_argument('--relax-steps', type=int, default=300, help='Max FIRE steps per eval') + p.add_argument('--optimizer', default='fire2', choices=['fire', 'fire2'], + help='FIRE variant (fire2 typically converges in fewer steps)') + p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') + p.add_argument('--steps', type=int, default=5_000_000, help='Number of GCMC steps') + p.add_argument('--mu-Ag', type=float, default=None, + help='Chemical potential of Ag (eV); default: derived from the potential') + p.add_argument('--mu-O', type=float, default=None, + help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') + p.add_argument('--delta-mu-O', type=float, default=-0.3, help='Shift applied to mu_O (eV)') + p.add_argument('--slab-size', type=int, nargs=3, default=[16, 16, 4], + help='fcc111 supercell (Nx Ny Nlayers); default 16 16 4 = 1024 atoms') + p.add_argument('--vacuum', type=float, default=10.0, help='Vacuum above/below slab (Å)') + p.add_argument('--cell-height', type=float, default=7.0, help='O insertion region height (Å)') + p.add_argument('--fix-layers', type=int, default=2, help='Number of bottom layers to freeze') + p.add_argument('--min-insert', type=float, default=0.5, + help='Min insertion distance to existing atoms (Å)') + p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') + p.add_argument('--write-interval', type=int, default=10, + help='Outfile/trajectory write interval') + p.add_argument('--outdir', default='.', help='Output directory') + return p.parse_args() + + +def build_slab(args): + nx, ny, nlayers = args.slab_size + atoms = fcc111('Ag', a=AG_LATTICE, size=(nx, ny, nlayers), + periodic=True, vacuum=args.vacuum) + # tag=1 is top layer, tag=nlayers is bottom layer (ASE convention) + fix_tag_threshold = nlayers - args.fix_layers + 1 + bottom = [a.index for a in atoms if a.tag >= fix_tag_threshold] + atoms.set_constraint(FixAtoms(indices=bottom)) + return atoms + + +def build_calculator(args): + enable_cueq = not args.no_cueq + compile_model = not args.no_compile + if args.no_relax: + return AlchemiCalculator( + checkpoint=args.checkpoint, + device=args.device, + enable_cueq=enable_cueq, + compile_model=compile_model, + ) + return AlchemiFCalculator( + checkpoint=args.checkpoint, + steps=args.relax_steps, + fmax=args.fmax, + device=args.device, + enable_cueq=enable_cueq, + compile_model=compile_model, + optimizer=args.optimizer, + ) + + +def main(): + args = parse_args() + os.makedirs(args.outdir, exist_ok=True) + + ss = np.random.SeedSequence(args.seed) + seed_del, seed_ins, seed_cell, seed_sel, seed_ens = ( + int(s) for s in ss.generate_state(5, dtype=np.uint32) + ) + + atoms = build_slab(args) + + # One interlayer below the top layer, so the region covers the subsurface + # site plus the vacuum above (see R_O_AG). + z_top = float(atoms.positions[atoms.get_tags() == 1, 2].max()) + cell_bottom_z = z_top - R_O_AG + + cell = CustomCell( + atoms, + custom_height=args.cell_height, + bottom_z=cell_bottom_z, + species_radii=SPECIES_RADII, + seed=seed_cell, + ) + + calculator = build_calculator(args) + + # mu references come from the running potential: hardcoded values silently + # change meaning when the checkpoint changes. + mu_ag = (derive_mu_bulk(calculator, 'Ag', a=AG_LATTICE) + MU_AG_CORRECTION + if args.mu_Ag is None else args.mu_Ag) + mu_o = derive_mu_gas(calculator, 'O2') if args.mu_O is None else args.mu_O + + import torch + gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A' + logger.info('Calculator device: %s | CUDA available: %s | GPU: %s', + args.device, torch.cuda.is_available(), gpu_name) + + species = ['O'] + move_selector = MoveSelector( + [1, 1], + [DeletionMove(cell, species=species, seed=seed_del), + InsertionMove(cell, species=species, min_insert=args.min_insert, seed=seed_ins)], + seed=seed_sel, + ) + + mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} + + nx, ny, nlayers = args.slab_size + mode = 'nrelax' if args.no_relax else f'fire_fmax{args.fmax}' + tag = f'Ag{nx * ny * nlayers}_111_dmuO_{args.delta_mu_O}_{mode}' + + gcmc = GrandCanonicalEnsemble( + atoms=atoms, + cells=[cell], + calculator=calculator, + mu=mus, + units_type='metal', + species=species, + temperature=args.T, + move_selector=move_selector, + random_seed=seed_ens, + outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), + traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), + trajectory_write_interval=args.write_interval, + outfile_write_interval=args.write_interval, + ) + + gcmc.run(args.steps) + + +if __name__ == '__main__': + main() diff --git a/examples/gcmc_big_alchemi.py b/examples/gcmc_big_alchemi.py deleted file mode 100644 index 57527b4..0000000 --- a/examples/gcmc_big_alchemi.py +++ /dev/null @@ -1,131 +0,0 @@ -"""GCMC of O adsorption on a large supported Ag nanoparticle (nvalchemi GPU stack). - -Larger-system companion to gcmc_nano_supported.py: - - Ag711 octahedron (Octahedron('Ag', 11, 4)) on a 10x10 Al2O3 slab (~3.7k atoms) - - DomeCell hemispherical insertion region over the supported particle - - AlchemiFCalculator: FIRE geometry relaxation on GPU at every energy evaluation - -A GCMC "step" performs ``move_selector.n_moves`` trial moves (here 2: one deletion, -one insertion), and each trial move triggers one relaxed energy evaluation. - -Requirements: - pip install 'nvalchemi-toolkit[mace]' - -Typical usage: - python gcmc_big_alchemi.py --support Al2O3.poscar --delta-mu-O -0.3 \ - --steps 300 --seed 7 --outdir out --device cuda - -GPU memory on long runs: the atom count changes every accepted move, which -fragments the CUDA caching allocator, so reserved GPU memory drifts up over a -long run (allocator fragmentation, not a leak). Launch with -``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True``. -""" -import argparse -import logging -import os - -import numpy as np -from ase.cluster import Octahedron -from ase.io import read -from ase.constraints import FixAtoms - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.calculators import AlchemiFCalculator # noqa: E402 -from mcpy.cell import DomeCell # noqa: E402 -from mcpy.utils.utils import get_p_at_support # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--support', default='Al2O3.poscar', help='POSCAR of the support') - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=300, help='Number of GCMC steps') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.5, help='Shift applied to mu_O (eV)') - p.add_argument('--relax-steps', type=int, default=100, help='Max FIRE steps per energy eval') - p.add_argument('--fmax', type=float, default=0.05, help='FIRE force threshold (eV/A)') - p.add_argument('--optimizer', default='fire', choices=['fire', 'fire2'], - help='FIRE variant used for relaxation') - p.add_argument('--max-neighbors', type=int, default=128, help='Neighbor-list cap (avoids OOM)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--device', default='cuda', help='Torch device') - p.add_argument('--outdir', default='.', help='Output directory') - p.add_argument('--debug', action='store_true', - help='DEBUG logging (per-relaxation FIRE step counts)') - return p.parse_args() - - -def main(): - args = parse_args() - if args.debug: - configure_logging(level=logging.DEBUG) - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seed_del, seed_ins = (int(s) for s in ss.generate_state(2, dtype=np.uint32)) - - support = read(args.support).repeat((10, 10, 1)) - support.center(vacuum=10.0, axis=2) - z = support.positions[:, 2] - z_half = z.min() + 0.5 * (z.max() - z.min()) - mask = z < z_half - support.set_constraint(FixAtoms(mask=mask)) - - surface_z = float(np.max(support.positions[:, 2])) - particle = Octahedron('Ag', 11, 4) - atoms = get_p_at_support(support, particle, contact_surface='100', gap=2.0) - - scell = DomeCell(atoms, particle_species='Ag', bottom_z=surface_z, vacuum=3.0, - species_radii={'Ag': 2.068, 'O': 0, 'Al': 3}) - - calculator = AlchemiFCalculator( - checkpoint='medium-mpa-0', - steps=args.relax_steps, - fmax=args.fmax, - device=args.device, - compile_model=True, - max_neighbors=args.max_neighbors, - optimizer=args.optimizer, - ) - - species = ['O'] - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=species, seed=seed_del), - InsertionMove(scell, species=species, min_insert=0.5, seed=seed_ins)], - ) - - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} - - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}_{args.optimizer}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=args.T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - ) - - gcmc.run(args.steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_bimet_nano.py b/examples/gcmc_bimet_nano.py deleted file mode 100644 index bb081f7..0000000 --- a/examples/gcmc_bimet_nano.py +++ /dev/null @@ -1,86 +0,0 @@ -"""GCMC of H adsorption on a Pt/Au bimetallic octahedral nanoparticle. - -Also performs Pt<->Au permutation moves on the metal framework. -""" -import argparse -import os - -import numpy as np -from ase.cluster import Octahedron -from ase.build import molecule -from mace.calculators import mace_mp - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.calculators import BaseCalculator # noqa: E402 -from mcpy.moves import DeletionMove, InsertionMove, PermutationMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.cell import SphericalCell # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=1_000_000, help='Number of GCMC steps') - p.add_argument('--delta-mu-H', type=float, default=-0.5, help='Shift to mu_H = E(H2)/2 (eV)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--device', default='cuda', help='Torch device for MACE') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seed_del, seed_ins, seed_perm = (int(s) for s in ss.generate_state(3, dtype=np.uint32)) - - atoms = Octahedron('Pt', 7, 2) - half = len(atoms) // 2 - atoms.symbols = ['Pt'] * half + ['Au'] * half + ['Pt'] - atoms.set_pbc(False) - - scell = SphericalCell(atoms, vacuum=3, species_radii={'Pt': 2, 'Au': 2.5, 'H': 0}, - mc_sample_points=100_000) - - # BaseCalculator relaxes with LBFGS before each energy evaluation; a bare - # mace_mp would return unrelaxed energies. - calculator = BaseCalculator(mace_mp(device=args.device), steps=100, fmax=0.05) - - species = ['H'] - move_selector = MoveSelector( - [1, 1, 1], - [DeletionMove(scell, species=species, seed=seed_del), - InsertionMove(scell, species=species, min_insert=0.5, seed=seed_ins), - PermutationMove(species=['Au', 'Pt'], seed=seed_perm)], - ) - - e_h2 = calculator.get_potential_energy(molecule('H2')) - - mus = {'H': e_h2 / 2 + args.delta_mu_H} - - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_H}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=args.T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - trajectory_write_interval=1, - outfile_write_interval=1, - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - ) - - gcmc.run(args.steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_custom_cell.py b/examples/gcmc_custom_cell.py deleted file mode 100644 index d754513..0000000 --- a/examples/gcmc_custom_cell.py +++ /dev/null @@ -1,90 +0,0 @@ -"""GCMC of Ag and O on an fcc(111) Ag slab with two CustomCell insertion regions.""" -import argparse -import os - -import numpy as np -from ase.build import fcc111, bulk, molecule -from ase.constraints import FixAtoms -from mace.calculators import mace_mp - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.calculators import BaseCalculator # noqa: E402 -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.cell import CustomCell # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=1_000_000, help='Number of GCMC steps') - p.add_argument('--delta-mu-O', type=float, default=-0.5, help='Shift to mu_O = E(O2)/2 (eV)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--device', default='cuda', help='Torch device for MACE') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seeds = [int(s) for s in ss.generate_state(4, dtype=np.uint32)] - - atoms = fcc111('Ag', a=4.165, size=(4, 4, 3), periodic=True, vacuum=8) - bottom_layer = [a.index for a in atoms if a.tag == 3] - atoms.set_constraint(FixAtoms(indices=bottom_layer)) - - cell_ag_ag = CustomCell(atoms, custom_height=7, bottom_z=12.8 - 2.068, - species_radii={'Ag': 2.947, 'O': 0}) - cell_ag_o = CustomCell(atoms, custom_height=7, bottom_z=12.8 - 2.068, - species_radii={'Ag': 2.068, 'O': 0}) - - mace = mace_mp(device=args.device) - - species = ['Ag', 'O'] - move_selector = MoveSelector( - [25, 25, 25, 25], - [DeletionMove(cell_ag_ag, species=['Ag'], seed=seeds[0]), - DeletionMove(cell_ag_o, species=['O'], seed=seeds[1]), - InsertionMove(cell_ag_ag, species=['Ag'], min_insert=0.5, seed=seeds[2]), - InsertionMove(cell_ag_o, species=['O'], min_insert=0.5, seed=seeds[3])], - ) - - # BaseCalculator relaxes with LBFGS before each energy evaluation; a bare - # mace_mp would return unrelaxed energies. Reference mu values use a - # stricter relaxation than the GCMC production settings below. - ref_calculator = BaseCalculator(mace, steps=100, fmax=0.05) - e_o2 = ref_calculator.get_potential_energy(molecule('O2')) - e_ag = ref_calculator.get_potential_energy(bulk('Ag', a=4.165)) - - mus = {'Ag': e_ag - 0.176, 'O': e_o2 / 2 + args.delta_mu_O} - - calculator = BaseCalculator(mace, steps=20, fmax=0.1) - - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[cell_ag_ag, cell_ag_o], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=args.T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - trajectory_write_interval=1, - outfile_write_interval=1, - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - ) - - gcmc.run(args.steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_molecule.py b/examples/gcmc_molecule.py deleted file mode 100644 index fd4ffbd..0000000 --- a/examples/gcmc_molecule.py +++ /dev/null @@ -1,81 +0,0 @@ -"""GCMC insertion/deletion of a rigid O2 molecule in a periodic box (LJ units). - -Runs without torch/mace: uses ASE's LennardJones calculator directly. -Demonstrates the molecular-move API; see -docs/gcmc_acceptance_convention.rst for the acceptance convention -(mu is the full molecular chemical potential). -""" -import argparse - -from ase import Atoms -from ase.calculators.lj import LennardJones - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import (MoleculeDeletionMove, MoleculeInsertionMove, # noqa: E402 - MoveSelector) -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.cell import Cell # noqa: E402 - - -class LJEnergy: - """Minimal calculator adapter: plain LJ energy, no relaxation.""" - - def __init__(self): - self.lj = LennardJones(sigma=1.0, epsilon=1.0, rc=3.0) - - def get_potential_energy(self, atoms): - atoms.calc = self.lj - return atoms.get_potential_energy() - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--box', type=float, default=8.0, help='Cubic box edge') - p.add_argument('--T', type=float, default=1.0, help='Temperature (LJ units)') - p.add_argument('--mu', type=float, default=-2.0, - help='Full molecular chemical potential of O2 (LJ units)') - p.add_argument('--bond', type=float, default=1.1, help='O-O bond length') - p.add_argument('--min-insert', type=float, default=0.9, - help='Min insertion distance to existing atoms') - p.add_argument('--gcmc-steps', type=int, default=500) - p.add_argument('--seed', type=int, default=42) - return p.parse_args() - - -def main(): - args = parse_args() - - o2 = Atoms('O2', positions=[[0, 0, 0], [0, 0, args.bond]]) - atoms = Atoms(cell=[args.box] * 3, pbc=True) - cell = Cell(atoms, species_radii={'O': 0.0}, seed=args.seed) - - move_selector = MoveSelector( - [1, 1], - [MoleculeInsertionMove(cell, o2, 'O2', seed=args.seed + 1, - min_insert=args.min_insert), - MoleculeDeletionMove(cell, o2, 'O2', seed=args.seed + 2)], - seed=args.seed + 3, - ) - - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[cell], - units_type='LJ', - calculator=LJEnergy(), - mu={'O2': args.mu}, - species=[], - temperature=args.T, - move_selector=move_selector, - molecules={'O2': o2}, - random_seed=args.seed, - outfile='gcmc_molecule.out', - traj_file='gcmc_molecule.xyz', - ) - gcmc.run(args.gcmc_steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_molecule_mace.py b/examples/gcmc_molecule_mace.py deleted file mode 100644 index 0cb846d..0000000 --- a/examples/gcmc_molecule_mace.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Molecular GCMC with a MACE potential: rigid-molecule insertion/deletion of -O2 (or another ASE-buildable molecule) above an fcc(111) Ag slab. - -The chemical potential is the FULL molecular chemical potential -(mu = E_MACE(relaxed molecule) + delta_mu); orientations are sampled -uniformly, so the rotational partition function is absorbed into mu -(see docs/gcmc_acceptance_convention.rst, "Molecular moves"). -""" -import argparse -import os - -import numpy as np -from ase.build import fcc111, molecule -from ase.constraints import FixAtoms -from mace.calculators import mace_mp - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.calculators import BaseCalculator # noqa: E402 -from mcpy.moves import MoleculeDeletionMove, MoleculeInsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.cell import CustomCell # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--molecule', default='O2', - help='ASE g2 molecule name used as the rigid template') - p.add_argument('--T', type=float, default=400.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=1000, help='Number of GCMC steps') - p.add_argument('--delta-mu', type=float, default=-0.65, - help='Shift added to mu = E_MACE(relaxed molecule) (eV). ' - 'For O2/Ag(111) at 400 K: -0.3 saturates the region, ' - '-0.8 strips it bare, -0.65 equilibrates with two-way ' - 'exchange at ~19 molecules') - p.add_argument('--min-insert', type=float, default=1.5, - help='Min distance of inserted atoms to existing atoms (A)') - p.add_argument('--cell-height', type=float, default=7.0, - help='Insertion region height above the surface (A)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--device', default='cuda', help='Torch device for MACE') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seeds = [int(s) for s in ss.generate_state(3, dtype=np.uint32)] - - atoms = fcc111('Ag', a=4.165, size=(4, 4, 3), periodic=True, vacuum=8) - bottom_layer = [a.index for a in atoms if a.tag == 3] - atoms.set_constraint(FixAtoms(indices=bottom_layer)) - surface_z = atoms.positions[:, 2].max() - - template = molecule(args.molecule) - name = args.molecule - # Radii: exclude overlap volume around slab atoms; molecule species get 0 - # so the free-volume estimate only subtracts substrate atoms. - radii = {'Ag': 2.068} - radii.update({s: 0.0 for s in set(template.get_chemical_symbols())}) - cell = CustomCell(atoms, custom_height=args.cell_height, - bottom_z=surface_z, species_radii=radii) - - mace = mace_mp(device=args.device) - - # BaseCalculator relaxes before each energy evaluation; the mu reference - # uses a stricter relaxation than the GCMC production settings below. - ref_calculator = BaseCalculator(mace, steps=100, fmax=0.05) - e_mol = ref_calculator.get_potential_energy(template.copy()) - mus = {name: e_mol + args.delta_mu} - - move_selector = MoveSelector( - [1, 1], - [MoleculeInsertionMove(cell, template, name, seed=seeds[0], - min_insert=args.min_insert), - MoleculeDeletionMove(cell, template, name, seed=seeds[1])], - ) - - calculator = BaseCalculator(mace, steps=20, fmax=0.1) - - tag = f'{atoms.get_chemical_formula()}_{name}_dmu_{args.delta_mu}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[cell], - calculator=calculator, - mu=mus, - units_type='metal', - species=[], - molecules={name: template}, - temperature=args.T, - move_selector=move_selector, - random_seed=seeds[2], - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - trajectory_write_interval=1, - outfile_write_interval=1, - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - ) - - gcmc.run(args.steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_nano.py b/examples/gcmc_nano.py deleted file mode 100644 index 5752957..0000000 --- a/examples/gcmc_nano.py +++ /dev/null @@ -1,78 +0,0 @@ -"""GCMC of O adsorption on an Ag octahedral nanoparticle in vacuum.""" -import argparse -import os - -import numpy as np -from ase.cluster import Octahedron -from mace.calculators import mace_mp - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.cell import SphericalCell # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=1_000_000, help='Number of GCMC steps') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.5, help='Shift applied to mu_O (eV)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--device', default='cuda', help='Torch device for MACE') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seed_del, seed_ins = (int(s) for s in ss.generate_state(2, dtype=np.uint32)) - - atoms = Octahedron('Ag', 6, 1) - - scell = SphericalCell(atoms, vacuum=3, species_radii={'Ag': 2.947, 'O': 0}, - mc_sample_points=100_000) - - calculator = mace_mp(device=args.device) - - species = ['O'] - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=species, seed=seed_del), - InsertionMove(scell, species=species, min_insert=0.5, seed=seed_ins)], - ) - - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} - - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=args.T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - ) - - gcmc.run(args.steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_nano_alchemi.py b/examples/gcmc_nano_alchemi.py deleted file mode 100644 index 8b61cf5..0000000 --- a/examples/gcmc_nano_alchemi.py +++ /dev/null @@ -1,143 +0,0 @@ -"""GCMC of O adsorption on an Ag octahedral nanoparticle using the nvalchemi GPU stack. - -Mirrors gcmc_nano.py but replaces the ASE MACE path with AlchemiFCalculator -(FIRE geometry relaxation on GPU) or AlchemiCalculator (energy-only, no relax). - -Requirements: - pip install 'nvalchemi-toolkit[mace]' - -Typical usage: - # With FIRE relaxation (recommended for >=100 atoms): - python gcmc_nano_alchemi.py --checkpoint medium-mpa-0 --device cuda - - # Energy-only (fastest; no geometry relaxation): - python gcmc_nano_alchemi.py --checkpoint medium-mpa-0 --no-relax - - # Reuse a shared model across restart runs (pass a .pt path): - python gcmc_nano_alchemi.py --checkpoint /path/to/model.pt - -GPU memory on long runs: the atom count changes every accepted move, which -fragments the CUDA caching allocator, so reserved GPU memory drifts up over a -long run (allocator fragmentation, not a leak). Launch with -``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True``. -""" -import argparse -import os - -import numpy as np -from ase.cluster import Octahedron - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.calculators import AlchemiCalculator, AlchemiFCalculator # noqa: E402 -from mcpy.cell import SphericalCell # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--checkpoint', default='medium-mpa-0', - help='Named checkpoint (e.g. medium-mpa-0) or path to a .pt file') - p.add_argument('--no-relax', action='store_true', - help='Use AlchemiCalculator (energy-only); default uses AlchemiFCalculator ' - '(FIRE relax)') - p.add_argument('--device', default='cuda', help='Torch device (cuda or cpu)') - p.add_argument('--no-cueq', action='store_true', - help='Disable cuEquivariance kernel fusion') - p.add_argument('--no-compile', action='store_true', - help='Disable torch.compile (faster startup, slower inference)') - p.add_argument('--fmax', type=float, default=0.05, - help='FIRE force convergence threshold (eV/Å) [relax mode only]') - p.add_argument('--relax-steps', type=int, default=500, - help='Max FIRE steps per energy evaluation [relax mode only]') - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=1_000_000, help='Number of GCMC steps') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.5, help='Shift applied to mu_O (eV)') - p.add_argument('--min-insert', type=float, default=0.5, - help='Min insertion distance to existing atoms (Å)') - p.add_argument('--vacuum', type=float, default=3.0, help='Vacuum around nanoparticle (Å)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--outdir', default='.', help='Output directory') - p.add_argument('--write-interval', type=int, default=1, - help='Outfile/trajectory write interval') - return p.parse_args() - - -def build_calculator(args): - enable_cueq = not args.no_cueq - compile_model = not args.no_compile - - if args.no_relax: - return AlchemiCalculator( - checkpoint=args.checkpoint, - device=args.device, - enable_cueq=enable_cueq, - compile_model=compile_model, - ) - return AlchemiFCalculator( - checkpoint=args.checkpoint, - steps=args.relax_steps, - fmax=args.fmax, - device=args.device, - enable_cueq=enable_cueq, - compile_model=compile_model, - ) - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seed_del, seed_ins = (int(s) for s in ss.generate_state(2, dtype=np.uint32)) - - atoms = Octahedron('Ag', 6, 1) - - scell = SphericalCell(atoms, vacuum=args.vacuum, species_radii={'Ag': 2.947, 'O': 0}, - mc_sample_points=100_000) - - calculator = build_calculator(args) - - species = ['O'] - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=species, seed=seed_del), - InsertionMove(scell, species=species, min_insert=args.min_insert, seed=seed_ins)], - ) - - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} - - mode = 'nrelax' if args.no_relax else f'fire_fmax{args.fmax}' - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}_{mode}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=args.T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - trajectory_write_interval=args.write_interval, - outfile_write_interval=args.write_interval, - ) - - gcmc.run(args.steps) - - -if __name__ == '__main__': - main() diff --git a/examples/gcmc_nano_supported.py b/examples/gcmc_nano_supported.py deleted file mode 100644 index 947f260..0000000 --- a/examples/gcmc_nano_supported.py +++ /dev/null @@ -1,91 +0,0 @@ -"""GCMC of O adsorption on an Ag nanoparticle supported on Al2O3.""" -import argparse -import os - -import numpy as np -from ase.cluster import Octahedron -from ase.io import read -from ase.constraints import FixAtoms -from mace.calculators import mace_mp - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.cell import CustomCell # noqa: E402 -from mcpy.utils.utils import get_p_at_support # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--support', default='Al2O3.poscar', help='POSCAR of the support') - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=1_000_000, help='Number of GCMC steps') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.5, help='Shift applied to mu_O (eV)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--device', default='cuda', help='Torch device for MACE') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seed_del, seed_ins = (int(s) for s in ss.generate_state(2, dtype=np.uint32)) - - support = read(args.support).repeat((4, 4, 1)) - support.center(vacuum=10.0, axis=2) - z = support.positions[:, 2] - z_half = z.min() + 0.5 * (z.max() - z.min()) - mask = z < z_half - support.set_constraint(FixAtoms(mask=mask)) - - surface_z = float(np.max(support.positions[:, 2])) - particle = Octahedron('Ag', 5, 2) - atoms = get_p_at_support(support, particle, contact_surface='100', gap=2.0) - - scell = CustomCell(atoms, custom_height=20, bottom_z=surface_z, - species_radii={'Ag': 2.068, 'O': 0, 'Al': 3}) - - calculator = mace_mp(device=args.device) - - species = ['O'] - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=species, seed=seed_del), - InsertionMove(scell, species=species, min_insert=0.5, seed=seed_ins)], - ) - - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} - - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=args.T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - ) - - gcmc.run(args.steps) - - -if __name__ == '__main__': - main() diff --git a/examples/plot_supported_np.py b/examples/plot_supported_np.py deleted file mode 100644 index 7eabb2e..0000000 --- a/examples/plot_supported_np.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Render a supported-nanoparticle GCMC trajectory frame as a showcase figure. - -A single 3/4 view with depth shading (atoms farther from the viewer are -darkened) so the particle reads as sitting on the support instead of pasted -onto it. Adsorbed atoms are identified as those appended after the initial -structure: the first frame of the trajectory is the clean (pre-adsorption) -system, so any atom whose index is >= ``len(frame0)`` was inserted by GCMC. -This is robust to deletions because insertions are appended at the end and the -support is never deleted. - -Needs matplotlib and ase (e.g. ``conda run -n base python ...``). - -Usage:: - - python examples/plot_supported_np.py traj.xyz --out figure.png - python examples/plot_supported_np.py traj.xyz --particle-species Ag --frame -1 -""" -import argparse - -import numpy as np -import matplotlib -matplotlib.use('Agg') -import matplotlib.pyplot as plt # noqa: E402 -import matplotlib.patches as mpatches # noqa: E402 -from ase.io import read # noqa: E402 -from ase.io.utils import rotate # noqa: E402 -from ase.visualize.plot import plot_atoms # noqa: E402 - -# role -> (base rgb, draw radius) -AG = (0.74, 0.77, 0.83) -SUPPORT_METAL = (0.82, 0.73, 0.55) -SUPPORT_O = (0.93, 0.80, 0.80) -ADSORBED = (0.90, 0.07, 0.07) - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('trajectory', help='Extended-XYZ GCMC trajectory') - p.add_argument('--out', default='supported_np.png', help='Output PNG path') - p.add_argument('--particle-species', default='Ag', - help='Symbol of the nanoparticle metal (coloured silver)') - p.add_argument('--frame', type=int, default=-1, help='Frame index to render') - p.add_argument('--rotation', default='-80x,-15y,0z', help='ASE view rotation') - p.add_argument('--dpi', type=int, default=200) - return p.parse_args() - - -def atom_styles(atoms, n_initial, particle_species): - """Per-atom base colours and draw radii, keyed on role.""" - symbols = np.array(atoms.get_chemical_symbols()) - n = len(atoms) - colors = np.zeros((n, 3)) - radii = np.zeros(n) - for i in range(n): - if i >= n_initial: # appended by GCMC = adsorbed - colors[i], radii[i] = ADSORBED, 0.66 - elif symbols[i] == particle_species: - colors[i], radii[i] = AG, 1.44 - elif symbols[i] == 'O': - colors[i], radii[i] = SUPPORT_O, 0.40 - else: - colors[i], radii[i] = SUPPORT_METAL, 0.66 - return colors, radii - - -def depth_shade(colors, atoms, rotation): - """Darken atoms by their depth along the view direction (fake volume).""" - depth = atoms.positions @ rotate(rotation) - z = depth[:, 2] - norm = (z - z.min()) / (np.ptp(z) + 1e-9) # 0 = far, 1 = near - return colors * (0.40 + 0.60 * norm)[:, None], np.argsort(z) - - -def main(): - args = parse_args() - n_initial = len(read(args.trajectory, index=0)) - atoms = read(args.trajectory, index=args.frame) - n_ads = len(atoms) - n_initial - - colors, radii = atom_styles(atoms, n_initial, args.particle_species) - colors, order = depth_shade(colors, atoms, args.rotation) # far drawn first - - fig, ax = plt.subplots(figsize=(8.5, 8)) - plot_atoms(atoms[order], ax, colors=colors[order], radii=radii[order], - rotation=args.rotation, show_unit_cell=0) - for patch in ax.patches: # soften the cartoon outlines - fc = patch.get_facecolor() - patch.set_edgecolor((fc[0] * 0.55, fc[1] * 0.55, fc[2] * 0.55)) - patch.set_linewidth(0.25) - ax.set_axis_off() - ax.set_aspect('equal') - ax.set_title(f'O on Al$_2$O$_3$-supported {args.particle_species} nanoparticle ' - f'| {n_ads} adsorbed', fontsize=13) - - handles = [ - mpatches.Patch(color=AG, label=f'{args.particle_species} (nanoparticle)'), - mpatches.Patch(color=SUPPORT_METAL, label='support metal'), - mpatches.Patch(color=SUPPORT_O, label='support O'), - mpatches.Patch(color=ADSORBED, label='adsorbed'), - ] - ax.legend(handles=handles, loc='upper right', frameon=False, fontsize=10) - fig.tight_layout() - fig.savefig(args.out, dpi=args.dpi, facecolor='white') - print(f'saved {args.out} ({n_ads} adsorbed atoms, frame {args.frame})') - - -if __name__ == '__main__': - main() diff --git a/examples/re_gcmc.py b/examples/re_gcmc.py deleted file mode 100644 index c1966d9..0000000 --- a/examples/re_gcmc.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Replica-exchange GCMC of Ag and O on an fcc(111) Ag slab. - -Requires mpi4py and must be launched with `mpirun -n python examples/re_gcmc.py`, -with one MPI rank per temperature. -""" -import argparse -import os - -import numpy as np -from ase.build import fcc111 -from ase.constraints import FixAtoms -from mace.calculators import mace_mp - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.calculators import BaseCalculator # noqa: E402 -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.cell import CustomCell as Cell # noqa: E402 -from mcpy.ensembles import ReplicaExchange # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__) - p.add_argument('--temperatures', type=float, nargs='+', - default=[250, 300, 350, 400, 450, 500], - help='Replica temperatures (K) — one per MPI rank') - p.add_argument('--gcmc-steps', type=int, default=200, help='GCMC steps between exchanges') - p.add_argument('--exchange-interval', type=int, default=10, help='Steps between exchanges') - p.add_argument('--write-interval', type=int, default=1, help='Outfile/traj write interval') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.5, help='Shift applied to mu_O (eV)') - p.add_argument('--rel-max-steps', type=int, default=40, help='LBFGS relaxation steps') - p.add_argument('--rel-fmax', type=float, default=0.1, help='LBFGS force convergence (eV/Å)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--device', default='cuda', help='Torch device for MACE') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - seeds = [int(s) for s in ss.generate_state(5, dtype=np.uint32)] - np.savetxt(os.path.join(args.outdir, 'moves_seeds.txt'), seeds) - - atoms = fcc111('Ag', a=4.1592, size=(4, 4, 3), periodic=True, vacuum=8) - bottom_layer = [a.index for a in atoms if a.tag == 3] - atoms.set_constraint(FixAtoms(indices=bottom_layer)) - - cell_ag_ag = Cell(atoms, custom_height=5.5, bottom_z=12.8 - 2.11, - species_radii={'Ag': 2.75, 'O': 0}) - cell_ag_o = Cell(atoms, custom_height=5.5, bottom_z=12.8 - 2.11, - species_radii={'Ag': 2.11, 'O': 0}) - - # BaseCalculator wraps the bare ASE calculator with an LBFGS relaxation - # before each energy evaluation (a bare mace_mp would never relax). - calculator = BaseCalculator(mace_mp(device=args.device), - steps=args.rel_max_steps, fmax=args.rel_fmax) - - species = ['Ag', 'O'] - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} - - move_selector = MoveSelector( - [25, 25, 25, 25], - [DeletionMove(cell_ag_ag, species=['Ag'], seed=seeds[0]), - DeletionMove(cell_ag_o, species=['O'], seed=seeds[1]), - InsertionMove(cell_ag_ag, species=['Ag'], min_insert=0.5, seed=seeds[2]), - InsertionMove(cell_ag_o, species=['O'], min_insert=0.5, seed=seeds[3])], - ) - - def gcmc_factory(T, rank=0): - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}_rank{rank}' - return GrandCanonicalEnsemble( - atoms=atoms, - cells=[cell_ag_ag, cell_ag_o], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_relax_{tag}.out'), - trajectory_write_interval=args.write_interval, - outfile_write_interval=args.write_interval, - traj_file=os.path.join(args.outdir, f'gcmc_relax_{tag}.xyz'), - ) - - pt_gcmc = ReplicaExchange( - gcmc_factory, - temperatures=args.temperatures, - gcmc_steps=args.gcmc_steps, - exchange_interval=args.exchange_interval, - write_out_interval=args.write_interval, - seed=seeds[4], - ) - pt_gcmc.run() - - -if __name__ == '__main__': - main() diff --git a/examples/re_gcmc_batched_chunked.py b/examples/re_gcmc_batched_chunked.py deleted file mode 100644 index 845aa3c..0000000 --- a/examples/re_gcmc_batched_chunked.py +++ /dev/null @@ -1,169 +0,0 @@ -"""Batched replica-exchange GCMC with memory-capping chunks (single GPU). - -All replicas live in one process and share one ``AlchemiFCalculator``. Each MC -step relaxes and scores every replica's trial move in a batched FIRE pass. The -``chunk_size`` argument splits that batch into sub-batches, so peak GPU memory is -set by the largest chunk (``chunk_size`` x atoms-per-replica) instead of by the -replica count. That lets you run many replicas, or large nanoparticles, that -would otherwise exceed device memory at the cost of more (smaller) forward -passes. - -Picking ``chunk_size``: - * ``None`` -> whole batch in one pass (fastest, most memory). - * ``k`` -> ceil(n_replicas / k) passes, peak memory ~ k replicas. - * ``1`` -> one replica per pass (lowest memory, slowest). -Each forward carries a fixed ~17 ms overhead, so prefer the largest ``k`` whose -peak memory fits your card rather than ``1``. - -(For a no-relaxation run, ``AlchemiCalculator`` also takes ``energy_only=True``, -which skips force autograd for a further ~12% memory saving.) - -Requirements: - pip install -e .[alchemi] - -Run:: - - python examples/re_gcmc_batched_chunked.py --temperatures 300 400 500 600 \ - --np-length 10 --chunk-size 1 - -No mpirun. One GPU. Each replica builds its OWN cell and move_selector -(BatchedReplicaExchange would otherwise share state across replicas). - -GPU memory on long runs: the atom count changes every accepted move, which -fragments the CUDA caching allocator, so reserved GPU memory drifts up over a -long run (allocator fragmentation, not a leak). Launch with -``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True``. -""" -import argparse -import os - -import numpy as np -from ase.cluster import Octahedron - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.calculators import AlchemiFCalculator # noqa: E402 -from mcpy.cell import SphericalCell # noqa: E402 -from mcpy.ensembles import BatchedReplicaExchange # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--temperatures', type=float, nargs='+', - default=[300, 400, 500, 600, 700, 800], - help='Replica temperatures (K); one replica each') - p.add_argument('--np-length', type=int, default=10, - help='Octahedron edge length (larger = bigger nanoparticle)') - p.add_argument('--chunk-size', type=int, default=1, - help='Replicas per forward pass. Caps peak GPU memory at ' - 'chunk_size x atoms-per-replica. 0/negative -> whole batch.') - p.add_argument('--gcmc-steps', type=int, default=200) - p.add_argument('--exchange-interval', type=int, default=10) - p.add_argument('--write-interval', type=int, default=1) - p.add_argument('--checkpoint', default='medium-mpa-0') - p.add_argument('--device', default='cuda') - p.add_argument('--no-cueq', action='store_true') - p.add_argument('--no-compile', action='store_true') - p.add_argument('--fmax', type=float, default=0.05) - p.add_argument('--relax-steps', type=int, default=500) - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.5) - p.add_argument('--min-insert', type=float, default=0.5) - p.add_argument('--vacuum', type=float, default=3.0) - p.add_argument('--seed', type=int, default=None) - p.add_argument('--outdir', default='.') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - chunk_size = args.chunk_size if args.chunk_size and args.chunk_size > 0 else None - - ss = np.random.SeedSequence(args.seed) - n_replicas = len(args.temperatures) - # 2 move seeds per replica (deletion + insertion) + 1 master seed for RE. - seeds = ss.generate_state(2 * n_replicas + 1, dtype=np.uint32) - move_seeds = [int(s) for s in seeds[:2 * n_replicas]] - master_seed = int(seeds[-1]) - - base_atoms = Octahedron('Ag', args.np_length, 1) - print(f'{n_replicas} replicas x {len(base_atoms)} atoms; ' - f'chunk_size={chunk_size} -> peak memory ~ ' - f'{(chunk_size or n_replicas) * len(base_atoms)} atoms/pass') - - # ONE model + calculator, shared across replicas. chunk_size caps peak memory. - calculator = AlchemiFCalculator( - checkpoint=args.checkpoint, - steps=args.relax_steps, - fmax=args.fmax, - device=args.device, - enable_cueq=not args.no_cueq, - compile_model=not args.no_compile, - chunk_size=chunk_size, - ) - - species = ['O'] - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} - - def gcmc_factory(T, rank): - atoms = base_atoms.copy() - scell = SphericalCell(atoms, vacuum=args.vacuum, - species_radii={'Ag': 2.947, 'O': 0}, - mc_sample_points=100_000) - s = move_seeds[2 * rank:2 * (rank + 1)] - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=species, seed=s[0]), - InsertionMove(scell, species=species, min_insert=args.min_insert, seed=s[1])], - ) - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}_rank{rank}' - return GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_chunked_{tag}.out'), - trajectory_write_interval=args.write_interval, - outfile_write_interval=args.write_interval, - traj_file=os.path.join(args.outdir, f'gcmc_chunked_{tag}.xyz'), - ) - - pt = BatchedReplicaExchange( - gcmc_factory, - calculator=calculator, - temperatures=args.temperatures, - gcmc_steps=args.gcmc_steps, - exchange_interval=args.exchange_interval, - write_out_interval=args.write_interval, - seed=master_seed, - outfile=os.path.join(args.outdir, 'replica_exchange_chunked.log'), - ) - - if args.device == 'cuda': - import torch - torch.cuda.reset_peak_memory_stats() - pt.run() - print(f'peak GPU memory: {torch.cuda.max_memory_allocated() / 1024**2:.0f} MB') - else: - pt.run() - - -if __name__ == '__main__': - main() diff --git a/examples/re_gcmc_co_cupd_batched.py b/examples/re_gcmc_co_cupd_batched.py deleted file mode 100644 index a820bd9..0000000 --- a/examples/re_gcmc_co_cupd_batched.py +++ /dev/null @@ -1,317 +0,0 @@ -"""CO adsorption phase diagram on a CuPd nanoparticle via batched-RE GCMC. - -One replica per CO chemical potential (mu ladder), all replicas evaluated in -batched relaxed-energy passes on one GPU (AlchemiFCalculator). Moves: rigid CO -molecule insertion/deletion (mu = E(CO) + delta_mu, full molecular chemical -potential) plus Cu<->Pd permutation so the alloy can segregate under CO. - -Outputs per replica outfiles/trajectories plus a coverage-vs-delta_mu phase -diagram (phase_diagram_co_cupd.png). - -Run (alchemi env):: - - python examples/re_gcmc_co_cupd_batched.py --gcmc-steps 80 -""" -import argparse -import os - -import numpy as np -from ase.build import molecule -from ase.cluster import Octahedron - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import (MoleculeDeletionMove, MoleculeDisplacementMove, # noqa: E402 - MoleculeInsertionMove, PermutationMove) -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.ensembles import BatchedReplicaExchange # noqa: E402 -from mcpy.calculators import AlchemiFCalculator # noqa: E402 -from mcpy.cell import SphericalCell # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--delta-mus', type=float, nargs='+', - default=[-1.5, -1.2, -0.9, -0.6, -0.3], - help='CO delta_mu ladder, one replica each (eV)') - p.add_argument('--T', type=float, default=400.0) - p.add_argument('--n-pd', type=int, default=5, help='Pd atoms in the Cu NP') - p.add_argument('--octa-length', type=int, default=4, - help='Octahedron length (4,1 -> 38 atoms; 9,3 -> 405)') - p.add_argument('--octa-cutoff', type=int, default=1) - p.add_argument('--gcmc-steps', type=int, default=80, - help='GCMC steps per replica') - p.add_argument('--exchange-interval', type=int, default=8, - help='Steps between replica-exchange attempts. Frequent ' - 'exchange is fine: a correct swap rule leaves each ' - "rung's distribution intact. (The old advice to use " - '>= 50 here was working around the mu-ladder swap ' - 'bug, which accepted every swap and so did flatten ' - 'the coverage curve the more often it ran.)') - p.add_argument('--checkpoint', default='medium-mpa-0') - p.add_argument('--rel-steps', type=int, default=100) - p.add_argument('--rel-fmax', type=float, default=0.05) - p.add_argument('--min-insert', type=float, default=1.3) - p.add_argument('--mol-disp-weight', type=int, default=0, - help='Weight of the rigid CO translate+rotate move ' - '(0 = disabled, matching pre-move behaviour)') - p.add_argument('--mol-disp-max', type=float, default=1.0, - help='Max COM displacement of the rigid move (A)') - p.add_argument('--mol-disp-angle', type=float, default=None, - help='Max rotation angle of the rigid move (rad); ' - 'None = full uniform rotation') - p.add_argument('--chunk-size', type=int, default=None, - help='Relax at most this many replicas per forward pass. ' - 'Peak GPU memory follows the largest chunk rather ' - 'than the replica count, which is what makes a wide ' - 'ladder fit: acceptance-equalized spacing needs ~29 ' - 'rungs here, and 29 x ~460 atoms is several times the ' - 'whole-batch ceiling. None relaxes the whole batch.') - p.add_argument('--no-compile', action='store_true') - p.add_argument('--seed', type=int, default=7, - help='Master seed for moves, RE, and the Pd placement') - p.add_argument('--outdir', - default=os.path.expanduser('~/mcpy_tmp_runs/co_cupd'), - help='Output directory (kept outside the git repo)') - p.add_argument('--init-dir', default=None, - help='Seed each replica from the last trajectory frame of a ' - 'previous run in this directory (same delta-mus) to ' - 'continue toward convergence') - return p.parse_args() - - -def build_nanoparticle(n_pd, seed, length=4, cutoff=1): - atoms = Octahedron('Cu', length, cutoff=cutoff) - rng = np.random.default_rng(seed) - pd_idx = rng.choice(len(atoms), size=n_pd, replace=False) - symbols = np.array(atoms.get_chemical_symbols()) - symbols[pd_idx] = 'Pd' - atoms.set_chemical_symbols(symbols.tolist()) - atoms.center(vacuum=10.0) - atoms.pbc = False - return atoms - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - n_rep = len(args.delta_mus) - seeds = ss.generate_state(3 * n_rep + 2, dtype=np.uint32) - move_seeds = [int(s) for s in seeds[:3 * n_rep]] - master_seed, np_seed = int(seeds[-1]), int(seeds[-2]) - - base_atoms = build_nanoparticle(args.n_pd, np_seed, - args.octa_length, args.octa_cutoff) - n_metal = len(base_atoms) - print(f'nanoparticle: {base_atoms.get_chemical_formula()} ({n_metal} atoms)') - - co = molecule('CO') - - calculator = AlchemiFCalculator( - checkpoint=args.checkpoint, - steps=args.rel_steps, - fmax=args.rel_fmax, - compile_model=not args.no_compile, - chunk_size=args.chunk_size, - ) - - e_co = calculator.get_potential_energy( - molecule('CO', cell=[20.0, 20.0, 20.0], pbc=False)) - print(f'E(CO, relaxed, isolated) = {e_co:.4f} eV') - mus = [{'CO': e_co + d} for d in args.delta_mus] - - def gcmc_factory(mu, rank): - atoms = base_atoms.copy() - if args.init_dir: - import ase.io - d = args.delta_mus[rank] - prev = os.path.join( - args.init_dir, - f'gcmc_{base_atoms.get_chemical_formula()}_CO_dmu_{d}.xyz') - atoms = ase.io.read(prev, index=-1) # molecule_id round-trips - atoms.pbc = False - cell = SphericalCell( - atoms, vacuum=3.5, - species_radii={'Cu': 2.4, 'Pd': 2.5, 'C': 0.0, 'O': 0.0}, - seed=move_seeds[3 * rank] + 1, - ) - s = move_seeds[3 * rank:3 * (rank + 1)] - weights = [2, 2, 1] - moves = [MoleculeInsertionMove(cell, co, 'CO', seed=s[0], - min_insert=args.min_insert), - MoleculeDeletionMove(cell, co, 'CO', seed=s[1]), - PermutationMove(species=['Cu', 'Pd'], seed=s[2])] - if args.mol_disp_weight > 0: - weights.append(args.mol_disp_weight) - moves.append(MoleculeDisplacementMove( - cell, co, 'CO', seed=s[0] + 101, - max_displacement=args.mol_disp_max, - max_angle=args.mol_disp_angle)) - move_selector = MoveSelector(weights, moves, n_moves=3) - d = args.delta_mus[rank] - tag = f'{base_atoms.get_chemical_formula()}_CO_dmu_{d}' - return GrandCanonicalEnsemble( - atoms=atoms, - cells=[cell], - calculator=calculator, - mu=mu, - units_type='metal', - species=['Cu', 'Pd'], - molecules={'CO': co}, - temperature=args.T, - move_selector=move_selector, - random_seed=s[0] + 17, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - outfile_write_interval=1, - trajectory_write_interval=1, - ) - - pt = BatchedReplicaExchange( - gcmc_factory, - calculator=calculator, - mus=mus, - gcmc_steps=args.gcmc_steps, - exchange_interval=args.exchange_interval, - write_out_interval=args.exchange_interval, - seed=master_seed, - outfile=os.path.join(args.outdir, 'replica_exchange_co_cupd.log'), - global_minimum_file=os.path.join(args.outdir, 'global_minimum.xyz'), - ) - pt.run() - - plot_coverage(args, n_metal) - render_snapshots(args, base_atoms.get_chemical_formula()) - try: - library_phase_diagram(args, base_atoms.get_chemical_formula(), e_co) - except ValueError as exc: - # Deep continuation chains may hold no adsorbate-free reference frame; - # pool the chain's trajectories by hand for the final diagram then. - print(f'library phase diagram skipped: {exc}') - - -def library_phase_diagram(args, formula, e_co): - """The proper thermodynamic phase diagram via mcpy's own analysis: - per-stoichiometry free-energy lines, stable envelope, transitions, and - per-phase structure thumbnails. CO molecules are counted via their single - C atom; mu_ref = E(CO) matches the full-molecular-mu convention.""" - import ase.io - from mcpy.utils.phase_diagram import plot_phase_diagram - - frames = [] - for d in args.delta_mus: - frames.append(ase.io.read( - os.path.join(args.outdir, f'gcmc_{formula}_CO_dmu_{d}.xyz'), - index=':')) - # Continuation runs start CO-loaded; pull the adsorbate-free - # reference frames from the run they were seeded from. - if args.init_dir: - frames.append(ase.io.read( - os.path.join(args.init_dir, f'gcmc_{formula}_CO_dmu_{d}.xyz'), - index=':')) - pad = 0.15 - out = os.path.join(args.outdir, 'phase_diagram_co_cupd.png') - res = plot_phase_diagram( - frames, - adsorbate='C', - metal_symbols=('Cu', 'Pd'), - mu_ref=e_co, - kind='nano', - T=args.T, - dmu_range=(min(args.delta_mus) - pad, max(args.delta_mus) + pad), - rotation='12z,-75x', - gamma_in_ev=True, - outfile=out, - adsorbate_label='CO', - atoms_per_reservoir_molecule=1, - ) - print(f'phase diagram (library): {out}') - print('stable CO stoichiometries:', - [res['stoich'][i] for i in res['phase_order']]) - print('transitions at delta_mu:', np.round(res['transitions'], 3).tolist()) - return res - - -def plot_coverage(args, n_metal): - """Coverage vs delta_mu from the per-replica outfiles (last 50%).""" - import matplotlib - matplotlib.use('Agg') - import matplotlib.pyplot as plt - - cov, err = [], [] - for d in args.delta_mus: - n_series = [] - path = None - for f in os.listdir(args.outdir): - if f.endswith(f'_dmu_{d}.out'): - path = os.path.join(args.outdir, f) - with open(path) as fh: - for line in fh: - s = line.split() - if s and s[0].isdigit(): - n_series.append((int(s[1]) - n_metal) / 2.0) - tail = np.array(n_series[len(n_series) // 2:]) - nb = max(2, min(5, len(tail))) - blocks = [b.mean() for b in np.array_split(tail, nb) if len(b)] - cov.append(np.mean(blocks)) - err.append(np.std(blocks, ddof=1) / np.sqrt(len(blocks))) - # Stationarity check: first vs second half of the post-burn tail must - # agree within 2x the combined block error, else more steps needed. - h1, h2 = np.array_split(tail, 2) - s_comb = np.hypot(h1.std(ddof=1) / np.sqrt(max(len(h1), 2)), - h2.std(ddof=1) / np.sqrt(max(len(h2), 2))) - drift = abs(h1.mean() - h2.mean()) - status = 'CONVERGED' if drift < 2 * max(s_comb, 1e-9) else 'NOT CONVERGED' - print(f' dmu={d}: tail halves {h1.mean():.1f} vs {h2.mean():.1f} ' - f'(drift {drift:.1f}, 2s={2 * s_comb:.1f}) {status}') - - fig, ax = plt.subplots(figsize=(5, 3.4), constrained_layout=True) - ax.errorbar(args.delta_mus, cov, yerr=err, fmt='o-', capsize=4) - ax.set_xlabel(r'$\Delta\mu_\mathrm{CO}$ (eV)') - ax.set_ylabel(r'$\langle N_\mathrm{CO} \rangle$ adsorbed') - ax.set_title(f'CO on CuPd nanoparticle, T = {args.T:.0f} K') - out = os.path.join(args.outdir, 'coverage_co_cupd.png') - fig.savefig(out, dpi=150) - print(f'coverage isotherm: {out}') - print('coverage:', {d: f'{c:.1f}±{e:.1f}' for d, c, e in - zip(args.delta_mus, cov, err)}) - - -def render_snapshots(args, formula): - """One rendered snapshot per replica (final trajectory frame) plus the - cross-replica global minimum, on a single row.""" - import ase.io - import matplotlib - matplotlib.use('Agg') - import matplotlib.pyplot as plt - from ase.visualize.plot import plot_atoms - - n = len(args.delta_mus) - fig, axes = plt.subplots(1, n + 1, figsize=(2.6 * (n + 1), 3.0), - constrained_layout=True) - for ax, d in zip(axes[:n], args.delta_mus): - atoms = ase.io.read( - os.path.join(args.outdir, f'gcmc_{formula}_CO_dmu_{d}.xyz'), - index=-1) - n_co = (len(atoms) - len([s for s in atoms.get_chemical_symbols() - if s in ('Cu', 'Pd')])) // 2 - plot_atoms(atoms, ax, rotation='12z,-75x', show_unit_cell=0) - ax.set_title(f'$\\Delta\\mu$ = {d} eV\n{n_co} CO', fontsize=9) - ax.set_axis_off() - gm = ase.io.read(os.path.join(args.outdir, 'global_minimum.xyz')) - plot_atoms(gm, axes[n], rotation='12z,-75x', show_unit_cell=0) - axes[n].set_title(f'global minimum\n{gm.get_chemical_formula()}', fontsize=9) - axes[n].set_axis_off() - out = os.path.join(args.outdir, 'snapshots_co_cupd.png') - fig.savefig(out, dpi=180) - print(f'snapshots: {out}') - - -if __name__ == '__main__': - main() diff --git a/examples/re_gcmc_nano_alchemi.py b/examples/re_gcmc_nano_alchemi.py deleted file mode 100644 index c4e43d9..0000000 --- a/examples/re_gcmc_nano_alchemi.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Batched replica-exchange GCMC of O adsorption on an Ag octahedral nanoparticle. - -Replica-exchange counterpart of ``gcmc_nano_alchemi.py``, built on the -single-GPU batched pattern of ``re_gcmc_batched.py``. All replicas live in one -Python process and share one ``AlchemiFCalculator``; every replica's trial move -is relaxed with FIRE and scored in a single batched pass per MC step (one -optimizer / one model forward per FIRE step, with converged replicas retired -from the active batch). - -Requirements: - pip install 'nvalchemi-toolkit[mace]' - -Run:: - - python examples/re_gcmc_nano_alchemi.py --temperatures 300 400 500 600 - -No mpirun. One GPU. Each replica builds its OWN cell and move_selector -(BatchedReplicaExchange would otherwise share state across replicas). - -GPU memory on long runs: the atom count changes every accepted move, which -fragments the CUDA caching allocator, so reserved GPU memory drifts up over a -long run (allocator fragmentation, not a leak). Launch with -``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True``. -""" -import argparse -import os - -import numpy as np -from ase.cluster import Octahedron - -from mcpy.utils.logging import configure as configure_logging - -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.calculators import AlchemiFCalculator # noqa: E402 -from mcpy.cell import SphericalCell # noqa: E402 -from mcpy.ensembles import BatchedReplicaExchange # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--temperatures', type=float, nargs='+', - default=[300, 400, 500, 600, 700, 800], - help='Replica temperatures (K)') - p.add_argument('--gcmc-steps', type=int, default=200) - p.add_argument('--exchange-interval', type=int, default=10) - p.add_argument('--write-interval', type=int, default=1) - p.add_argument('--checkpoint', default='medium-mpa-0', - help='Named checkpoint (e.g. medium-mpa-0) or path to a .pt file') - p.add_argument('--device', default='cuda') - p.add_argument('--no-cueq', action='store_true', - help='Disable cuEquivariance kernel fusion') - p.add_argument('--no-compile', action='store_true', - help='Disable torch.compile (faster startup, slower inference)') - p.add_argument('--fmax', type=float, default=0.05, - help='FIRE force convergence threshold (eV/Å)') - p.add_argument('--relax-steps', type=int, default=500, - help='Max FIRE steps per batched energy evaluation') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.5, help='Shift applied to mu_O (eV)') - p.add_argument('--min-insert', type=float, default=0.5, - help='Min insertion distance to existing atoms (Å)') - p.add_argument('--vacuum', type=float, default=3.0, help='Vacuum around nanoparticle (Å)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def main(): - args = parse_args() - os.makedirs(args.outdir, exist_ok=True) - - ss = np.random.SeedSequence(args.seed) - n_replicas = len(args.temperatures) - # 2 move seeds per replica (deletion + insertion) + 1 master seed for RE. - seeds = ss.generate_state(2 * n_replicas + 1, dtype=np.uint32) - move_seeds = [int(s) for s in seeds[:2 * n_replicas]] - master_seed = int(seeds[-1]) - - # Shared template -- each factory call copies it and builds its own cell/moves. - base_atoms = Octahedron('Ag', 9, 3) - - # ONE model, shared across replicas. Batched FIRE relax reuses this calculator. - calculator = AlchemiFCalculator( - checkpoint=args.checkpoint, - steps=args.relax_steps, - fmax=args.fmax, - device=args.device, - enable_cueq=not args.no_cueq, - compile_model=not args.no_compile, - ) - - species = ['O'] - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} - - def gcmc_factory(T, rank): - atoms = base_atoms.copy() - - scell = SphericalCell(atoms, vacuum=args.vacuum, species_radii={'Ag': 2.947, 'O': 0}, - mc_sample_points=100_000) - - s = move_seeds[2 * rank:2 * (rank + 1)] - move_selector = MoveSelector( - [25, 25], - [DeletionMove(scell, species=species, seed=s[0]), - InsertionMove(scell, species=species, min_insert=args.min_insert, seed=s[1])], - ) - - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}_rank{rank}' - return GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=species, - temperature=T, - move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_re_nano_{tag}.out'), - trajectory_write_interval=args.write_interval, - outfile_write_interval=args.write_interval, - traj_file=os.path.join(args.outdir, f'gcmc_re_nano_{tag}.xyz'), - ) - - pt = BatchedReplicaExchange( - gcmc_factory, - calculator=calculator, - temperatures=args.temperatures, - gcmc_steps=args.gcmc_steps, - exchange_interval=args.exchange_interval, - write_out_interval=args.write_interval, - seed=master_seed, - outfile=os.path.join(args.outdir, 'replica_exchange_nano.log'), - ) - pt.run() - - -if __name__ == '__main__': - main() From 48bf5428397d2ef7feb1ba4f709dbbe3f0c2d439 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:14:36 +0200 Subject: [PATCH 06/20] Prune docs to the two remaining examples Deletes the five example pages that documented parked scripts and drops them from the examples index; the four surviving pages carry their code inline and never named a file. Retargets the live references: the installation page and the cluster tutorial now point at the two tracked scripts, and the molecular adsorbates page and the README keep only the notebook pointer. Two cross-references broke with the deleted pages (docs/cells.rst pointing at the dome walk-through, and the NVT example pointing at the replica-exchange page); both are removed or redirected, and the Sphinx build is warning-free. CHANGELOG.md, docs/superpowers/ and docs/research/ keep their references on purpose: they are historical records of when those scripts existed. The notebooks still reference three parked scripts and are left untouched, since editing executed notebook JSON is a separate job. --- README.md | 5 +- docs/cells.rst | 2 - docs/examples.rst | 6 - docs/examples/gcmc_bimet_nano.rst | 118 ---------------- docs/examples/gcmc_dome_supported.rst | 123 ----------------- docs/examples/gcmc_nano.rst | 104 -------------- docs/examples/gcmc_nano_supported.rst | 114 --------------- docs/examples/mc_simulations.rst | 4 +- .../examples/replica_exchange_simulations.rst | 130 ------------------ docs/installation.rst | 5 +- docs/molecular_adsorbates.rst | 4 - docs/tutorials/cluster_install_iqtc.rst | 4 +- examples/gcmc_ag111.py | 4 +- tests/test_ag111_gcmc_golden.py | 4 +- 14 files changed, 12 insertions(+), 615 deletions(-) delete mode 100644 docs/examples/gcmc_bimet_nano.rst delete mode 100644 docs/examples/gcmc_dome_supported.rst delete mode 100644 docs/examples/gcmc_nano.rst delete mode 100644 docs/examples/gcmc_nano_supported.rst delete mode 100644 docs/examples/replica_exchange_simulations.rst diff --git a/README.md b/README.md index e64ba60..807526d 100644 --- a/README.md +++ b/README.md @@ -197,10 +197,7 @@ gcmc = GrandCanonicalEnsemble( move_selector=moves) ``` -Runnable examples: `examples/gcmc_molecule_mace.py` (O2 on Ag(111); any g2 molecule via `--molecule`) and -`examples/re_gcmc_co_cupd_batched.py` (CO on a CuPd nanoparticle with batched -replica exchange and a phase diagram). The notebooks in `notebooks/` walk -through the machinery step by step. +The notebooks in `notebooks/` walk through the machinery step by step. ## Contributing diff --git a/docs/cells.rst b/docs/cells.rst index 2a575d8..b29a669 100644 --- a/docs/cells.rst +++ b/docs/cells.rst @@ -166,8 +166,6 @@ Key parameters: ) dome_cell.calculate_volume(atoms) -A worked end-to-end run is in :doc:`examples/gcmc_dome_supported`. - Choosing the right cell ----------------------- diff --git a/docs/examples.rst b/docs/examples.rst index 7eacabb..3d2d864 100644 --- a/docs/examples.rst +++ b/docs/examples.rst @@ -59,7 +59,6 @@ Phase diagrams from GCMC ensembles The formation-energy construction, the lower-envelope logic, and the post-processing utilities are described in :doc:`phase_diagrams`. -The replica-exchange examples below end in exactly that analysis. Examples by workflow @@ -80,10 +79,6 @@ Grand Canonical Monte Carlo (GCMC) :maxdepth: 1 examples/gcmc_custom_cell.rst - examples/gcmc_nano.rst - examples/gcmc_bimet_nano.rst - examples/gcmc_nano_supported.rst - examples/gcmc_dome_supported.rst examples/gcmc_simulations.rst Replica-Exchange GCMC (RE-GCMC) @@ -92,5 +87,4 @@ Replica-Exchange GCMC (RE-GCMC) .. toctree:: :maxdepth: 1 - examples/replica_exchange_simulations.rst examples/phase_diagram_analysis.rst diff --git a/docs/examples/gcmc_bimet_nano.rst b/docs/examples/gcmc_bimet_nano.rst deleted file mode 100644 index c18a456..0000000 --- a/docs/examples/gcmc_bimet_nano.rst +++ /dev/null @@ -1,118 +0,0 @@ -GCMC of H on a Pt/Au bimetallic nanoparticle with metal permutations -==================================================================== - -Hydrogen GCMC on a Pt/Au truncated-octahedral nanoparticle, combined with -Pt↔Au :class:`PermutationMove` on the metal framework. Demonstrates how to -mix variable-particle-number moves with chemical-ordering moves in a -single ensemble. - -Mirrors ``examples/gcmc_bimet_nano.py``. - -Goal ----- - -At fixed :math:`T = 500\,\mathrm{K}` and :math:`\mu_{\mathrm{H}}` derived -from :math:`\tfrac{1}{2}E(\mathrm{H}_2)`, sample equilibrium H coverage on -the cluster while letting the Pt/Au shell composition equilibrate. - -Prerequisites -------------- - -- An ASE calculator that can evaluate H\ :sub:`2` (MACE in the script, - but any ASE calculator works). - -Code ----- - -.. code-block:: python - - import numpy as np - from ase.build import molecule - from ase.cluster import Octahedron - from mace.calculators import mace_mp - - from mcpy.calculators import BaseCalculator - from mcpy.cell import SphericalCell - from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble - from mcpy.moves import DeletionMove, InsertionMove, PermutationMove - from mcpy.moves.move_selector import MoveSelector - - T = 500.0 - delta_mu_H = -0.5 - - ss = np.random.SeedSequence(0) - seed_del, seed_ins, seed_perm = ( - int(s) for s in ss.generate_state(3, dtype=np.uint32) - ) - - atoms = Octahedron('Pt', 7, 2) - half = len(atoms) // 2 - atoms.symbols = ['Pt'] * half + ['Au'] * half + ['Pt'] - atoms.set_pbc(False) - - scell = SphericalCell( - atoms, - vacuum=3, - species_radii={'Pt': 2, 'Au': 2.5, 'H': 0}, - mc_sample_points=100_000, - ) - - # BaseCalculator relaxes with LBFGS before each energy evaluation; - # a bare mace_mp would return unrelaxed energies. - calculator = BaseCalculator(mace_mp(device='cuda'), steps=100, fmax=0.05) - - move_selector = MoveSelector( - [1, 1, 1], - [DeletionMove(scell, species=['H'], seed=seed_del), - InsertionMove(scell, species=['H'], min_insert=0.5, seed=seed_ins), - PermutationMove(species=['Au', 'Pt'], seed=seed_perm)], - ) - - # Reference mu_H from molecular H2. - e_h2 = calculator.get_potential_energy(molecule('H2')) - mus = {'H': e_h2 / 2 + delta_mu_H} - - tag = f'{atoms.get_chemical_formula()}_dmu_{delta_mu_H}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=['H'], - temperature=T, - move_selector=move_selector, - outfile=f'gcmc_{tag}.out', - trajectory_write_interval=1, - outfile_write_interval=1, - traj_file=f'gcmc_{tag}.xyz', - ) - gcmc.run(1_000_000) - -Outputs -------- - -- ``gcmc__dmu_.out`` — step log with per-move acceptance for - H insertion, H deletion, and Pt↔Au permutation. -- ``gcmc__dmu_.xyz`` — extended XYZ trajectory. - -Interpretation --------------- - -- ``species=['H']`` declares the *variable-number* species. Pt and Au - counts are conserved by :class:`PermutationMove`, which only swaps - identities of existing atoms. -- The ``species_radii`` entry for ``'H'`` is 0 so that hydrogen never - blocks itself or the metal scaffold during the free-volume estimate. -- Reference energies use a stricter relaxation (``steps=100``, - ``fmax=0.05``) than the production loop; restore the production - tolerance before calling ``gcmc.run`` if you reuse the same calculator - object. - -Next steps ----------- - -- Drop the permutation move to recover a pure GCMC run: see - :doc:`gcmc_nano`. -- Replace H with O and the metal symbols with Ag to study oxidation on - a single-metal cluster. diff --git a/docs/examples/gcmc_dome_supported.rst b/docs/examples/gcmc_dome_supported.rst deleted file mode 100644 index 6e38289..0000000 --- a/docs/examples/gcmc_dome_supported.rst +++ /dev/null @@ -1,123 +0,0 @@ -GCMC of O on a supported Ag nanoparticle (dome region) -====================================================== - -Grand-canonical sampling of oxygen on an Ag truncated-octahedral -nanoparticle placed on an Al\ :sub:`2`\ O\ :sub:`3` support, using a -:class:`DomeCell` to confine insertion/deletion to a hemispherical region -centred on the particle. Compared with :doc:`gcmc_nano_supported` (which -uses a slab-spanning :class:`CustomCell`), the dome keeps trial insertions -on and around the particle and its metal-support contact rim instead of -spreading them across the bare support. - -Adapted from ``examples/gcmc_nano_supported.py`` by swapping -:class:`CustomCell` for :class:`DomeCell`. - -Goal ----- - -At fixed :math:`T = 500\,\mathrm{K}` and chosen -:math:`\Delta\mu_{\mathrm{O}}`, sample equilibrium O coverage on the -supported cluster, with insertions localised on the nanoparticle. - -Prerequisites -------------- - -- A POSCAR of the support (``Al2O3.poscar`` in the script). -- A MACE checkpoint or other ASE calculator that can describe Ag, Al, - and O. - -Code ----- - -.. code-block:: python - - import numpy as np - from ase.cluster import Octahedron - from ase.constraints import FixAtoms - from ase.io import read - from mace.calculators import mace_mp - - from mcpy.cell import DomeCell - from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble - from mcpy.moves import DeletionMove, InsertionMove - from mcpy.moves.move_selector import MoveSelector - from mcpy.utils.utils import get_p_at_support - - T = 500.0 - delta_mu_O = -0.5 - mus = {'Ag': -2.99, 'O': -4.91 + delta_mu_O} - - ss = np.random.SeedSequence(0) - seed_del, seed_ins = (int(s) for s in ss.generate_state(2, dtype=np.uint32)) - - support = read('Al2O3.poscar').repeat((4, 4, 1)) - support.center(vacuum=10.0, axis=2) - z = support.positions[:, 2] - z_half = z.min() + 0.5 * (z.max() - z.min()) - support.set_constraint(FixAtoms(mask=(z < z_half))) - - surface_z = float(np.max(support.positions[:, 2])) - particle = Octahedron('Ag', 5, 2) - atoms = get_p_at_support(support, particle, contact_surface='100', gap=2.0) - - scell = DomeCell( - atoms, - particle_species='Ag', - bottom_z=surface_z, - vacuum=3.0, - species_radii={'Ag': 2.068, 'O': 0, 'Al': 3}, - ) - - calculator = mace_mp(device='cuda') - - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=['O'], seed=seed_del), - InsertionMove(scell, species=['O'], min_insert=0.5, seed=seed_ins)], - ) - - tag = f'{atoms.get_chemical_formula()}_dmu_{delta_mu_O}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=['O'], - temperature=T, - move_selector=move_selector, - outfile=f'gcmc_{tag}.out', - traj_file=f'gcmc_{tag}.xyz', - ) - gcmc.run(1_000_000) - -Outputs -------- - -- ``gcmc__dmu_.out`` -- step log with per-move acceptance. -- ``gcmc__dmu_.xyz`` -- extended XYZ trajectory. - -Interpretation --------------- - -- ``particle_species='Ag'`` tells the dome which atoms are the - nanoparticle: the region is centred on the Ag centroid and its radius - is the farthest Ag atom plus ``vacuum``. The support is never used to - place or size the dome. -- ``bottom_z=surface_z`` clips the dome at the topmost support atom, so - sampled points stay above the surface; the part of the ball that would - dip into the support is discarded. -- Only O is exchanged. Al and Ag enter ``species_radii`` solely so the - free-volume estimator excludes overlaps with them. -- Because insertions are confined to the dome, O accumulates on the - particle far more efficiently than with a slab-spanning cell, where - many trial insertions land on the bare support. - -Next steps ----------- - -- Compare against the slab-cell version: :doc:`gcmc_nano_supported`. -- Visualise the loaded particle with - ``examples/plot_supported_np.py``. -- Sweep :math:`\Delta\mu_{\mathrm{O}}` and post-process with - :doc:`phase_diagram_analysis`. diff --git a/docs/examples/gcmc_nano.rst b/docs/examples/gcmc_nano.rst deleted file mode 100644 index f91718d..0000000 --- a/docs/examples/gcmc_nano.rst +++ /dev/null @@ -1,104 +0,0 @@ -GCMC of O on an Ag octahedral nanoparticle -========================================== - -Grand-canonical sampling of oxygen on a truncated-octahedral Ag -nanoparticle in vacuum. A :class:`SphericalCell` confines insertions and -deletions to a shell around the cluster. - -Mirrors ``examples/gcmc_nano.py``. - -Goal ----- - -At fixed :math:`T = 500\,\mathrm{K}` and a chosen :math:`\mu_{\mathrm{O}}`, -sample equilibrium O coverage on the nanoparticle. Repeat across -:math:`\Delta\mu_{\mathrm{O}}` to construct a coverage–chemical-potential -curve. - -Prerequisites -------------- - -- A MACE checkpoint, or any ASE calculator (the script defaults to - ``mace_mp`` on CUDA — swap for EMT or another calculator for a quick - smoke test). -- See :doc:`../getting_started` for the minimal API walkthrough. - -Code ----- - -.. code-block:: python - - import numpy as np - from ase.cluster import Octahedron - from mace.calculators import mace_mp - - from mcpy.cell import SphericalCell - from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble - from mcpy.moves import DeletionMove, InsertionMove - from mcpy.moves.move_selector import MoveSelector - - T = 500.0 - delta_mu_O = -0.5 - mus = {'Ag': -2.99, 'O': -4.91 + delta_mu_O} - - ss = np.random.SeedSequence(0) - seed_del, seed_ins = (int(s) for s in ss.generate_state(2, dtype=np.uint32)) - - atoms = Octahedron('Ag', 6, 1) - - scell = SphericalCell( - atoms, - vacuum=3, - species_radii={'Ag': 2.947, 'O': 0}, - mc_sample_points=100_000, - ) - - calculator = mace_mp(device='cuda') - - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=['O'], seed=seed_del), - InsertionMove(scell, species=['O'], min_insert=0.5, seed=seed_ins)], - ) - - tag = f'{atoms.get_chemical_formula()}_dmu_{delta_mu_O}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=['O'], - temperature=T, - move_selector=move_selector, - outfile=f'gcmc_{tag}.out', - traj_file=f'gcmc_{tag}.xyz', - ) - gcmc.run(1_000_000) - -Outputs -------- - -- ``gcmc__dmu_.out`` — step log with per-move acceptance. -- ``gcmc__dmu_.xyz`` — extended XYZ trajectory. - -Interpretation --------------- - -- :class:`SphericalCell` recomputes its free volume after every accepted - move; ``species_radii`` set per-element exclusion radii that prevent - trial insertions from overlapping existing atoms. -- Healthy insertion/deletion acceptance is typically 5–30 %. Adjust - ``min_insert`` and the ``vacuum`` padding if acceptance collapses. -- Run a short warm-up and discard the leading frames before computing - averages; the trajectory is written from step 0. - -Next steps ----------- - -- Sweep :math:`\Delta\mu_{\mathrm{O}}` and post-process the trajectories - with :doc:`phase_diagram_analysis`. -- Replace the single sphere with a supported-particle setup in - :doc:`gcmc_nano_supported`. -- Add a metal permutation move on a bimetallic cluster in - :doc:`gcmc_bimet_nano`. diff --git a/docs/examples/gcmc_nano_supported.rst b/docs/examples/gcmc_nano_supported.rst deleted file mode 100644 index ff09867..0000000 --- a/docs/examples/gcmc_nano_supported.rst +++ /dev/null @@ -1,114 +0,0 @@ -GCMC of O on a supported Ag nanoparticle -======================================== - -Grand-canonical sampling of oxygen on an Ag truncated-octahedral -nanoparticle placed on an Al\ :sub:`2`\ O\ :sub:`3` support. A -:class:`CustomCell` confines insertion/deletion to a slab-shaped region -above the support, and the lower half of the support is held fixed with -ASE constraints. - -Mirrors ``examples/gcmc_nano_supported.py``. - -Goal ----- - -At fixed :math:`T = 500\,\mathrm{K}` and chosen -:math:`\Delta\mu_{\mathrm{O}}`, sample equilibrium O coverage on the -supported cluster. - -Prerequisites -------------- - -- A POSCAR of the support (``Al2O3.poscar`` in the script). -- A MACE checkpoint or other ASE calculator that can describe Ag, Al, - and O. - -Code ----- - -.. code-block:: python - - import numpy as np - from ase.cluster import Octahedron - from ase.constraints import FixAtoms - from ase.io import read - from mace.calculators import mace_mp - - from mcpy.cell import CustomCell - from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble - from mcpy.moves import DeletionMove, InsertionMove - from mcpy.moves.move_selector import MoveSelector - from mcpy.utils.utils import get_p_at_support - - T = 500.0 - delta_mu_O = -0.5 - mus = {'Ag': -2.99, 'O': -4.91 + delta_mu_O} - - ss = np.random.SeedSequence(0) - seed_del, seed_ins = (int(s) for s in ss.generate_state(2, dtype=np.uint32)) - - support = read('Al2O3.poscar').repeat((4, 4, 1)) - support.center(vacuum=10.0, axis=2) - z = support.positions[:, 2] - z_half = z.min() + 0.5 * (z.max() - z.min()) - support.set_constraint(FixAtoms(mask=(z < z_half))) - - surface_z = float(np.max(support.positions[:, 2])) - particle = Octahedron('Ag', 5, 2) - atoms = get_p_at_support(support, particle, contact_surface='100', gap=2.0) - - scell = CustomCell( - atoms, - custom_height=20, - bottom_z=surface_z, - species_radii={'Ag': 2.068, 'O': 0, 'Al': 3}, - ) - - calculator = mace_mp(device='cuda') - - move_selector = MoveSelector( - [1, 1], - [DeletionMove(scell, species=['O'], seed=seed_del), - InsertionMove(scell, species=['O'], min_insert=0.5, seed=seed_ins)], - ) - - tag = f'{atoms.get_chemical_formula()}_dmu_{delta_mu_O}' - gcmc = GrandCanonicalEnsemble( - atoms=atoms, - cells=[scell], - calculator=calculator, - mu=mus, - units_type='metal', - species=['O'], - temperature=T, - move_selector=move_selector, - outfile=f'gcmc_{tag}.out', - traj_file=f'gcmc_{tag}.xyz', - ) - gcmc.run(1_000_000) - -Outputs -------- - -- ``gcmc__dmu_.out`` — step log with per-move acceptance. -- ``gcmc__dmu_.xyz`` — extended XYZ trajectory. - -Interpretation --------------- - -- ``bottom_z=surface_z`` anchors the MC region just above the topmost - support atom; ``custom_height=20`` opens a 20 Å vertical region for - insertions. Both values must match your support's vacuum padding. -- The Al support enters ``species_radii`` only so the free-volume - estimator excludes overlaps with it. Al atoms are *not* in - ``species`` and so are neither inserted nor deleted. -- ``get_p_at_support`` from :mod:`mcpy.utils.utils` positions the - cluster's chosen facet against the support with a fixed gap. - -Next steps ----------- - -- Strip the support to recover the in-vacuum cluster: see - :doc:`gcmc_nano`. -- Sweep :math:`\Delta\mu_{\mathrm{O}}` and post-process with - :doc:`phase_diagram_analysis`. diff --git a/docs/examples/mc_simulations.rst b/docs/examples/mc_simulations.rst index ab61997..61b93e1 100644 --- a/docs/examples/mc_simulations.rst +++ b/docs/examples/mc_simulations.rst @@ -74,6 +74,6 @@ Next steps - For variable composition, switch to :doc:`gcmc_simulations`. - For temperature-ladder parallel tempering of the chemical ordering (high-T replicas escape poor orderings and feed good ones down to low-T replicas), - drive several ``CanonicalEnsemble`` replicas with - :doc:`replica_exchange_simulations`. + drive several ``CanonicalEnsemble`` replicas with ``ReplicaExchange`` (see + :doc:`../ensembles`). - See :doc:`../moves` for the full move set. diff --git a/docs/examples/replica_exchange_simulations.rst b/docs/examples/replica_exchange_simulations.rst deleted file mode 100644 index a2bee64..0000000 --- a/docs/examples/replica_exchange_simulations.rst +++ /dev/null @@ -1,130 +0,0 @@ -Replica-Exchange GCMC across a temperature ladder -================================================= - -Parallel tempering over six replicas of an Ag(111) + O system, exchanging -configurations between neighbouring temperatures every ``exchange_interval`` -GCMC steps. One MPI rank corresponds to one replica. - -Mirrors ``examples/re_gcmc.py``. - -Goal ----- - -Run RE-GCMC at :math:`T \in \{250, 300, 350, 400, 450, 500\}\,\mathrm{K}` -with shared :math:`\mu_{\mathrm{Ag}}, \mu_{\mathrm{O}}` to escape local -basins that a single GCMC trajectory would not leave. - -Prerequisites -------------- - -- ``mpi4py`` and a working system MPI (see :doc:`../installation`). -- A MACE checkpoint, or any ASE calculator that you wrap in - ``BaseCalculator``. - -Code ----- - -.. code-block:: python - - import numpy as np - from ase.build import fcc111 - from ase.constraints import FixAtoms - from mace.calculators import mace_mp - - from mcpy.calculators import BaseCalculator - from mcpy.cell import CustomCell - from mcpy.ensembles import ReplicaExchange - from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble - from mcpy.moves import InsertionMove, DeletionMove - from mcpy.moves.move_selector import MoveSelector - - temperatures = [250, 300, 350, 400, 450, 500] - gcmc_steps = 200 - exchange_interval = 10 - delta_mu_O = -0.5 - - ss = np.random.SeedSequence(0) - seeds = [int(s) for s in ss.generate_state(5, dtype=np.uint32)] - - atoms = fcc111('Ag', a=4.1592, size=(4, 4, 3), periodic=True, vacuum=8.0) - atoms.set_constraint(FixAtoms(indices=[a.index for a in atoms if a.tag == 3])) - - cell_ag = CustomCell(atoms, custom_height=5.5, bottom_z=12.8 - 2.11, - species_radii={'Ag': 2.75, 'O': 0.0}) - cell_o = CustomCell(atoms, custom_height=5.5, bottom_z=12.8 - 2.11, - species_radii={'Ag': 2.11, 'O': 0.0}) - - # BaseCalculator relaxes with LBFGS before each energy evaluation; - # a bare mace_mp would return unrelaxed energies. - calculator = BaseCalculator(mace_mp(device='cuda'), steps=40, fmax=0.1) - - mus = {'Ag': -2.99, 'O': -4.91 + delta_mu_O} - - move_selector = MoveSelector( - [25, 25, 25, 25], - [DeletionMove(cell_ag, species=['Ag'], seed=seeds[0]), - DeletionMove(cell_o, species=['O'], seed=seeds[1]), - InsertionMove(cell_ag, species=['Ag'], min_insert=0.5, seed=seeds[2]), - InsertionMove(cell_o, species=['O'], min_insert=0.5, seed=seeds[3])], - ) - - def gcmc_factory(T, rank=0): - tag = f'rank{rank}_T{int(T)}' - return GrandCanonicalEnsemble( - atoms=atoms, - cells=[cell_ag, cell_o], - calculator=calculator, - mu=mus, - units_type='metal', - species=['Ag', 'O'], - temperature=T, - move_selector=move_selector, - outfile=f'gcmc_{tag}.out', - traj_file=f'gcmc_{tag}.xyz', - trajectory_write_interval=1, - outfile_write_interval=1, - ) - - rex = ReplicaExchange( - gcmc_factory, - temperatures=temperatures, - gcmc_steps=gcmc_steps, - exchange_interval=exchange_interval, - seed=seeds[4], - ) - rex.run() - -Run it ------- - -.. code-block:: bash - - mpirun -n 6 python re_gcmc.py - -The number of MPI ranks must equal ``len(temperatures)``. - -Outputs -------- - -- One ``gcmc_rank*_T*.out`` log and matching ``.xyz`` trajectory per rank. -- ``replica_exchange.log`` — the cross-replica exchange log. - -Interpretation --------------- - -- ``gcmc_factory`` must produce *per-rank* output filenames, otherwise all - ranks race on the same files. -- A healthy exchange acceptance is typically 20–40 %. If it is too low, - reduce the temperature gap between neighbouring replicas; if too high, - widen it. -- For a chemical-potential ladder, pass ``mus=[mu_0, mu_1, ...]`` instead - of ``temperatures``; the corresponding acceptance criterion is selected - automatically. - -Next steps ----------- - -- See ``ReplicaExchange`` in :doc:`../ensembles` for the parallel-tempering - acceptance rule. -- Combine the per-rank trajectories into a phase diagram with - :doc:`phase_diagram_analysis`. diff --git a/docs/installation.rst b/docs/installation.rst index f73c608..1dcd539 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -90,8 +90,9 @@ Install it with conda before running RE-GCMC: conda install mpi4py -Then launch a replica-exchange simulation with one MPI rank per replica: +Then launch a replica-exchange simulation with one MPI rank per replica. On a +single GPU prefer ``BatchedReplicaExchange``, which needs no MPI: .. code-block:: bash - mpirun -n python examples/re_gcmc.py + python examples/re_gcmc_batched.py diff --git a/docs/molecular_adsorbates.rst b/docs/molecular_adsorbates.rst index e15e287..d55f9e6 100644 --- a/docs/molecular_adsorbates.rst +++ b/docs/molecular_adsorbates.rst @@ -174,9 +174,5 @@ Worked examples Cu\ :sub:`33`\ Pd\ :sub:`5` particle against the CO chemical potential (see :doc:`phase_diagrams`). -- ``examples/gcmc_molecule_mace.py`` adsorbs any ASE-buildable molecule on - Ag(111) with MACE. -- ``examples/re_gcmc_co_cupd_batched.py`` runs CO on a CuPd nanoparticle - across a mu ladder with batched replica exchange. - The molecular GCMC notebook walks the machinery step by step with executed output (see :doc:`notebooks`). diff --git a/docs/tutorials/cluster_install_iqtc.rst b/docs/tutorials/cluster_install_iqtc.rst index 69f992c..d36f21b 100644 --- a/docs/tutorials/cluster_install_iqtc.rst +++ b/docs/tutorials/cluster_install_iqtc.rst @@ -187,7 +187,7 @@ Save as ``submit_gcmc_alchemi.sh``: set -eo pipefail PYTHON=$HOME/.conda/envs/alchemi128/bin/python - SCRIPT=$HOME/bin/mcpy/examples/gcmc_nano_alchemi.py + SCRIPT=$HOME/bin/mcpy/examples/gcmc_ag111.py RESULTS_DIR=$SLURM_SUBMIT_DIR/results/job_${SLURM_JOB_ID:-local} export PYTHONNOUSERSITE=1 @@ -237,7 +237,7 @@ The script above runs a single replica. ``mcpy`` also runs **replica exchange on one GPU**: all replicas live in a single Python process, share one ``AlchemiCalculator``, and have their trial-move energies evaluated in a single batched forward pass per MC step. This is the GPU path the package is built -around. Unlike the CPU ``re_gcmc.py`` (one MPI rank per replica), there is **no +around. Unlike the MPI ``ReplicaExchange`` (one rank per replica), there is **no mpirun** and still only ``--gres=gpu:1``. The script is ``examples/re_gcmc_batched.py``. Replica temperatures are passed diff --git a/examples/gcmc_ag111.py b/examples/gcmc_ag111.py index 4aa2dd0..aef9f7c 100644 --- a/examples/gcmc_ag111.py +++ b/examples/gcmc_ag111.py @@ -46,7 +46,7 @@ from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 # Calibrated set for O on Ag(111) with the small-density-agnesi MACE model, -# taken from examples/gcmc_custom_cell.py. Fixed on purpose: this is fixture +# Fixed on purpose: this is fixture # geometry, not a per-checkpoint number. Only mu is derived at run time. AG_LATTICE = 4.165 # Ag fcc lattice constant (Angstrom) # Relaxed O-Ag pair distance: the free-volume exclusion radius, and also how @@ -54,7 +54,7 @@ # oxidation path). Anchoring at the top layer makes the region pure vacuum. R_O_AG = 2.068 SPECIES_RADII = {'Ag': R_O_AG, 'O': 0.0} -# Bulk Ag reference correction (examples/gcmc_custom_cell.py). +# Bulk Ag reference correction. MU_AG_CORRECTION = -0.176 diff --git a/tests/test_ag111_gcmc_golden.py b/tests/test_ag111_gcmc_golden.py index 5a6e482..a2999a4 100644 --- a/tests/test_ag111_gcmc_golden.py +++ b/tests/test_ag111_gcmc_golden.py @@ -15,8 +15,8 @@ overlap rejection, de Broglie acceptance, arrays+constraints rollback, trajectory writing) reproduces exactly when seeded. -Parameters come from ``examples/gcmc_custom_cell.py``, the calibrated set for -this checkpoint: a=4.165, the O-Ag exclusion radius 2.068 used both as the +Parameters come from ``examples/gcmc_ag111.py``, the calibrated set for this +checkpoint: a=4.165, the O-Ag exclusion radius 2.068 used both as the radius and as the region depth below the top layer, height 7, min_insert 0.5, and the -0.176 bulk Ag correction. """ From 10abab38774e2804e94ae38b08f7a736558adf08 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:20:53 +0200 Subject: [PATCH 07/20] Delete dead code found by the audit Findings 1, 7, 11 and 15, all zero-behaviour-change: - mcpy/moves/go_moves.py, whole file. Three of its classes were stale copy-based twins of the exported moves; the other four had no caller, no export and no test. - mcpy/calculators/mace_calculator.py. No code in the repo used it: MACE_F_Calculator builds the upstream mace.calculators.MACECalculator, not this one, which hand-rolled the AtomicData and batch dicts instead. - Six unused helpers in mcpy/utils/utils.py (find_surface_indices, sphere_volume, total_volume, total_volume_with_overlap, overlap_volume, get_volume) and the module logger they were the only users of. The analytic overlap volume was superseded by the cells' Monte Carlo free-volume sampling. get_p_at_support stays: it is a real helper, currently used only by parked examples. - Duplicate get_species() in SphericalCell and CustomCell, plus the species_radii reassignment those two and DomeCell repeated after Cell.__init__ had already stored it. The audit claimed DomeCell also duplicated get_species; it did not. Two tests went with them: the overlap_volume pair, which existed only to keep that function alive, and one assertion in test_logging that required mcpy.utils.utils to expose a module logger, which it no longer needs because nothing in it logs. Sphinx build is warning-free and the torch-free suite passes (200 tests). --- docs/installation.rst | 2 +- docs/moves.rst | 11 -- docs/reference/calculators.rst | 20 +-- docs/reference/moves.rst | 9 -- mcpy/calculators/__init__.py | 3 +- mcpy/calculators/alchemi_calculator.py | 2 +- mcpy/calculators/mace_calculator.py | 58 -------- mcpy/cell/custom_cell.py | 7 - mcpy/cell/dome_cell.py | 1 - mcpy/cell/spherical_cell.py | 7 - mcpy/moves/go_moves.py | 187 ------------------------- mcpy/utils/utils.py | 74 ---------- pyproject.toml | 2 +- tests/test_audit_regressions.py | 20 --- tests/test_logging.py | 2 - 15 files changed, 7 insertions(+), 398 deletions(-) delete mode 100644 mcpy/calculators/mace_calculator.py delete mode 100644 mcpy/moves/go_moves.py diff --git a/docs/installation.rst b/docs/installation.rst index 1dcd539..bf6078e 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -34,7 +34,7 @@ MLIP backends ------------- Install only the backends you intend to use. For MACE -(``MACECalculator`` / ``MACE_F_Calculator`` and the bundled examples): +(``MACE_F_Calculator`` and the bundled examples): .. code-block:: bash diff --git a/docs/moves.rst b/docs/moves.rst index 53e1288..7201e61 100644 --- a/docs/moves.rst +++ b/docs/moves.rst @@ -103,14 +103,3 @@ as a relaxation-like perturbation that explores the local basin more thoroughly single random displacement. -Experimental moves ------------------- - -`mcpy.moves.go_moves` contains additional specialised moves used in exploratory workflows: - -- `BallMove`, `ShellMove`, -- `BondMove`, -- `HighEnergyAtomsMove`, -- alternative `PermutationMove`, `ShakeMove`, and `BrownianMove` variants. - -These are not part of the default `mcpy.moves` export list and their interfaces may change. diff --git a/docs/reference/calculators.rst b/docs/reference/calculators.rst index 03c3740..bda02cc 100644 --- a/docs/reference/calculators.rst +++ b/docs/reference/calculators.rst @@ -8,7 +8,7 @@ background, including the relaxation-inside-energy design, is in Each wrapper exposes ``get_potential_energy(atoms) -> float``. The Alchemi classes add ``get_potential_energies(atoms_list, chunk_size=None) -> ndarray`` for batched evaluation, where ``chunk_size`` caps peak GPU memory at one chunk. -``MACECalculator``, ``MACE_F_Calculator``, and ``BaseCalculator`` import +``MACE_F_Calculator`` and ``BaseCalculator`` import unconditionally; the Alchemi classes import only when ``nvalchemi-toolkit`` is installed. @@ -29,20 +29,6 @@ relaxed energy. - ``fmax`` (float): force tolerance in eV/Å. -MACECalculator --------------- - -.. code-block:: python - - MACECalculator(model_paths, device='cpu') - -Single-point MACE evaluation with no relaxation. -``get_potential_energy(atoms)`` returns one forward pass. - -- ``model_paths`` (str): path to a MACE model file. -- ``device`` (str): ``'cpu'`` or ``'cuda'``. - - MACE_F_Calculator ----------------- @@ -54,8 +40,8 @@ MACE_F_Calculator Relax-then-energy MACE wrapper used in most GCMC runs. Records ``last_relax_steps`` and ``total_relax_steps`` after each call. -- ``model_paths`` (str | MACECalculator): model file path, or a built - ``MACECalculator`` to reuse. +- ``model_paths`` (str | mace.calculators.MACECalculator): model file path, or + a built upstream MACE calculator to reuse. - ``steps`` (int): maximum relaxation steps. - ``fmax`` (float): force tolerance in eV/Å. - ``device`` (str): ``'cpu'`` or ``'cuda'``. diff --git a/docs/reference/moves.rst b/docs/reference/moves.rst index a7e3f4f..5eef688 100644 --- a/docs/reference/moves.rst +++ b/docs/reference/moves.rst @@ -143,12 +143,3 @@ AlchemiBrownianMove GPU-native Langevin Brownian move. Runs a short NVT Langevin trajectory through ``calculator.run_md`` (an ``AlchemiCalculator`` or ``AlchemiFCalculator``), reusing its model. Honors ``FixAtoms``. - - -Experimental moves ------------------- - -``mcpy.moves.go_moves`` contains specialised exploratory moves (``BallMove``, -``ShellMove``, ``BondMove``, ``HighEnergyAtomsMove``, and variant -permutation/shake/Brownian moves). They are not part of the ``mcpy.moves`` -export list and their interfaces may change. diff --git a/mcpy/calculators/__init__.py b/mcpy/calculators/__init__.py index ec3830f..5a82df6 100644 --- a/mcpy/calculators/__init__.py +++ b/mcpy/calculators/__init__.py @@ -3,9 +3,8 @@ __all__ = ['BaseCalculator'] try: - from .mace_calculator import MACECalculator from .mace_f_calculator import MACE_F_Calculator - __all__ += ['MACECalculator', 'MACE_F_Calculator'] + __all__ += ['MACE_F_Calculator'] except ImportError: # mace-torch is an optional backend; install with: pip install -e .[mace] pass diff --git a/mcpy/calculators/alchemi_calculator.py b/mcpy/calculators/alchemi_calculator.py index 3150f4c..41508cf 100644 --- a/mcpy/calculators/alchemi_calculator.py +++ b/mcpy/calculators/alchemi_calculator.py @@ -24,7 +24,7 @@ class AlchemiCalculator: """ Energy-only Alchemi calculator (no geometry relaxation). - Equivalent to MACECalculator but uses the nvalchemi GPU-native stack. + Single-point MACE evaluation on the nvalchemi GPU-native stack. Use this when mcpy handles its own MC moves and only needs E(atoms). Parameters diff --git a/mcpy/calculators/mace_calculator.py b/mcpy/calculators/mace_calculator.py deleted file mode 100644 index 6e7ec81..0000000 --- a/mcpy/calculators/mace_calculator.py +++ /dev/null @@ -1,58 +0,0 @@ -import torch -from mace.tools import torch_tools, utils -from mace import data -from ase import Atoms - - -class MACECalculator: - def __init__(self, model_paths: str, device: str = 'cpu') -> None: - """ - Initialize the Calculator with the given model path and device. - - :param model_paths: Path to the model file. - :param device: Device to load the model on ('cpu' or 'cuda'). - """ - torch_tools.set_default_dtype('float64') - self.device = device - # weights_only=False: MACE checkpoints are pickled full models, which - # torch >= 2.6 refuses to load under its weights_only default. - self.model = torch.load(model_paths, map_location=device, - weights_only=False).to(device) - for param in self.model.parameters(): - param.requires_grad = False - self.z_table = utils.AtomicNumberTable([int(z) for z in self.model.atomic_numbers]) - - def get_atom_data(self, atoms: Atoms) -> dict: - """ - Convert ASE Atoms object to the required atomic data format. - - :param atoms: ASE Atoms object. - :return: Dictionary containing atomic data. - """ - atom_data = data.AtomicData.from_config( - data.config_from_atoms(atoms), z_table=self.z_table, - cutoff=float(self.model.r_max), heads=None - ).to_dict() - atom_data.update({ - 'batch': torch.zeros(len(atoms), dtype=torch.long), - 'head': torch.tensor([0], dtype=torch.long), - 'ptr': torch.tensor([0, len(atoms)]) - }) - # AtomicData builds on CPU; the model may live on CUDA. - return { - k: v.to(self.device) if isinstance(v, torch.Tensor) else v - for k, v in atom_data.items() - } - - def get_potential_energy(self, atoms: Atoms) -> float: - """ - Calculate the potential energy of the given atoms. - - :param atoms: ASE Atoms object. - :return: Potential energy as a float. - """ - atom_data = self.get_atom_data(atoms) - with torch.no_grad(): - energy = self.model(atom_data, compute_stress=False, compute_force=False)['energy'] - atoms.info['energy'] = energy.item() - return energy.item() diff --git a/mcpy/cell/custom_cell.py b/mcpy/cell/custom_cell.py index 1b0bf96..b36c704 100644 --- a/mcpy/cell/custom_cell.py +++ b/mcpy/cell/custom_cell.py @@ -23,7 +23,6 @@ def __init__(self, atoms, custom_height=None, bottom_z=None, species_radii=None, """ super().__init__(atoms, species_radii=species_radii, seed=seed) self.dimensions, self.offset = self.get_custom_height_cell(custom_height, bottom_z) - self.species_radii = species_radii if species_radii else {} self.cell_volume = float(abs(np.linalg.det(self.dimensions))) self.mc_sample_points = int(mc_sample_points) self._dim_inv = np.linalg.inv(self.dimensions) @@ -154,9 +153,3 @@ def is_point_exchangeable(self, point): frac = (point - self.offset) @ self._dim_inv return bool(np.all(frac[:2] >= 0.0) and np.all(frac[:2] < 1.0) and frac[2] >= 0.0) - - def get_species(self): - """ - Get the species present in the custom cell. - """ - return list(self.species_radii.keys()) diff --git a/mcpy/cell/dome_cell.py b/mcpy/cell/dome_cell.py index 3795089..d51d735 100644 --- a/mcpy/cell/dome_cell.py +++ b/mcpy/cell/dome_cell.py @@ -57,7 +57,6 @@ def __init__(self, atoms, particle_species, bottom_z, vacuum, self.radius = float( np.linalg.norm(particle_pos - self.center, axis=1).max() + vacuum ) - self.species_radii = species_radii self.mc_sample_points = int(mc_sample_points) self.sphere_volume = (4.0 / 3.0) * np.pi * (self.radius ** 3) self.volume = self.sphere_volume diff --git a/mcpy/cell/spherical_cell.py b/mcpy/cell/spherical_cell.py index a7af832..94231b8 100644 --- a/mcpy/cell/spherical_cell.py +++ b/mcpy/cell/spherical_cell.py @@ -32,7 +32,6 @@ def __init__(self, atoms, vacuum, species_radii, mc_sample_points=100_000, seed= self.radius = float( np.linalg.norm(atoms.positions - self.center, axis=1).max() + vacuum ) - self.species_radii = species_radii self.mc_sample_points = int(mc_sample_points) self.sphere_volume = (4.0 / 3.0) * np.pi * (self.radius ** 3) self.volume = self.sphere_volume @@ -79,12 +78,6 @@ def get_atoms_specie_inside_cell(self, atoms, specie): inside = np.einsum('ij,ij->i', diff, diff) <= self.radius ** 2 return np.where(species_mask & inside)[0] - def get_species(self): - """ - Get the species present in the custom cell. - """ - return list(self.species_radii.keys()) - def _sample_sphere_points(self, n): """ Uniform random points inside the sphere. Rejection sample from the diff --git a/mcpy/moves/go_moves.py b/mcpy/moves/go_moves.py deleted file mode 100644 index c011521..0000000 --- a/mcpy/moves/go_moves.py +++ /dev/null @@ -1,187 +0,0 @@ -import numpy as np -from ase import Atoms -from ase.md.velocitydistribution import MaxwellBoltzmannDistribution -from ase.md.verlet import VelocityVerlet -from ase.neighborlist import neighbor_list -from ase.units import fs -from .base_move import BaseMove -from ..cell import NullCell - - -def get_surface_atoms(atoms, cutoff=3.2, threshold=8): - indices = neighbor_list('i', atoms, cutoff) - counts = np.bincount(indices, minlength=len(atoms)) - return np.where(counts < threshold)[0] - - -class PermutationMove(BaseMove): - """ - Swap the chemical symbols of two atoms of different species. - Useful for optimizing chemical order in nanoalloys. - """ - def __init__(self, species: list[str], seed: int) -> None: - cell = NullCell() - super().__init__(cell, species, seed) - - def do_trial_move(self, atoms) -> Atoms: - atoms_new = atoms.copy() - species_pair = self.rng.random.sample(self.species, 2) - indices_symbol_a = [atom.index for atom in atoms_new if atom.symbol == species_pair[0]] - indices_symbol_b = [atom.index for atom in atoms_new if atom.symbol == species_pair[1]] - if len(indices_symbol_a) == 0 or len(indices_symbol_b) == 0: - return False, 0, 'X' - i = self.rng.random.choice(indices_symbol_a) - j = self.rng.random.choice(indices_symbol_b) - atoms_new[i].symbol, atoms_new[j].symbol = atoms_new[j].symbol, atoms_new[i].symbol - return atoms_new, 0, 'X' - - -class ShakeMove(BaseMove): - """ - Randomly displace all atoms within a sphere of radius r_max. - Very efficient for shape optimization. - """ - def __init__(self, r_max: float, seed: int) -> None: - cell = NullCell() - super().__init__(cell, species=['X'], seed=seed) - self.r_max = r_max - self.name = 'Shake' - self.rng = np.random.default_rng(seed) - - def do_trial_move(self, atoms: Atoms) -> Atoms: - atoms_new = atoms.copy() - n = len(atoms) - # Uniform directions on S^2 (Gaussian normalized; a cube-uniform - # vector would bias toward diagonals). - displacements = self.rng.standard_normal(size=(n, 3)) - directions = displacements / np.linalg.norm(displacements, axis=1, keepdims=True) - # Uniform-in-ball radius: r_max * u^(1/3), u ~ U(0,1). - radii = self.r_max * np.cbrt(self.rng.random(size=(n, 1))) - new_positions = atoms_new.get_positions() + directions * radii - atoms_new.set_positions(new_positions) - return atoms_new, 0, 'X' - - -class BrownianMove(BaseMove): - """ - Apply short high-temperature molecular dynamics to simulate Brownian motion. - Useful for large cluster relaxation. - """ - def __init__(self, temperature: float, calculator, steps: int, d_t: float, seed: int) -> None: - cell = NullCell() - super().__init__(cell, species=['X'], seed=seed) - self.temperature = temperature - self.calculator = calculator - self.steps = steps - self.d_t = d_t - self.name = 'Brownian' - - def do_trial_move(self, atoms: Atoms) -> Atoms: - atoms_new = atoms.copy() - MaxwellBoltzmannDistribution(atoms_new, temperature_K=self.temperature) - atoms_new.calc = self.calculator - dyn = VelocityVerlet(atoms_new, self.d_t * fs, logfile=None) - dyn.run(steps=self.steps) - return atoms_new, 0, 'X' - - -class BallMove(BaseMove): - """ - Randomly displace a single atom within the cluster volume. - """ - def __init__(self, radius: float, seed: int) -> None: - cell = NullCell() - super().__init__(cell, species=['X'], seed=seed) - self.radius = radius - self.name = 'Ball' - self.rng = np.random.default_rng(seed) - - def do_trial_move(self, atoms: Atoms) -> Atoms: - atoms_new = atoms.copy() - i = self.rng.integers(0, len(atoms)) - # Uniform inside a ball of radius `self.radius`: uniform direction - # times r = radius * u^(1/3). - direction = self.rng.normal(size=3) - direction /= np.linalg.norm(direction) - r = self.radius * np.cbrt(self.rng.random()) - atoms_new[i].position += direction * r - return atoms_new, 0, 'X' - - -class ShellMove(BaseMove): - """ - Move a surface atom within a spherical shell, refining surface arrangement. - """ - def __init__(self, r_shell: float, seed: int) -> None: - cell = NullCell() - super().__init__(cell, species=['X'], seed=seed) - self.r_shell = r_shell - self.name = 'Shell' - self.rng = np.random.default_rng(seed) - - def do_trial_move(self, atoms: Atoms) -> Atoms: - atoms_new = atoms.copy() - surface_indices = get_surface_atoms(atoms_new) - if len(surface_indices) == 0: - return False, 0, 'X' - i = self.rng.choice(surface_indices) - direction = self.rng.normal(size=3) - direction /= np.linalg.norm(direction) - atoms_new[i].position += direction * self.r_shell - return atoms_new, 0, 'X' - - -class BondMove(BaseMove): - """ - Move a low-coordination (weakly bonded) surface atom. - """ - def __init__(self, r_max: float, seed: int) -> None: - cell = NullCell() - super().__init__(cell, species=['X'], seed=seed) - self.r_max = r_max - self.name = 'Bond' - self.rng = np.random.default_rng(seed) - - def do_trial_move(self, atoms: Atoms) -> Atoms: - atoms_new = atoms.copy() - surface_indices = get_surface_atoms(atoms_new) - if len(surface_indices) == 0: - return False, 0, 'X' - i = self.rng.choice(surface_indices) - displacement = self.rng.normal(size=3) - displacement *= self.r_max / np.linalg.norm(displacement) - atoms_new[i].position += displacement - return atoms_new, 0, 'X' - - -class HighEnergyAtomsMove(BaseMove): - """ - Displace atoms depending on local energy: - - atoms below threshold move slightly (shake) - - atoms above threshold move more (bond-like) - """ - def __init__(self, calculator, energy_threshold: float, r_max: float, seed: int) -> None: - cell = NullCell() - super().__init__(cell, species=['X'], seed=seed) - self.energy_threshold = energy_threshold - self.calculator = calculator - self.r_max = r_max - self.rng = np.random.default_rng(seed) - self.name = 'HighEnergy' - - def do_trial_move(self, atoms: Atoms) -> Atoms: - atoms_new = atoms.copy() - try: - local_energies = self.calculator.calculator.results['node_energy'] - except Exception: - return False, 0, 'X' - - if local_energies is None: - return False, 0, 'X' - - for i, e in enumerate(local_energies): - if e < self.energy_threshold: - atoms_new[i].position += self.rng.normal(size=3) * self.r_max * 0.2 - else: - atoms_new[i].position += self.rng.normal(size=3) * self.r_max - return atoms_new, 0, 'X' diff --git a/mcpy/utils/utils.py b/mcpy/utils/utils.py index 0cc8051..c47c8cf 100644 --- a/mcpy/utils/utils.py +++ b/mcpy/utils/utils.py @@ -1,81 +1,7 @@ -import logging -import math - from ase import Atoms -from ase.neighborlist import NeighborList, natural_cutoffs -from ase.cell import Cell import numpy as np -logger = logging.getLogger(__name__) - - -def get_volume(box: list) -> float: - """ - Get the volume of a box. - """ - box = Cell(box) - return box.volume - - -def find_surface_indices(atoms: Atoms, tolerance: float = 0.1) -> list: - """ - Find the surface indices of an atoms object. - - Parameters: - atoms (Atoms): The atoms object. - tolerance (float): The tolerance for identifying surface atoms. - - Returns: - list: A list of surface atom indices. - """ - # Create a neighbor list - cutoffs = natural_cutoffs(atoms) - neighbor_list = NeighborList(cutoffs, self_interaction=False, bothways=True) - neighbor_list.update(atoms) - - surface_indices = [] - for i in range(len(atoms)): - indices, offsets = neighbor_list.get_neighbors(i) - if len(indices) < 12: - surface_indices.append(i) - - return surface_indices - - -def sphere_volume(radius): - return (4/3) * math.pi * radius**3 - - -def total_volume(spheres): - total_vol = sum(sphere_volume(radius) for radius in spheres) - return total_vol - - -def overlap_volume(r1, r2, d): - """Calculate the volume of overlap between two spheres of radius r1, r2 and distance d.""" - if d >= r1 + r2: - return 0 # No overlap - if d <= abs(r1 - r2): - return (4/3) * math.pi * min(r1, r2)**3 - - return (math.pi * (r1 + r2 - d)**2 - * (d**2 + 2*d*(r1 + r2) - 3*(r1 - r2)**2)) / (12 * d) - - -def total_volume_with_overlap(spheres, positions): - total_vol = 0 - num_spheres = len(spheres) - for i in range(num_spheres): - total_vol += sphere_volume(spheres[i]) - logger.debug('Total volume of spheres: %s', total_vol) - - for i in range(num_spheres): - for j in range(i + 1, num_spheres): - d = math.dist(positions[i], positions[j]) # Distance between sphere centers - total_vol -= overlap_volume(spheres[i], spheres[j], d) - return total_vol - def get_p_at_support(support: Atoms, particle: Atoms, contact_surface: str = '100', gap: float = 2.0, diff --git a/pyproject.toml b/pyproject.toml index 2d41b0f..87acba4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ ] [project.optional-dependencies] -# MACE-torch backend for MACECalculator / MACE_F_Calculator: +# MACE-torch backend for MACE_F_Calculator: # pip install -e .[mace] mace = [ "mace-torch>=0.3.9", diff --git a/tests/test_audit_regressions.py b/tests/test_audit_regressions.py index abf670d..5f5a666 100644 --- a/tests/test_audit_regressions.py +++ b/tests/test_audit_regressions.py @@ -23,7 +23,6 @@ InsertionMove, MoveSelector, PermutationMove) from mcpy.moves.base_move import BaseMove from mcpy.utils.set_unit_constant import SetUnits -from mcpy.utils.utils import overlap_volume class StubCalc: @@ -79,25 +78,6 @@ def test_gcmc_runs_in_lj_units(): assert np.isfinite(g.E_old) -# -------------------------------------------------------------------------- -# overlap_volume analytic value (bug: returned exactly twice the lens volume) -# -------------------------------------------------------------------------- - -def test_overlap_volume_matches_analytic(): - # Equal spheres radius r at distance d: V = pi*(4r + d)*(2r - d)^2 / 12 - r, d = 1.0, 1.0 - exact = math.pi * (4 * r + d) * (2 * r - d) ** 2 / 12.0 - assert overlap_volume(r, r, d) == pytest.approx(exact) - - -def test_overlap_volume_limits(): - assert overlap_volume(1.0, 1.0, 2.5) == 0 # disjoint - engulfed = overlap_volume(2.0, 1.0, 0.5) # small sphere inside big one - assert engulfed == pytest.approx((4 / 3) * math.pi) - # Touching from inside: d -> r1 + r2 gives vanishing overlap. - assert overlap_volume(1.0, 1.0, 1.999999) == pytest.approx(0.0, abs=1e-9) - - # -------------------------------------------------------------------------- # NullCell honors the Cell interface (bug: TypeError in calculate_volume) # -------------------------------------------------------------------------- diff --git a/tests/test_logging.py b/tests/test_logging.py index d60fb1b..72550f6 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -78,14 +78,12 @@ def test_library_loggers_use_module_name(): """Each module exposes ``logger = logging.getLogger(__name__)``.""" from mcpy.ensembles import base_ensemble, canonical_ensemble from mcpy.ensembles import grand_canonical_ensemble - from mcpy.utils import utils as utils_mod assert base_ensemble.logger.name == "mcpy.ensembles.base_ensemble" assert canonical_ensemble.logger.name == "mcpy.ensembles.canonical_ensemble" assert grand_canonical_ensemble.logger.name == ( "mcpy.ensembles.grand_canonical_ensemble" ) - assert utils_mod.logger.name == "mcpy.utils.utils" # --------------------------------------------------------------------------- From 2df7712a6e11b9e17373b8866460d4ce7fb11a1e Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:22:40 +0200 Subject: [PATCH 08/20] Remove dead parameters and unreachable MoveSelector methods Finding 13. Signatures that promised something the code never did: - BaseEnsemble(units_type=, user_tag=). units_type appeared only in the signature and was never read; user_tag was stored and never read. GrandCanonicalEnsemble keeps its own units_type, which is real: it builds SetUnits from it. It just stops forwarding it to a base class that discards it. - CanonicalEnsemble(units_type='metal'). Accepted and ignored: the class computes beta from ase.units.kB directly, so passing units_type='LJ' silently got you metal units. Removing it turns a wrong answer into a TypeError. - RandomNumberGenerator(warm_up=0). The docstring already said the Mersenne Twister does not need warm-up and the default was 0; no caller set it. - MoveSelector.get_operator and .calculate_volumes: no caller anywhere. - MoveSelector.get_acceptance_ration, the deprecated typo'd alias, and the one test that existed to assert it still warned. get_acceptance_ratio stays. Reference docs follow: units_type moves from the shared BaseEnsemble parameter list to the GrandCanonicalEnsemble section where it belongs. Two of these are caller-visible changes (CanonicalEnsemble's units_type and the typo'd alias). No caller in the repo passed either. Sphinx build warning-free, 199 tests pass. --- docs/reference/ensembles.rst | 7 ++++--- docs/reference/utils.rst | 4 +--- mcpy/ensembles/base_ensemble.py | 3 --- mcpy/ensembles/canonical_ensemble.py | 2 -- mcpy/ensembles/grand_canonical_ensemble.py | 1 - mcpy/moves/move_selector.py | 19 +------------------ mcpy/utils/random_number_generator.py | 9 +-------- tests/test_audit_regressions.py | 7 +------ 8 files changed, 8 insertions(+), 44 deletions(-) diff --git a/docs/reference/ensembles.rst b/docs/reference/ensembles.rst index 81d56bc..5070da0 100644 --- a/docs/reference/ensembles.rst +++ b/docs/reference/ensembles.rst @@ -17,7 +17,7 @@ BaseEnsemble .. code-block:: python - BaseEnsemble(atoms, cells, units_type, calculator, user_tag=None, + BaseEnsemble(atoms, cells, calculator, random_seed=None, traj_file='trajectory.xyz', traj_mode='w', trajectory_write_interval=1, outfile='outfile.out', outfile_mode='w', outfile_write_interval=1, @@ -30,7 +30,6 @@ Shared parameters (inherited by the concrete ensembles): - ``atoms`` (ase.Atoms): starting configuration. - ``cells`` (list): cell objects providing free volume and insertion points. -- ``units_type`` (str): ``'metal'`` or ``'LJ'`` (see :doc:`../ensembles`). - ``calculator``: an mcpy calculator wrapper or any object exposing ``get_potential_energy(atoms)``. - ``random_seed`` (int, optional): seed source for derived generators. @@ -57,7 +56,7 @@ CanonicalEnsemble .. code-block:: python - CanonicalEnsemble(atoms, calculator, cells=None, units_type='metal', + CanonicalEnsemble(atoms, calculator, cells=None, random_seed=None, optimizer=None, fmax=0.1, temperature=300, move_selector=None, constraints=None, traj_file='trajectory.xyz', traj_mode='w', @@ -100,6 +99,8 @@ de Broglie criterion (see :doc:`../ensembles`). Parameters beyond the shared set: +- ``units_type`` (str): ``'metal'`` or ``'LJ'``; selects the constants the + de Broglie factors are built from (see :doc:`../ensembles`). - ``mu`` (dict): species to reservoir chemical potential in eV. - ``species`` (list[str]): symbols the move set may insert or delete. - ``temperature`` (float): temperature in K. diff --git a/docs/reference/utils.rst b/docs/reference/utils.rst index 091a148..7c3c6f6 100644 --- a/docs/reference/utils.rst +++ b/docs/reference/utils.rst @@ -38,14 +38,12 @@ RandomNumberGenerator .. code-block:: python - RandomNumberGenerator(seed=None, warm_up=0) + RandomNumberGenerator(seed=None) Seeded wrapper over Python's ``random.Random``. Each move and the acceptance test carries its own instance, so the global ``random`` state is never touched. - ``seed`` (int, optional): seed for reproducibility. -- ``warm_up`` (int): discarded draws after seeding. The default ``0`` is - sufficient for the Mersenne Twister. Methods: ``get_uniform(a=0.0, b=1.0)``, ``get_gaussian(mu=0.0, sigma=1.0)``. diff --git a/mcpy/ensembles/base_ensemble.py b/mcpy/ensembles/base_ensemble.py index 70cc7ad..37aaf42 100644 --- a/mcpy/ensembles/base_ensemble.py +++ b/mcpy/ensembles/base_ensemble.py @@ -18,9 +18,7 @@ class BaseEnsemble(ABC): def __init__(self, atoms: Atoms, cells: List[Cell], - units_type: str, calculator: Calculator, - user_tag: Optional[str] = None, random_seed: Optional[int] = None, traj_file: Optional[str] = 'trajectory.xyz', traj_mode: str = 'w', @@ -53,7 +51,6 @@ def __init__(self, self._atoms = atoms self._cells = cells self._calculator = calculator - self._user_tag = user_tag self._trajectory_write_interval = trajectory_write_interval self._outfile_write_interval = outfile_write_interval diff --git a/mcpy/ensembles/canonical_ensemble.py b/mcpy/ensembles/canonical_ensemble.py index 5cd3a2b..3dfa27a 100644 --- a/mcpy/ensembles/canonical_ensemble.py +++ b/mcpy/ensembles/canonical_ensemble.py @@ -25,7 +25,6 @@ def __init__(self, atoms, calculator, cells=None, - units_type: str = 'metal', random_seed=None, optimizer=None, fmax=0.1, @@ -42,7 +41,6 @@ def __init__(self, minima_mode: str = 'a') -> None: super().__init__(atoms=atoms, cells=cells if cells is not None else [], - units_type=units_type, calculator=calculator, random_seed=random_seed, traj_file=traj_file, diff --git a/mcpy/ensembles/grand_canonical_ensemble.py b/mcpy/ensembles/grand_canonical_ensemble.py index 024b5f9..ad23ebb 100644 --- a/mcpy/ensembles/grand_canonical_ensemble.py +++ b/mcpy/ensembles/grand_canonical_ensemble.py @@ -41,7 +41,6 @@ def __init__(self, super().__init__(atoms=atoms, cells=cells, - units_type=units_type, calculator=calculator, random_seed=random_seed, traj_file=traj_file, diff --git a/mcpy/moves/move_selector.py b/mcpy/moves/move_selector.py index b4db54e..9d06fda 100644 --- a/mcpy/moves/move_selector.py +++ b/mcpy/moves/move_selector.py @@ -1,5 +1,4 @@ import re -import warnings import numpy as np from mcpy.utils import RandomNumberGenerator @@ -113,14 +112,6 @@ def get_exchange_count(self): def calculate_volume(self, atoms): return self.move_list[self.to_use].calculate_volume(atoms) - def calculate_volumes(self, atoms): - for move in self.move_list: - move.calculate_volume(atoms) - - def get_operator(self): - to_use = self.__get_index__() - return self.move_list[to_use] - def acceptance_counter(self): self.move_acceptance[self.to_use] += 1 self.move_acceptance_total[self.to_use] += 1 @@ -146,18 +137,10 @@ def total_ratios(self): self.move_counter_total, self.move_failed_counter_total) - # Backwards-compatible aliases. + # Backwards-compatible alias. def get_acceptance_ratio(self): return self.interval_ratios() - def get_acceptance_ration(self): # legacy typo'd name - warnings.warn( - "MoveSelector.get_acceptance_ration is a deprecated typo'd alias; " - "use interval_ratios() (or get_acceptance_ratio()).", - DeprecationWarning, stacklevel=2, - ) - return self.interval_ratios() - def reset_counters(self): n = len(self.move_list) self.move_counter = [0] * n diff --git a/mcpy/utils/random_number_generator.py b/mcpy/utils/random_number_generator.py index 66097e9..157fdf1 100644 --- a/mcpy/utils/random_number_generator.py +++ b/mcpy/utils/random_number_generator.py @@ -4,22 +4,15 @@ class RandomNumberGenerator: """Class for generating random numbers with optional seeding.""" - def __init__(self, seed=None, warm_up=0) -> None: + def __init__(self, seed=None) -> None: """ Initialize the random number generator. Parameters: seed (int, optional): Seed for reproducibility. Default is None. - warm_up (int): Discarded-draw count after seeding. Python's Mersenne - Twister does not need warm-up; default is 0. Set to a positive - integer to reproduce legacy behavior. """ self.random = Random(seed) - for _ in range(warm_up): - self.get_uniform() - self.get_gaussian() - def get_uniform(self, a=0.0, b=1.0) -> float: """ Generate a random number uniformly distributed between a and b. diff --git a/tests/test_audit_regressions.py b/tests/test_audit_regressions.py index 5f5a666..d472f73 100644 --- a/tests/test_audit_regressions.py +++ b/tests/test_audit_regressions.py @@ -196,17 +196,12 @@ def test_moveselector_explicit_trials_per_step_overrides_weights(): assert ms.n_moves == 2 -def test_moveselector_legacy_typo_alias_warns(): - ms = MoveSelector([1], [DisplacementMove(species=['H'], seed=1)], seed=2) - with pytest.warns(DeprecationWarning): - ms.get_acceptance_ration() - - # -------------------------------------------------------------------------- # set_state refreshes cell free volumes (bug: acceptance after an accepted # replica swap used the previous configuration's volume) # -------------------------------------------------------------------------- + class _RecordingCell: def __init__(self): self.calls = 0 From cd5b456c34eb0cfa9d364753c69d18cfbcda9096 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:28:56 +0200 Subject: [PATCH 09/20] Generate the API reference with autodoc Finding 2. The five reference pages restated by hand what the docstrings already say, so they drifted: the BaseEnsemble signature there still listed units_type and user_tag until the previous commit removed them, and MACECalculator had its own section until it was deleted. Replaced the hand-typed signature and parameter blocks with autoclass/automodule, keeping the narrative intros that carry information no docstring holds (the two cell predicates, the move return contract, which backends import conditionally). 688 lines of reference prose become 208 lines of directives, while the rendered pages get larger, not smaller: full method documentation, inherited signatures and [source] links, all generated from the code so they cannot drift again. Enabling this needed three things in conf.py: autodoc_mock_imports for the optional backends (importing a real torch on the docs builder is unnecessary, and in a misconfigured env fatal, SIGBUS rather than ImportError), autoclass_content='both' because the parameters are documented in __init__, and napoleon for the Args: blocks. Also fixes the SetUnits docstring, whose "Parameters:" block was malformed RST (unindented items followed by an indented continuation) and which autodoc surfaced as a docutils error the moment it was rendered. It is now a napoleon Args: block, and documents temperature, which it had never mentioned. Sphinx build is warning-free and every previously documented class still appears on its page. --- docs/conf.py | 12 ++- docs/reference/calculators.rst | 117 +++------------------- docs/reference/cells.rst | 118 +++++----------------- docs/reference/ensembles.rst | 170 ++++---------------------------- docs/reference/moves.rst | 159 ++++++++--------------------- docs/reference/utils.rst | 122 ++++------------------- mcpy/utils/set_unit_constant.py | 17 ++-- 7 files changed, 135 insertions(+), 580 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 214586b..c14033f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -20,7 +20,17 @@ # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration -extensions = ['sphinx.ext.autodoc', 'sphinx.ext.todo', 'sphinx.ext.viewcode'] +extensions = ['sphinx.ext.autodoc', 'sphinx.ext.napoleon', + 'sphinx.ext.todo', 'sphinx.ext.viewcode'] + +# The optional backends are not installed on the docs builder, and importing a +# real torch here is both slow and, in a misconfigured env, fatal (SIGBUS). +# Mocking lets autodoc read the wrappers' signatures and docstrings anyway. +autodoc_mock_imports = ['torch', 'mace', 'nvalchemi', 'cuequivariance', 'mpi4py'] +autodoc_member_order = 'bysource' +# Parameters are documented in __init__, so render both docstrings. +autoclass_content = 'both' +autodoc_typehints = 'description' templates_path = ['_templates'] exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] diff --git a/docs/reference/calculators.rst b/docs/reference/calculators.rst index bda02cc..6e5d39e 100644 --- a/docs/reference/calculators.rst +++ b/docs/reference/calculators.rst @@ -8,114 +8,19 @@ background, including the relaxation-inside-energy design, is in Each wrapper exposes ``get_potential_energy(atoms) -> float``. The Alchemi classes add ``get_potential_energies(atoms_list, chunk_size=None) -> ndarray`` for batched evaluation, where ``chunk_size`` caps peak GPU memory at one chunk. -``MACE_F_Calculator`` and ``BaseCalculator`` import -unconditionally; the Alchemi classes import only when ``nvalchemi-toolkit`` is -installed. +``BaseCalculator`` imports unconditionally; ``MACE_F_Calculator`` needs +``mace-torch`` and the Alchemi classes need ``nvalchemi-toolkit``, and each is +exported only when its backend is installed. -BaseCalculator --------------- +.. autoclass:: mcpy.calculators.BaseCalculator + :members: -.. code-block:: python +.. autoclass:: mcpy.calculators.mace_f_calculator.MACE_F_Calculator + :members: - BaseCalculator(calculator, steps, fmax) +.. autoclass:: mcpy.calculators.alchemi_calculator.AlchemiCalculator + :members: -Adapter for any ASE calculator. ``get_potential_energy(atoms)`` relaxes -``atoms`` with LBFGS up to ``steps`` iterations or ``fmax``, then returns the -relaxed energy. - -- ``calculator``: a constructed ASE calculator. -- ``steps`` (int): maximum LBFGS steps. -- ``fmax`` (float): force tolerance in eV/Å. - - -MACE_F_Calculator ------------------ - -.. code-block:: python - - MACE_F_Calculator(model_paths, steps, fmax, device='cpu', cueq=False, - optimizer='lbfgs') - -Relax-then-energy MACE wrapper used in most GCMC runs. Records -``last_relax_steps`` and ``total_relax_steps`` after each call. - -- ``model_paths`` (str | mace.calculators.MACECalculator): model file path, or - a built upstream MACE calculator to reuse. -- ``steps`` (int): maximum relaxation steps. -- ``fmax`` (float): force tolerance in eV/Å. -- ``device`` (str): ``'cpu'`` or ``'cuda'``. -- ``cueq`` (bool): enable cuEquivariance kernels. -- ``optimizer`` (str): ``'lbfgs'`` or ``'fire'``. Raises ``ValueError`` for any - other value. - - -AlchemiCalculator ------------------ - -.. code-block:: python - - AlchemiCalculator(checkpoint='medium-mpa-0', device='cuda', - dtype=torch.float32, enable_cueq=True, compile_model=True, - max_neighbors=None, chunk_size=None, energy_only=False) - -GPU-native MACE evaluation with no relaxation (optional backend). Provides -``get_potential_energy(atoms)``, ``get_potential_energies(atoms_list, -chunk_size=None)``, and ``run_md(...)``. - -- ``checkpoint`` (str | MACEWrapper): named checkpoint, local ``.pt`` path, or a - shared ``MACEWrapper``. -- ``enable_cueq`` (bool): fused equivariance kernels. -- ``compile_model`` (bool): ``torch.compile`` the model. Keep ``True`` (the - default): after a one-time warmup and one recompile on the first atom-count - change, dynamic shapes handle GCMC's varying N at compiled speed. Set - ``False`` only for short smoke tests where the warmup dominates. -- ``max_neighbors`` (int, optional): neighbor-list cap. -- ``chunk_size`` (int, optional): default sub-batch size for - ``get_potential_energies``; caps peak GPU memory at one chunk. ``None`` - evaluates the whole batch in one pass. See :doc:`../calculators`. -- ``energy_only`` (bool): drop force computation (no autograd graph) for a - ~12% memory saving. Energy is unchanged; forces are unavailable. The switch - lives on the model's own config, so it cannot be scoped to one calculator: - combining it with a pre-loaded ``MACEWrapper`` (which other calculators may - share, including a relaxing ``AlchemiFCalculator``) raises ``ValueError``. - Pass a checkpoint path instead so this calculator loads its own model. - - -AlchemiFCalculator ------------------- - -.. code-block:: python - - AlchemiFCalculator(checkpoint='medium-mpa-0', steps=500, fmax=0.05, - device='cuda', dtype=torch.float32, enable_cueq=True, - compile_model=True, dt=1.0, optimizer='fire', - max_neighbors=None, chunk_size=None) - -GPU counterpart of ``MACE_F_Calculator``: FIRE relaxation then energy, honoring -``FixAtoms`` (optional backend). Provides ``get_potential_energy(atoms)``, -``get_potential_energies(atoms_list, chunk_size=None)``, and ``run_md(...)``. - -- ``steps`` (int): maximum FIRE steps. -- ``fmax`` (float): force tolerance in eV/Å. -- ``dt`` (float): FIRE initial timestep. -- ``optimizer`` (str): ``'fire'`` or ``'fire2'``. Raises ``ValueError`` - otherwise. -- ``compile_model`` (bool): keep ``True`` (see ``AlchemiCalculator`` above); - measured 1.3-1.4x on GCMC after the one-time warmup. -- ``chunk_size`` (int, optional): default sub-batch size for - ``get_potential_energies``; caps peak GPU memory at one chunk. Chunked - relaxation reaches the same minimum to within ``fmax`` but is not bit - identical. See :doc:`../calculators`. - - -run_md ------- - -.. code-block:: python - - run_md(atoms, temperature, friction=0.01, dt=2.0, steps=100, seed=42) - -Method on both Alchemi classes. Runs NVT Langevin MD in place on ``atoms``, -reusing the loaded model. ``temperature`` in K, ``friction`` in 1/fs, ``dt`` in -fs. Honors ``FixAtoms``. Backs ``AlchemiBrownianMove``. +.. autoclass:: mcpy.calculators.alchemi_f_calculator.AlchemiFCalculator + :members: diff --git a/docs/reference/cells.rst b/docs/reference/cells.rst index 4e966e1..02a4571 100644 --- a/docs/reference/cells.rst +++ b/docs/reference/cells.rst @@ -9,103 +9,37 @@ Every non-trivial cell exposes ``calculate_volume(atoms)`` (refresh the cached free volume), ``get_volume()``, ``get_random_point()``, and ``get_atoms_specie_inside_cell(atoms, species)``. +Every cell also exposes two point-membership predicates. +``is_point_inside(point)`` bounds the *proposal* region -- the one +``get_random_point`` samples. ``is_point_exchangeable(point)`` bounds the region +the reservoir may take molecules back from, and is what the molecule moves use +to pick candidates. They coincide for every cell except :class:`CustomCell`; +see :doc:`../gcmc_acceptance_convention`. -BaseCell --------- +Every ``species_radii`` mapping must cover every species the system can hold, +including species inserted during the run; a missing entry raises +``ValueError`` from ``calculate_volume``. -.. code-block:: python - BaseCell() +.. autoclass:: mcpy.cell.base_cell.BaseCell + :members: -Abstract interface. Declares ``calculate_volume``, ``get_random_point``, and -``get_volume``. Not instantiated directly. +.. autoclass:: mcpy.cell.Cell + :members: + :show-inheritance: +.. autoclass:: mcpy.cell.CustomCell + :members: + :show-inheritance: -Cell ----- +.. autoclass:: mcpy.cell.SphericalCell + :members: + :show-inheritance: -.. code-block:: python +.. autoclass:: mcpy.cell.DomeCell + :members: + :show-inheritance: - Cell(atoms, species_radii=None, seed=None) - -The full ASE simulation box as the active region. Random points are sampled -uniformly in fractional coordinates. Its volume is the fixed box volume, so -``calculate_volume`` does no free-volume sampling. - -- ``species_radii`` (dict, optional): per-element exclusion radii. Every - species the system can hold needs one, including species inserted during the - run; a missing entry raises ``ValueError`` from ``calculate_volume``. -- ``seed`` (int, optional): cell-local RNG seed. - -Every cell exposes two point-membership predicates. ``is_point_inside(point)`` -bounds the *proposal* region -- the one ``get_random_point`` samples. -``is_point_exchangeable(point)`` bounds the region the reservoir may take -molecules back from, and is what the molecule moves use to pick candidates. -They coincide for every cell except ``CustomCell``; see -:doc:`../gcmc_acceptance_convention`. - - -CustomCell ----------- - -.. code-block:: python - - CustomCell(atoms, custom_height=None, bottom_z=None, species_radii=None, - mc_sample_points=100_000, seed=None) - -A rectangular sub-slab spanning the full x and y extent with finite thickness -along z. Standard choice for surface slabs. - -- ``custom_height`` (float): thickness along z. Must not exceed the box height. -- ``bottom_z`` (float): z of the lower face. -- ``species_radii`` (dict): per-element exclusion radii. -- ``mc_sample_points`` (int): random points for the free-volume estimate. - - -SphericalCell -------------- - -.. code-block:: python - - SphericalCell(atoms, vacuum, species_radii, mc_sample_points=100_000, - seed=None) - -A sphere around an isolated cluster. On construction it translates the atoms so -their center of mass sits at the origin, then sets the radius to the outermost -atom distance plus ``vacuum``. - -- ``vacuum`` (float): padding added to the bounding radius. -- ``species_radii`` (dict): per-element exclusion radii. -- ``mc_sample_points`` (int): random points for the free-volume estimate. - - -DomeCell --------- - -.. code-block:: python - - DomeCell(atoms, particle_species, bottom_z, vacuum, species_radii, - mc_sample_points=100_000, seed=None) - -A hemispherical region for a supported nanoparticle. Centers on the -``particle_species`` centroid and clips the ball at the support surface -(``z >= bottom_z``). Unlike ``SphericalCell`` it does not translate the atoms. - -- ``particle_species`` (str | list[str]): symbol(s) identifying the particle. - Raises ``ValueError`` if none are present. -- ``bottom_z`` (float): z of the support surface. -- ``vacuum`` (float): padding added to the particle bounding radius. -- ``species_radii`` (dict): per-element exclusion radii. -- ``mc_sample_points`` (int): random points for the free-volume estimate. - - -NullCell --------- - -.. code-block:: python - - NullCell() - -A no-op cell. ``get_volume()`` returns ``0`` and it provides no insertion -points. Use it as a placeholder where the API needs a cell but no insertion -region exists, such as a pure-displacement canonical run. +.. autoclass:: mcpy.cell.NullCell + :members: + :show-inheritance: diff --git a/docs/reference/ensembles.rst b/docs/reference/ensembles.rst index 5070da0..5af0226 100644 --- a/docs/reference/ensembles.rst +++ b/docs/reference/ensembles.rst @@ -6,162 +6,30 @@ from its module, for example ``from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble``. Conceptual background is in :doc:`../ensembles`. -All concrete ensembles share the file and output parameters listed under -:ref:`BaseEnsemble `. Each one runs through ``run(steps)``. +All concrete ensembles share the file and output parameters of +:class:`~mcpy.ensembles.base_ensemble.BaseEnsemble`. Each one runs through +``run(steps)``. +``GrandCanonicalEnsemble`` takes ``units_type`` (``'metal'`` or ``'LJ'``), which +selects the constants its de Broglie factors are built from; see +:doc:`../ensembles`. -.. _ref-base-ensemble: -BaseEnsemble ------------- +.. autoclass:: mcpy.ensembles.base_ensemble.BaseEnsemble + :members: -.. code-block:: python +.. autoclass:: mcpy.ensembles.grand_canonical_ensemble.GrandCanonicalEnsemble + :members: + :show-inheritance: - BaseEnsemble(atoms, cells, calculator, - random_seed=None, traj_file='trajectory.xyz', traj_mode='w', - trajectory_write_interval=1, outfile='outfile.out', - outfile_mode='w', outfile_write_interval=1, - minima_file=None, minima_mode='a') +.. autoclass:: mcpy.ensembles.canonical_ensemble.CanonicalEnsemble + :members: + :show-inheritance: -Abstract base for every ensemble. Owns the atoms, cells, calculator, file -handles, and step counter. Not instantiated directly. +.. autoclass:: mcpy.ensembles.replica_exchange.ReplicaExchange + :members: -Shared parameters (inherited by the concrete ensembles): +.. autoclass:: mcpy.ensembles.batched_replica_exchange.BatchedReplicaExchange + :members: -- ``atoms`` (ase.Atoms): starting configuration. -- ``cells`` (list): cell objects providing free volume and insertion points. -- ``calculator``: an mcpy calculator wrapper or any object exposing - ``get_potential_energy(atoms)``. -- ``random_seed`` (int, optional): seed source for derived generators. -- ``traj_file`` / ``traj_mode`` / ``trajectory_write_interval``: sampling - trajectory path, file mode, and write stride. ``traj_file=None`` disables it. -- ``outfile`` / ``outfile_mode`` / ``outfile_write_interval``: log path, mode, - and write stride. ``outfile=None`` disables it. -- ``minima_file`` / ``minima_mode``: basin-hopping output, off by default. - ``'a'`` appends every new minimum; ``'w'`` keeps only the running best. - -Methods: - -- ``run(steps)``: open files, run ``steps`` ensemble steps, finalize. -- ``initialize_run()`` / ``finalize_run()``: open and close handles. Idempotent. - Also available through the context-manager protocol. -- ``atoms`` / ``cells`` / ``step``: properties. - -Abstract: ``_run()`` performs one ensemble step and is implemented by -subclasses. - - -CanonicalEnsemble ------------------ - -.. code-block:: python - - CanonicalEnsemble(atoms, calculator, cells=None, - random_seed=None, optimizer=None, fmax=0.1, - temperature=300, move_selector=None, constraints=None, - traj_file='trajectory.xyz', traj_mode='w', - outfile='outfile.out', outfile_mode='w', - outfile_write_interval=10, trajectory_write_interval=1, - minima_file=None, minima_mode='a') - -Fixed-composition (NVT) Metropolis sampler. Each step applies one move from -``move_selector``, relaxes it with ``optimizer`` to ``fmax``, and accepts by the -Metropolis rule at ``temperature``. - -Parameters beyond the shared set: - -- ``optimizer``: an ASE optimizer class (e.g. ``ase.optimize.LBFGS``) used to - relax each trial. -- ``fmax`` (float): force tolerance for the relaxation. -- ``temperature`` (float): canonical temperature in K. -- ``move_selector`` (MoveSelector): the trial-move sampler. -- ``constraints``: ASE constraint applied to each mutated structure. - -Methods: ``run(steps)``, plus ``get_state()`` / ``set_state(state)`` for use as a -temperature replica inside ``ReplicaExchange``. - - -GrandCanonicalEnsemble ----------------------- - -.. code-block:: python - - GrandCanonicalEnsemble(atoms, cells, units_type, calculator, mu, species, - temperature, move_selector, random_seed=None, - traj_file='trajectory.xyz', traj_mode='w', - trajectory_write_interval=1, outfile='outfile.out', - outfile_mode='w', outfile_write_interval=1, - minima_file=None, minima_mode='a') - -Variable-composition (μVT) sampler. Insertion and deletion moves exchange -``species`` with a reservoir at chemical potentials ``mu``; acceptance uses the -de Broglie criterion (see :doc:`../ensembles`). - -Parameters beyond the shared set: - -- ``units_type`` (str): ``'metal'`` or ``'LJ'``; selects the constants the - de Broglie factors are built from (see :doc:`../ensembles`). -- ``mu`` (dict): species to reservoir chemical potential in eV. -- ``species`` (list[str]): symbols the move set may insert or delete. -- ``temperature`` (float): temperature in K. -- ``move_selector`` (MoveSelector): the trial-move sampler. - -Methods: ``run(steps)``, plus ``get_state()`` / ``set_state(state)`` used by the -replica-exchange drivers. The running-minimum score is the grand potential -``E - Σ μ_i N_i``. - - -ReplicaExchange ---------------- - -.. code-block:: python - - ReplicaExchange(gcmc_factory, temperatures=None, mus=None, gcmc_steps=100, - exchange_interval=10, outfile='replica_exchange.log', - write_out_interval=20, seed=31, - global_minimum_file='global_minimum.xyz') - -MPI parallel-tempering wrapper, one replica per rank. Requires ``mpi4py``; -constructing it without ``mpi4py`` installed raises ``ImportError``. - -Parameters: - -- ``gcmc_factory`` (callable): builds the rank-local ensemble. Called as - ``gcmc_factory(T=..., rank=i)`` for a temperature ladder or - ``gcmc_factory(mu=..., rank=i)`` for a chemical-potential ladder. Must produce - per-rank output paths so ranks do not race on the same files. -- ``temperatures`` (list, optional): one temperature per rank. -- ``mus`` (list, optional): one ``mu`` dict per rank. Pass exactly one of - ``temperatures`` or ``mus``. -- ``gcmc_steps`` (int): GCMC steps per replica. -- ``exchange_interval`` (int): steps between exchange attempts. -- ``write_out_interval`` (int): steps between summary writes. -- ``seed`` (int): RNG seed for exchange decisions. -- ``global_minimum_file`` (str): path for the lowest-score frame across ranks. - ``None`` disables it. - -Methods: ``run()`` (no step argument; the step count is fixed at construction). - - -BatchedReplicaExchange ----------------------- - -.. code-block:: python - - BatchedReplicaExchange(gcmc_factory, calculator, temperatures=None, mus=None, - gcmc_steps=100, exchange_interval=10, - outfile='replica_exchange.log', write_out_interval=20, - seed=31, global_minimum_file='global_minimum.xyz') - -Single-process replica exchange that batches every replica's trial energy into -one calculator call. The right choice on a single GPU. - -Parameters beyond those of ``ReplicaExchange``: - -- ``calculator``: shared calculator. Must implement - ``get_potential_energies(atoms_list)``. Passing one without it raises - ``TypeError``. ``AlchemiCalculator`` provides the method. - -The factory must return a fresh ensemble per replica with its own cells and -move selector; sharing them across replicas corrupts per-replica state. Methods: -``run()``. +.. autofunction:: mcpy.ensembles.base_ensemble.write_xyz diff --git a/docs/reference/moves.rst b/docs/reference/moves.rst index 5eef688..3cc1835 100644 --- a/docs/reference/moves.rst +++ b/docs/reference/moves.rst @@ -11,135 +11,60 @@ or ``0``, and ``species`` is the affected symbol. Moves mutate ``atoms`` in place; the ensemble rolls back on rejection. -MoveSelector ------------- +.. autoclass:: mcpy.moves.MoveSelector + :members: -.. code-block:: python +.. autoclass:: mcpy.moves.base_move.BaseMove + :members: - MoveSelector(probabilities, move_list, seed=None) +Particle exchange +----------------- -Samples one move per step from ``move_list`` with weights ``probabilities`` -(weights need not sum to one). Tracks per-interval and cumulative acceptance -counters, excluding non-viable proposals from the denominator. +.. autoclass:: mcpy.moves.InsertionMove + :members: + :show-inheritance: -Parameters: +.. autoclass:: mcpy.moves.DeletionMove + :members: + :show-inheritance: -- ``probabilities`` (list): relative weight of each move. Their sum sets - ``n_moves``, the number of trial moves per GCMC step. -- ``move_list`` (list): the move instances to sample from. -- ``seed`` (int, optional): RNG seed. +Molecular exchange +------------------ -Methods: ``do_trial_move(atoms)``, ``acceptance_counter()``, ``get_volume()``, -``interval_ratios()``, ``total_ratios()``, ``reset_counters()``. +.. autoclass:: mcpy.moves.MoleculeInsertionMove + :members: + :show-inheritance: +.. autoclass:: mcpy.moves.MoleculeDeletionMove + :members: + :show-inheritance: -BaseMove --------- +.. autoclass:: mcpy.moves.MoleculeDisplacementMove + :members: + :show-inheritance: -.. code-block:: python +.. automodule:: mcpy.moves.molecule_utils + :members: - BaseMove(cell, species, seed) +Displacement and reordering +--------------------------- -Abstract base. Holds the attached ``cell``, the ``species`` list, and a seeded -RNG. Provides ``get_volume()`` and ``calculate_volume(atoms)`` over the cell. -Subclasses implement ``do_trial_move(atoms)``. +.. autoclass:: mcpy.moves.DisplacementMove + :members: + :show-inheritance: +.. autoclass:: mcpy.moves.PermutationMove + :members: + :show-inheritance: -InsertionMove -------------- +.. autoclass:: mcpy.moves.ShakeMove + :members: + :show-inheritance: -.. code-block:: python +.. autoclass:: mcpy.moves.BrownianMove + :members: + :show-inheritance: - InsertionMove(cell, species, seed, min_insert=None, max_atoms=None) - -Inserts one atom of a random selected species at a random point in ``cell``. -Sets ``delta_particles = +1``. - -- ``min_insert`` (float, optional): minimum distance to existing cell atoms. The - move retries up to an internal cap and reports a failed move if it cannot - place the atom without a closer contact. -- ``max_atoms`` (int, optional): if the structure already contains at least this - many atoms of the selected species, the move is skipped without mutating - ``atoms`` (recorded as a failed move, no energy evaluation). - - -DeletionMove ------------- - -.. code-block:: python - - DeletionMove(cell, species, seed, min_atoms=None) - -Deletes a random atom of the selected species lying inside ``cell``. Sets -``delta_particles = -1``. Returns a falsy result when no candidate atom exists, -when ``min_atoms`` would be violated, or recorded as a failed move rather than a -rejection. - -- ``min_atoms`` (int, optional): if the structure contains at most this many - atoms of the selected species, the move is skipped without mutating ``atoms``. - - -DisplacementMove ----------------- - -.. code-block:: python - - DisplacementMove(species, seed, constraints=None, max_displacement=0.1, - n_steps=1) - -Displaces ``n_steps`` distinct atoms by random vectors of magnitude up to -``max_displacement``. Particle count is unchanged. Uses a ``NullCell`` -internally (no insertion region). - -- ``constraints`` (list, optional): indices held fixed. -- ``max_displacement`` (float): maximum per-atom step. -- ``n_steps`` (int): atoms moved per trial. Exceeding the movable-atom count - raises ``ValueError``. - - -PermutationMove ---------------- - -.. code-block:: python - - PermutationMove(species, seed, n_swaps=1) - -Swaps the chemical identities of ``n_swaps`` atom pairs drawn from different -species groups, in a single trial. Particle count is unchanged. Returns a falsy -result if a requested species is absent. - - -ShakeMove ---------- - -.. code-block:: python - - ShakeMove(r_max, seed) - -Displaces every atom by an independent vector drawn uniformly inside a ball of -radius ``r_max``. A global perturbation, usually paired with a relaxing -calculator. - - -BrownianMove ------------- - -.. code-block:: python - - BrownianMove(temperature, calculator, steps, d_t, seed) - -Runs ``steps`` of Velocity-Verlet MD at ``temperature`` (timestep ``d_t`` in fs) -as the trial move, from a Maxwell-Boltzmann velocity draw. - - -AlchemiBrownianMove -------------------- - -.. code-block:: python - - AlchemiBrownianMove(calculator, temperature, friction=0.01, steps=100, - dt=2.0, seed=0) - -GPU-native Langevin Brownian move. Runs a short NVT Langevin trajectory through -``calculator.run_md`` (an ``AlchemiCalculator`` or ``AlchemiFCalculator``), -reusing its model. Honors ``FixAtoms``. +.. autoclass:: mcpy.moves.AlchemiBrownianMove + :members: + :show-inheritance: diff --git a/docs/reference/utils.rst b/docs/reference/utils.rst index 7c3c6f6..93152e9 100644 --- a/docs/reference/utils.rst +++ b/docs/reference/utils.rst @@ -2,119 +2,31 @@ Utilities ========= ``mcpy.utils`` holds the thermodynamic constants, the seeded RNG, logging setup, -and phase-diagram post-processing. +reference chemical potentials, and phase-diagram post-processing. +``phase_diagram`` is imported lazily, so the Monte Carlo core never pulls in +matplotlib. -SetUnits --------- -.. code-block:: python +.. autoclass:: mcpy.utils.set_unit_constant.SetUnits + :members: - SetUnits(unit_type, temperature, species) +.. autoclass:: mcpy.utils.RandomNumberGenerator + :members: -Computes the constants the GCMC acceptance rules need. Built automatically by -``GrandCanonicalEnsemble`` from its ``units_type``, ``temperature``, and -``species``. Conceptual background is in :doc:`../ensembles` (section -"Setting a unit system"). +.. automodule:: mcpy.utils.reference_mu + :members: -- ``unit_type`` (str): ``'metal'`` or ``'LJ'``. Any other value raises - ``ValueError``. -- ``temperature`` (float): temperature in K. -- ``species`` (list[str]): symbols to precompute masses and wavelengths for. +.. automodule:: mcpy.utils.chunking + :members: -Attributes: ``beta``, ``lambda_dbs`` (per-species thermal de Broglie -wavelength), ``masses``, ``BOLTZMANN_CONSTANT``, ``PLANCK_CONSTANT``. +.. autofunction:: mcpy.utils.logging.configure -Methods: +Phase diagrams +-------------- -- ``de_broglie_insertion(volume, n_atoms, specie)``: the insertion prefactor - ``V / ((n_atoms + 1) Λ³)``. Raises ``ValueError`` if ``n_atoms < 0``. -- ``de_broglie_deletion(volume, n_atoms, specie)``: the deletion prefactor - ``n_atoms Λ³ / V``. +A worked example is in :doc:`../examples/phase_diagram_analysis`. +.. autofunction:: mcpy.utils.phase_diagram.plot_phase_diagram -RandomNumberGenerator ---------------------- - -.. code-block:: python - - RandomNumberGenerator(seed=None) - -Seeded wrapper over Python's ``random.Random``. Each move and the acceptance -test carries its own instance, so the global ``random`` state is never touched. - -- ``seed`` (int, optional): seed for reproducibility. - -Methods: ``get_uniform(a=0.0, b=1.0)``, ``get_gaussian(mu=0.0, sigma=1.0)``. - - -configure ---------- - -.. code-block:: python - - configure(level=logging.INFO, file=None, mpi_rank=None, stream=True) - -Attaches handlers to the top-level ``mcpy`` logger. Importing `mcpy` attaches -no handlers and mutates no global state; call this once at startup to opt in to -log output. Importable from ``mcpy.utils.logging``. - -- ``level`` (int): log level for the ``mcpy`` logger. -- ``file`` (str, optional): log file path. Pass ``f"mcpy_rank_{rank}.log"`` for - per-rank files. -- ``mpi_rank`` (int, optional): included in the format string when set. -- ``stream`` (bool): also emit to stderr. - -Returns the configured ``mcpy`` logger. Re-calling is idempotent: it drops the -handlers it previously attached. - - -plot_phase_diagram ------------------- - -.. code-block:: python - - plot_phase_diagram(frames, adsorbate='H', metal_symbols=('Pd', 'Ag'), - mu_ref=-3.05, kind='surface', T=300.0, - dmu_range=(-1.0, 0.0), n_bins=400, ...) - -Builds a grand-canonical adsorbate phase diagram for one system configuration -from a list of relaxed frames (or a list of trajectories). The lowest-energy -adsorbate-free frame is the reference. Returns a dict of arrays and metadata and -optionally saves the figure. A worked example is in -:doc:`../examples/phase_diagram_analysis`. - -Key parameters: - -- ``frames`` (list[ase.Atoms] | list[list]): frames for one configuration. -- ``adsorbate`` (str): species whose count sets the stoichiometry. -- ``kind`` (str): ``'surface'`` normalizes by cell area (meV/Ų); ``'nano'`` by - metal-atom count (meV/atom). Any other value raises ``ValueError``. -- ``mu_ref`` (float): reference chemical potential of the adsorbate in eV. -- ``adsorbate_count_fn`` (callable, optional): custom ``atoms -> int`` count for - adsorbates that also occur in an inert sublattice. - -See the function docstring for the full plotting and normalization parameters. - - -analyze_phase_diagram_results ------------------------------ - -.. code-block:: python - - analyze_phase_diagram_results(trajectory_path='relaxed_structures.xyz', - host_symbol='Ag', oxygen_symbol='O', - idx_ref=2400, ...) - -Reads a relaxed-structure trajectory from disk and builds a surface oxidation -phase diagram against the oxygen chemical potential, returning a dict of arrays -and metadata. The module is also runnable as a script through ``main()``. See -the function docstring for the full set of thermodynamic-reference parameters. - - -compute_radii -------------- - -``scripts/compute_radii.py`` is a calibration script, not a library API. It -automates the ``species_radii`` workflow for FCC(111) hosts. Configuration and -usage are documented in :doc:`../species_radii`. +.. autofunction:: mcpy.utils.phase_diagram.analyze_phase_diagram_results diff --git a/mcpy/utils/set_unit_constant.py b/mcpy/utils/set_unit_constant.py index dc684e1..e2d2370 100644 --- a/mcpy/utils/set_unit_constant.py +++ b/mcpy/utils/set_unit_constant.py @@ -14,14 +14,15 @@ def __init__(self, """ Initialize the SetUnits class with a specific unit type. - Parameters: - unit_type (str): The type of units to set. Possible values are "LJ" or "metal". - species (List): List of atomic species for which to set units. - molecules (dict, optional): Mapping of molecular species name to an - ASE Atoms template. The name is the key used in ``mu`` and - ``lambda_dbs``; the mass is the sum of the template's atomic - masses. Molecular species are identified by composition, so two - templates with the same composition cannot coexist. + Args: + unit_type: The type of units to set, ``"LJ"`` or ``"metal"``. + temperature: Temperature in K. + species: List of atomic species for which to set units. + molecules: Mapping of molecular species name to an ASE Atoms + template. The name is the key used in ``mu`` and + ``lambda_dbs``; the mass is the sum of the template's atomic + masses. Molecular species are identified by composition, so two + templates with the same composition cannot coexist. """ self.unit_type = unit_type self.species = species From 458943cc223a8b2b363027aac26023f80c024999 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:32:26 +0200 Subject: [PATCH 10/20] Stop suppressing E402 in the examples and tests Finding X1. Both examples called configure_logging() between their imports, which forced a `# noqa: E402` on every mcpy import below it. The call belongs in main(): configure() is idempotent, and library modules only ever call getLogger, so nothing needs logging configured at import time. Thirteen noqa comments go away and the imports sort normally. Two unrelated E402 sources, found while checking the first: - tests/test_molecule_moves.py grew nine sectional mid-file imports. Hoisted into one sorted block; the sections still read fine without them. - tests/test_phase_diagram.py must call matplotlib.use('Agg') before anything imports pyplot, so its E402 is correct code. Marked with noqa and a comment saying why, rather than "fixed". tests/test_calculator_gcmc_equivalence.py keeps its noqa markers: they sit below pytest.importorskip calls, which have to run before the imports they guard. flake8 is now clean across mcpy/, tests/ and examples/. 199 tests pass, and both examples still respond to --help in the alchemi env. --- examples/gcmc_ag111.py | 21 +++++++++--------- examples/re_gcmc_batched.py | 18 +++++++--------- tests/test_molecule_moves.py | 41 +++++++++--------------------------- tests/test_phase_diagram.py | 6 +++--- 4 files changed, 31 insertions(+), 55 deletions(-) diff --git a/examples/gcmc_ag111.py b/examples/gcmc_ag111.py index aef9f7c..6a6bfae 100644 --- a/examples/gcmc_ag111.py +++ b/examples/gcmc_ag111.py @@ -33,21 +33,19 @@ import logging +from mcpy.calculators import AlchemiCalculator, AlchemiFCalculator +from mcpy.cell import CustomCell +from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble +from mcpy.moves import DeletionMove, InsertionMove +from mcpy.moves.move_selector import MoveSelector +from mcpy.utils import derive_mu_bulk, derive_mu_gas from mcpy.utils.logging import configure as configure_logging -configure_logging() logger = logging.getLogger('mcpy') -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.calculators import AlchemiCalculator, AlchemiFCalculator # noqa: E402 -from mcpy.cell import CustomCell # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - -# Calibrated set for O on Ag(111) with the small-density-agnesi MACE model, -# Fixed on purpose: this is fixture -# geometry, not a per-checkpoint number. Only mu is derived at run time. +# Calibrated set for O on Ag(111) with the small-density-agnesi MACE model. +# Fixed on purpose: this is fixture geometry, not a per-checkpoint number. +# Only mu is derived at run time. AG_LATTICE = 4.165 # Ag fcc lattice constant (Angstrom) # Relaxed O-Ag pair distance: the free-volume exclusion radius, and also how # far below the top layer the region starts so O can go subsurface (the @@ -127,6 +125,7 @@ def build_calculator(args): def main(): args = parse_args() + configure_logging() os.makedirs(args.outdir, exist_ok=True) ss = np.random.SeedSequence(args.seed) diff --git a/examples/re_gcmc_batched.py b/examples/re_gcmc_batched.py index 7dc4ae5..2edd510 100644 --- a/examples/re_gcmc_batched.py +++ b/examples/re_gcmc_batched.py @@ -18,18 +18,15 @@ from ase.build import fcc111 from ase.constraints import FixAtoms +from mcpy.calculators import AlchemiFCalculator +from mcpy.cell import CustomCell as Cell +from mcpy.ensembles import BatchedReplicaExchange +from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble +from mcpy.moves import DeletionMove, InsertionMove +from mcpy.moves.move_selector import MoveSelector +from mcpy.utils import derive_mu_bulk, derive_mu_gas from mcpy.utils.logging import configure as configure_logging -configure_logging() - -from mcpy.moves import DeletionMove, InsertionMove # noqa: E402 -from mcpy.moves.move_selector import MoveSelector # noqa: E402 -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble # noqa: E402 -from mcpy.calculators import AlchemiFCalculator # noqa: E402 -from mcpy.cell import CustomCell as Cell # noqa: E402 -from mcpy.ensembles import BatchedReplicaExchange # noqa: E402 -from mcpy.utils import derive_mu_bulk, derive_mu_gas # noqa: E402 - def parse_args(): p = argparse.ArgumentParser(description=__doc__, @@ -56,6 +53,7 @@ def parse_args(): def main(): args = parse_args() + configure_logging() os.makedirs(args.outdir, exist_ok=True) ss = np.random.SeedSequence(args.seed) diff --git a/tests/test_molecule_moves.py b/tests/test_molecule_moves.py index 1218111..03d02f9 100644 --- a/tests/test_molecule_moves.py +++ b/tests/test_molecule_moves.py @@ -10,8 +10,18 @@ import numpy as np import pytest from ase import Atoms +from ase.build import molecule as build_molecule from mcpy.cell import Cell +from mcpy.ensembles.batched_replica_exchange import BatchedReplicaExchange +from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble +from mcpy.moves import (DeletionMove, InsertionMove, MoleculeDeletionMove, + MoleculeDisplacementMove, MoleculeInsertionMove, + MoveSelector) +from mcpy.moves.molecule_utils import (find_molecules, molecule_com, + random_rotation_matrix) +from mcpy.utils import RandomNumberGenerator +from mcpy.utils.set_unit_constant import SetUnits def _box_atoms(): @@ -24,11 +34,6 @@ def test_box_cell_is_point_inside_always_true(): assert cell.is_point_inside(np.array([-100.0, 0.0, 1e6])) -from mcpy.moves.molecule_utils import (find_molecules, molecule_com, - random_rotation_matrix) -from mcpy.utils import RandomNumberGenerator - - def _water_box(): """Box with one H2O (id 0), one O2 (id 1), and one lone H (id -1).""" atoms = Atoms('OH2O2H', @@ -117,11 +122,6 @@ def test_random_rotation_matrix_is_rotation(): assert not np.allclose(r1, r2) -from ase.build import molecule as build_molecule - -from mcpy.utils.set_unit_constant import SetUnits - - def test_setunits_molecular_mass_and_lambda(): water = build_molecule('H2O') u = SetUnits('metal', temperature=300.0, species=['Ag'], @@ -152,9 +152,6 @@ def test_setunits_rejects_isomer_compositions(): molecules={'co2_linear': a, 'co2_alt': b}) -from mcpy.moves import MoleculeInsertionMove - - def _water_template(): return build_molecule('H2O') @@ -225,9 +222,6 @@ def test_molecule_insertion_min_insert_unplaceable(): assert len(atoms) == 1 -from mcpy.moves import MoleculeDeletionMove - - def test_molecule_deletion_removes_whole_molecule(): atoms = _water_box() move = MoleculeDeletionMove(Cell(atoms), _water_template(), 'H2O', seed=1) @@ -268,9 +262,6 @@ def test_molecule_deletion_ignores_other_compositions(): assert sorted(atoms.get_chemical_symbols()) == ['H', 'H', 'H', 'O'] -from mcpy.moves import DeletionMove, MoveSelector - - def test_move_selector_exchange_count(): atoms = _water_box() cell = Cell(atoms) @@ -286,9 +277,6 @@ def test_move_selector_exchange_count(): assert ms2.get_exchange_count() is None -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble - - class _StubCalc: def get_potential_energy(self, atoms): return -1.0 * len(atoms) @@ -392,9 +380,6 @@ def factory(T=None, rank=None): assert isinstance(re.gcmc, _FakeCanonical) -from mcpy.ensembles.batched_replica_exchange import BatchedReplicaExchange - - def test_batched_re_grand_potential_counts_molecules(): """BatchedReplicaExchange._grand_potential delegates to _minimum_score, which counts whole molecules via find_molecules; Counter(get_chemical_ @@ -501,9 +486,6 @@ def _two_o2_box(): assert deleted_ids == {0, 1} -from mcpy.moves import InsertionMove # noqa: E402 - - def test_atomic_insertion_alongside_molecules_stays_free(): # ASE ``extend`` zero-pads arrays missing from the fragment; without the # explicit -1 tag the inserted atom would join molecule id 0. @@ -534,9 +516,6 @@ def test_atomic_deletion_never_picks_molecule_members(): assert again[0] is False -from mcpy.moves import MoleculeDisplacementMove # noqa: E402 - - def test_molecule_displacement_rigid_and_bounded(): atoms = _water_box() d_before = atoms.get_all_distances(mic=True)[np.ix_([0, 1, 2], [0, 1, 2])] diff --git a/tests/test_phase_diagram.py b/tests/test_phase_diagram.py index fe8e001..80cf730 100644 --- a/tests/test_phase_diagram.py +++ b/tests/test_phase_diagram.py @@ -6,11 +6,11 @@ Run with: python -m pytest tests/test_phase_diagram.py -v """ import matplotlib -matplotlib.use('Agg') +matplotlib.use('Agg') # must precede any pyplot import; hence the E402 below -from ase import Atoms +from ase import Atoms # noqa: E402 -from mcpy.utils.phase_diagram import plot_phase_diagram +from mcpy.utils.phase_diagram import plot_phase_diagram # noqa: E402 def _frame(n_co, energy): From bb3734c49f2c4e1b62617e4d52a502181bbb5b08 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:34:26 +0200 Subject: [PATCH 11/20] Share the LAMMPS parity benchmark scaffolding Finding B1. The two parity scripts each carried their own run_lammps, parse_thermo, block_stats, compare and ZeroCalc; mace_gcmc_parity.py even said so in a comment ("shared infra (mirrors benchmark/lammps_gcmc_parity.py)"). block_stats was byte-identical, and the rest differed only by a field width, a dropped docstring, and one environment variable. They now live in benchmark/_parity_common.py, with the two real differences as parameters: compare(..., width=) sizes the mean columns, since metal-units energies need more room than LJ ones, and run_lammps(..., extra_env=) carries the libtorch loader path the MACE deck needs. The MACE side's compare returned a bare bool while the LJ side returned (ok, row); the shared one returns the tuple and the MACE call sites unpack it. The module has a demo() self-check under a __main__ guard, since parse_thermo and block_stats are the kind of parsing and averaging that fails quietly. Neither benchmark had any test before. --- benchmark/_parity_common.py | 116 ++++++++++++++++++++++++++++++++ benchmark/lammps_gcmc_parity.py | 69 +------------------ benchmark/mace_gcmc_parity.py | 84 +++-------------------- 3 files changed, 126 insertions(+), 143 deletions(-) create mode 100644 benchmark/_parity_common.py diff --git a/benchmark/_parity_common.py b/benchmark/_parity_common.py new file mode 100644 index 0000000..d0ab427 --- /dev/null +++ b/benchmark/_parity_common.py @@ -0,0 +1,116 @@ +"""Shared scaffolding for the LAMMPS parity benchmarks. + +``lammps_gcmc_parity.py`` (LJ units) and ``mace_gcmc_parity.py`` (metal units, +MACE) run the same comparison against LAMMPS ``fix gcmc``, so they need the same +four things: drive LAMMPS on a generated deck, parse its thermo output, block +average a series, and print a pass/fail line. Both had their own copy; +``block_stats`` was byte-identical and the rest differed only by a field width, +a dropped docstring, and one environment variable. +""" +import math +import os +import subprocess + +import numpy as np + + +class ZeroCalc: + """Zero energy, for the ideal-gas stage where only the prefactor matters.""" + + def get_potential_energy(self, atoms): + return 0.0 + + +def run_lammps(deck, outdir, tag, lmp, extra_env=None): + """Write ``deck``, run ``lmp`` on it, and return the parsed thermo output. + + ``extra_env`` is merged over ``os.environ`` for the child; the MACE + benchmark uses it to put libtorch on ``LD_LIBRARY_PATH``. + """ + deckfile = os.path.join(outdir, f'in.{tag}') + logfile = os.path.join(outdir, f'log.{tag}') + with open(deckfile, 'w') as f: + f.write(deck) + env = None + if extra_env: + env = dict(os.environ) + env.update(extra_env) + subprocess.run([lmp, '-in', os.path.abspath(deckfile), + '-log', os.path.abspath(logfile), '-screen', 'none'], + check=True, cwd=outdir, env=env) + return parse_thermo(logfile) + + +def parse_thermo(logfile): + """Collect thermo rows from every run section. Returns dict of arrays + keyed by lower-cased column names (step, atoms, pe, ...).""" + cols, rows = None, [] + with open(logfile) as f: + in_section = False + for line in f: + s = line.split() + if not s: + continue + if s[0] == 'Step': + cols = [c.lower() for c in s] + in_section = True + continue + if in_section: + if s[0] == 'Loop' or line.startswith('WARNING'): + in_section = False + continue + try: + rows.append([float(v) for v in s]) + except ValueError: + in_section = False + if cols is None or not rows: + raise RuntimeError(f'no thermo data parsed from {logfile}') + rows = [r for r in rows if len(r) == len(cols)] + data = np.array(rows) + out = {c: data[:, i] for i, c in enumerate(cols)} + if 'poteng' in out: # thermo header prints "PotEng" for pe + out['pe'] = out['poteng'] + return out + + +def block_stats(series, burn_frac=0.4, nblocks=15): + """Block-averaged mean and standard error after discarding ``burn_frac``.""" + x = np.asarray(series, dtype=float) + x = x[int(len(x) * burn_frac):] + blocks = np.array_split(x, nblocks) + means = np.array([b.mean() for b in blocks]) + return means.mean(), means.std(ddof=1) / math.sqrt(nblocks) + + +def compare(label, m_mcpy, s_mcpy, m_lmp, s_lmp, nsig, width=10): + """Print one pass/fail line and return ``(ok, row)``. + + ``ok`` is ``|difference| < nsig * sigma_combined``; ``row`` is the tuple the + CSV writer expects. ``width`` sizes the mean columns (metal-units energies + need more room than LJ ones). + """ + diff = abs(m_mcpy - m_lmp) + scomb = math.sqrt(s_mcpy ** 2 + s_lmp ** 2) + ok = diff < nsig * scomb + sigmas = diff / scomb if scomb > 0 else float('inf') + print(f' {label:6s} mcpy {m_mcpy:{width}.4f} ± {s_mcpy:.4f} ' + f'lammps {m_lmp:{width}.4f} ± {s_lmp:.4f} ' + f'|d|={diff:.4f} ({sigmas:.2f} sigma) ' + f'{"PASS" if ok else "FAIL"}') + return ok, (m_mcpy, s_mcpy, m_lmp, s_lmp, diff, scomb) + + +def demo(): + """Self-check: block_stats on a known series, compare's verdict logic.""" + mean, err = block_stats(list(range(100)) * 2, burn_frac=0.0, nblocks=10) + assert abs(mean - 49.5) < 1e-9, mean + assert err > 0 + ok, row = compare('demo', 1.0, 0.1, 1.05, 0.1, 2) + assert ok and len(row) == 6 + ok, _ = compare('demo', 1.0, 0.01, 2.0, 0.01, 2) + assert not ok + print('ok') + + +if __name__ == '__main__': + demo() diff --git a/benchmark/lammps_gcmc_parity.py b/benchmark/lammps_gcmc_parity.py index 10f9f87..185d508 100644 --- a/benchmark/lammps_gcmc_parity.py +++ b/benchmark/lammps_gcmc_parity.py @@ -16,12 +16,12 @@ import argparse import math import os -import subprocess import numpy as np from ase import Atoms from ase.calculators.lj import LennardJones +from _parity_common import ZeroCalc, block_stats, compare, run_lammps from mcpy.cell import Cell from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble from mcpy.moves import (DeletionMove, DisplacementMove, InsertionMove, @@ -52,57 +52,6 @@ def mu_map(mu_mcpy, lam): # shared infra # -------------------------------------------------------------------------- -def run_lammps(deck, outdir, tag, lmp): - deckfile = os.path.join(outdir, f'in.{tag}') - logfile = os.path.join(outdir, f'log.{tag}') - with open(deckfile, 'w') as f: - f.write(deck) - subprocess.run([lmp, '-in', os.path.abspath(deckfile), - '-log', os.path.abspath(logfile), '-screen', 'none'], - check=True, cwd=outdir) - return parse_thermo(logfile) - - -def parse_thermo(logfile): - """Collect thermo rows from every run section. Returns dict of arrays - keyed by lower-cased column names (step, atoms, pe, ...).""" - cols, rows = None, [] - with open(logfile) as f: - in_section = False - for line in f: - s = line.split() - if not s: - continue - if s[0] == 'Step': - cols = [c.lower() for c in s] - in_section = True - continue - if in_section: - if s[0] == 'Loop' or line.startswith('WARNING'): - in_section = False - continue - try: - rows.append([float(v) for v in s]) - except ValueError: - in_section = False - if cols is None or not rows: - raise RuntimeError(f'no thermo data parsed from {logfile}') - rows = [r for r in rows if len(r) == len(cols)] - data = np.array(rows) - out = {c: data[:, i] for i, c in enumerate(cols)} - if 'poteng' in out: # thermo header prints "PotEng" for pe - out['pe'] = out['poteng'] - return out - - -def block_stats(series, burn_frac=0.4, nblocks=15): - x = np.asarray(series, dtype=float) - x = x[int(len(x) * burn_frac):] - blocks = np.array_split(x, nblocks) - means = np.array([b.mean() for b in blocks]) - return means.mean(), means.std(ddof=1) / math.sqrt(nblocks) - - def write_csv(outdir, stage, results): """One row per (mu, observable): means, stderrs, diff, verdict.""" path = os.path.join(outdir, f'stage{stage}_results.csv') @@ -145,26 +94,10 @@ def plot_isotherms(outdir, s1_results, s2_results): return path -def compare(label, m_mcpy, s_mcpy, m_lmp, s_lmp, nsig): - diff = abs(m_mcpy - m_lmp) - scomb = math.sqrt(s_mcpy ** 2 + s_lmp ** 2) - ok = diff < nsig * scomb - print(f' {label:6s} mcpy {m_mcpy:10.4f} ± {s_mcpy:.4f} ' - f'lammps {m_lmp:10.4f} ± {s_lmp:.4f} ' - f'|d|={diff:.4f} ({diff / scomb if scomb > 0 else float("inf"):.2f} sigma) ' - f'{"PASS" if ok else "FAIL"}') - return ok, (m_mcpy, s_mcpy, m_lmp, s_lmp, diff, scomb) - - # -------------------------------------------------------------------------- # mcpy side # -------------------------------------------------------------------------- -class ZeroCalc: - def get_potential_energy(self, atoms): - return 0.0 - - class LJCalc: """Vectorized truncated-and-shifted LJ (sigma=eps=1), minimum image in a cubic box of side L. Verified against ASE LennardJones by diff --git a/benchmark/mace_gcmc_parity.py b/benchmark/mace_gcmc_parity.py index b99acce..a70b611 100644 --- a/benchmark/mace_gcmc_parity.py +++ b/benchmark/mace_gcmc_parity.py @@ -22,12 +22,12 @@ import argparse import math import os -import subprocess import numpy as np from ase import Atoms from ase.build import fcc111, molecule +from _parity_common import ZeroCalc, block_stats, compare, run_lammps from mcpy.cell import Cell, CustomCell from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble from mcpy.moves import (MoleculeDeletionMove, MoleculeInsertionMove, @@ -45,81 +45,15 @@ '~/miniconda3/envs/mcpy/lib/python3.14/site-packages/torch/lib')) -# -------------------------------------------------------------------------- -# shared infra (mirrors benchmark/lammps_gcmc_parity.py) -# -------------------------------------------------------------------------- - -def run_lammps(deck, outdir, tag, lmp): - deckfile = os.path.join(outdir, f'in.{tag}') - logfile = os.path.join(outdir, f'log.{tag}') - with open(deckfile, 'w') as f: - f.write(deck) - env = dict(os.environ) - env['LD_LIBRARY_PATH'] = TORCH_LIB + ':' + env.get('LD_LIBRARY_PATH', '') - subprocess.run([lmp, '-in', os.path.abspath(deckfile), - '-log', os.path.abspath(logfile), '-screen', 'none'], - check=True, cwd=outdir, env=env) - return parse_thermo(logfile) - - -def parse_thermo(logfile): - cols, rows = None, [] - with open(logfile) as f: - in_section = False - for line in f: - s = line.split() - if not s: - continue - if s[0] == 'Step': - cols = [c.lower() for c in s] - in_section = True - continue - if in_section: - if s[0] == 'Loop' or line.startswith('WARNING'): - in_section = False - continue - try: - rows.append([float(v) for v in s]) - except ValueError: - in_section = False - if cols is None or not rows: - raise RuntimeError(f'no thermo data parsed from {logfile}') - rows = [r for r in rows if len(r) == len(cols)] - data = np.array(rows) - out = {c: data[:, i] for i, c in enumerate(cols)} - if 'poteng' in out: - out['pe'] = out['poteng'] - return out - - -def block_stats(series, burn_frac=0.4, nblocks=15): - x = np.asarray(series, dtype=float) - x = x[int(len(x) * burn_frac):] - blocks = np.array_split(x, nblocks) - means = np.array([b.mean() for b in blocks]) - return means.mean(), means.std(ddof=1) / math.sqrt(nblocks) - - -def compare(label, m_mcpy, s_mcpy, m_lmp, s_lmp, nsig): - diff = abs(m_mcpy - m_lmp) - scomb = math.sqrt(s_mcpy ** 2 + s_lmp ** 2) - ok = diff < nsig * scomb - print(f' {label:6s} mcpy {m_mcpy:12.4f} ± {s_mcpy:.4f} ' - f'lammps {m_lmp:12.4f} ± {s_lmp:.4f} ' - f'|d|={diff:.4f} ({diff / scomb if scomb > 0 else float("inf"):.2f} sigma) ' - f'{"PASS" if ok else "FAIL"}') - return ok +def _libtorch_env(): + """LAMMPS pair_style mace needs libtorch on the loader path.""" + return {'LD_LIBRARY_PATH': TORCH_LIB + ':' + os.environ.get('LD_LIBRARY_PATH', '')} # -------------------------------------------------------------------------- # mcpy side # -------------------------------------------------------------------------- -class ZeroCalc: - def get_potential_energy(self, atoms): - return 0.0 - - class MaceSinglePoint: """Bare single-point MACE energy (no relaxation), float32 on GPU.""" @@ -224,7 +158,7 @@ def gate(outdir, lmp, model, py_calc): thermo_style custom step atoms pe run 0 """ - th = run_lammps(deck, outdir, 'gate', lmp) + th = run_lammps(deck, outdir, 'gate', lmp, extra_env=_libtorch_env()) e_l = th['pe'][-1] diff = abs(e_py - e_l) ok = diff < 5e-4 @@ -276,7 +210,7 @@ def stage1(outdir, lmp, steps_mcpy, lmp_steps): thermo_style custom step atoms pe run {lmp_steps} """ - th = run_lammps(deck, outdir, 'stage1', lmp) + th = run_lammps(deck, outdir, 'stage1', lmp, extra_env=_libtorch_env()) ml, sl = block_stats(th['atoms'] / 2.0) ok_lmp = abs(ml - n_exact) < 2 * sl print(f' lammps = {ml:.3f} ± {sl:.3f} ' @@ -338,12 +272,12 @@ def stage2(outdir, lmp, model, py_calc, steps_mcpy, lmp_attempts, deltas): thermo_style custom step atoms pe run {nsteps} """ - th = run_lammps(deck, outdir, f'stage2_d{delta}', lmp) + th = run_lammps(deck, outdir, f'stage2_d{delta}', lmp, extra_env=_libtorch_env()) mln, sln = block_stats((th['atoms'] - n_slab) / 2.0) mle, sle = block_stats(th['pe']) - okn = compare('', mn, sn, mln, sln, 2) - oke = compare('', me, se, mle, sle, 3) + okn, _ = compare('', mn, sn, mln, sln, 2, width=12) + oke, _ = compare('', me, se, mle, sle, 3, width=12) all_ok &= okn and oke print(f' Stage II: {"PASS" if all_ok else "FAIL"}') return all_ok From 73f697f490d2eb52823677ae8cfcfae7f6aec735 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 12:43:22 +0200 Subject: [PATCH 12/20] Share one accept/reject step between the serial and batched GCMC loops Finding E1, the audit's biggest duplication. do_gcmc_step and BatchedReplicaExchange._batched_single_move were the same sixty-five lines, comments included: the arrays+constraints snapshot, the "is False or is None" sentinel with its identity-check explanation, the "returned a different Atoms object" guard, the de Broglie count fallback, and the whole accept branch (wrap, n_atoms, E_old, counter, cell volumes, _record_minimum). Keeping them in step was manual, and the repo carried a paired regression test for each loop precisely because they could drift. Split into GrandCanonicalEnsemble._propose() and ._commit_or_rollback(), with _restore() shared by both. The split is what the batched loop needs anyway: every replica must have proposed before their energies can go into one batched forward pass, and each replica's MoveSelector still holds its own pending volume and exchange count when it is scored. One behaviour change, deliberate: the batched loop had a defensive branch for a move returning a non-tuple ("result if isinstance(result, tuple) else (result, 0, None)"). The serial loop never had it and every move returns a three-tuple, so it is gone; such a move now raises instead of being silently read as delta_particles=0. test_batched_rejected_deletion_restores_constraint used to hand-build a SimpleNamespace replica, which stubbed out exactly the code the batched loop now delegates to. It uses a real GrandCanonicalEnsemble replica instead, so it tests the shared path rather than a mock of it. Verified against the GPU golden test: 14/14 pass unchanged, both the single-replica fixture and the three-rung mu ladder, so every accept/reject decision and swap outcome is identical. 199 torch-free tests pass. --- mcpy/ensembles/batched_replica_exchange.py | 82 ++++-------- mcpy/ensembles/grand_canonical_ensemble.py | 147 ++++++++++++--------- tests/test_constraint_rollback.py | 23 +--- 3 files changed, 115 insertions(+), 137 deletions(-) diff --git a/mcpy/ensembles/batched_replica_exchange.py b/mcpy/ensembles/batched_replica_exchange.py index 28a1ca4..a7a3a9c 100644 --- a/mcpy/ensembles/batched_replica_exchange.py +++ b/mcpy/ensembles/batched_replica_exchange.py @@ -235,71 +235,35 @@ def _batched_gcmc_step(self) -> None: def _batched_single_move(self, active: List[int]) -> None: """ One trial move on each active replica: - 1. propose a trial move on each replica (in place, with snapshot) - 2. batched energy eval over the replicas with viable trials - 3. per-replica Metropolis accept/reject + 1. propose on each replica (in place, with a rollback snapshot) + 2. one batched energy evaluation over the replicas that proposed + 3. per-replica accept/reject + + Steps 1 and 3 are ``GrandCanonicalEnsemble._propose`` and + ``._commit_or_rollback``, the same methods the serial loop uses, so the + two cannot drift apart. Splitting them is what lets the energies be + evaluated together: every replica must have proposed before the batch + can be built. + + ``get_volume`` and ``get_exchange_count`` are stateful on each + replica's own MoveSelector, so they still describe that replica's + pending proposal when it is scored later in this method. """ - snapshots: Dict[int, Dict] = {} - constraint_snapshots: Dict[int, list] = {} - trial_meta: Dict[int, tuple] = {} - + pending = {} for i in active: - r = self.replicas[i] - snapshots[i] = {k: v.copy() for k, v in r.atoms.arrays.items()} - # ``del atoms[i]`` in a deletion move remaps FixAtoms indices in - # place; a rejected deletion must restore the constraint too or - # the fixed set drifts (see GrandCanonicalEnsemble.do_gcmc_step). - constraint_snapshots[i] = [c.copy() for c in r.atoms.constraints] - result = r.move_selector.do_trial_move(r.atoms) - atoms_new, delta_particles, species = ( - result if isinstance(result, tuple) else (result, 0, None) - ) - if atoms_new is False or atoms_new is None: - # Move couldn't propose. Identity check, not truthiness: an - # empty Atoms (last atom deleted) is falsy but is a real - # proposal. Restore rather than trust the sentinel to mean - # "nothing was touched" (see GrandCanonicalEnsemble). - r.atoms.arrays = snapshots[i] - r.atoms.set_constraint(constraint_snapshots[i]) - continue - if atoms_new is not r.atoms: - raise RuntimeError( - f"move '{r.move_selector.get_name()}' returned a " - "different Atoms object; GCMC moves must mutate the " - "passed atoms in place" - ) - trial_meta[i] = (delta_particles, species) + proposal = self.replicas[i]._propose() + if proposal is not None: + pending[i] = proposal - viable = [i for i in active if i in trial_meta] - if not viable: + if not pending: return - atoms_list = [self.replicas[i].atoms for i in viable] - energies = self.calculator.get_potential_energies(atoms_list) + viable = list(pending) + energies = self.calculator.get_potential_energies( + [self.replicas[i].atoms for i in viable]) for i, E_new in zip(viable, energies): - r = self.replicas[i] - delta_particles, species = trial_meta[i] - delta_E = float(E_new) - r.E_old - volume = r.move_selector.get_volume() - # de Broglie particle count: molecule moves report their in-cell - # molecule count, atomic moves fall back to the pre-move total - # atom count (``r.n_atoms`` is updated only on acceptance). See - # docs/gcmc_acceptance_convention.rst. - n_exchange = r.move_selector.get_exchange_count() - if n_exchange is None: - n_exchange = r.n_atoms - if r._acceptance_condition(delta_E, delta_particles, volume, species, - n_exchange): - if r._wrap_on_accept: - r.atoms.wrap() - r.n_atoms = len(r.atoms) - r.E_old = float(E_new) - r.move_selector.acceptance_counter() - r.calculate_cells_volume(r.atoms) - r._record_minimum(r.atoms, r.E_old) - else: - r.atoms.arrays = snapshots[i] - r.atoms.set_constraint(constraint_snapshots[i]) + snapshot, meta = pending[i] + self.replicas[i]._commit_or_rollback(float(E_new), snapshot, meta) # --------------------------------------------------------- exchange diff --git a/mcpy/ensembles/grand_canonical_ensemble.py b/mcpy/ensembles/grand_canonical_ensemble.py index ad23ebb..5e4e590 100644 --- a/mcpy/ensembles/grand_canonical_ensemble.py +++ b/mcpy/ensembles/grand_canonical_ensemble.py @@ -240,71 +240,94 @@ def _acceptance_condition(self, return True return p > self.rng_acceptance.get_uniform() - def do_gcmc_step(self) -> None: + def _restore(self, snapshot) -> None: + """Put the configuration back as it was before the trial move.""" + arrays, constraints = snapshot + self.atoms.arrays = arrays + self.atoms.set_constraint(constraints) + + def _propose(self): + """Run one trial move in place. Returns ``(snapshot, meta)``, or + ``None`` when the move could not propose (already rolled back). + + Split out of :meth:`do_gcmc_step` so ``BatchedReplicaExchange`` runs + the identical proposal step: it has to propose on every replica before + it can score any of them in one batched forward pass. + + The snapshot holds ``atoms.arrays`` (positions, numbers, ...) plus the + constraints. Restoring it on rejection also undoes any in-place + relaxation ``compute_energy`` did. The constraints are part of it + because ``del atoms[i]`` remaps FixAtoms indices in place (ASE + ``delete_atoms``), so a rejected deletion would otherwise leave the + restored configuration with shifted fixed indices -- wrong atoms frozen + from then on. + """ atoms = self.atoms + snapshot = ({k: v.copy() for k, v in atoms.arrays.items()}, + [c.copy() for c in atoms.constraints]) + + atoms_new, delta_particles, species = self.move_selector.do_trial_move(atoms) + + if atoms_new is False or atoms_new is None: + # Move couldn't be proposed (e.g. empty cell). MoveSelector already + # recorded the failure so it won't depress the acceptance ratio. + # Identity check, not truthiness: an empty Atoms (last atom + # deleted) is falsy but is a real proposal that must be scored. + # + # Restore rather than trust the sentinel to mean "nothing was + # touched": a move that mutates before bailing would otherwise + # leave the configuration changed while ``E_old`` still describes + # the previous one. + self._restore(snapshot) + return None + + if atoms_new is not atoms: + raise RuntimeError( + f"move '{self.move_selector.get_name()}' returned a " + "different Atoms object; GCMC moves must mutate the " + "passed atoms in place (copy-based moves are for " + "CanonicalEnsemble only)" + ) + return snapshot, (delta_particles, species) + + def _commit_or_rollback(self, E_new: float, snapshot, meta) -> bool: + """Score the pending proposal and either keep it or restore the + snapshot. Returns whether it was accepted.""" + delta_particles, species = meta + delta_E = E_new - self.E_old + volume = self.move_selector.get_volume() + # de Broglie particle count. Molecule moves report their in-cell + # molecule count via ``get_exchange_count`` (textbook convention); + # atomic moves return None and fall back to the total atom count + # before the move (``self.n_atoms`` is updated only on acceptance). + # See docs/gcmc_acceptance_convention.rst. + n_exchange = self.move_selector.get_exchange_count() + if n_exchange is None: + n_exchange = self.n_atoms + if not self._acceptance_condition(delta_E, delta_particles, volume, + species, n_exchange): + self._restore(snapshot) + return False + + if self._wrap_on_accept: + self.atoms.wrap() + self.n_atoms = len(self.atoms) + self.E_old = E_new + self.move_selector.acceptance_counter() + self.calculate_cells_volume(self.atoms) + self._record_minimum(self.atoms, self.E_old) + self.logger.debug("Volume: %.3f, Delta_particles: %d, Species: %s", + volume, delta_particles, species) + return True + + def do_gcmc_step(self) -> None: for _ in range(self.move_selector.n_moves): - # Snapshot arrays (positions, numbers, ...) before the trial move. - # On rejection we restore from this snapshot, which also undoes - # any in-place relaxation done by ``compute_energy``. This - # replaces the previous ``atoms.copy()`` per trial in the moves. - # Constraints too: ``del atoms[i]`` remaps FixAtoms indices in - # place (ASE ``delete_atoms``), so a rejected deletion would - # otherwise leave the restored configuration with shifted fixed - # indices — wrong atoms frozen from then on. - saved_arrays = {k: v.copy() for k, v in atoms.arrays.items()} - saved_constraints = [c.copy() for c in atoms.constraints] - - atoms_new, delta_particles, species = self.move_selector.do_trial_move(atoms) - - if atoms_new is False or atoms_new is None: - # Move couldn't be proposed (e.g. empty cell). MoveSelector - # already recorded the failure so it won't depress the - # acceptance ratio. Identity check, not truthiness: an empty - # Atoms (last atom deleted) is falsy but is a real proposal - # that must be scored. - # - # Restore the snapshot rather than trusting the sentinel to - # mean "nothing was touched": a move that mutates before - # bailing would otherwise leave the configuration changed - # while ``E_old`` still describes the previous one. - atoms.arrays = saved_arrays - atoms.set_constraint(saved_constraints) + proposal = self._propose() + if proposal is None: continue - - if atoms_new is not atoms: - raise RuntimeError( - f"move '{self.move_selector.get_name()}' returned a " - "different Atoms object; GCMC moves must mutate the " - "passed atoms in place (copy-based moves are for " - "CanonicalEnsemble only)" - ) - - E_new = self.compute_energy(atoms) - delta_E = E_new - self.E_old - volume = self.move_selector.get_volume() - # de Broglie particle count. Molecule moves report their in-cell - # molecule count via ``get_exchange_count`` (textbook convention); - # atomic moves return None and fall back to the total atom count - # before the move (``self.n_atoms`` is updated only on acceptance). - # See docs/gcmc_acceptance_convention.rst. - n_exchange = self.move_selector.get_exchange_count() - if n_exchange is None: - n_exchange = self.n_atoms - if self._acceptance_condition(delta_E, delta_particles, volume, - species, n_exchange): - if self._wrap_on_accept: - atoms.wrap() - self.n_atoms = len(atoms) - self.E_old = E_new - self.move_selector.acceptance_counter() - self.calculate_cells_volume(atoms) - self._record_minimum(atoms, self.E_old) - - self.logger.debug("Volume: %.3f, Delta_particles: %d, Species: %s", - volume, delta_particles, species) - else: - atoms.arrays = saved_arrays - atoms.set_constraint(saved_constraints) + snapshot, meta = proposal + self._commit_or_rollback(self.compute_energy(self.atoms), + snapshot, meta) def initialize_run(self) -> None: """Open files, write headers, log start, write initial state.""" diff --git a/tests/test_constraint_rollback.py b/tests/test_constraint_rollback.py index ccc1017..78b7126 100644 --- a/tests/test_constraint_rollback.py +++ b/tests/test_constraint_rollback.py @@ -89,25 +89,16 @@ def test_accepted_deletion_keeps_ase_remap(): def test_batched_rejected_deletion_restores_constraint(): + """The batched loop must roll back exactly like the serial one. + + The replica is a real GrandCanonicalEnsemble, not a stand-in: the batched + loop delegates the snapshot and the accept/reject to the ensemble's own + _propose / _commit_or_rollback, so stubbing those would test nothing. + """ atoms, move = _make_atoms_and_move() positions_before = atoms.positions.copy() - replica = types.SimpleNamespace( - atoms=atoms, - move_selector=types.SimpleNamespace( - do_trial_move=move.do_trial_move, - get_volume=lambda: 100.0, - get_name=lambda: 'deletion', - acceptance_counter=lambda: None, - get_exchange_count=lambda: None, - ), - E_old=0.0, - n_atoms=len(atoms), - _acceptance_condition=lambda *a, **k: False, - _wrap_on_accept=False, - calculate_cells_volume=lambda a: None, - _record_minimum=lambda a, en: None, - ) + replica = _gcmc(atoms, move, accept=False) re = BatchedReplicaExchange.__new__(BatchedReplicaExchange) re.replicas = [replica] From 2ddab33d623eae21e3722692ebd720d8d818f2bc Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 13:36:26 +0200 Subject: [PATCH 13/20] Deduplicate the Alchemi calculator internals Findings A1, A2, A3, A6, A7. A5 was already applied with the float64 fix. A1, the one that mattered: get_potential_energy re-implemented the batched relaxation for a single structure, and the two had drifted apart. This path called opt.run, which steps the whole batch until every graph has converged, while the batched path retires each graph at its own first convergence. One structure is a batch of one, so it now delegates: get_potential_energies ([atoms])[0]. A2: the bootstrap ritual (pre-allocate the tensors compute() writes into, build the neighbor list, compute initial forces, zero the frozen rows) appeared three times with the same three-line comment. Now _prepare_batch and _bootstrap. A3: run_md was byte-identical in both calculators, docstring included. Now a _MDMixin in the shared module. The energy_only guard reads through getattr, so the F variant, which has no such attribute and rejects a forces-less model in its constructor instead, inherits it harmlessly. A6: _make_batch was exactly _make_multi_batch([atoms]); one function now takes either. A7: _per_graph_energies carried six lines of comment describing a scatter-reduced fallback above a branch that only raises. A4 is NOT applied: the audit called the fire2 optimizer uncalled, but its only caller then lived in the ignored scratch/ directory. It is now the default in examples/gcmc_ag111.py and what the golden test uses, so the registry stays. Verified three ways: the GPU golden test passes 14/14 with unchanged goldens, so the single-graph compacted path is equivalent to opt.run as expected; benchmark/internal/verify_compact_parity.py passes (max |dE| 0.0006 eV against the serial path, fresh max|F| within fmax, FixAtoms rows bit-identical); and the 199 torch-free tests pass. That guard needed its import updated for A6. --- mcpy/calculators/_alchemi_common.py | 90 ++++++++++++++++------ mcpy/calculators/alchemi_calculator.py | 39 +--------- mcpy/calculators/alchemi_f_calculator.py | 95 ++++-------------------- 3 files changed, 85 insertions(+), 139 deletions(-) diff --git a/mcpy/calculators/_alchemi_common.py b/mcpy/calculators/_alchemi_common.py index 938486e..403319f 100644 --- a/mcpy/calculators/_alchemi_common.py +++ b/mcpy/calculators/_alchemi_common.py @@ -117,12 +117,10 @@ def _load_model( ) -def _make_batch(atoms: Atoms, device: str, dtype: torch.dtype) -> Batch: - data = AtomicData.from_atoms(atoms, device=device, dtype=dtype) - return Batch.from_data_list([data], device=device) - - -def _make_multi_batch(atoms_list: List[Atoms], device: str, dtype: torch.dtype) -> Batch: +def _make_batch(atoms_list, device: str, dtype: torch.dtype) -> Batch: + """Batch one ``Atoms`` or a list of them; a single structure is a batch of one.""" + if isinstance(atoms_list, Atoms): + atoms_list = [atoms_list] data = [AtomicData.from_atoms(a, device=device, dtype=dtype) for a in atoms_list] return Batch.from_data_list(data, device=device) @@ -132,9 +130,6 @@ def _per_graph_energies(out_energy: torch.Tensor, n_graphs: int) -> np.ndarray: e = out_energy.detach().to('cpu') if e.numel() == n_graphs: return e.view(-1).numpy() - # Fallback: scatter-reduced layout — sum each graph's atomic contributions. - # If this branch is hit, the model returned per-atom energies; for now we - # surface the mismatch loudly so the caller can adapt. raise RuntimeError( f"Unexpected energy tensor shape {tuple(out_energy.shape)} for {n_graphs} graphs. " "Batched eval expects one energy per graph." @@ -146,6 +141,27 @@ def _build_nl(batch: Batch, nl_hook: NeighborListHook) -> None: nl_hook(ctx, DynamicsStage.BEFORE_COMPUTE) +def _prepare_batch(batch: Batch, n_graphs: int, device: str, + dtype: torch.dtype) -> None: + """Pre-allocate the tensors ``compute()`` writes into via ``copy_()``.""" + batch.forces = torch.zeros_like(batch.positions) + batch.energy = torch.zeros(n_graphs, 1, device=device, dtype=dtype) + + +def _bootstrap(opt, batch: Batch, nl_hook: NeighborListHook, + fixed: List[int]) -> None: + """Build the neighbor list and the initial forces before the loop starts. + + ``FreezeAtomsHook`` zeros frozen forces only at ``AFTER_POST_UPDATE``, and + this ``compute()`` runs outside the hook loop, so step 0 would otherwise + displace the frozen rows once. + """ + _build_nl(batch, nl_hook) + opt.compute(batch) + if fixed: + batch.forces[fixed] = 0.0 + + def _write_back_positions(atoms: Atoms, batch: Batch) -> None: """Copy relaxed positions from batch back to atoms, skipping FixAtoms-constrained indices.""" # float64: the batch is float32, and assigning the float64 frozen rows @@ -208,10 +224,7 @@ def _run_langevin_md( converts to internal units). Mutates ``atoms`` in place; returns nothing. """ batch = _make_batch(atoms, device, dtype) - - # compute() writes forces/energy via copy_() — pre-allocate the targets. - batch.forces = torch.zeros_like(batch.positions) - batch.energy = torch.zeros(1, 1, device=device, dtype=dtype) + _prepare_batch(batch, 1, device, dtype) batch.velocities = torch.zeros_like(batch.positions) temp = torch.full((batch.num_graphs,), float(temperature), device=device, dtype=dtype) initialize_velocities( @@ -235,14 +248,47 @@ def _run_langevin_md( ) 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 - + _bootstrap(opt, batch, nl_hook, fixed) opt.run(batch) _write_back_positions(atoms, batch) + + +class _MDMixin: + """``run_md`` for both Alchemi calculators. + + Identical in each, docstring included, so it lives here. The energy_only + guard applies only to :class:`AlchemiCalculator`; the F variant has no such + attribute and rejects a forces-less model in its constructor instead. + """ + + def run_md( + self, + atoms: Atoms, + *, + temperature: float, + friction: float = 0.01, + dt: float = 2.0, + steps: int = 100, + seed: int = 42, + ) -> None: + """Run NVT Langevin MD in place on ``atoms`` (Maxwell-Boltzmann IC at T). + + Reuses this calculator's model and neighbor-list config, so no second + model is loaded. ``friction`` is in 1/fs, ``dt`` in fs, ``temperature`` + in K. ``FixAtoms`` constraints are honored. + """ + if getattr(self, 'energy_only', False): + # 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, + seed=seed, device=self.device, dtype=self.dtype, + max_neighbors=self.max_neighbors, + ) diff --git a/mcpy/calculators/alchemi_calculator.py b/mcpy/calculators/alchemi_calculator.py index 41508cf..9004f63 100644 --- a/mcpy/calculators/alchemi_calculator.py +++ b/mcpy/calculators/alchemi_calculator.py @@ -14,13 +14,12 @@ _build_nl, _load_model, _make_batch, - _make_multi_batch, _per_graph_energies, - _run_langevin_md, + _MDMixin, ) -class AlchemiCalculator: +class AlchemiCalculator(_MDMixin): """ Energy-only Alchemi calculator (no geometry relaxation). @@ -147,41 +146,9 @@ def get_potential_energies( out: List[np.ndarray] = [] 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 = _make_batch(chunk, self.device, self.dtype) nl_hook = NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) _build_nl(batch, nl_hook) model_out = self._forward(batch) out.append(_per_graph_energies(model_out['energy'], len(chunk))) return np.concatenate(out) - - def run_md( - self, - atoms: Atoms, - *, - temperature: float, - friction: float = 0.01, - dt: float = 2.0, - steps: int = 100, - seed: int = 42, - ) -> None: - """Run NVT Langevin MD in place on ``atoms`` (Maxwell-Boltzmann IC at T). - - Reuses this calculator's model and neighbor-list config, so no second - 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, - seed=seed, device=self.device, dtype=self.dtype, - max_neighbors=self.max_neighbors, - ) diff --git a/mcpy/calculators/alchemi_f_calculator.py b/mcpy/calculators/alchemi_f_calculator.py index 2900f92..17313d6 100644 --- a/mcpy/calculators/alchemi_f_calculator.py +++ b/mcpy/calculators/alchemi_f_calculator.py @@ -15,15 +15,14 @@ from ..utils.chunking import chunk_ranges from ._alchemi_common import ( - _build_nl, + _bootstrap, _fixed_indices, _freeze_hook_for, _load_model, _make_batch, - _make_multi_batch, _per_graph_energies, - _run_langevin_md, - _write_back_positions, + _prepare_batch, + _MDMixin, ) @@ -32,7 +31,7 @@ _ALCHEMI_OPTIMIZERS = {'fire': AlchemiFIRE, 'fire2': AlchemiFIRE2} -class AlchemiFCalculator: +class AlchemiFCalculator(_MDMixin): """ Alchemi calculator with FIRE geometry relaxation. @@ -117,6 +116,12 @@ def get_potential_energy(self, atoms: Atoms) -> float: """ Relax with Alchemi FIRE, then return the relaxed potential energy. + One structure is a batch of one: this delegates to the same batched + relaxation the multi-structure path uses, so the two cannot use + different FIRE strategies. (They did: this path called ``opt.run``, + which steps the whole batch until every graph converges, while the + batched path retires each graph at its own first convergence.) + Parameters ---------- atoms : ase.Atoms @@ -129,57 +134,13 @@ def get_potential_energy(self, atoms: Atoms) -> float: float Relaxed potential energy in eV. """ - batch = _make_batch(atoms, self.device, self.dtype) - - # compute() writes via copy_() — pre-allocate target tensors - batch.forces = torch.zeros_like(batch.positions) - batch.energy = torch.zeros(1, 1, device=self.device, dtype=self.dtype) - - 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( - model=self.model, - dt=self.dt, - convergence_hook=ConvergenceHook.from_fmax(self.fmax), - n_steps=self.steps, - # 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) - return float(batch.energy.sum().item()) + return float(self.get_potential_energies([atoms])[0]) def _relax_batch(self, atoms_list: List[Atoms]) -> tuple: """Batched FIRE relaxation of one (sub-)batch. Mutates positions in place; returns ``(per_graph_energies, step_count)``.""" - n_graphs = len(atoms_list) - batch = _make_multi_batch(atoms_list, self.device, self.dtype) - batch.forces = torch.zeros_like(batch.positions) - batch.energy = torch.zeros(n_graphs, 1, device=self.device, dtype=self.dtype) + batch = _make_batch(atoms_list, self.device, self.dtype) + _prepare_batch(batch, len(atoms_list), self.device, self.dtype) # Map each graph's FixAtoms indices into the concatenated (sub-)batch. fixed: List[int] = [] @@ -198,12 +159,7 @@ def _relax_batch(self, atoms_list: List[Atoms]) -> tuple: 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 - + _bootstrap(opt, batch, nl_hook, fixed) return self._run_compacted(opt, batch, atoms_list) def _run_compacted(self, opt, batch, atoms_list: List[Atoms]) -> tuple: @@ -320,26 +276,3 @@ def get_potential_energies( len(atoms_list), len(steps), self.last_relax_steps, self.steps, self.fmax) return np.concatenate(out) - - def run_md( - self, - atoms: Atoms, - *, - temperature: float, - friction: float = 0.01, - dt: float = 2.0, - steps: int = 100, - seed: int = 42, - ) -> None: - """Run NVT Langevin MD in place on ``atoms`` (Maxwell-Boltzmann IC at T). - - Reuses this calculator's model and neighbor-list config, so no second - model is loaded. ``friction`` is in 1/fs, ``dt`` in fs, ``temperature`` - in K. ``FixAtoms`` constraints are honored. - """ - _run_langevin_md( - self.model, self._nl_config, atoms, - temperature=temperature, friction=friction, dt=dt, steps=steps, - seed=seed, device=self.device, dtype=self.dtype, - max_neighbors=self.max_neighbors, - ) From e390e9064a615f5746f0857e17dee795d9cf5a5a Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 13:49:19 +0200 Subject: [PATCH 14/20] Deduplicate the ensembles' bookkeeping Findings E2 to E8, minus two the audit got wrong. E2: the seven-column replica table was typed four times (header and rows, in each replica-exchange driver), and ReplicaExchange printed its whole parameter banner twice, once through logger.info and once to the outfile. Now re_header, re_row, RE_RULE and re_banner. The two drivers label their first column differently (MPI numbers ranks, the batched driver numbers replicas), so that column is a parameter and both outputs are byte-identical to before. E3: CanonicalEnsemble kept the current energy in three places at once: _current_energy, atoms.info['key_value_pairs']['potential_energy'] (an ASE-GA convention nothing else here reads), and lowest_energy beside the base class's _best_energy. relax() now returns the energy instead of stashing it in atoms.info, trial_step compares against _current_energy, and lowest_energy is gone in favour of the base _best_energy. set_state no longer has to reconstruct the info dict so a swapped configuration carries a baseline. E4: _write_global_minimum existed twice; only picking the winner differed (MPI gather vs local min), so the open/write/except tail moves to base_ensemble.write_global_minimum. E5: _write_initial_row re-typed write_outfile's format string to emit N/A placeholders. Both call _row now. E6: consolidate_logging reached for the replicas' logger by hardcoded dotted path. It asks BaseEnsemble.__module__ instead, so moving the file cannot silently break the silencing. The audit suggested dropping the feature and logging at DEBUG, which would change what every serial run prints; not worth it for a string. E7: BaseEnsemble.__del__, a third cleanup path after finalize_run and __exit__, which swallowed every exception. E8: CanonicalEnsemble(optimizer=None, move_selector=None) were optional in the signature and mandatory in fact, either default reaching a TypeError inside relax() or do_mutation(). Now required. Callers all pass them by keyword. NOT applied: the audit also called _exchange_pairs a one-caller helper. It has a test (test_re_exchange_pairs_alternate_by_offset), which the audit missed by grepping only mcpy/; it stays. 199 torch-free tests pass, the GPU golden test passes 14/14 with unchanged goldens, and the Sphinx build is warning-free. --- mcpy/ensembles/base_ensemble.py | 22 ++-- mcpy/ensembles/batched_replica_exchange.py | 42 +++---- mcpy/ensembles/canonical_ensemble.py | 42 +++---- mcpy/ensembles/grand_canonical_ensemble.py | 17 ++- mcpy/ensembles/replica_exchange.py | 139 ++++++++++----------- tests/test_canonical_nvt.py | 4 +- 6 files changed, 117 insertions(+), 149 deletions(-) diff --git a/mcpy/ensembles/base_ensemble.py b/mcpy/ensembles/base_ensemble.py index 37aaf42..9c68dd7 100644 --- a/mcpy/ensembles/base_ensemble.py +++ b/mcpy/ensembles/base_ensemble.py @@ -176,14 +176,6 @@ def close_files(self) -> None: self.logger.exception("Error closing minima file %s", self._minima_file) self._minima_handle = None - def __del__(self): - # Safety net only. Deterministic cleanup goes through finalize_run / - # the context manager. - try: - self.close_files() - except Exception: - pass - def compute_energy(self, atoms: Atoms) -> float: return self._calculator.get_potential_energy(atoms) @@ -261,6 +253,20 @@ def run(self, steps: int) -> None: self.finalize_run() +def write_global_minimum(path, atoms, energy, score, rank, logger): + """Write one XYZ frame for the best configuration seen across replicas. + + Shared by ``ReplicaExchange`` (which gathers over MPI ranks) and + ``BatchedReplicaExchange`` (which picks the local minimum); only finding + the winner differs between them. + """ + try: + with open(path, 'w') as fh: + write_xyz(atoms, energy, fh, extra=f'rank={rank} score={score:.6f}') + except OSError: + logger.exception("Error writing global minimum to %s", path) + + def write_xyz(atoms, energy, file_or_path, extra: str = ''): """ Write an XYZ frame to an open file handle or a path (append mode). diff --git a/mcpy/ensembles/batched_replica_exchange.py b/mcpy/ensembles/batched_replica_exchange.py index a7a3a9c..ea6bce9 100644 --- a/mcpy/ensembles/batched_replica_exchange.py +++ b/mcpy/ensembles/batched_replica_exchange.py @@ -32,7 +32,8 @@ import numpy as np from ..utils.random_number_generator import RandomNumberGenerator -from .base_ensemble import write_xyz +from .base_ensemble import BaseEnsemble, write_global_minimum +from .replica_exchange import RE_RULE, re_header, re_row logger = logging.getLogger(__name__) @@ -131,11 +132,13 @@ def __init__( # ------------------------------------------------------------------ run def run(self) -> None: - import logging as _logging - quiet = _logging.getLogger('mcpy.ensembles.base_ensemble') + # The replicas' lifecycle logs all go to their base class's module + # logger; ask the class for it rather than repeating its dotted path, + # which a file move would silently break. + quiet = logging.getLogger(BaseEnsemble.__module__) prior_level = quiet.level if self._consolidate_logging: - quiet.setLevel(_logging.WARNING) + quiet.setLevel(logging.WARNING) for r in self.replicas: r.initialize_run() @@ -182,17 +185,9 @@ def _write_global_minimum(self) -> None: if not candidates: return best = min(candidates, key=lambda r: r._best_score) - rank = self.replicas.index(best) - try: - with open(self._global_minimum_file, 'w') as fh: - write_xyz( - best._best_atoms, best._best_energy, fh, - extra=f'rank={rank} score={best._best_score:.6f}', - ) - except OSError: - self.logger.exception( - "Error writing global minimum to %s", self._global_minimum_file - ) + write_global_minimum(self._global_minimum_file, best._best_atoms, + best._best_energy, best._best_score, + self.replicas.index(best), self.logger) def _rebatch_initial_energies(self) -> None: """Re-evaluate replica energies as one batch so all replicas start @@ -391,12 +386,8 @@ def _initialize_outfile(self) -> None: if self.mus is not None: for k, mu in enumerate(self.mus): fh.write(f"Chemical potentials replica {k}: {mu}\n") - fh.write("{:<8} {:<10} {:<25} {:<15} {:<35} {:<20} {:<25}\n".format( - "Replica", "Step", "Atom Count (by species)", "Energy (eV)", - "Chemical Potentials (eV)", "Temperature (K)", - "Accepted Exchange (%)", - )) - fh.write("-" * 140 + "\n") + fh.write(re_header("Replica", 8) + "\n") + fh.write(RE_RULE + "\n") self._outfile_handle = open(self._outfile, 'a') except OSError: self.logger.exception("Error opening output file %s", self._outfile) @@ -496,12 +487,9 @@ def _write_outfile(self, step: int) -> None: mu_str = ', '.join(f"{s}: {v:.3f}" for s, v in r._mu.items()) attempts = self.exchange_attempts[i] pct = (self.exchange_successes[i] / attempts * 100) if attempts else 0.0 - self._outfile_handle.write( - "{:<8} {:<10} {:<25} {:<15.6f} {:<35} {:<20} {:<25.2f}\n".format( - i, r._step, count_str, r.E_old, mu_str, - r._temperature, pct, - ) - ) + self._outfile_handle.write(re_row( + i, r._step, count_str, r.E_old, mu_str, r._temperature, + pct, width=8) + "\n") self._outfile_handle.flush() except (OSError, AttributeError): self.logger.exception("Error writing to %s", self._outfile) diff --git a/mcpy/ensembles/canonical_ensemble.py b/mcpy/ensembles/canonical_ensemble.py index 3dfa27a..1345932 100644 --- a/mcpy/ensembles/canonical_ensemble.py +++ b/mcpy/ensembles/canonical_ensemble.py @@ -2,7 +2,6 @@ import time import numpy as np -from ase import Atoms from ase.units import kB as boltzmann_constant from .base_ensemble import BaseEnsemble @@ -24,12 +23,12 @@ class CanonicalEnsemble(BaseEnsemble): def __init__(self, atoms, calculator, + optimizer, + move_selector, cells=None, random_seed=None, - optimizer=None, fmax=0.1, temperature=300, - move_selector=None, constraints=None, traj_file: str = 'trajectory.xyz', traj_mode: str = 'w', @@ -57,7 +56,6 @@ def __init__(self, # promises as much). self._rng_acceptance = RandomNumberGenerator(seed=self._random_seed + 2) - self.lowest_energy = float('inf') self._current_energy = None self.atoms = atoms self.constraints = constraints @@ -88,11 +86,10 @@ def set_state(self, state): # (and thus _beta) across an exchange; restoring them would collapse # the temperature ladder. self.atoms = state["atoms"] + # The energy baseline the next trial_step compares against. It travels + # in the state dict rather than in atoms.info, which would have to + # survive being pickled across MPI. self._current_energy = state["energy"] - # Guarantee the next trial_step reads a valid energy baseline after a - # 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"] # Restored only when present, so the restart round-trip is lossless; # ReplicaExchange strips these slot-bound keys before a swap -- see # GrandCanonicalEnsemble.set_state. @@ -110,15 +107,12 @@ def _acceptance_condition(self, potential_diff: float) -> bool: p = np.exp(-potential_diff / (boltzmann_constant * self._temperature)) return p > self._rng_acceptance.get_uniform() - def relax(self, atoms) -> Atoms: - atoms.info['key_value_pairs'] = {} + def relax(self, atoms) -> float: + """Relax ``atoms`` in place and return its relaxed potential energy.""" atoms.calc = self._calculator opt = self._optimizer(atoms, logfile=None) opt.run(fmax=self._fmax) - - Epot = atoms.get_potential_energy() - atoms.info['key_value_pairs']['potential_energy'] = Epot - return atoms + return atoms.get_potential_energy() def do_mutation(self): new_atoms = self.atoms.copy() @@ -135,15 +129,10 @@ def trial_step(self): if new_atoms is None: return 0 - new_atoms = self.relax(new_atoms) - - potential_i = self.atoms.info['key_value_pairs']['potential_energy'] - potential_f = new_atoms.info['key_value_pairs']['potential_energy'] - potential_diff = potential_f - potential_i + potential_f = self.relax(new_atoms) + potential_diff = potential_f - self._current_energy if self._acceptance_condition(potential_diff): - if potential_f < self.lowest_energy: - self.lowest_energy = potential_f self.atoms = new_atoms self._current_energy = potential_f self.move_selector.acceptance_counter() @@ -159,9 +148,7 @@ def initialize_run(self) -> None: self.logger.info("Canonical Ensemble Monte Carlo starting " "(T=%s K, outfile=%s)", self._temperature, self._outfile) - self.relax(self.atoms) - self._current_energy = self.atoms.get_potential_energy() - self.lowest_energy = self._current_energy + self._current_energy = self.relax(self.atoms) self.write_coordinates(self.atoms, self._current_energy) self.write_outfile(self._step, self._current_energy) self._record_minimum(self.atoms, self._current_energy) @@ -175,9 +162,10 @@ def _run(self) -> None: if self._step % self._outfile_write_interval == 0: self.write_outfile(self._step, self._current_energy) - self.logger.debug("step=%d E=%s lowest_E=%s accepted=%d t=%.2fs", - self._step, self._current_energy, self.lowest_energy, - self._accepted_trials, self._last_step_seconds) + self.logger.debug("step=%d E=%s best_E=%s accepted=%d t=%.2fs", + self._step, self._current_energy, + self._best_energy, self._accepted_trials, + self._last_step_seconds) if self._step % self._trajectory_write_interval == 0: # Current accepted configuration (unchanged on a rejected step). diff --git a/mcpy/ensembles/grand_canonical_ensemble.py b/mcpy/ensembles/grand_canonical_ensemble.py index 5e4e590..ef79b46 100644 --- a/mcpy/ensembles/grand_canonical_ensemble.py +++ b/mcpy/ensembles/grand_canonical_ensemble.py @@ -139,6 +139,11 @@ def get_outfile_metadata(self) -> str: + "\n" + "-" * self._table_width + "\n" ) + def _row(self, ratio_str: str) -> str: + return "{:<10} {:<10} {:<15.6f} {:<{w}}".format( + self._step, self.n_atoms, self.E_old, ratio_str, + w=self._ratio_col_width) + def write_outfile(self, step: int = None, energy: float = None) -> None: """Write one row: step, N, energy, per-interval acceptance ratios. @@ -162,13 +167,7 @@ def write_outfile(self, step: int = None, energy: float = None) -> None: for r in acceptance_ratios ) try: - self._outfile_handle.write("{:<10} {:<10} {:<15.6f} {:<{w}}\n".format( - self._step, - self.n_atoms, - self.E_old, - ratio_str, - w=self._ratio_col_width, - )) + self._outfile_handle.write(self._row(ratio_str) + "\n") self._outfile_handle.flush() self._last_logged_step = self._step except (OSError, AttributeError): @@ -350,9 +349,7 @@ def _write_initial_row(self) -> None: return try: placeholder = ", ".join("N/A" for _ in self.move_selector.move_list) - self._outfile_handle.write("{:<10} {:<10} {:<15.6f} {:<{w}}\n".format( - 0, self.n_atoms, self.E_old, placeholder, w=self._ratio_col_width, - )) + self._outfile_handle.write(self._row(placeholder) + "\n") self._outfile_handle.flush() self._last_logged_step = 0 except (OSError, AttributeError): diff --git a/mcpy/ensembles/replica_exchange.py b/mcpy/ensembles/replica_exchange.py index 5412ef5..0c34415 100644 --- a/mcpy/ensembles/replica_exchange.py +++ b/mcpy/ensembles/replica_exchange.py @@ -8,10 +8,54 @@ MPI = None from ..utils import RandomNumberGenerator -from .base_ensemble import write_xyz +from .base_ensemble import write_global_minimum logger = logging.getLogger(__name__) +# One replica table, printed by the console log and the outfile of both +# replica-exchange drivers. The first column differs between them (MPI numbers +# ranks, the batched driver numbers replicas), so it is a parameter; everything +# after it is shared. +RE_COLUMNS = ("Step", "Atom Count (by species)", "Energy (eV)", + "Chemical Potentials (eV)", "Temperature (K)", + "Accepted Exchange (%)") +_RE_TAIL_HEADER = "{:<10} {:<25} {:<15} {:<35} {:<20} {:<25}" +_RE_TAIL_ROW = "{:<10} {:<25} {:<15.6f} {:<35} {:<20} {:<25.2f}" +RE_RULE = "-" * 140 + + +def re_header(label="Rank", width=5): + return f"{label:<{width}} " + _RE_TAIL_HEADER.format(*RE_COLUMNS) + + +def re_row(index, *values, width=5): + return f"{index:<{width}} " + _RE_TAIL_ROW.format(*values) + + +def re_banner(gcmc_steps, exchange_interval, atom_count, temperatures, mus, + n_ranks): + """The parameter block both drivers print, as a list of lines.""" + lines = [ + "+-------------------------------------------------+", + "| Replica Exchange Monte Carlo Simulation |", + "+-------------------------------------------------+", + "Simulation Parameters:", + f"Total GCMC steps: {gcmc_steps}", + f"Exchange interval (steps): {exchange_interval}", + f"Number of atoms in initial configuration: {atom_count}", + ] + if temperatures is not None: + lines.append(f"Temperatures (K): {temperatures}") + else: + lines.append("Temperatures: Not specified (default)") + if mus is not None: + lines += [f"Chemical potentials: Rank {i} - {mu}" + for i, mu in enumerate(mus)] + else: + lines.append("Chemical potentials: Not specified (default)") + lines += [f"Number of MPI ranks: {n_ranks}", re_header(), RE_RULE] + return lines + class ReplicaExchange: def __init__(self, @@ -281,41 +325,20 @@ def _write_global_minimum(self) -> None: if not candidates: return rank, (score, energy, atoms) = min(candidates, key=lambda kv: kv[1][0]) - try: - with open(self._global_minimum_file, 'w') as fh: - write_xyz(atoms, energy, fh, - extra=f'rank={rank} score={score:.6f}') - except OSError: - self.logger.exception( - "Error writing global minimum to %s", self._global_minimum_file - ) + write_global_minimum(self._global_minimum_file, atoms, energy, score, + rank, self.logger) def initialize_run(self): - self.logger.info("+-------------------------------------------------+") - self.logger.info("| Replica Exchange Monte Carlo Simulation |") - self.logger.info("+-------------------------------------------------+") - self.logger.info("Simulation Parameters:") - self.logger.info(f"Total GCMC steps: {self.gcmc_steps}") - self.logger.info(f"Exchange interval (steps): {self.exchange_interval}") - - atom_count = len(self.gcmc.get_state()['atoms']) - self.logger.info(f"Number of atoms in initial configuration: {atom_count}") - - if hasattr(self, 'temperatures') and self.temperatures is not None: - self.logger.info(f"Temperatures (K): {self.temperatures}") - else: - self.logger.info("Temperatures: Not specified (default)") - if hasattr(self, 'mus') and self.mus is not None: - for i, mus in enumerate(self.mus): - self.logger.info(f"Chemical potentials: Rank {i} - {mus}") - else: - self.logger.info("Chemical potentials: Not specified (default)") - self.logger.info(f"Number of MPI ranks: {self.size}") - self.logger.info("{:<5} {:<10} {:<25} {:<15} {:<35} {:<20} {:<25}".format( - "Rank", "Step", "Atom Count (by species)", "Energy (eV)", "Chemical Potentials (eV)", - "Temperature (K)", "Accepted Exchange (%)" - )) - self.logger.info("-" * 140) + for line in self._banner(): + self.logger.info(line) + + def _banner(self): + return re_banner( + self.gcmc_steps, self.exchange_interval, + len(self.gcmc.get_state()['atoms']), + getattr(self, 'temperatures', None), getattr(self, 'mus', None), + self.size, + ) def summarize_states(self, step): states = self.comm.gather(self.gcmc.get_state(), root=0) @@ -360,10 +383,9 @@ def summarize_states(self, step): "temperature": temperature, "accepted_percentage": accepted_percentage } - self.logger.info("{:<5} {:<10} {:<25} {:<15.6f} {:<35} {:<20} {:<25.2f}".format( - i, gcmc_step, atom_count_by_species, energy, chemical_potential_str, - temperature, accepted_percentage - )) + self.logger.info(re_row( + i, gcmc_step, atom_count_by_species, energy, + chemical_potential_str, temperature, accepted_percentage)) return summary def initialize_outfile(self): @@ -372,31 +394,7 @@ def initialize_outfile(self): return try: with open(self._outfile, 'w') as outfile: - outfile.write("+-------------------------------------------------+\n") - outfile.write("| Replica Exchange Monte Carlo Simulation |\n") - outfile.write("+-------------------------------------------------+\n") - outfile.write("Simulation Parameters:\n") - outfile.write(f"Total GCMC steps: {self.gcmc_steps}\n") - outfile.write(f"Exchange interval (steps): {self.exchange_interval}\n") - - atom_count = len(self.gcmc.get_state()['atoms']) - outfile.write(f"Number of atoms in initial configuration: {atom_count}\n") - - if hasattr(self, 'temperatures') and self.temperatures is not None: - outfile.write(f"Temperatures (K): {self.temperatures}\n") - else: - outfile.write("Temperatures: Not specified (default)\n") - if hasattr(self, 'mus') and self.mus is not None: - for i, mus in enumerate(self.mus): - outfile.write(f"Chemical potentials: Rank {i} - {mus}\n") - else: - outfile.write("Chemical potentials: Not specified (default)\n") - outfile.write(f"Number of MPI ranks: {self.size}\n") - outfile.write("{:<5} {:<10} {:<25} {:<15} {:<35} {:<20} {:<25}\n".format( - "Rank", "Step", "Atom Count (by species)", "Energy (eV)", - "Chemical Potentials (eV)", "Temperature (K)", "Accepted Exchange (%)" - )) - outfile.write("-" * 140 + "\n") + outfile.write('\n'.join(self._banner()) + '\n') self._outfile_handle = open(self._outfile, 'a') except OSError: self.logger.exception("Error opening output file %s", self._outfile) @@ -409,17 +407,10 @@ def write_outfile(self, step: int) -> None: return try: for rank, state in summary.items(): - self._outfile_handle.write( - "{:<5} {:<10} {:<25} {:<15.6f} " - "{:<35} {:<20} {:<25.2f}\n".format( - rank, state["step"], - state["atom_count_by_species"], - state["energy"], - state["chemical_potential"], - state["temperature"], - state["accepted_percentage"] - ) - ) + self._outfile_handle.write(re_row( + rank, state["step"], state["atom_count_by_species"], + state["energy"], state["chemical_potential"], + state["temperature"], state["accepted_percentage"]) + "\n") self._outfile_handle.flush() except (OSError, AttributeError): self.logger.exception("Error writing summary to file %s", self._outfile) diff --git a/tests/test_canonical_nvt.py b/tests/test_canonical_nvt.py index 81a2bec..6b7a3b7 100644 --- a/tests/test_canonical_nvt.py +++ b/tests/test_canonical_nvt.py @@ -106,7 +106,7 @@ def test_minima_file_append_only_records_improvements(tmp_path): # any further frames must each correspond to a strict improvement. assert n_frames >= 1 assert mc._best_atoms is not None - assert mc._best_score == pytest.approx(mc.lowest_energy) + assert mc._best_score == pytest.approx(mc._best_energy) def test_minima_file_overwrite_keeps_single_frame(tmp_path): @@ -141,7 +141,5 @@ def test_get_set_state_roundtrip(tmp_path): mc.set_state(state) assert mc._current_energy == saved_energy - kvp = mc.atoms.info['key_value_pairs']['potential_energy'] - assert kvp == saved_energy finally: mc.finalize_run() From 935af966e944887fb41a572df8f2797563e2394a Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 14:03:36 +0200 Subject: [PATCH 15/20] Share the cells' region tests and use ase's constants Findings 6, 9, 16 and 18. 6: the free-volume sampler's covered-point test (one cKDTree per distinct radius, nearest atom within it) was written three times. Now Cell._covered_mask, with positions and radii overridable for the periodic-image expansion CustomCell needs. The cells still sample their own points: that is where they genuinely differ, and it is what consumes their RNG, so the draw order is untouched. 9: get_atoms_specie_inside_cell repeated the same symbol-mask and str/list normalisation in four cells, differing only in the region predicate. One base implementation now calls a per-cell _inside_mask(atoms); the box cell answers "everywhere" and each region cell keeps its own test, comments included. 16: the metal-unit constants were written out by hand (Planck, Boltzmann, amu->kg, sqrt(e)) at 6 to 7 significant digits. They come from ase.units now, which is already a hard dependency. The values shift by about 1e-6 relative, which moves Lambda^3 by 4e-6, so the acceptance prefactor changes in its sixth digit; the golden fixture's decisions are unaffected. 18: the weighted move pick was cumsum + searchsorted + an index clamp. It is random.choices(cum_weights=...), which bisects the same cumulative weights against one random() draw. Verified equivalent over 20k draws: identical selections and identical final RNG state, so no run changes. Three findings in this group are NOT applied, on evidence: - 17, replacing the hand-built quaternion rotation with scipy.spatial.transform.Rotation.random: that takes a numpy Generator and rejects a python Random ("SeedSequence expects int or sequence of ints"), while the move draws from its own seeded RandomNumberGenerator. Applying it would need a second generator per move, so a single --seed would no longer determine the run. The twelve hand-written lines buy reproducibility. - 10, replacing write_xyz with ase.io extxyz: ase writes a different comment line (Properties= and pbc= present, energy last, different float formatting), so every trajectory this package has ever produced would change shape. That is a user-visible output format change for 40 lines, and the writer has tests. Not worth it. - 20, retiring NullCell by making BaseMove.cell optional: NullCell is exported, tested, and reads as a deliberate null object; replacing it spreads None checks through five moves to save 35 lines. 199 torch-free tests pass, the GPU golden test passes 14/14 with unchanged goldens. --- mcpy/cell/cell.py | 51 +++++++++++++++++++++++++++++---- mcpy/cell/custom_cell.py | 30 +++---------------- mcpy/cell/dome_cell.py | 36 ++++------------------- mcpy/cell/spherical_cell.py | 36 +++-------------------- mcpy/moves/move_selector.py | 19 ++++++------ mcpy/utils/set_unit_constant.py | 17 +++++++---- 6 files changed, 80 insertions(+), 109 deletions(-) diff --git a/mcpy/cell/cell.py b/mcpy/cell/cell.py index e0c45a4..26ee99d 100644 --- a/mcpy/cell/cell.py +++ b/mcpy/cell/cell.py @@ -1,4 +1,5 @@ import numpy as np +from scipy.spatial import cKDTree from .base_cell import BaseCell @@ -48,15 +49,27 @@ def get_volume(self): """ return self.volume - def get_atoms_specie_inside_cell(self, atoms, species): + def get_atoms_specie_inside_cell(self, atoms, specie): + """Indices of atoms whose symbol is in ``specie`` and that lie inside + the exchangeable region. + + ``specie`` may be a single symbol or a list. The region test is + :meth:`_inside_mask`, which the box cell answers with "everywhere"; + the region cells override it. """ - Get the indices of atoms of a specific species inside the cell. + if len(atoms) == 0: + return np.empty(0, dtype=int) + symbols = np.asarray(atoms.get_chemical_symbols()) + species_list = [specie] if isinstance(specie, str) else list(specie) + species_mask = np.isin(symbols, species_list) + return np.where(species_mask & self._inside_mask(atoms))[0] - :param atoms: ASE Atoms object containing the atomic configuration. - :param species: List of species to filter. - :return: Indices of atoms of the specified species inside the cell. + def _inside_mask(self, atoms): + """Per-atom boolean: is this atom in the exchangeable region? + + The box cell spans the whole periodic cell, so every atom qualifies. """ - return np.where(np.isin(atoms.get_chemical_symbols(), species))[0] + return np.ones(len(atoms), dtype=bool) def get_species(self): """ @@ -88,6 +101,32 @@ def _radii_for(self, atoms): return np.fromiter((self.species_radii[s] for s in symbols), dtype=float, count=len(symbols)) + def _covered_mask(self, points, atoms, positions=None, radii=None): + """Which of ``points`` fall inside some atom's exclusion sphere. + + One cKDTree per distinct radius, querying the nearest atom of that + radius: a point is covered when that nearest atom is within it. Avoids + the O(N_points * N_atoms) broadcast. Shared by every cell that + estimates a free volume; the cells differ in how they *sample* the + points, not in how they test them. + + ``positions`` and ``radii`` override the ones taken from ``atoms``, for + the periodic-image expansion :class:`CustomCell` needs. + """ + if radii is None: + radii = self._radii_for(atoms) + if positions is None: + positions = atoms.positions + covered = np.zeros(len(points), dtype=bool) + for r in np.unique(radii): + if r <= 0.0: + continue + tree = cKDTree(positions[radii == r]) + # distance_upper_bound returns inf above r; that is the rejection. + dists, _ = tree.query(points, k=1, distance_upper_bound=float(r)) + covered |= np.isfinite(dists) + return covered + def is_point_inside(self, point): """The box cell spans the whole periodic cell: every point is inside. diff --git a/mcpy/cell/custom_cell.py b/mcpy/cell/custom_cell.py index b36c704..93b958c 100644 --- a/mcpy/cell/custom_cell.py +++ b/mcpy/cell/custom_cell.py @@ -1,5 +1,4 @@ import numpy as np -from scipy.spatial import cKDTree from .cell import Cell @@ -75,16 +74,7 @@ def calculate_volume(self, atoms) -> float: positions = self._periodic_images(atoms) radii = np.tile(self._radii_for(atoms), len(positions) // n_atoms) - - covered = np.zeros(self.mc_sample_points, dtype=bool) - for r in np.unique(radii): - if r <= 0.0: - continue - mask = radii == r - tree = cKDTree(positions[mask]) - dists, _ = tree.query(cart_coords, k=1, distance_upper_bound=float(r)) - covered |= np.isfinite(dists) - + covered = self._covered_mask(cart_coords, atoms, positions, radii) occupied_fraction = float(np.count_nonzero(covered)) / self.mc_sample_points self.volume = self._clamp_free_volume( self.cell_volume * (1.0 - occupied_fraction), self.cell_volume) @@ -105,10 +95,8 @@ def _periodic_images(self, atoms): ] return np.concatenate(images) - def get_atoms_specie_inside_cell(self, atoms, specie): - """ - Vectorized: indices of atoms with symbol in ``specie`` that are - exchangeable with the reservoir. + def _inside_mask(self, atoms): + """Exchangeable with the reservoir. In xy an atom must lie within the cell footprint. In z the upper bound is intentionally dropped: atoms at or above the cell floor count as @@ -118,19 +106,9 @@ def get_atoms_specie_inside_cell(self, atoms, specie): Atoms below the floor (absorbed into the subsurface layers) stay excluded so they are kept. """ - if len(atoms) == 0: - return np.empty(0, dtype=int) - symbols = np.asarray(atoms.get_chemical_symbols()) - if isinstance(specie, str): - species_list = [specie] - else: - species_list = list(specie) - species_mask = np.isin(symbols, species_list) frac = (atoms.positions - self.offset) @ self._dim_inv in_xy = np.all((frac[:, :2] >= 0.0) & (frac[:, :2] < 1.0), axis=1) - above_floor = frac[:, 2] >= 0.0 - inside = in_xy & above_floor - return np.where(species_mask & inside)[0] + return in_xy & (frac[:, 2] >= 0.0) def is_point_inside(self, point): """ diff --git a/mcpy/cell/dome_cell.py b/mcpy/cell/dome_cell.py index d51d735..c7aa5db 100644 --- a/mcpy/cell/dome_cell.py +++ b/mcpy/cell/dome_cell.py @@ -1,5 +1,4 @@ import numpy as np -from scipy.spatial import cKDTree from .cell import Cell from .spherical_cell import SphericalCell @@ -82,25 +81,13 @@ def get_random_point(self): if point[2] >= self.bottom_z: return point - def get_atoms_specie_inside_cell(self, atoms, specie): - """ - Vectorized: indices of atoms whose symbol is in ``specie`` and whose - position is inside the dome (within the ball and above the surface). - Atoms that drifted below the surface (absorbed into the support) are - excluded, so they are never selected for deletion. - """ - if len(atoms) == 0: - return np.empty(0, dtype=int) - symbols = np.asarray(atoms.get_chemical_symbols()) - if isinstance(specie, str): - species_list = [specie] - else: - species_list = list(specie) - species_mask = np.isin(symbols, species_list) + def _inside_mask(self, atoms): + """Within the ball and at or above the support surface. Atoms that + drifted below it (absorbed into the support) are excluded, so they are + never selected for deletion.""" diff = atoms.positions - self.center within = np.einsum('ij,ij->i', diff, diff) <= self.radius ** 2 - above = atoms.positions[:, 2] >= self.bottom_z - return np.where(species_mask & within & above)[0] + return within & (atoms.positions[:, 2] >= self.bottom_z) def calculate_volume(self, atoms): """ @@ -120,18 +107,7 @@ def calculate_volume(self, atoms): self.volume = dome_fraction * self.sphere_volume return - radii = self._radii_for(atoms) - positions = atoms.positions - - covered = np.zeros(len(pts), dtype=bool) - for r in np.unique(radii): - if r <= 0.0: - continue - mask = radii == r - tree = cKDTree(positions[mask]) - dists, _ = tree.query(pts, k=1, distance_upper_bound=float(r)) - covered |= np.isfinite(dists) - + covered = self._covered_mask(pts, atoms) free_fraction = float(np.count_nonzero(~covered)) / self.mc_sample_points self.volume = self._clamp_free_volume( free_fraction * self.sphere_volume, self.sphere_volume) diff --git a/mcpy/cell/spherical_cell.py b/mcpy/cell/spherical_cell.py index 94231b8..9537383 100644 --- a/mcpy/cell/spherical_cell.py +++ b/mcpy/cell/spherical_cell.py @@ -1,5 +1,4 @@ import numpy as np -from scipy.spatial import cKDTree from .cell import Cell @@ -61,22 +60,10 @@ def get_random_point(self): r = self.radius * (self._rng.random() ** (1.0 / 3.0)) return self.center + r * direction - def get_atoms_specie_inside_cell(self, atoms, specie): - """ - Vectorized: indices of atoms whose symbol is in ``specie`` and whose - position is inside the cell. - """ - if len(atoms) == 0: - return np.empty(0, dtype=int) - symbols = np.asarray(atoms.get_chemical_symbols()) - if isinstance(specie, str): - species_list = [specie] - else: - species_list = list(specie) - species_mask = np.isin(symbols, species_list) + def _inside_mask(self, atoms): + """Inside the sphere.""" diff = atoms.positions - self.center - inside = np.einsum('ij,ij->i', diff, diff) <= self.radius ** 2 - return np.where(species_mask & inside)[0] + return np.einsum('ij,ij->i', diff, diff) <= self.radius ** 2 def _sample_sphere_points(self, n): """ @@ -114,22 +101,7 @@ def calculate_volume(self, atoms): pts = self._sample_sphere_points(self.mc_sample_points) - radii = self._radii_for(atoms) - positions = atoms.positions - - covered = np.zeros(self.mc_sample_points, dtype=bool) - # One kdtree per unique radius; query nearest atom of that radius. - # A point is covered if the nearest such atom is within r. - for r in np.unique(radii): - if r <= 0.0: - continue - mask = radii == r - sub_pos = positions[mask] - tree = cKDTree(sub_pos) - # distance_upper_bound returns inf above r; that's the rejection. - dists, _ = tree.query(pts, k=1, distance_upper_bound=float(r)) - covered |= np.isfinite(dists) - + covered = self._covered_mask(pts, atoms) free_fraction = float(np.count_nonzero(~covered)) / self.mc_sample_points self.volume = self._clamp_free_volume( free_fraction * self.sphere_volume, self.sphere_volume) diff --git a/mcpy/moves/move_selector.py b/mcpy/moves/move_selector.py index 9d06fda..47fa07a 100644 --- a/mcpy/moves/move_selector.py +++ b/mcpy/moves/move_selector.py @@ -57,9 +57,10 @@ def __init__(self, probabilities, move_list, seed=None, n_moves=None): if n_moves < 1: raise ValueError(f"n_moves must be >= 1, got {n_moves}") self.n_moves = int(n_moves) - self.rho = np.asarray(np.cumsum(probabilities), dtype=float) - self._rho_total = float(self.rho[-1]) - self._n_moves_idx = len(self.rho) + # Cumulative weights for random.choices; it bisects them against one + # random() draw, which is exactly what the previous cumsum + + # searchsorted did, on the same stream. + self.rho = list(np.cumsum(probabilities, dtype=float)) self.rng = RandomNumberGenerator(seed=seed) n = len(probabilities) self.move_counter = [0] * n @@ -70,18 +71,16 @@ def __init__(self, probabilities, move_list, seed=None, n_moves=None): self.move_failed_counter_total = [0] * n self.to_use = 0 - def __get_index__(self): - v = self.rng.get_uniform() * self._rho_total - i = int(np.searchsorted(self.rho, v, side='right')) - if i >= self._n_moves_idx: - i = self._n_moves_idx - 1 - return i + def _pick_index(self): + """Index of the next move, sampled by weight.""" + return self.rng.random.choices(range(len(self.rho)), + cum_weights=self.rho, k=1)[0] def do_trial_move(self, atoms): """Choose a move and apply it. Tracks attempt and (if the move returned the ``False`` sentinel) records the failure so ratios use viable attempts.""" - self.to_use = self.__get_index__() + self.to_use = self._pick_index() self.move_counter[self.to_use] += 1 self.move_counter_total[self.to_use] += 1 result = self.move_list[self.to_use].do_trial_move(atoms) diff --git a/mcpy/utils/set_unit_constant.py b/mcpy/utils/set_unit_constant.py index e2d2370..173f149 100644 --- a/mcpy/utils/set_unit_constant.py +++ b/mcpy/utils/set_unit_constant.py @@ -1,5 +1,6 @@ import numpy as np from typing import List +from ase import units from ase.data import atomic_masses, atomic_numbers @@ -63,11 +64,17 @@ def _set_lj_units(self) -> None: self.lambda_dbs = {specie: 1 for specie in names} def _set_metal_units(self) -> None: - """Set units for metal potential.""" - self.PLANCK_CONSTANT = 4.13567e-15 # eV/s - self.BOLTZMANN_CONSTANT = 8.617333e-5 # eV/K - self.mass_conversion_factor = 1.66053906660e-27 # amu to kg - self.lambda_conversion_factor = np.sqrt(1.60218e-19) * 1e10 + """Set units for metal potential. + + The constants come from ``ase.units`` rather than being written out + here: ase is already a hard dependency, and its CODATA values keep the + de Broglie wavelengths consistent with every other energy in the run. + """ + self.PLANCK_CONSTANT = units._hplanck / units._e # eV s + self.BOLTZMANN_CONSTANT = units.kB # eV/K + self.mass_conversion_factor = units._amu # amu to kg + # sqrt(J -> eV) folded with m -> Angstrom, so lambda comes out in A. + self.lambda_conversion_factor = np.sqrt(units._e) * 1e10 self.beta = 1/(self.temperature*self.BOLTZMANN_CONSTANT) self.masses = {specie: atomic_masses[atomic_numbers[specie]] for specie in self.species} From 966eb72c8c74e956e447b6b2c6fcbd66089ff0f9 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 14:23:49 +0200 Subject: [PATCH 16/20] Replace the Alchemi _common module with a base class _alchemi_common.py exported twelve underscore-prefixed names, of which only four were used by both calculators (_load_model, _make_batch, _per_graph_energies, _MDMixin). Five belonged solely to AlchemiFCalculator and one solely to AlchemiCalculator. "Common" was a location, not a concept, and things drifted into it: a private module that two sibling modules and two benchmark scripts all imported from. What the two calculators actually shared was state and a template, which free functions cannot absorb. Their __init__ bodies agreed on 19 of 29/39 lines (the whole signature plus every assignment) and their get_potential_energies on 18 (the entire chunking loop, differing only in what happens per chunk). Both survived the previous dedupe pass for exactly that reason. So: one module, mcpy/calculators/alchemi.py, with _AlchemiBase holding the shared state, the batching and neighbour-list helpers, the bootstrap, run_md, get_potential_energy and the chunked get_potential_energies template. Each subclass now implements only _evaluate_chunk: one forward pass, or a batched FIRE relaxation. 726 lines across three files become 646 in one, and the 26 lines of cross-module private imports are gone. The subclasses keep their full explicit signatures rather than *args/**kwargs: the parameters are the public API, positional callers must keep working, and autodoc has to be able to render them. Two guards kept asymmetric on purpose, since a base class makes it tempting to unify them: AlchemiCalculator refuses energy_only with a shared pre-loaded wrapper and overrides run_md to reject a forces-less model, while AlchemiFCalculator rejects one in its constructor instead. Verified: GPU golden test 14/14 with unchanged goldens, verify_compact_parity PASS (max |dE| 0.0003 eV, fresh max|F| within fmax), 199 torch-free tests, Sphinx warning-free. The reference page now documents the base class too. benchmark/internal/verify_compact_parity.py and verify_compile_parity.py import the moved names; both updated, though that directory is untracked. --- docs/reference/calculators.rst | 9 +- mcpy/calculators/__init__.py | 3 +- mcpy/calculators/_alchemi_common.py | 294 ----------- mcpy/calculators/alchemi.py | 646 +++++++++++++++++++++++ mcpy/calculators/alchemi_calculator.py | 154 ------ mcpy/calculators/alchemi_f_calculator.py | 278 ---------- 6 files changed, 654 insertions(+), 730 deletions(-) delete mode 100644 mcpy/calculators/_alchemi_common.py create mode 100644 mcpy/calculators/alchemi.py delete mode 100644 mcpy/calculators/alchemi_calculator.py delete mode 100644 mcpy/calculators/alchemi_f_calculator.py diff --git a/docs/reference/calculators.rst b/docs/reference/calculators.rst index 6e5d39e..0fda577 100644 --- a/docs/reference/calculators.rst +++ b/docs/reference/calculators.rst @@ -19,8 +19,13 @@ exported only when its backend is installed. .. autoclass:: mcpy.calculators.mace_f_calculator.MACE_F_Calculator :members: -.. autoclass:: mcpy.calculators.alchemi_calculator.AlchemiCalculator +.. autoclass:: mcpy.calculators.alchemi._AlchemiBase :members: -.. autoclass:: mcpy.calculators.alchemi_f_calculator.AlchemiFCalculator +.. autoclass:: mcpy.calculators.alchemi.AlchemiCalculator :members: + :show-inheritance: + +.. autoclass:: mcpy.calculators.alchemi.AlchemiFCalculator + :members: + :show-inheritance: diff --git a/mcpy/calculators/__init__.py b/mcpy/calculators/__init__.py index 5a82df6..1052496 100644 --- a/mcpy/calculators/__init__.py +++ b/mcpy/calculators/__init__.py @@ -10,8 +10,7 @@ pass try: - from .alchemi_calculator import AlchemiCalculator - from .alchemi_f_calculator import AlchemiFCalculator + from .alchemi import AlchemiCalculator, AlchemiFCalculator __all__ += ['AlchemiCalculator', 'AlchemiFCalculator'] except ImportError: # nvalchemi is an optional backend; skip if not installed. diff --git a/mcpy/calculators/_alchemi_common.py b/mcpy/calculators/_alchemi_common.py deleted file mode 100644 index 403319f..0000000 --- a/mcpy/calculators/_alchemi_common.py +++ /dev/null @@ -1,294 +0,0 @@ -from __future__ import annotations - -import os -from typing import List, Union - -import numpy as np -import torch -from ase import Atoms -from ase.constraints import FixAtoms - -from nvalchemi._typing import AtomCategory -from nvalchemi.data import AtomicData -from nvalchemi.data.batch import Batch -from nvalchemi.models.mace import MACEWrapper -from nvalchemi.hooks.neighbor_list import NeighborListHook -from nvalchemi.hooks import DynamicsContext -from nvalchemi.dynamics import NVTLangevin, initialize_velocities -from nvalchemi.dynamics.base import DynamicsStage -from nvalchemi.dynamics.hooks import FreezeAtomsHook, NaNDetectorHook - - -class _HeadMACEWrapper(MACEWrapper): - """MACEWrapper that pins a multihead MACE model to one head. - - nvalchemi's batch carries no ``head`` field, so the inner MACE model falls - back to head 0. Fine-tuned models often keep the pretrain head at index 0 - and the fine-tuned head elsewhere, so head 0 is the wrong potential. This - injects a fixed head index into the MACE input. - """ - - def __init__(self, model: torch.nn.Module, head_index: int) -> None: - super().__init__(model) - self._head_index = head_index - - def adapt_input(self, data, **kwargs): - d = super().adapt_input(data, **kwargs) - d['head'] = torch.full((data.num_graphs,), self._head_index, - dtype=torch.long, device=data.positions.device) - return d - - -def _load_model( - checkpoint: Union[str, MACEWrapper], - device: str, - dtype: torch.dtype, - enable_cueq: bool, - compile_model: bool, - 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 - # 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. - if isinstance(checkpoint, (str, os.PathLike)) and os.path.exists(checkpoint): - raw = torch.load(checkpoint, map_location=device, weights_only=False).to(dtype) - # Resolve the head index before any cuEq conversion, while ``raw.heads`` - # is guaranteed present. - idx = None if head is None else ( - head if isinstance(head, int) else list(raw.heads).index(head) - ) - if enable_cueq: - try: - import cuequivariance # noqa: F401 - from mace.cli.convert_e3nn_cueq import run as _to_cueq - except ImportError as exc: - raise ImportError( - "enable_cueq=True requires the 'cuequivariance' package; " - "install with: pip install 'nvalchemi-toolkit[mace]'" - ) from exc - raw = _to_cueq(raw, return_model=True, device=device) - raw = raw.to(device) - if compile_model: - # Mirror MACEWrapper.from_checkpoint step 3: the model is - # inference-only after this. The e3nn patch is a private nvalchemi - # helper; benchmark/verify_compile_parity.py pins it so an upgrade - # that renames it fails loudly. - from nvalchemi.models.mace import _patch_e3nn_irrep_len_for_compile - _patch_e3nn_irrep_len_for_compile() - raw.eval() - for param in raw.parameters(): - param.requires_grad = False - raw = torch.compile(raw) - 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; " - f"got alias checkpoint {checkpoint!r}" - ) - return MACEWrapper.from_checkpoint( - checkpoint, - device=torch.device(device), - dtype=dtype, - enable_cueq=enable_cueq, - compile_model=compile_model, - ) - - -def _make_batch(atoms_list, device: str, dtype: torch.dtype) -> Batch: - """Batch one ``Atoms`` or a list of them; a single structure is a batch of one.""" - if isinstance(atoms_list, Atoms): - atoms_list = [atoms_list] - data = [AtomicData.from_atoms(a, device=device, dtype=dtype) for a in atoms_list] - return Batch.from_data_list(data, device=device) - - -def _per_graph_energies(out_energy: torch.Tensor, n_graphs: int) -> np.ndarray: - """Reduce model output to a (n_graphs,) numpy array regardless of layout.""" - e = out_energy.detach().to('cpu') - if e.numel() == n_graphs: - return e.view(-1).numpy() - raise RuntimeError( - f"Unexpected energy tensor shape {tuple(out_energy.shape)} for {n_graphs} graphs. " - "Batched eval expects one energy per graph." - ) - - -def _build_nl(batch: Batch, nl_hook: NeighborListHook) -> None: - ctx = DynamicsContext(batch=batch, step_count=0) - nl_hook(ctx, DynamicsStage.BEFORE_COMPUTE) - - -def _prepare_batch(batch: Batch, n_graphs: int, device: str, - dtype: torch.dtype) -> None: - """Pre-allocate the tensors ``compute()`` writes into via ``copy_()``.""" - batch.forces = torch.zeros_like(batch.positions) - batch.energy = torch.zeros(n_graphs, 1, device=device, dtype=dtype) - - -def _bootstrap(opt, batch: Batch, nl_hook: NeighborListHook, - fixed: List[int]) -> None: - """Build the neighbor list and the initial forces before the loop starts. - - ``FreezeAtomsHook`` zeros frozen forces only at ``AFTER_POST_UPDATE``, and - this ``compute()`` runs outside the hook loop, so step 0 would otherwise - displace the frozen rows once. - """ - _build_nl(batch, nl_hook) - opt.compute(batch) - if fixed: - batch.forces[fixed] = 0.0 - - -def _write_back_positions(atoms: Atoms, batch: Batch) -> None: - """Copy relaxed positions from batch back to atoms, skipping FixAtoms-constrained indices.""" - # float64: the batch is float32, and assigning the float64 frozen rows - # into a float32 array would truncate them, leaving the "restored" atoms - # about 5e-7 A off their pinned positions at slab-sized coordinates. - relaxed = batch.positions.detach().cpu().numpy().astype(np.float64) - fixed = _fixed_indices(atoms) - if fixed: - relaxed[fixed] = atoms.positions[fixed] - atoms.positions = relaxed - - -def _fixed_indices(atoms: Atoms) -> List[int]: - """Indices held by ASE ``FixAtoms`` constraints, sorted and de-duplicated.""" - idx: set[int] = set() - for c in atoms.constraints: - if isinstance(c, FixAtoms): - idx.update(int(i) for i in c.index) - return sorted(idx) - - -def _freeze_hook_for(batch: Batch, fixed: List[int]) -> List[FreezeAtomsHook]: - """Tag ``fixed`` batch rows as SPECIAL and return the hook that holds them. - - nvalchemi's FIRE has no notion of ASE ``FixAtoms``; without this the - "fixed" atoms relax freely and are only snapped back afterward, leaving the - returned energy inconsistent with the stored geometry. Marking them - ``AtomCategory.SPECIAL`` and registering :class:`FreezeAtomsHook` zeros their - forces/velocities and restores positions every FIRE step. - """ - if not fixed: - return None - batch.atom_categories[fixed] = AtomCategory.SPECIAL.value - return [FreezeAtomsHook()] - - -def _run_langevin_md( - model: MACEWrapper, - nl_config, - atoms: Atoms, - *, - temperature: float, - friction: float, - dt: float, - steps: int, - seed: int, - device: str, - dtype: torch.dtype, - max_neighbors: int | None, -) -> None: - """Run NVT Langevin (BAOAB) MD in place on ``atoms`` using ``model``. - - Mirrors the FIRE bootstrap in :meth:`AlchemiFCalculator.get_potential_energy`: - pre-allocate the tensors ``compute()`` writes into, seed velocities from a - Maxwell-Boltzmann draw at ``temperature``, build the neighbor list once, then - run ``steps`` of Langevin dynamics. ``FixAtoms``-constrained atoms are held - every step by :class:`FreezeAtomsHook` and restored on write-back. - - ``dt`` is in fs, ``friction`` in 1/fs, ``temperature`` in K (the integrator - converts to internal units). Mutates ``atoms`` in place; returns nothing. - """ - batch = _make_batch(atoms, device, dtype) - _prepare_batch(batch, 1, device, dtype) - batch.velocities = torch.zeros_like(batch.positions) - temp = torch.full((batch.num_graphs,), float(temperature), device=device, dtype=dtype) - initialize_velocities( - batch.velocities, batch.atomic_masses, temp, batch.batch_idx.int(), - random_seed=seed, - ) - - fixed = _fixed_indices(atoms) - freeze_hooks = _freeze_hook_for(batch, fixed) - nl_hook = NeighborListHook(nl_config, max_neighbors=max_neighbors) - opt = NVTLangevin( - model=model, - dt=dt, - temperature=temperature, - friction=friction, - random_seed=seed, - n_steps=steps, - # 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(opt, batch, nl_hook, fixed) - opt.run(batch) - _write_back_positions(atoms, batch) - - -class _MDMixin: - """``run_md`` for both Alchemi calculators. - - Identical in each, docstring included, so it lives here. The energy_only - guard applies only to :class:`AlchemiCalculator`; the F variant has no such - attribute and rejects a forces-less model in its constructor instead. - """ - - def run_md( - self, - atoms: Atoms, - *, - temperature: float, - friction: float = 0.01, - dt: float = 2.0, - steps: int = 100, - seed: int = 42, - ) -> None: - """Run NVT Langevin MD in place on ``atoms`` (Maxwell-Boltzmann IC at T). - - Reuses this calculator's model and neighbor-list config, so no second - model is loaded. ``friction`` is in 1/fs, ``dt`` in fs, ``temperature`` - in K. ``FixAtoms`` constraints are honored. - """ - if getattr(self, 'energy_only', False): - # 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, - seed=seed, device=self.device, dtype=self.dtype, - max_neighbors=self.max_neighbors, - ) diff --git a/mcpy/calculators/alchemi.py b/mcpy/calculators/alchemi.py new file mode 100644 index 0000000..9bdbf2a --- /dev/null +++ b/mcpy/calculators/alchemi.py @@ -0,0 +1,646 @@ +"""GPU-native MACE calculators on the nvalchemi stack. + +Two calculators over one base: :class:`AlchemiCalculator` evaluates energies +with a single forward pass, :class:`AlchemiFCalculator` relaxes with FIRE first. +Everything they share -- loading and configuring the model, batching, the +neighbour-list hook, the chunked batched-evaluation template, Langevin MD -- +lives on :class:`_AlchemiBase`, so a subclass only says what is different about +its per-chunk evaluation. + +Both are optional: importing this module needs ``nvalchemi-toolkit[mace]``. +""" +from __future__ import annotations + +import logging +import os +from typing import List, Union + +import numpy as np +import torch +from ase import Atoms +from ase.constraints import FixAtoms + +from nvalchemi._typing import AtomCategory +from nvalchemi.data import AtomicData +from nvalchemi.data.batch import Batch +from nvalchemi.dynamics import (FIRE as AlchemiFIRE, FIRE2 as AlchemiFIRE2, + ConvergenceHook, NVTLangevin, + initialize_velocities) +from nvalchemi.dynamics.base import DynamicsStage +from nvalchemi.dynamics.hooks import FreezeAtomsHook, NaNDetectorHook +from nvalchemi.hooks import DynamicsContext +from nvalchemi.hooks.neighbor_list import NeighborListHook +from nvalchemi.models.mace import MACEWrapper + +from ..utils.chunking import chunk_ranges + +logger = logging.getLogger(__name__) + +_ALCHEMI_OPTIMIZERS = {'fire': AlchemiFIRE, 'fire2': AlchemiFIRE2} + + +# --------------------------------------------------------------- model loading + +class _HeadMACEWrapper(MACEWrapper): + """MACEWrapper that pins a multihead MACE model to one head. + + nvalchemi's batch carries no ``head`` field, so the inner MACE model falls + back to head 0. Fine-tuned models often keep the pretrain head at index 0 + and the fine-tuned head elsewhere, so head 0 is the wrong potential. This + injects a fixed head index into the MACE input. + """ + + def __init__(self, model: torch.nn.Module, head_index: int) -> None: + super().__init__(model) + self._head_index = head_index + + def adapt_input(self, data, **kwargs): + d = super().adapt_input(data, **kwargs) + d['head'] = torch.full((data.num_graphs,), self._head_index, + dtype=torch.long, device=data.positions.device) + return d + + +def _load_model( + checkpoint: Union[str, MACEWrapper], + device: str, + dtype: torch.dtype, + enable_cueq: bool, + compile_model: bool, + 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 + # 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. + if isinstance(checkpoint, (str, os.PathLike)) and os.path.exists(checkpoint): + raw = torch.load(checkpoint, map_location=device, weights_only=False).to(dtype) + # Resolve the head index before any cuEq conversion, while ``raw.heads`` + # is guaranteed present. + idx = None if head is None else ( + head if isinstance(head, int) else list(raw.heads).index(head) + ) + if enable_cueq: + try: + import cuequivariance # noqa: F401 + from mace.cli.convert_e3nn_cueq import run as _to_cueq + except ImportError as exc: + raise ImportError( + "enable_cueq=True requires the 'cuequivariance' package; " + "install with: pip install 'nvalchemi-toolkit[mace]'" + ) from exc + raw = _to_cueq(raw, return_model=True, device=device) + raw = raw.to(device) + if compile_model: + # Mirror MACEWrapper.from_checkpoint step 3: the model is + # inference-only after this. The e3nn patch is a private nvalchemi + # helper; benchmark/verify_compile_parity.py pins it so an upgrade + # that renames it fails loudly. + from nvalchemi.models.mace import _patch_e3nn_irrep_len_for_compile + _patch_e3nn_irrep_len_for_compile() + raw.eval() + for param in raw.parameters(): + param.requires_grad = False + raw = torch.compile(raw) + 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; " + f"got alias checkpoint {checkpoint!r}" + ) + return MACEWrapper.from_checkpoint( + checkpoint, + device=torch.device(device), + dtype=dtype, + enable_cueq=enable_cueq, + compile_model=compile_model, + ) + + +# ------------------------------------------------------------ batch/atoms glue + +def _make_batch(atoms_list, device: str, dtype: torch.dtype) -> Batch: + """Batch one ``Atoms`` or a list of them; a single structure is a batch of one.""" + if isinstance(atoms_list, Atoms): + atoms_list = [atoms_list] + data = [AtomicData.from_atoms(a, device=device, dtype=dtype) for a in atoms_list] + return Batch.from_data_list(data, device=device) + + +def _per_graph_energies(out_energy: torch.Tensor, n_graphs: int) -> np.ndarray: + """Reduce model output to a (n_graphs,) numpy array regardless of layout.""" + e = out_energy.detach().to('cpu') + if e.numel() == n_graphs: + return e.view(-1).numpy() + raise RuntimeError( + f"Unexpected energy tensor shape {tuple(out_energy.shape)} for {n_graphs} graphs. " + "Batched eval expects one energy per graph." + ) + + +def _build_nl(batch: Batch, nl_hook: NeighborListHook) -> None: + """Run the neighbour-list hook once, outside any dynamics loop.""" + ctx = DynamicsContext(batch=batch, step_count=0) + nl_hook(ctx, DynamicsStage.BEFORE_COMPUTE) + + +def _fixed_indices(atoms: Atoms) -> List[int]: + """Indices held by ASE ``FixAtoms`` constraints, sorted and de-duplicated.""" + idx: set[int] = set() + for c in atoms.constraints: + if isinstance(c, FixAtoms): + idx.update(int(i) for i in c.index) + return sorted(idx) + + +def _freeze_hook_for(batch: Batch, fixed: List[int]) -> List[FreezeAtomsHook]: + """Tag ``fixed`` batch rows as SPECIAL and return the hook that holds them. + + nvalchemi's FIRE has no notion of ASE ``FixAtoms``; without this the + "fixed" atoms relax freely and are only snapped back afterward, leaving the + returned energy inconsistent with the stored geometry. Marking them + ``AtomCategory.SPECIAL`` and registering :class:`FreezeAtomsHook` zeros their + forces/velocities and restores positions every step. + """ + if not fixed: + return None + batch.atom_categories[fixed] = AtomCategory.SPECIAL.value + return [FreezeAtomsHook()] + + +def _write_back_positions(atoms: Atoms, batch: Batch) -> None: + """Copy positions from batch back to atoms, restoring FixAtoms rows.""" + # float64: the batch is float32, and assigning the float64 frozen rows + # into a float32 array would truncate them, leaving the "restored" atoms + # about 5e-7 A off their pinned positions at slab-sized coordinates. + relaxed = batch.positions.detach().cpu().numpy().astype(np.float64) + fixed = _fixed_indices(atoms) + if fixed: + relaxed[fixed] = atoms.positions[fixed] + atoms.positions = relaxed + + +# ------------------------------------------------------------------- base class + +class _AlchemiBase: + """Shared machinery for the Alchemi calculators. + + Holds the model and the device/dtype/chunking state, the batching and + neighbour-list helpers, the chunked ``get_potential_energies`` template, and + Langevin MD. A subclass implements :meth:`_evaluate_chunk`, which is the + only thing that differs between evaluating and relaxing. + + Parameters + ---------- + checkpoint : str | MACEWrapper + Named checkpoint (e.g. 'medium-mpa-0'), local .pt path, or a pre-loaded + MACEWrapper to share across calculators. + device : str + 'cuda' or 'cpu'. + dtype : torch.dtype + Model and data dtype. float32 recommended for speed. + enable_cueq : bool + cuEquivariance kernel fusion — significant speedup on GPU. + compile_model : bool + torch.compile the model. Adds ~30s warmup, then faster. + max_neighbors : int | None + Neighbour-list cap handed to ``NeighborListHook``. + chunk_size : int | None + Default sub-batch size for :meth:`get_potential_energies`. ``None`` + evaluates the whole list at once. + head : str | int | None + Head of a multihead model, by name or index. Local paths only. + """ + + def __init__( + self, + checkpoint: Union[str, MACEWrapper] = 'medium-mpa-0', + device: str = 'cuda', + dtype: torch.dtype = torch.float32, + enable_cueq: bool = True, + compile_model: bool = True, + max_neighbors: int | None = None, + chunk_size: int | None = None, + head: Union[str, int, None] = None, + ) -> None: + self.device = device + self.dtype = dtype + self.max_neighbors = max_neighbors + self.chunk_size = chunk_size + self.model = _load_model(checkpoint, device, dtype, enable_cueq, + compile_model, head) + self._nl_config = self.model.model_config.neighbor_config + + # -- helpers a subclass builds its evaluation from + + def _batch(self, atoms_or_list) -> Batch: + return _make_batch(atoms_or_list, self.device, self.dtype) + + def _nl_hook(self) -> NeighborListHook: + return NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) + + def _prepare(self, batch: Batch, n_graphs: int) -> None: + """Pre-allocate the tensors ``compute()`` writes into via ``copy_()``.""" + batch.forces = torch.zeros_like(batch.positions) + batch.energy = torch.zeros(n_graphs, 1, device=self.device, dtype=self.dtype) + + @staticmethod + def _bootstrap(opt, batch: Batch, nl_hook: NeighborListHook, + fixed: List[int]) -> None: + """Build the neighbour list and initial forces before the loop starts. + + ``FreezeAtomsHook`` zeros frozen forces only at ``AFTER_POST_UPDATE``, + and this ``compute()`` runs outside the hook loop, so step 0 would + otherwise displace the frozen rows once. + """ + _build_nl(batch, nl_hook) + opt.compute(batch) + if fixed: + batch.forces[fixed] = 0.0 + + # -- the public evaluation surface + + def _evaluate_chunk(self, chunk: List[Atoms]) -> np.ndarray: + """Per-graph energies for one sub-batch. Implemented by subclasses.""" + raise NotImplementedError + + def get_potential_energy(self, atoms: Atoms) -> float: + """Energy of one structure, as a batch of one. + + Delegating keeps the single-structure and batched paths from drifting + into different strategies, which is what happened before: the relaxing + calculator used to step the whole batch to full convergence here and + retire graphs individually in the batched path. + """ + return float(self.get_potential_energies([atoms])[0]) + + def get_potential_energies( + self, atoms_list: List[Atoms], chunk_size: int | None = None + ) -> np.ndarray: + """ + Per-graph energies over multiple (possibly differently sized) structures. + + One CUDA kernel launch per layer instead of N — the win behind batched + replica exchange on a single GPU. + + ``chunk_size`` splits ``atoms_list`` into consecutive sub-batches of at + most that many structures, evaluated one at a time, capping peak GPU + memory at the largest chunk instead of the whole batch. ``None`` (the + default) falls back to the instance ``chunk_size``; if that is also + ``None`` the whole list goes in one pass. Results are unaffected by + chunking (MACE message passing and FIRE convergence are both per-graph), + up to GPU run-to-run noise. + + Parameters + ---------- + atoms_list : list[ase.Atoms] + Structures to evaluate. Lengths may differ. Whether they are + mutated depends on the subclass: relaxation writes positions back. + chunk_size : int | None + Per-call override of the instance ``chunk_size``. + + Returns + ------- + np.ndarray + Potential energies in eV, shape (len(atoms_list),). + """ + if not atoms_list: + return np.empty(0, dtype=np.float64) + cs = self.chunk_size if chunk_size is None else chunk_size + out = [self._evaluate_chunk(atoms_list[start:stop]) + for start, stop in chunk_ranges(len(atoms_list), cs)] + return np.concatenate(out) + + def run_md( + self, + atoms: Atoms, + *, + temperature: float, + friction: float = 0.01, + dt: float = 2.0, + steps: int = 100, + seed: int = 42, + ) -> None: + """Run NVT Langevin (BAOAB) MD in place on ``atoms``. + + Velocities start from a Maxwell-Boltzmann draw at ``temperature``. + Reuses this calculator's model and neighbour-list config, so no second + model is loaded. ``friction`` is in 1/fs, ``dt`` in fs, ``temperature`` + in K (the integrator converts to internal units). ``FixAtoms`` + constraints are held every step and restored on write-back. + """ + batch = self._batch(atoms) + self._prepare(batch, 1) + batch.velocities = torch.zeros_like(batch.positions) + temp = torch.full((batch.num_graphs,), float(temperature), + device=self.device, dtype=self.dtype) + initialize_velocities( + batch.velocities, batch.atomic_masses, temp, batch.batch_idx.int(), + random_seed=seed, + ) + + fixed = _fixed_indices(atoms) + freeze_hooks = _freeze_hook_for(batch, fixed) + nl_hook = self._nl_hook() + opt = NVTLangevin( + model=self.model, + dt=dt, + temperature=temperature, + friction=friction, + random_seed=seed, + n_steps=steps, + # 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) + self._bootstrap(opt, batch, nl_hook, fixed) + opt.run(batch) + _write_back_positions(atoms, batch) + + +# ------------------------------------------------------------------ energy only + +class AlchemiCalculator(_AlchemiBase): + """ + Energy-only Alchemi calculator (no geometry relaxation). + + Single-point MACE evaluation on the nvalchemi GPU-native stack. Use this + when mcpy handles its own MC moves and only needs E(atoms). + + Takes the shared parameters of :class:`_AlchemiBase`, plus: + + Parameters + ---------- + energy_only : bool + Drop ``forces`` from the model's active outputs. MC energy evaluation + never uses them, and without them the MACE forward builds no autograd + graph: lower peak memory, energy unchanged up to fp32 rounding. + """ + + def __init__( + self, + checkpoint: Union[str, MACEWrapper] = 'medium-mpa-0', + device: str = 'cuda', + dtype: torch.dtype = torch.float32, + enable_cueq: bool = True, + compile_model: bool = True, + max_neighbors: int | None = None, + chunk_size: int | None = None, + energy_only: bool = False, + head: Union[str, int, None] = None, + ) -> 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 -- + # 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 + # silently and order-dependently. + raise ValueError( + "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 instead, so this " + "calculator loads its own model." + ) + super().__init__(checkpoint, device, dtype, enable_cueq, compile_model, + max_neighbors, chunk_size, head) + self.energy_only = energy_only + if energy_only: + 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 _evaluate_chunk(self, chunk: List[Atoms]) -> np.ndarray: + """One forward pass; ``chunk`` is not mutated.""" + batch = self._batch(chunk) + nl_hook = self._nl_hook() + _build_nl(batch, nl_hook) + out = self._forward(batch) + return _per_graph_energies(out['energy'], len(chunk)) + + def run_md(self, atoms: Atoms, **kwargs) -> None: + 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.' + ) + return super().run_md(atoms, **kwargs) + + +# -------------------------------------------------------------- FIRE relaxation + +class AlchemiFCalculator(_AlchemiBase): + """ + Alchemi calculator with FIRE geometry relaxation. + + Mirrors MACE_F_Calculator: relax atoms with FIRE then return energy. Uses + the fully GPU-resident nvalchemi dynamics stack — significantly faster than + ASE FIRE above ~100 atoms. + + Takes the shared parameters of :class:`_AlchemiBase`, plus: + + Parameters + ---------- + steps : int + Maximum FIRE steps before returning. + fmax : float + Force convergence threshold in eV/Å. + dt : float + 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). + """ + + def __init__( + self, + checkpoint: Union[str, MACEWrapper] = 'medium-mpa-0', + steps: int = 500, + fmax: float = 0.05, + device: str = 'cuda', + dtype: torch.dtype = torch.float32, + enable_cueq: bool = True, + compile_model: bool = True, + dt: float = 1.0, + optimizer: str = 'fire', + max_neighbors: int | None = None, + chunk_size: int | None = None, + head: Union[str, int, None] = None, + ) -> None: + self.steps = steps + self.fmax = fmax + self.dt = dt + self.last_relax_steps = 0 + self.total_relax_steps = 0 + if optimizer not in _ALCHEMI_OPTIMIZERS: + raise ValueError( + f"optimizer must be one of {list(_ALCHEMI_OPTIMIZERS)}, got {optimizer!r}" + ) + self._optimizer_cls = _ALCHEMI_OPTIMIZERS[optimizer] + self.optimizer_name = optimizer + super().__init__(checkpoint, device, dtype, enable_cueq, compile_model, + max_neighbors, chunk_size, 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.' + ) + + def get_potential_energies( + self, atoms_list: List[Atoms], chunk_size: int | None = None + ) -> np.ndarray: + """Relax every structure with batched FIRE, then return the energies. + + Each structure in ``atoms_list`` is mutated in place: positions are + updated to the relaxed configuration, respecting FixAtoms. Adds the + relax-step accounting to the inherited chunking template. + """ + self._chunk_steps: List[int] = [] + energies = super().get_potential_energies(atoms_list, chunk_size) + if self._chunk_steps: + # last = deepest chunk relaxation; total accumulates across chunks. + self.last_relax_steps = max(self._chunk_steps) + self.total_relax_steps += sum(self._chunk_steps) + logger.debug("FIRE relaxation (batched, %d graphs, %d chunk(s)): " + "max %d/%d steps (fmax=%.3g eV/A)", + len(atoms_list), len(self._chunk_steps), + self.last_relax_steps, self.steps, self.fmax) + return energies + + def _evaluate_chunk(self, chunk: List[Atoms]) -> np.ndarray: + """Batched FIRE relaxation of one sub-batch, mutating positions.""" + batch = self._batch(chunk) + self._prepare(batch, len(chunk)) + + # Map each graph's FixAtoms indices into the concatenated sub-batch. + fixed: List[int] = [] + offset = 0 + for a in chunk: + fixed.extend(offset + i for i in _fixed_indices(a)) + offset += len(a) + freeze_hooks = _freeze_hook_for(batch, fixed) + + nl_hook = self._nl_hook() + opt = self._optimizer_cls( + model=self.model, + dt=self.dt, + convergence_hook=ConvergenceHook.from_fmax(self.fmax), + n_steps=self.steps, + # 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) + self._bootstrap(opt, batch, nl_hook, fixed) + energies, n_steps = self._run_compacted(opt, batch, chunk) + self._chunk_steps.append(n_steps) + return energies + + def _run_compacted(self, opt, batch, atoms_list: List[Atoms]) -> tuple: + """Step FIRE manually, retiring each graph from the batch at its first + convergence so subsequent steps only compute the still-active graphs. + + ``opt.run`` computes the FULL batch every step and stops only when all + graphs are converged at the same step; with mixed trial moves the fast + graphs pay for the slowest one. + + Batch rows are removed with ``Batch.index_select`` and the optimizer's + per-graph FIRE state is shrunk in lockstep with + ``_sync_state_to_batch`` — the same primitives nvalchemi's inflight + refill machinery uses. Positions and energy are harvested per graph at + retirement; returns ``(per_graph_energies, step_count)``. + """ + alive = np.arange(len(atoms_list)) # batch row -> atoms_list index + energies = np.zeros(len(atoms_list)) + n_steps = 0 + opt._open_hooks() + try: + while alive.size and n_steps < self.steps: + batch, conv = opt.step(batch) + n_steps += 1 + if conv is None or conv.numel() == 0: + continue + rows = np.unique(conv.detach().cpu().numpy()) + self._harvest(batch, rows, alive, energies, atoms_list) + keep = np.setdiff1d(np.arange(alive.size), rows) + alive = alive[keep] + if alive.size: + keep_t = torch.as_tensor(keep, device=batch.device) + opt._sync_state_to_batch(keep_t, 0, batch) + batch = batch.index_select(keep_t) + # _last_converged indexes the pre-shrink batch; clear it + # 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: + opt._close_hooks() + return energies, n_steps + + @staticmethod + def _harvest(batch, rows, alive, energies, atoms_list) -> None: + """Write energy + relaxed positions of ``rows`` (batch row indices) + back to their originating Atoms objects, restoring FixAtoms rows.""" + e = _per_graph_energies(batch.energy, int(batch.num_graphs)) + # float64 before restoring frozen rows: see _write_back_positions. + pos = batch.positions.detach().cpu().numpy().astype(np.float64) + batch_idx = batch.batch_idx.detach().cpu().numpy() + for row in rows: + g = int(alive[row]) + atoms = atoms_list[g] + energies[g] = e[row] + new_pos = pos[batch_idx == row].copy() + for j in _fixed_indices(atoms): + new_pos[j] = atoms.positions[j] + atoms.positions = new_pos diff --git a/mcpy/calculators/alchemi_calculator.py b/mcpy/calculators/alchemi_calculator.py deleted file mode 100644 index 9004f63..0000000 --- a/mcpy/calculators/alchemi_calculator.py +++ /dev/null @@ -1,154 +0,0 @@ -from __future__ import annotations - -from typing import List, Union - -import numpy as np -import torch -from ase import Atoms - -from nvalchemi.models.mace import MACEWrapper -from nvalchemi.hooks.neighbor_list import NeighborListHook - -from ..utils.chunking import chunk_ranges -from ._alchemi_common import ( - _build_nl, - _load_model, - _make_batch, - _per_graph_energies, - _MDMixin, -) - - -class AlchemiCalculator(_MDMixin): - """ - Energy-only Alchemi calculator (no geometry relaxation). - - Single-point MACE evaluation on the nvalchemi GPU-native stack. - Use this when mcpy handles its own MC moves and only needs E(atoms). - - Parameters - ---------- - checkpoint : str | MACEWrapper - Named checkpoint (e.g. 'medium-mpa-0'), local .pt path, - or a pre-loaded MACEWrapper to share across calculators. - device : str - 'cuda' or 'cpu'. - dtype : torch.dtype - Model and data dtype. float32 recommended for speed. - enable_cueq : bool - cuEquivariance kernel fusion — significant speedup on GPU. - compile_model : bool - torch.compile the model. Adds ~30s warmup, then faster. - """ - - def __init__( - self, - checkpoint: Union[str, MACEWrapper] = 'medium-mpa-0', - device: str = 'cuda', - dtype: torch.dtype = torch.float32, - enable_cueq: bool = True, - compile_model: bool = True, - max_neighbors: int | None = None, - chunk_size: int | None = None, - energy_only: bool = False, - head: Union[str, int, None] = None, - ) -> 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 -- - # 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 - # silently and order-dependently. - raise ValueError( - "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 instead, so this " - "calculator loads its own model." - ) - self.device = device - self.dtype = dtype - self.max_neighbors = max_neighbors - self.chunk_size = chunk_size - self.energy_only = energy_only - self.model = _load_model(checkpoint, device, dtype, enable_cueq, compile_model, head) - self._nl_config = self.model.model_config.neighbor_config - if energy_only: - # MC energy evaluation never uses forces. Dropping 'forces' from - # active_outputs sets compute_force=False in the MACE forward, so no - # autograd graph is built — lower peak memory, energy unchanged up to - # 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. - - Parameters - ---------- - atoms : ase.Atoms - Structure to evaluate. Not mutated. - - Returns - ------- - float - Potential energy in eV. - """ - batch = _make_batch(atoms, self.device, self.dtype) - nl_hook = NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) - _build_nl(batch, nl_hook) - out = self._forward(batch) - return float(out['energy'].sum().item()) - - def get_potential_energies( - self, atoms_list: List[Atoms], chunk_size: int | None = None - ) -> np.ndarray: - """ - Batched forward pass over multiple (possibly differently sized) structures. - - One CUDA kernel launch per layer instead of N — the win behind batched - replica exchange on a single GPU. - - ``chunk_size`` splits ``atoms_list`` into consecutive sub-batches of at - most that many structures, evaluated one forward pass each, capping peak - GPU memory at the largest chunk instead of the whole batch. ``None`` (the - default) falls back to the instance ``chunk_size``; if that is also - ``None`` the whole list is evaluated in a single pass. Per-graph energies - are unaffected by chunking (MACE message passing is per-graph), up to - GPU run-to-run noise. - - Parameters - ---------- - atoms_list : list[ase.Atoms] - Structures to evaluate. Not mutated. Lengths may differ. - chunk_size : int | None - Per-call override of the instance ``chunk_size``. - - Returns - ------- - np.ndarray - Potential energies in eV, shape (len(atoms_list),). - """ - if not atoms_list: - return np.empty(0, dtype=np.float64) - cs = self.chunk_size if chunk_size is None else chunk_size - out: List[np.ndarray] = [] - for start, stop in chunk_ranges(len(atoms_list), cs): - chunk = atoms_list[start:stop] - batch = _make_batch(chunk, self.device, self.dtype) - nl_hook = NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) - _build_nl(batch, nl_hook) - 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 deleted file mode 100644 index 17313d6..0000000 --- a/mcpy/calculators/alchemi_f_calculator.py +++ /dev/null @@ -1,278 +0,0 @@ -from __future__ import annotations - -import logging -from typing import List, Union - -import numpy as np -import torch -from ase import Atoms - -from nvalchemi.models.mace import MACEWrapper -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 ( - _bootstrap, - _fixed_indices, - _freeze_hook_for, - _load_model, - _make_batch, - _per_graph_energies, - _prepare_batch, - _MDMixin, -) - - -logger = logging.getLogger(__name__) - -_ALCHEMI_OPTIMIZERS = {'fire': AlchemiFIRE, 'fire2': AlchemiFIRE2} - - -class AlchemiFCalculator(_MDMixin): - """ - Alchemi calculator with FIRE geometry relaxation. - - Mirrors MACE_F_Calculator: relax atoms with FIRE then return energy. - Uses the fully GPU-resident nvalchemi dynamics stack — - significantly faster than ASE FIRE above ~100 atoms. - - Parameters - ---------- - checkpoint : str | MACEWrapper - Named checkpoint, local .pt path, or pre-loaded MACEWrapper. - steps : int - Maximum FIRE steps before returning. - fmax : float - Force convergence threshold in eV/Å. - device : str - 'cuda' or 'cpu'. - dtype : torch.dtype - float32 recommended; float64 for higher accuracy. - enable_cueq : bool - cuEquivariance kernel fusion. - compile_model : bool - torch.compile. Best to pre-warm in __init__ if reusing. - dt : float - 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). - chunk_size : int | None - Default sub-batch size for batched relaxation (see - ``get_potential_energies``). ``None`` relaxes the whole batch at once. - """ - - def __init__( - self, - checkpoint: Union[str, MACEWrapper] = 'medium-mpa-0', - steps: int = 500, - fmax: float = 0.05, - device: str = 'cuda', - dtype: torch.dtype = torch.float32, - enable_cueq: bool = True, - compile_model: bool = True, - dt: float = 1.0, - optimizer: str = 'fire', - max_neighbors: int | None = None, - chunk_size: int | None = None, - head: Union[str, int, None] = None, - ) -> None: - self.steps = steps - self.fmax = fmax - self.device = device - self.dtype = dtype - self.dt = dt - self.max_neighbors = max_neighbors - self.chunk_size = chunk_size - self.last_relax_steps = 0 - self.total_relax_steps = 0 - if optimizer not in _ALCHEMI_OPTIMIZERS: - raise ValueError( - f"optimizer must be one of {list(_ALCHEMI_OPTIMIZERS)}, got {optimizer!r}" - ) - 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: - """ - Relax with Alchemi FIRE, then return the relaxed potential energy. - - One structure is a batch of one: this delegates to the same batched - relaxation the multi-structure path uses, so the two cannot use - different FIRE strategies. (They did: this path called ``opt.run``, - which steps the whole batch until every graph converges, while the - batched path retires each graph at its own first convergence.) - - Parameters - ---------- - atoms : ase.Atoms - Starting geometry. Mutated in place: positions are updated to - the relaxed configuration (FixAtoms rows restored). The GCMC - ensembles rely on this write-back. - - Returns - ------- - float - Relaxed potential energy in eV. - """ - return float(self.get_potential_energies([atoms])[0]) - - def _relax_batch(self, atoms_list: List[Atoms]) -> tuple: - """Batched FIRE relaxation of one (sub-)batch. Mutates positions in - place; returns ``(per_graph_energies, step_count)``.""" - batch = _make_batch(atoms_list, self.device, self.dtype) - _prepare_batch(batch, len(atoms_list), self.device, self.dtype) - - # Map each graph's FixAtoms indices into the concatenated (sub-)batch. - fixed: List[int] = [] - offset = 0 - for a in atoms_list: - fixed.extend(offset + i for i in _fixed_indices(a)) - offset += len(a) - freeze_hooks = _freeze_hook_for(batch, fixed) - - nl_hook = NeighborListHook(self._nl_config, max_neighbors=self.max_neighbors) - opt = self._optimizer_cls( - model=self.model, - dt=self.dt, - convergence_hook=ConvergenceHook.from_fmax(self.fmax), - n_steps=self.steps, - hooks=(freeze_hooks or []) + [NaNDetectorHook()], - ) - opt.register_hook(nl_hook, stage=DynamicsStage.BEFORE_COMPUTE) - _bootstrap(opt, batch, nl_hook, fixed) - return self._run_compacted(opt, batch, atoms_list) - - def _run_compacted(self, opt, batch, atoms_list: List[Atoms]) -> tuple: - """Step FIRE manually, retiring each graph from the batch at its first - convergence so subsequent steps only compute the still-active graphs. - - ``opt.run`` computes the FULL batch every step and stops only when all - graphs are converged at the same step; with mixed trial moves the fast - graphs pay for the slowest one. Retiring at first convergence also - matches the single-graph ``get_potential_energy`` semantics (a lone - graph stops at its own first convergence). - - Batch rows are removed with ``Batch.index_select`` and the optimizer's - per-graph FIRE state is shrunk in lockstep with - ``_sync_state_to_batch`` — the same primitives nvalchemi's inflight - refill machinery uses. Positions and energy are harvested per graph at - retirement; returns ``(per_graph_energies, step_count)``. - """ - alive = np.arange(len(atoms_list)) # batch row -> atoms_list index - energies = np.zeros(len(atoms_list)) - n_steps = 0 - opt._open_hooks() - try: - while alive.size and n_steps < self.steps: - batch, conv = opt.step(batch) - n_steps += 1 - if conv is None or conv.numel() == 0: - continue - rows = np.unique(conv.detach().cpu().numpy()) - self._harvest(batch, rows, alive, energies, atoms_list) - keep = np.setdiff1d(np.arange(alive.size), rows) - alive = alive[keep] - if alive.size: - keep_t = torch.as_tensor(keep, device=batch.device) - opt._sync_state_to_batch(keep_t, 0, batch) - batch = batch.index_select(keep_t) - # _last_converged indexes the pre-shrink batch; clear it - # 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: - opt._close_hooks() - return energies, n_steps - - @staticmethod - def _harvest(batch, rows, alive, energies, atoms_list) -> None: - """Write energy + relaxed positions of ``rows`` (batch row indices) - back to their originating Atoms objects, restoring FixAtoms rows.""" - e = _per_graph_energies(batch.energy, int(batch.num_graphs)) - # float64 before restoring frozen rows: see _write_back_positions. - pos = batch.positions.detach().cpu().numpy().astype(np.float64) - batch_idx = batch.batch_idx.detach().cpu().numpy() - for row in rows: - g = int(alive[row]) - atoms = atoms_list[g] - energies[g] = e[row] - new_pos = pos[batch_idx == row].copy() - for j in _fixed_indices(atoms): - new_pos[j] = atoms.positions[j] - atoms.positions = new_pos - - def get_potential_energies( - self, atoms_list: List[Atoms], chunk_size: int | None = None - ) -> np.ndarray: - """ - Batched FIRE relaxation over multiple structures, then per-graph energies. - - All graphs in a (sub-)batch share one optimizer / one model forward pass - per FIRE step. Each graph is removed from the batch at its first - convergence, so already-converged graphs stop paying for the forward - pass. Returns when every graph has converged or ``steps`` is reached. - - ``chunk_size`` splits ``atoms_list`` into consecutive sub-batches of at - most that many structures, each relaxed independently, capping peak GPU - memory at the largest chunk. ``None`` (the default) falls back to the - instance ``chunk_size``; if that is also ``None`` the whole list is - relaxed at once. FIRE dynamics are per-graph (forces and the convergence - criterion are per-graph), so a graph's relaxed energy is independent of - which others share its sub-batch, up to GPU run-to-run noise. - - Parameters - ---------- - atoms_list : list[ase.Atoms] - Starting geometries. Each is mutated in place: positions are - updated to the relaxed configuration (respecting FixAtoms). - chunk_size : int | None - Per-call override of the instance ``chunk_size``. - - Returns - ------- - np.ndarray - Relaxed potential energies in eV, shape (len(atoms_list),). - """ - if not atoms_list: - return np.empty(0, dtype=np.float64) - cs = self.chunk_size if chunk_size is None else chunk_size - out: List[np.ndarray] = [] - steps: List[int] = [] - for start, stop in chunk_ranges(len(atoms_list), cs): - energies, step_count = self._relax_batch(atoms_list[start:stop]) - out.append(energies) - steps.append(step_count) - # last = deepest chunk relaxation; total accumulates work across chunks. - self.last_relax_steps = max(steps) - self.total_relax_steps += sum(steps) - logger.debug("FIRE relaxation (batched, %d graphs, %d chunk(s)): " - "max %d/%d steps (fmax=%.3g eV/A)", - len(atoms_list), len(steps), self.last_relax_steps, - self.steps, self.fmax) - return np.concatenate(out) From 560a604d811ecacb894c2f09dc81b363a7a524d3 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 14:58:28 +0200 Subject: [PATCH 17/20] Add three CPU-only examples plus an examples index The two existing examples both need a GPU and the nvalchemi backend, so a new user could not run anything without one. These three run on any CPU with nothing but ASE installed, and cover the parts of the library the GPU examples do not: the canonical ensemble, basin hopping, and the molecular moves. They deliberately use toy potentials. EMT has no usable oxygen chemistry (O on Ag(111) either piles up without bound or never adsorbs, depending on delta_mu by 0.1 eV), so the GCMC examples use the Lennard-Jones system that the LAMMPS parity benchmark already validates, where the expected averages are known. EMT is kept for the canonical example, where Cu and Au are parameters it actually has. --- examples/00_hello_gcmc.py | 85 ++++++++++++++++++++++++++++++++++ examples/01_canonical_mc.py | 79 ++++++++++++++++++++++++++++++++ examples/02_molecule_gcmc.py | 89 ++++++++++++++++++++++++++++++++++++ examples/README.md | 22 +++++++++ 4 files changed, 275 insertions(+) create mode 100644 examples/00_hello_gcmc.py create mode 100644 examples/01_canonical_mc.py create mode 100644 examples/02_molecule_gcmc.py create mode 100644 examples/README.md diff --git a/examples/00_hello_gcmc.py b/examples/00_hello_gcmc.py new file mode 100644 index 0000000..5d8d4b0 --- /dev/null +++ b/examples/00_hello_gcmc.py @@ -0,0 +1,85 @@ +"""Smallest possible GCMC run: a Lennard-Jones fluid in a periodic box. + +Runs on a laptop CPU in a few seconds. No GPU, no MACE, no downloads. +This is the textbook GCMC test case, and the same system mcpy is validated +against LAMMPS ``fix gcmc`` with (see ``benchmark/README.md``), so the average +particle number it settles at is a real, checkable number. + + python examples/00_hello_gcmc.py + +Writes ``hello_gcmc.out`` (step / N / energy / acceptance ratios) and +``hello_gcmc.xyz`` (trajectory). + +Once this makes sense, read ``examples/gcmc_ag111.py`` for the same loop +driven by a machine-learned potential on a real surface. +""" +from ase import Atoms +from ase.calculators.lj import LennardJones + +from mcpy.cell import Cell +from mcpy.ensembles import GrandCanonicalEnsemble +from mcpy.moves import DeletionMove, DisplacementMove, InsertionMove, MoveSelector +from mcpy.utils.logging import configure as configure_logging + +# LJ reduced units: sigma = epsilon = 1. mcpy's 'LJ' unit system sets the +# de Broglie wavelength to 1, so mu is the reduced chemical potential. +BOX = 9.0 # box side, in sigma +CUTOFF = 3.0 # pair cutoff, in sigma +TEMPERATURE = 2.0 # reduced temperature +MU = -4.0 # reduced chemical potential of the reservoir +STEPS = 3000 + +# One seed per stochastic component. Leaving any of them at None makes the +# run unrepeatable. +SEED_INS, SEED_DEL, SEED_DIS, SEED_SEL, SEED_ENS = 1, 2, 3, 4, 5 + + +def main(): + configure_logging() + + # 1. An empty periodic box. GCMC fills it from the reservoir. + atoms = Atoms(cell=[BOX, BOX, BOX], pbc=True) + + # 2. The region where particles may be inserted or deleted. ``Cell`` is + # the whole periodic box. + cell = Cell(atoms) + + # 3. The calculator. Any ASE calculator works directly; wrap it in + # ``mcpy.calculators.BaseCalculator`` instead if each trial structure + # should be relaxed before its energy is used. + calculator = LennardJones(sigma=1.0, epsilon=1.0, rc=CUTOFF, smooth=False) + + # 4. The trial moves, sampled with equal weight. Insertion and deletion + # change N; displacement equilibrates the particles already present. + move_selector = MoveSelector( + [1, 1, 1], + [InsertionMove(cell, species=['H'], seed=SEED_INS), + DeletionMove(cell, species=['H'], seed=SEED_DEL), + DisplacementMove(species=['H'], max_displacement=0.3, seed=SEED_DIS)], + seed=SEED_SEL, + ) + + gcmc = GrandCanonicalEnsemble( + atoms=atoms, + cells=[cell], + calculator=calculator, + mu={'H': MU}, + species=['H'], + units_type='LJ', + temperature=TEMPERATURE, + move_selector=move_selector, + random_seed=SEED_ENS, + outfile='hello_gcmc.out', + traj_file='hello_gcmc.xyz', + outfile_write_interval=100, + trajectory_write_interval=100, + ) + gcmc.run(STEPS) + + # The observable of a GCMC run is an average over the equilibrated part + # of the trajectory, never the last frame. + print(f'final N = {len(gcmc.atoms)}, final E = {gcmc.E_old:.3f}') + + +if __name__ == '__main__': + main() diff --git a/examples/01_canonical_mc.py b/examples/01_canonical_mc.py new file mode 100644 index 0000000..17f37c6 --- /dev/null +++ b/examples/01_canonical_mc.py @@ -0,0 +1,79 @@ +"""Canonical (NVT) Monte Carlo: chemical ordering in a CuAu nanoparticle. + +Runs on a laptop CPU in about a minute with ASE's built-in EMT potential, +which is a genuine (if simple) potential for Cu and Au. + +The particle count never changes here. Two moves explore the configuration +space instead: ``PermutationMove`` swaps a Cu and an Au atom, and +``DisplacementMove`` nudges atoms. Every trial is relaxed with LBFGS before +the Metropolis test, so this doubles as a basin-hopping search -- +``minima_file`` collects every configuration that beats the previous best. + + python examples/01_canonical_mc.py + +Writes ``canonical.out``, ``canonical.xyz`` and ``canonical_minima.xyz``. +""" +import numpy as np +from ase.calculators.emt import EMT +from ase.cluster import Icosahedron +from ase.optimize import LBFGS + +from mcpy.ensembles import CanonicalEnsemble +from mcpy.moves import DisplacementMove, MoveSelector, PermutationMove +from mcpy.utils.logging import configure as configure_logging + +TEMPERATURE = 500 +STEPS = 300 + +# One seed per stochastic component. Leaving any of them at None makes the +# run unrepeatable. +SEED_STRUCTURE, SEED_PERM, SEED_DIS, SEED_SEL, SEED_ENS = 0, 1, 2, 3, 4 + + +def build_particle(): + """A 55-atom icosahedron with half its atoms turned into Au at random.""" + atoms = Icosahedron('Cu', noshells=3) + rng = np.random.default_rng(SEED_STRUCTURE) + au = rng.choice(len(atoms), size=len(atoms) // 2, replace=False) + symbols = np.array(atoms.get_chemical_symbols()) + symbols[au] = 'Au' + atoms.set_chemical_symbols(symbols.tolist()) + return atoms + + +def main(): + configure_logging() + + atoms = build_particle() + + # PermutationMove changes the chemical order; DisplacementMove relaxes + # the geometry around it. The permutation is sampled three times as often + # because ordering is what this run is about. + move_selector = MoveSelector( + [3, 1], + [PermutationMove(species=['Cu', 'Au'], seed=SEED_PERM), + DisplacementMove(species=['Cu', 'Au'], max_displacement=0.1, seed=SEED_DIS)], + seed=SEED_SEL, + ) + + # CanonicalEnsemble relaxes each trial itself, so it takes a plain ASE + # calculator plus an ASE optimizer class (not an instance). + mc = CanonicalEnsemble( + atoms=atoms, + calculator=EMT(), + optimizer=LBFGS, + fmax=0.1, + move_selector=move_selector, + temperature=TEMPERATURE, + random_seed=SEED_ENS, + outfile='canonical.out', + traj_file='canonical.xyz', + minima_file='canonical_minima.xyz', + outfile_write_interval=10, + trajectory_write_interval=10, + ) + mc.run(STEPS) + + +if __name__ == '__main__': + main() diff --git a/examples/02_molecule_gcmc.py b/examples/02_molecule_gcmc.py new file mode 100644 index 0000000..414ce82 --- /dev/null +++ b/examples/02_molecule_gcmc.py @@ -0,0 +1,89 @@ +"""Molecular GCMC: rigid Lennard-Jones dimers exchanged with a reservoir. + +Same box and potential as ``00_hello_gcmc.py``, but the exchanged species is +a whole molecule instead of a single atom. This is the setup mcpy's molecular +moves are validated against LAMMPS ``fix gcmc`` with (``benchmark/README.md``). + +Three molecular moves appear here: + * ``MoleculeInsertionMove`` places the rigid template at a random position + with a random orientation, + * ``MoleculeDeletionMove`` removes a whole molecule at once, + * ``MoleculeDisplacementMove`` translates and rotates one molecule rigidly. + +Membership is tracked by a per-atom ``molecule_id`` array that survives +rejection and is written into the trajectory, so a run can restart from its +own ``.xyz`` with molecules intact. + + python examples/02_molecule_gcmc.py + +Writes ``molecule_gcmc.out`` and ``molecule_gcmc.xyz``. For a real adsorbate +(CO, H2O, O2) on a surface, swap the template for an ASE-built molecule and +the calculator for a machine-learned potential; see ``docs/molecular_adsorbates.rst``. +""" +from ase import Atoms +from ase.calculators.lj import LennardJones + +from mcpy.cell import Cell +from mcpy.ensembles import GrandCanonicalEnsemble +from mcpy.moves import (MoleculeDeletionMove, MoleculeDisplacementMove, + MoleculeInsertionMove, MoveSelector) +from mcpy.utils.logging import configure as configure_logging + +BOX = 9.0 # box side, in sigma +CUTOFF = 3.0 # pair cutoff, in sigma +TEMPERATURE = 2.0 # reduced temperature +# mu is the chemical potential of the WHOLE molecule: the energy of one +# molecule in the reservoir plus its translational, rotational and pressure +# contributions. It is not the per-atom value. +MU = -6.0 +STEPS = 1000 + +SEED_INS, SEED_DEL, SEED_DIS, SEED_SEL, SEED_ENS = 1, 2, 3, 4, 5 + + +def main(): + configure_logging() + + atoms = Atoms(cell=[BOX, BOX, BOX], pbc=True) + cell = Cell(atoms) + calculator = LennardJones(sigma=1.0, epsilon=1.0, rc=CUTOFF, smooth=False) + + # The template is centered on its center of mass when the move is built; + # its name ('H2') is the key used in ``mu`` and in ``molecules``. + template = Atoms('H2', positions=[[0, 0, 0], [0, 0, 1.0]]) + + move_selector = MoveSelector( + [1, 1, 1], + [MoleculeInsertionMove(cell, template, 'H2', seed=SEED_INS), + MoleculeDeletionMove(cell, template, 'H2', seed=SEED_DEL), + MoleculeDisplacementMove(cell, template, 'H2', max_displacement=0.3, + seed=SEED_DIS)], + seed=SEED_SEL, + ) + + gcmc = GrandCanonicalEnsemble( + atoms=atoms, + cells=[cell], + calculator=calculator, + mu={'H2': MU}, + # ``species`` lists ATOMIC exchangeable species only; leave it empty + # when every exchanged species is molecular. + species=[], + molecules={'H2': template}, + units_type='LJ', + temperature=TEMPERATURE, + move_selector=move_selector, + random_seed=SEED_ENS, + outfile='molecule_gcmc.out', + traj_file='molecule_gcmc.xyz', + outfile_write_interval=100, + trajectory_write_interval=100, + ) + gcmc.run(STEPS) + + print(f'final: {len(gcmc.atoms) // len(template)} molecules, ' + f'E = {gcmc.E_old:.3f}') + + +if __name__ == '__main__': + main() diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..546b9c4 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,22 @@ +# Examples + +Run them in order. The first three need nothing but `pip install -e .` and a CPU. + +| Example | What it shows | Needs | Runtime | +| --- | --- | --- | --- | +| `00_hello_gcmc.py` | The GCMC loop: cell, insertion/deletion/displacement moves, chemical potential, output files. Lennard-Jones fluid in a box. | ASE only | ~40 s | +| `01_canonical_mc.py` | Canonical (NVT) MC and basin hopping: permutation and displacement moves, LBFGS-relaxed trials, `minima_file`. CuAu nanoparticle with EMT. | ASE only | ~10 s | +| `02_molecule_gcmc.py` | Molecular adsorbates: rigid insertion, deletion and translate+rotate moves, `molecule_id` bookkeeping. Lennard-Jones dimers. | ASE only | ~20 s | +| `gcmc_ag111.py` | A production run: O on Ag(111) with a machine-learned potential, FIRE relaxation, chemical potentials derived from the potential itself. | `nvalchemi-toolkit[mace]`, GPU | hours | +| `re_gcmc_batched.py` | Replica exchange on one GPU, with every replica's trial energy evaluated in a single batched forward pass. | `nvalchemi-toolkit[mace]`, GPU | hours | + +The first three use toy potentials (Lennard-Jones, EMT) on purpose: they run +anywhere and their reference values are known. Real chemistry needs a real +potential, which is what the last two examples use. + +The same material with plots and derivations in between lives in `../notebooks`: +`gcmc_basics_lj.ipynb`, `molecular_gcmc.ipynb`, `grand_canonical_simulation_ag.ipynb`, +`co_on_cupd_replica_exchange.ipynb` and `phase_diagram.ipynb` (which builds a +surface phase diagram from the sweep in `notebooks/phase_diagram_demo`). +The scripts here are the same workflows without the narration, so they can be +run and diffed directly. From 29d0460d56387a42e00c68d73a6bfd8cd0019462 Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 15:03:07 +0200 Subject: [PATCH 18/20] Number the two GPU examples and give them explanatory names `gcmc_ag111.py` and `re_gcmc_batched.py` named their system and their implementation, not what a reader would learn from them, and they carried no position in the reading order the new CPU examples introduced. gcmc_ag111.py -> 03_gcmc_surface_mace.py re_gcmc_batched.py -> 04_replica_exchange_batched.py Also drops two references to `examples/gcmc.py` and `examples/re_gcmc.py`, which no longer exist. --- docs/installation.rst | 2 +- docs/tutorials/cluster_install_iqtc.rst | 6 +++--- examples/00_hello_gcmc.py | 2 +- examples/{gcmc_ag111.py => 03_gcmc_surface_mace.py} | 6 +++--- .../{re_gcmc_batched.py => 04_replica_exchange_batched.py} | 4 ++-- examples/README.md | 4 ++-- mcpy/ensembles/batched_replica_exchange.py | 2 +- tests/test_ag111_gcmc_golden.py | 2 +- 8 files changed, 14 insertions(+), 14 deletions(-) rename examples/{gcmc_ag111.py => 03_gcmc_surface_mace.py} (97%) rename examples/{re_gcmc_batched.py => 04_replica_exchange_batched.py} (96%) diff --git a/docs/installation.rst b/docs/installation.rst index bf6078e..81c1d82 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -95,4 +95,4 @@ single GPU prefer ``BatchedReplicaExchange``, which needs no MPI: .. code-block:: bash - python examples/re_gcmc_batched.py + python examples/04_replica_exchange_batched.py diff --git a/docs/tutorials/cluster_install_iqtc.rst b/docs/tutorials/cluster_install_iqtc.rst index d36f21b..259adfc 100644 --- a/docs/tutorials/cluster_install_iqtc.rst +++ b/docs/tutorials/cluster_install_iqtc.rst @@ -187,7 +187,7 @@ Save as ``submit_gcmc_alchemi.sh``: set -eo pipefail PYTHON=$HOME/.conda/envs/alchemi128/bin/python - SCRIPT=$HOME/bin/mcpy/examples/gcmc_ag111.py + SCRIPT=$HOME/bin/mcpy/examples/03_gcmc_surface_mace.py RESULTS_DIR=$SLURM_SUBMIT_DIR/results/job_${SLURM_JOB_ID:-local} export PYTHONNOUSERSITE=1 @@ -240,7 +240,7 @@ batched forward pass per MC step. This is the GPU path the package is built around. Unlike the MPI ``ReplicaExchange`` (one rank per replica), there is **no mpirun** and still only ``--gres=gpu:1``. -The script is ``examples/re_gcmc_batched.py``. Replica temperatures are passed +The script is ``examples/04_replica_exchange_batched.py``. Replica temperatures are passed with ``--temperatures``; the number of values is the replica count. Save as ``submit_re_gcmc_batched.sh``: @@ -259,7 +259,7 @@ Save as ``submit_re_gcmc_batched.sh``: set -eo pipefail PYTHON=$HOME/.conda/envs/alchemi128/bin/python - SCRIPT=$HOME/bin/mcpy/examples/re_gcmc_batched.py + SCRIPT=$HOME/bin/mcpy/examples/04_replica_exchange_batched.py RESULTS_DIR=$SLURM_SUBMIT_DIR/results/job_${SLURM_JOB_ID:-local} export PYTHONNOUSERSITE=1 diff --git a/examples/00_hello_gcmc.py b/examples/00_hello_gcmc.py index 5d8d4b0..f00c924 100644 --- a/examples/00_hello_gcmc.py +++ b/examples/00_hello_gcmc.py @@ -10,7 +10,7 @@ Writes ``hello_gcmc.out`` (step / N / energy / acceptance ratios) and ``hello_gcmc.xyz`` (trajectory). -Once this makes sense, read ``examples/gcmc_ag111.py`` for the same loop +Once this makes sense, read ``examples/03_gcmc_surface_mace.py`` for the same loop driven by a machine-learned potential on a real surface. """ from ase import Atoms diff --git a/examples/gcmc_ag111.py b/examples/03_gcmc_surface_mace.py similarity index 97% rename from examples/gcmc_ag111.py rename to examples/03_gcmc_surface_mace.py index 6a6bfae..6fa180f 100644 --- a/examples/gcmc_ag111.py +++ b/examples/03_gcmc_surface_mace.py @@ -16,13 +16,13 @@ Typical usage: # FIRE relaxation per step (recommended, GPU-native so fast enough): - python gcmc_ag111.py --checkpoint medium-mpa-0 --device cuda + python 03_gcmc_surface_mace.py --checkpoint medium-mpa-0 --device cuda # Energy-only (fastest, no geometry relaxation): - python gcmc_ag111.py --checkpoint medium-mpa-0 --no-relax + python 03_gcmc_surface_mace.py --checkpoint medium-mpa-0 --no-relax # Adjust chemical potential sweep: - python gcmc_ag111.py --checkpoint medium-mpa-0 --delta-mu-O -1.0 + python 03_gcmc_surface_mace.py --checkpoint medium-mpa-0 --delta-mu-O -1.0 """ import argparse import os diff --git a/examples/re_gcmc_batched.py b/examples/04_replica_exchange_batched.py similarity index 96% rename from examples/re_gcmc_batched.py rename to examples/04_replica_exchange_batched.py index 2edd510..56ec407 100644 --- a/examples/re_gcmc_batched.py +++ b/examples/04_replica_exchange_batched.py @@ -1,12 +1,12 @@ """Batched replica-exchange GCMC on a single GPU. -Single-process variant of ``re_gcmc.py``. All replicas live in one Python +Single-process variant of the MPI ``ReplicaExchange``. All replicas live in one Python process and share one ``AlchemiFCalculator`` instance; energies for trial moves are evaluated in a single batched forward pass per MC step. Run:: - python examples/re_gcmc_batched.py --temperatures 250 300 350 400 + python examples/04_replica_exchange_batched.py --temperatures 250 300 350 400 No mpirun. One GPU. Each replica must build its OWN cells and move_selector (BatchedReplicaExchange would otherwise share state across replicas). diff --git a/examples/README.md b/examples/README.md index 546b9c4..39dc2cb 100644 --- a/examples/README.md +++ b/examples/README.md @@ -7,8 +7,8 @@ Run them in order. The first three need nothing but `pip install -e .` and a CPU | `00_hello_gcmc.py` | The GCMC loop: cell, insertion/deletion/displacement moves, chemical potential, output files. Lennard-Jones fluid in a box. | ASE only | ~40 s | | `01_canonical_mc.py` | Canonical (NVT) MC and basin hopping: permutation and displacement moves, LBFGS-relaxed trials, `minima_file`. CuAu nanoparticle with EMT. | ASE only | ~10 s | | `02_molecule_gcmc.py` | Molecular adsorbates: rigid insertion, deletion and translate+rotate moves, `molecule_id` bookkeeping. Lennard-Jones dimers. | ASE only | ~20 s | -| `gcmc_ag111.py` | A production run: O on Ag(111) with a machine-learned potential, FIRE relaxation, chemical potentials derived from the potential itself. | `nvalchemi-toolkit[mace]`, GPU | hours | -| `re_gcmc_batched.py` | Replica exchange on one GPU, with every replica's trial energy evaluated in a single batched forward pass. | `nvalchemi-toolkit[mace]`, GPU | hours | +| `03_gcmc_surface_mace.py` | A production run: O on Ag(111) with a machine-learned potential, FIRE relaxation, chemical potentials derived from the potential itself. | `nvalchemi-toolkit[mace]`, GPU | hours | +| `04_replica_exchange_batched.py` | Replica exchange on one GPU, with every replica's trial energy evaluated in a single batched forward pass. | `nvalchemi-toolkit[mace]`, GPU | hours | The first three use toy potentials (Lennard-Jones, EMT) on purpose: they run anywhere and their reference values are known. Real chemistry needs a real diff --git a/mcpy/ensembles/batched_replica_exchange.py b/mcpy/ensembles/batched_replica_exchange.py index ea6bce9..18f8b98 100644 --- a/mcpy/ensembles/batched_replica_exchange.py +++ b/mcpy/ensembles/batched_replica_exchange.py @@ -17,7 +17,7 @@ ``gcmc_factory(T, rank)`` must return a fresh ``GrandCanonicalEnsemble`` with its own ``cells``, ``move_selector`` and atoms object. Sharing cells or a move_selector across replicas will corrupt per-replica state (cell volumes, -acceptance counters, RNG streams). See ``examples/re_gcmc_batched.py``. +acceptance counters, RNG streams). See ``examples/04_replica_exchange_batched.py``. The shared ``calculator`` must implement ``get_potential_energies(atoms_list)`` returning an ndarray of energies. ``AlchemiCalculator`` provides this. diff --git a/tests/test_ag111_gcmc_golden.py b/tests/test_ag111_gcmc_golden.py index a2999a4..7545620 100644 --- a/tests/test_ag111_gcmc_golden.py +++ b/tests/test_ag111_gcmc_golden.py @@ -15,7 +15,7 @@ overlap rejection, de Broglie acceptance, arrays+constraints rollback, trajectory writing) reproduces exactly when seeded. -Parameters come from ``examples/gcmc_ag111.py``, the calibrated set for this +Parameters come from ``examples/03_gcmc_surface_mace.py``, the calibrated set for this checkpoint: a=4.165, the O-Ag exclusion radius 2.068 used both as the radius and as the region depth below the top layer, height 7, min_insert 0.5, and the -0.176 bulk Ag correction. From bcf3ee2d5ffe797d4cb7a931684f9a49de07c4cb Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 15:24:07 +0200 Subject: [PATCH 19/20] Give example 03 its Ag exchange back and drop argparse from both GPU examples The Ag(111) example exchanged only O, so the lattice it started with was the lattice it ended with: O could decorate the surface but never reconstruct it into an oxide. The two-species setup that examples/gcmc.py carried before it was pruned is restored here, one cell per species (an incoming Ag needs a full Ag-Ag distance where an incoming O needs one O-Ag bond), with the mu_Ag the script already derived now actually driving moves. EXCHANGE_AG = False falls back to the O-only setup the golden test pins. Both GPU examples now read top to bottom as constants instead of an argparse block: an example is read more often than it is invoked, and the flags were duplicating every default twice. The cluster tutorial submits them without flags accordingly, changing into the results directory since outputs now land in the working directory. Also seeds the replica-exchange cells, selector and ensemble. They were the one unseeded path left in the examples, which silently made a --seed run unrepeatable. --- docs/tutorials/cluster_install_iqtc.rst | 55 ++--- examples/03_gcmc_surface_mace.py | 271 ++++++++++-------------- examples/04_replica_exchange_batched.py | 184 +++++++++------- examples/README.md | 5 +- tests/test_ag111_gcmc_golden.py | 3 +- 5 files changed, 248 insertions(+), 270 deletions(-) diff --git a/docs/tutorials/cluster_install_iqtc.rst b/docs/tutorials/cluster_install_iqtc.rst index 259adfc..2d1db53 100644 --- a/docs/tutorials/cluster_install_iqtc.rst +++ b/docs/tutorials/cluster_install_iqtc.rst @@ -199,17 +199,14 @@ Save as ``submit_gcmc_alchemi.sh``: nvidia-smi -L "$PYTHON" -c "import torch; print('torch', torch.__version__, 'cu_avail', torch.cuda.is_available())" - "$PYTHON" "$SCRIPT" \ - --checkpoint medium-mpa-0 \ - --device cuda \ - --T 500.0 \ - --steps 1000 \ - --delta-mu-O -0.5 \ - --fmax 0.05 \ - --relax-steps 500 \ - --seed 42 \ - --outdir "$RESULTS_DIR" \ - --write-interval 1 + cd "$RESULTS_DIR" + "$PYTHON" "$SCRIPT" + +The script takes no command-line flags: every setting is a constant at the +top of the file (``CHECKPOINT``, ``TEMPERATURE``, ``STEPS``, ``DELTA_MU_O``, +``FMAX``, ``SEED``, ...). Edit your copy before submitting, and note that the +output files land in the working directory, which is why the job changes into +``$RESULTS_DIR`` first. Submit with ``sbatch submit_gcmc_alchemi.sh``. @@ -227,7 +224,7 @@ Key SLURM-script rules: (after the first size change it switches to dynamic shapes; one-time warmup, then ~1.3-1.4x faster). On nodes with tight host-memory limits the compile guard-building phase can OOM at startup — if you hit that (see - troubleshooting below), pass ``--no-compile`` to the simulation script. + troubleshooting below), set ``COMPILE_MODEL = False`` in the script. Batched replica exchange on one GPU @@ -240,8 +237,9 @@ batched forward pass per MC step. This is the GPU path the package is built around. Unlike the MPI ``ReplicaExchange`` (one rank per replica), there is **no mpirun** and still only ``--gres=gpu:1``. -The script is ``examples/04_replica_exchange_batched.py``. Replica temperatures are passed -with ``--temperatures``; the number of values is the replica count. +The script is ``examples/04_replica_exchange_batched.py``. It takes no +command-line flags: the replica temperatures are the ``TEMPERATURES`` list at +the top of the file, and its length is the replica count. Save as ``submit_re_gcmc_batched.sh``: @@ -271,32 +269,23 @@ Save as ``submit_re_gcmc_batched.sh``: nvidia-smi -L "$PYTHON" -c "import torch; print('torch', torch.__version__, 'cu_avail', torch.cuda.is_available())" - "$PYTHON" "$SCRIPT" \ - --temperatures 250 300 350 400 450 500 \ - --checkpoint medium-mpa-0 \ - --device cuda \ - --gcmc-steps 200 \ - --exchange-interval 10 \ - --delta-mu-O -0.5 \ - --seed 42 \ - --outdir "$RESULTS_DIR" \ - --write-interval 1 + cd "$RESULTS_DIR" + "$PYTHON" "$SCRIPT" Submit with ``sbatch submit_re_gcmc_batched.sh``. Notes specific to the batched path: -- **One GPU, one process.** All six replicas (the six ``--temperatures`` - above) share a single ``AlchemiCalculator`` and are evaluated in one batched - forward pass. Keep ``--gres=gpu:1``; adding ranks or extra GPUs does not help - this path. +- **One GPU, one process.** All replicas in ``TEMPERATURES`` share a single + ``AlchemiCalculator`` and are evaluated in one batched forward pass. Keep + ``--gres=gpu:1``; adding ranks or extra GPUs does not help this path. - **``torch.compile`` stays on by default**, same as the single-replica run: it handles GCMC's dynamic atom count via dynamic shapes after a one-time - warmup. ``--no-compile`` is only the escape hatch if the compile + warmup. ``COMPILE_MODEL = False`` is only the escape hatch if the compile guard-building phase OOMs on a node with tight host-memory limits. -- **Replica count = number of ``--temperatures``.** More replicas means a - larger batch and more GPU memory; if startup OOMs, reduce the number of - temperatures before changing anything else. +- **Replica count = length of ``TEMPERATURES``.** More replicas means a + larger batch and more GPU memory; if startup OOMs, shorten that list + before changing anything else. - **Outputs** land in ``$RESULTS_DIR``: per-replica ``gcmc_batched_*.out`` and ``.xyz`` files, plus a combined ``replica_exchange_batched.log`` with the exchange history. @@ -338,7 +327,7 @@ Common failure modes ``Warp CUDA error 2: out of memory`` at startup - ``torch.compile`` tries to JIT-compile the model and the guard-building pass OOMs against the job's host-memory limit. - - Pass ``--no-compile`` to the script (sets ``compile_model=False``), or + - Set ``COMPILE_MODEL = False`` in the script (passes ``compile_model=False``), or request more host memory for the job. diff --git a/examples/03_gcmc_surface_mace.py b/examples/03_gcmc_surface_mace.py index 6fa180f..9263a7f 100644 --- a/examples/03_gcmc_surface_mace.py +++ b/examples/03_gcmc_surface_mace.py @@ -1,197 +1,158 @@ -"""GCMC of O adsorption on an Ag(111) surface using the nvalchemi GPU stack. +"""GCMC on an Ag(111) surface with a machine-learned potential (GPU). -The geometry constants below (lattice constant, O-Ag exclusion radius, region -height, the -0.176 bulk-Ag correction) are the calibrated set for the MACE -small-density-agnesi checkpoint; pass that checkpoint with ``--checkpoint`` to -reproduce ``tests/test_ag111_gcmc_golden.py``. Reference chemical potentials -are derived from whatever potential is passed, so another checkpoint stays -self-consistent, but its balance point for ``--delta-mu-O`` will differ. +A production run, in contrast to examples 00-02: a real potential, a real +surface, and both species exchanged with their reservoirs. -Builds a 16×16×4 Ag(111) slab (1024 atoms) by default, fixes the bottom two -layers, and runs GCMC inserting/deleting O in a 7 Å region that starts one -O-Ag distance *below* the top layer, so subsurface sites are reachable. + * O is inserted and deleted in a 7 A region that starts one O-Ag distance + *below* the top layer, so subsurface sites are reachable. + * Ag is inserted and deleted through its own cell, identical to the O one + except for the exclusion radius an insertion sees. Exchanging Ag is what + lets the surface reconstruct into an oxide instead of merely decorating + the lattice it started with. + * Reference chemical potentials are derived from the running potential, so + swapping the checkpoint keeps the setup self-consistent. Only DELTA_MU_O + is a physical knob, and its balance point differs per checkpoint. -Requirements: - pip install 'nvalchemi-toolkit[mace]' +Every setting is a constant below; edit and re-run. Set EXCHANGE_AG to False +for the O-only setup that ``tests/test_ag111_gcmc_golden.py`` pins. -Typical usage: - # FIRE relaxation per step (recommended, GPU-native so fast enough): - python 03_gcmc_surface_mace.py --checkpoint medium-mpa-0 --device cuda +Requirements: pip install 'nvalchemi-toolkit[mace]', and a CUDA GPU. - # Energy-only (fastest, no geometry relaxation): - python 03_gcmc_surface_mace.py --checkpoint medium-mpa-0 --no-relax - - # Adjust chemical potential sweep: - python 03_gcmc_surface_mace.py --checkpoint medium-mpa-0 --delta-mu-O -1.0 + python examples/03_gcmc_surface_mace.py """ -import argparse -import os - import numpy as np from ase.build import fcc111 from ase.constraints import FixAtoms -import logging - -from mcpy.calculators import AlchemiCalculator, AlchemiFCalculator +from mcpy.calculators import AlchemiFCalculator from mcpy.cell import CustomCell -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble -from mcpy.moves import DeletionMove, InsertionMove -from mcpy.moves.move_selector import MoveSelector +from mcpy.ensembles import GrandCanonicalEnsemble +from mcpy.moves import DeletionMove, InsertionMove, MoveSelector from mcpy.utils import derive_mu_bulk, derive_mu_gas from mcpy.utils.logging import configure as configure_logging -logger = logging.getLogger('mcpy') - -# Calibrated set for O on Ag(111) with the small-density-agnesi MACE model. -# Fixed on purpose: this is fixture geometry, not a per-checkpoint number. -# Only mu is derived at run time. -AG_LATTICE = 4.165 # Ag fcc lattice constant (Angstrom) -# Relaxed O-Ag pair distance: the free-volume exclusion radius, and also how -# far below the top layer the region starts so O can go subsurface (the -# oxidation path). Anchoring at the top layer makes the region pure vacuum. +# --- potential ------------------------------------------------------------- +# A named checkpoint (e.g. 'medium-mpa-0') or a path to a local .pt file. +CHECKPOINT = 'medium-mpa-0' +DEVICE = 'cuda' +ENABLE_CUEQ = True # cuEquivariance kernel fusion; verify per checkpoint +COMPILE_MODEL = True # torch.compile: ~2x faster forward pass +RELAX_STEPS = 300 # max FIRE steps per energy evaluation +FMAX = 0.05 # FIRE force threshold (eV/A) +OPTIMIZER = 'fire2' # 'fire2' usually converges in fewer steps than 'fire' + +# --- system ---------------------------------------------------------------- +# Calibrated for O on Ag(111) with the small-density-agnesi MACE checkpoint. +# Fixture geometry, not per-checkpoint numbers: only mu is derived at run time. +AG_LATTICE = 4.165 # Ag fcc lattice constant (A) +SLAB_SIZE = (16, 16, 4) # fcc111 supercell (Nx, Ny, Nlayers) = 1024 atoms +VACUUM = 10.0 # vacuum above and below the slab (A) +FIX_LAYERS = 2 # bottom layers frozen + +# Relaxed O-Ag pair distance. Used twice: as the exclusion radius an incoming +# O sees, and as how far below the top layer the exchange region starts. +# Anchoring the region at the top layer instead would make it pure vacuum. R_O_AG = 2.068 -SPECIES_RADII = {'Ag': R_O_AG, 'O': 0.0} -# Bulk Ag reference correction. -MU_AG_CORRECTION = -0.176 - - -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--checkpoint', default='medium-mpa-0', - help='Named checkpoint (e.g. medium-mpa-0) or path to a .pt file') - p.add_argument('--no-relax', action='store_true', - help='Use AlchemiCalculator (energy-only); default uses AlchemiFCalculator') - p.add_argument('--device', default='cuda', help='Torch device (cuda or cpu)') - p.add_argument('--no-cueq', action='store_true', help='Disable cuEquivariance kernel fusion') - p.add_argument('--no-compile', action='store_true', help='Disable torch.compile') - p.add_argument('--fmax', type=float, default=0.05, help='FIRE force threshold (eV/Å)') - p.add_argument('--relax-steps', type=int, default=300, help='Max FIRE steps per eval') - p.add_argument('--optimizer', default='fire2', choices=['fire', 'fire2'], - help='FIRE variant (fire2 typically converges in fewer steps)') - p.add_argument('--T', type=float, default=500.0, help='Temperature (K)') - p.add_argument('--steps', type=int, default=5_000_000, help='Number of GCMC steps') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.3, help='Shift applied to mu_O (eV)') - p.add_argument('--slab-size', type=int, nargs=3, default=[16, 16, 4], - help='fcc111 supercell (Nx Ny Nlayers); default 16 16 4 = 1024 atoms') - p.add_argument('--vacuum', type=float, default=10.0, help='Vacuum above/below slab (Å)') - p.add_argument('--cell-height', type=float, default=7.0, help='O insertion region height (Å)') - p.add_argument('--fix-layers', type=int, default=2, help='Number of bottom layers to freeze') - p.add_argument('--min-insert', type=float, default=0.5, - help='Min insertion distance to existing atoms (Å)') - p.add_argument('--seed', type=int, default=None, help='Master seed (random if unset)') - p.add_argument('--write-interval', type=int, default=10, - help='Outfile/trajectory write interval') - p.add_argument('--outdir', default='.', help='Output directory') - return p.parse_args() - - -def build_slab(args): - nx, ny, nlayers = args.slab_size +# An incoming Ag needs a full Ag-Ag distance instead, which is why each +# species gets its own cell over the same region. +R_AG_AG = 2.75 +CELL_HEIGHT = 7.0 # height of the exchange region (A) +MIN_INSERT = 0.5 # reject insertions closer than this to any atom (A) + +# --- thermodynamics -------------------------------------------------------- +TEMPERATURE = 500.0 +DELTA_MU_O = -0.3 # shift applied to mu_O = E(O2)/2 +MU_AG_CORRECTION = -0.176 # bulk Ag reference correction +EXCHANGE_AG = True # False = O only, the golden-test setup + +# --- run ------------------------------------------------------------------- +STEPS = 5_000_000 +WRITE_INTERVAL = 10 +SEED = 42 # one master seed; every RNG below derives from it + + +def build_slab(): + nx, ny, nlayers = SLAB_SIZE atoms = fcc111('Ag', a=AG_LATTICE, size=(nx, ny, nlayers), - periodic=True, vacuum=args.vacuum) - # tag=1 is top layer, tag=nlayers is bottom layer (ASE convention) - fix_tag_threshold = nlayers - args.fix_layers + 1 - bottom = [a.index for a in atoms if a.tag >= fix_tag_threshold] - atoms.set_constraint(FixAtoms(indices=bottom)) + periodic=True, vacuum=VACUUM) + # tag=1 is the top layer, tag=nlayers the bottom one (ASE convention). + fix_from = nlayers - FIX_LAYERS + 1 + atoms.set_constraint( + FixAtoms(indices=[a.index for a in atoms if a.tag >= fix_from])) return atoms -def build_calculator(args): - enable_cueq = not args.no_cueq - compile_model = not args.no_compile - if args.no_relax: - return AlchemiCalculator( - checkpoint=args.checkpoint, - device=args.device, - enable_cueq=enable_cueq, - compile_model=compile_model, - ) - return AlchemiFCalculator( - checkpoint=args.checkpoint, - steps=args.relax_steps, - fmax=args.fmax, - device=args.device, - enable_cueq=enable_cueq, - compile_model=compile_model, - optimizer=args.optimizer, - ) - - def main(): - args = parse_args() configure_logging() - os.makedirs(args.outdir, exist_ok=True) - ss = np.random.SeedSequence(args.seed) - seed_del, seed_ins, seed_cell, seed_sel, seed_ens = ( - int(s) for s in ss.generate_state(5, dtype=np.uint32) + # One master seed, one derived stream per stochastic component. Any of + # them left at None would make the run unrepeatable. + (seed_ins_o, seed_del_o, seed_cell_o, + seed_ins_ag, seed_del_ag, seed_cell_ag, + seed_sel, seed_ens) = ( + int(s) for s in np.random.SeedSequence(SEED).generate_state(8, dtype=np.uint32) ) - atoms = build_slab(args) + atoms = build_slab() - # One interlayer below the top layer, so the region covers the subsurface - # site plus the vacuum above (see R_O_AG). + # Both cells cover the same region and differ only in the radii, i.e. in + # how much of that region each species may actually be inserted into. z_top = float(atoms.positions[atoms.get_tags() == 1, 2].max()) - cell_bottom_z = z_top - R_O_AG - - cell = CustomCell( - atoms, - custom_height=args.cell_height, - bottom_z=cell_bottom_z, - species_radii=SPECIES_RADII, - seed=seed_cell, + bottom_z = z_top - R_O_AG + cell_o = CustomCell(atoms, custom_height=CELL_HEIGHT, bottom_z=bottom_z, + species_radii={'Ag': R_O_AG, 'O': 0.0}, seed=seed_cell_o) + cell_ag = CustomCell(atoms, custom_height=CELL_HEIGHT, bottom_z=bottom_z, + species_radii={'Ag': R_AG_AG, 'O': 0.0}, seed=seed_cell_ag) + + # AlchemiFCalculator relaxes each trial structure with GPU-native FIRE, + # honouring the FixAtoms constraint. Swap in AlchemiCalculator for + # energy-only evaluation (faster, but no relaxation of the trial). + calculator = AlchemiFCalculator( + checkpoint=CHECKPOINT, + steps=RELAX_STEPS, + fmax=FMAX, + device=DEVICE, + enable_cueq=ENABLE_CUEQ, + compile_model=COMPILE_MODEL, + optimizer=OPTIMIZER, ) - calculator = build_calculator(args) - - # mu references come from the running potential: hardcoded values silently - # change meaning when the checkpoint changes. - mu_ag = (derive_mu_bulk(calculator, 'Ag', a=AG_LATTICE) + MU_AG_CORRECTION - if args.mu_Ag is None else args.mu_Ag) - mu_o = derive_mu_gas(calculator, 'O2') if args.mu_O is None else args.mu_O - - import torch - gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A' - logger.info('Calculator device: %s | CUDA available: %s | GPU: %s', - args.device, torch.cuda.is_available(), gpu_name) + # Hardcoded reference energies silently change meaning when the + # checkpoint changes, so they are measured with the running potential. + mu_ag = derive_mu_bulk(calculator, 'Ag', a=AG_LATTICE) + MU_AG_CORRECTION + mu_o = derive_mu_gas(calculator, 'O2') + DELTA_MU_O + moves = [InsertionMove(cell_o, species=['O'], min_insert=MIN_INSERT, seed=seed_ins_o), + DeletionMove(cell_o, species=['O'], seed=seed_del_o)] + cells = [cell_o] species = ['O'] - move_selector = MoveSelector( - [1, 1], - [DeletionMove(cell, species=species, seed=seed_del), - InsertionMove(cell, species=species, min_insert=args.min_insert, seed=seed_ins)], - seed=seed_sel, - ) - - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} + if EXCHANGE_AG: + moves += [InsertionMove(cell_ag, species=['Ag'], min_insert=MIN_INSERT, + seed=seed_ins_ag), + DeletionMove(cell_ag, species=['Ag'], seed=seed_del_ag)] + cells.append(cell_ag) + species.insert(0, 'Ag') + move_selector = MoveSelector([1] * len(moves), moves, seed=seed_sel) - nx, ny, nlayers = args.slab_size - mode = 'nrelax' if args.no_relax else f'fire_fmax{args.fmax}' - tag = f'Ag{nx * ny * nlayers}_111_dmuO_{args.delta_mu_O}_{mode}' + nx, ny, nlayers = SLAB_SIZE + tag = f'Ag{nx * ny * nlayers}_111_dmuO_{DELTA_MU_O}_fire_fmax{FMAX}' gcmc = GrandCanonicalEnsemble( atoms=atoms, - cells=[cell], + cells=cells, calculator=calculator, - mu=mus, - units_type='metal', + mu={'Ag': mu_ag, 'O': mu_o}, species=species, - temperature=args.T, + units_type='metal', + temperature=TEMPERATURE, move_selector=move_selector, random_seed=seed_ens, - outfile=os.path.join(args.outdir, f'gcmc_{tag}.out'), - traj_file=os.path.join(args.outdir, f'gcmc_{tag}.xyz'), - trajectory_write_interval=args.write_interval, - outfile_write_interval=args.write_interval, + outfile=f'gcmc_{tag}.out', + traj_file=f'gcmc_{tag}.xyz', + outfile_write_interval=WRITE_INTERVAL, + trajectory_write_interval=WRITE_INTERVAL, ) - - gcmc.run(args.steps) + gcmc.run(STEPS) if __name__ == '__main__': diff --git a/examples/04_replica_exchange_batched.py b/examples/04_replica_exchange_batched.py index 56ec407..15b1a5e 100644 --- a/examples/04_replica_exchange_batched.py +++ b/examples/04_replica_exchange_batched.py @@ -1,131 +1,155 @@ -"""Batched replica-exchange GCMC on a single GPU. +"""Replica-exchange GCMC on a single GPU, with batched energy evaluation. -Single-process variant of the MPI ``ReplicaExchange``. All replicas live in one Python -process and share one ``AlchemiFCalculator`` instance; energies for trial -moves are evaluated in a single batched forward pass per MC step. +Parallel tempering without MPI. All replicas live in one Python process and +share one ``AlchemiFCalculator``, so every replica's trial energy for a step +is computed in a single batched forward pass. This is the supported +replica-exchange path on one GPU; the MPI ``ReplicaExchange`` needs one rank +per replica and cannot exchange molecules. -Run:: +Each replica must build its OWN cells and move selector: sharing them would +share their RNG streams and acceptance counters across the ladder. That is +what ``gcmc_factory`` below is for -- it is called once per replica. - python examples/04_replica_exchange_batched.py --temperatures 250 300 350 400 +Ladder spacing is the parameter that decides whether the run works at all. +See ``docs/replica_exchange_ladder_spacing.rst``; a ladder whose neighbours +never swap samples no better than independent runs. -No mpirun. One GPU. Each replica must build its OWN cells and move_selector -(BatchedReplicaExchange would otherwise share state across replicas). -""" -import argparse -import os +Every setting is a constant below; edit and re-run. + +Requirements: pip install 'nvalchemi-toolkit[mace]', and a CUDA GPU. + python examples/04_replica_exchange_batched.py +""" import numpy as np from ase.build import fcc111 from ase.constraints import FixAtoms from mcpy.calculators import AlchemiFCalculator -from mcpy.cell import CustomCell as Cell -from mcpy.ensembles import BatchedReplicaExchange -from mcpy.ensembles.grand_canonical_ensemble import GrandCanonicalEnsemble -from mcpy.moves import DeletionMove, InsertionMove -from mcpy.moves.move_selector import MoveSelector +from mcpy.cell import CustomCell +from mcpy.ensembles import BatchedReplicaExchange, GrandCanonicalEnsemble +from mcpy.moves import DeletionMove, InsertionMove, MoveSelector from mcpy.utils import derive_mu_bulk, derive_mu_gas from mcpy.utils.logging import configure as configure_logging +# --- the ladder ------------------------------------------------------------ +# A temperature ladder OR a mu ladder, never both. Neighbouring rungs must +# overlap in energy or no swap is ever accepted. +TEMPERATURES = [300, 400, 500, 600, 700, 800] +EXCHANGE_INTERVAL = 10 # GCMC steps between swap attempts +GCMC_STEPS = 200 # GCMC steps per replica, in total + +# --- potential ------------------------------------------------------------- +CHECKPOINT = 'mace-small-density-agnesi-stress.model' +DEVICE = 'cuda' +ENABLE_CUEQ = True # cuEquivariance kernel fusion; verify per checkpoint +COMPILE_MODEL = True # torch.compile: ~2x faster forward pass + +# --- system ---------------------------------------------------------------- +AG_LATTICE = 4.1592 +SLAB_SIZE = (4, 4, 3) +VACUUM = 8.0 +R_O_AG = 2.11 # exclusion radius an incoming O sees +R_AG_AG = 2.75 # exclusion radius an incoming Ag sees +CELL_HEIGHT = 7.0 +MIN_INSERT = 0.5 -def parse_args(): - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument('--temperatures', type=float, nargs='+', - default=[300, 400, 500, 600, 700, 800], - help='Replica temperatures (K)') - p.add_argument('--gcmc-steps', type=int, default=200) - p.add_argument('--exchange-interval', type=int, default=10) - p.add_argument('--write-interval', type=int, default=1) - p.add_argument('--checkpoint', default='mace-small-density-agnesi-stress.model') - p.add_argument('--device', default='cuda') - p.add_argument('--no-cueq', action='store_true') - p.add_argument('--no-compile', action='store_true') - p.add_argument('--mu-Ag', type=float, default=None, - help='Chemical potential of Ag (eV); default: derived from the potential') - p.add_argument('--mu-O', type=float, default=None, - help='Reference mu_O (eV); default: E(O2)/2 derived from the potential') - p.add_argument('--delta-mu-O', type=float, default=-0.3) - p.add_argument('--seed', type=int, default=None) - p.add_argument('--outdir', default='.') - return p.parse_args() +# --- thermodynamics -------------------------------------------------------- +DELTA_MU_O = -0.3 # shift applied to mu_O = E(O2)/2 + +# --- run ------------------------------------------------------------------- +WRITE_INTERVAL = 1 +SEED = 42 # one master seed; every RNG below derives from it + +# Per replica: 2 cells + 4 moves + 1 selector + 1 ensemble = 8 streams. +SEEDS_PER_REPLICA = 8 def main(): - args = parse_args() configure_logging() - os.makedirs(args.outdir, exist_ok=True) - ss = np.random.SeedSequence(args.seed) - n_replicas = len(args.temperatures) - # 4 move seeds per replica + 1 master seed for RE. - seeds = ss.generate_state(4 * n_replicas + 1, dtype=np.uint32) - move_seeds = [int(s) for s in seeds[:4 * n_replicas]] + n_replicas = len(TEMPERATURES) + seeds = np.random.SeedSequence(SEED).generate_state( + SEEDS_PER_REPLICA * n_replicas + 1, dtype=np.uint32) master_seed = int(seeds[-1]) - # Shared template — each factory call copies and builds its own cells/moves. - base_atoms = fcc111('Ag', a=4.1592, size=(4, 4, 3), periodic=True, vacuum=8) - bottom_layer = [a.index for a in base_atoms if a.tag == 3] + # Shared template. Every factory call copies it and builds its own cells. + base_atoms = fcc111('Ag', a=AG_LATTICE, size=SLAB_SIZE, + periodic=True, vacuum=VACUUM) + bottom_layer = [a.index for a in base_atoms if a.tag == SLAB_SIZE[2]] base_atoms.set_constraint(FixAtoms(indices=bottom_layer)) + bottom_z = float(base_atoms.positions[base_atoms.get_tags() == 1, 2].max()) - R_O_AG - # ONE model, shared across replicas. Batched eval reuses this calculator. + # ONE model, shared across replicas: this is what makes the batched + # forward pass possible in the first place. calculator = AlchemiFCalculator( - checkpoint=args.checkpoint, - device=args.device, - enable_cueq=not args.no_cueq, - compile_model=not args.no_compile, + checkpoint=CHECKPOINT, + device=DEVICE, + enable_cueq=ENABLE_CUEQ, + compile_model=COMPILE_MODEL, ) - species = ['Ag', 'O'] - mu_ag = args.mu_Ag if args.mu_Ag is not None else derive_mu_bulk(calculator, 'Ag') - mu_o = args.mu_O if args.mu_O is not None else derive_mu_gas(calculator) - mus = {'Ag': mu_ag, 'O': mu_o + args.delta_mu_O} + # Measured with the running potential rather than hardcoded, so swapping + # the checkpoint keeps the setup self-consistent. + mus = {'Ag': derive_mu_bulk(calculator, 'Ag', a=AG_LATTICE), + 'O': derive_mu_gas(calculator, 'O2') + DELTA_MU_O} def gcmc_factory(T, rank): + """Build replica ``rank``. Called once per rung of the ladder. + + BatchedReplicaExchange calls this as ``gcmc_factory(T=..., rank=...)`` + for a temperature ladder and ``gcmc_factory(mu=..., rank=...)`` for a + mu ladder, so the parameter name is part of the contract. + """ + s = [int(x) for x in + seeds[SEEDS_PER_REPLICA * rank:SEEDS_PER_REPLICA * (rank + 1)]] + atoms = base_atoms.copy() atoms.set_constraint(FixAtoms(indices=bottom_layer)) - cell_ag_ag = Cell(atoms, custom_height=7, bottom_z=12.8 - 2.11, - species_radii={'Ag': 2.75, 'O': 0}) - cell_ag_o = Cell(atoms, custom_height=7, bottom_z=12.8 - 2.11, - species_radii={'Ag': 2.11, 'O': 0}) + # Same region for both species, differing only in the radii, i.e. in + # how much of that region each species may be inserted into. + cell_ag = CustomCell(atoms, custom_height=CELL_HEIGHT, bottom_z=bottom_z, + species_radii={'Ag': R_AG_AG, 'O': 0.0}, seed=s[0]) + cell_o = CustomCell(atoms, custom_height=CELL_HEIGHT, bottom_z=bottom_z, + species_radii={'Ag': R_O_AG, 'O': 0.0}, seed=s[1]) - s = move_seeds[4 * rank:4 * (rank + 1)] move_selector = MoveSelector( - [5, 5, 5, 5], - [DeletionMove(cell_ag_ag, species=['Ag'], seed=s[0]), - DeletionMove(cell_ag_o, species=['O'], seed=s[1]), - InsertionMove(cell_ag_ag, species=['Ag'], min_insert=0.5, seed=s[2]), - InsertionMove(cell_ag_o, species=['O'], min_insert=0.5, seed=s[3])], + [1, 1, 1, 1], + [InsertionMove(cell_ag, species=['Ag'], min_insert=MIN_INSERT, seed=s[2]), + DeletionMove(cell_ag, species=['Ag'], seed=s[3]), + InsertionMove(cell_o, species=['O'], min_insert=MIN_INSERT, seed=s[4]), + DeletionMove(cell_o, species=['O'], seed=s[5])], + seed=s[6], ) - tag = f'{atoms.get_chemical_formula()}_dmu_{args.delta_mu_O}_rank{rank}' + tag = f'{atoms.get_chemical_formula()}_dmu_{DELTA_MU_O}_rank{rank}' return GrandCanonicalEnsemble( atoms=atoms, - cells=[cell_ag_ag, cell_ag_o], + cells=[cell_ag, cell_o], calculator=calculator, mu=mus, + species=['Ag', 'O'], units_type='metal', - species=species, temperature=T, move_selector=move_selector, - outfile=os.path.join(args.outdir, f'gcmc_batched_{tag}.out'), - trajectory_write_interval=args.write_interval, - outfile_write_interval=args.write_interval, - traj_file=os.path.join(args.outdir, f'gcmc_batched_{tag}.xyz'), + random_seed=s[7], + outfile=f'gcmc_batched_{tag}.out', + traj_file=f'gcmc_batched_{tag}.xyz', + outfile_write_interval=WRITE_INTERVAL, + trajectory_write_interval=WRITE_INTERVAL, ) - pt = BatchedReplicaExchange( + replica_exchange = BatchedReplicaExchange( gcmc_factory, calculator=calculator, - temperatures=args.temperatures, - gcmc_steps=args.gcmc_steps, - exchange_interval=args.exchange_interval, - write_out_interval=args.write_interval, + temperatures=TEMPERATURES, + gcmc_steps=GCMC_STEPS, + exchange_interval=EXCHANGE_INTERVAL, + write_out_interval=WRITE_INTERVAL, seed=master_seed, - outfile=os.path.join(args.outdir, 'replica_exchange_batched.log'), + outfile='replica_exchange_batched.log', ) - pt.run() + replica_exchange.run() if __name__ == '__main__': diff --git a/examples/README.md b/examples/README.md index 39dc2cb..f35a2b1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -2,12 +2,15 @@ Run them in order. The first three need nothing but `pip install -e .` and a CPU. +None of them take command-line flags: every setting is a named constant at the +top of the file. Edit and re-run. + | Example | What it shows | Needs | Runtime | | --- | --- | --- | --- | | `00_hello_gcmc.py` | The GCMC loop: cell, insertion/deletion/displacement moves, chemical potential, output files. Lennard-Jones fluid in a box. | ASE only | ~40 s | | `01_canonical_mc.py` | Canonical (NVT) MC and basin hopping: permutation and displacement moves, LBFGS-relaxed trials, `minima_file`. CuAu nanoparticle with EMT. | ASE only | ~10 s | | `02_molecule_gcmc.py` | Molecular adsorbates: rigid insertion, deletion and translate+rotate moves, `molecule_id` bookkeeping. Lennard-Jones dimers. | ASE only | ~20 s | -| `03_gcmc_surface_mace.py` | A production run: O on Ag(111) with a machine-learned potential, FIRE relaxation, chemical potentials derived from the potential itself. | `nvalchemi-toolkit[mace]`, GPU | hours | +| `03_gcmc_surface_mace.py` | A production run: O and Ag both exchanged on an Ag(111) surface (one cell per species) with a machine-learned potential, FIRE relaxation, and chemical potentials derived from the potential itself. | `nvalchemi-toolkit[mace]`, GPU | hours | | `04_replica_exchange_batched.py` | Replica exchange on one GPU, with every replica's trial energy evaluated in a single batched forward pass. | `nvalchemi-toolkit[mace]`, GPU | hours | The first three use toy potentials (Lennard-Jones, EMT) on purpose: they run diff --git a/tests/test_ag111_gcmc_golden.py b/tests/test_ag111_gcmc_golden.py index 7545620..1fae919 100644 --- a/tests/test_ag111_gcmc_golden.py +++ b/tests/test_ag111_gcmc_golden.py @@ -15,7 +15,8 @@ overlap rejection, de Broglie acceptance, arrays+constraints rollback, trajectory writing) reproduces exactly when seeded. -Parameters come from ``examples/03_gcmc_surface_mace.py``, the calibrated set for this +Parameters come from ``examples/03_gcmc_surface_mace.py --no-ag-moves`` (the +example exchanges Ag as well by default), the calibrated set for this checkpoint: a=4.165, the O-Ag exclusion radius 2.068 used both as the radius and as the region depth below the top layer, height 7, min_insert 0.5, and the -0.176 bulk Ag correction. From 155401ef412b50820cd7c6e943666d62558ef1ed Mon Sep 17 00:00:00 2001 From: Riccardo Farris Date: Mon, 24 Aug 2026 15:27:25 +0200 Subject: [PATCH 20/20] Drop the IQTC cluster tutorial Site-specific install and SLURM instructions for one university cluster, which every mcpy release has to keep in step with the examples it submits even though nobody outside that cluster can follow it. Its only remaining tutorial neighbour is an :orphan: redirect stub, so the Tutorials toctree goes with it. --- docs/index.rst | 6 - docs/tutorials/cluster_install_iqtc.rst | 343 ------------------------ 2 files changed, 349 deletions(-) delete mode 100644 docs/tutorials/cluster_install_iqtc.rst diff --git a/docs/index.rst b/docs/index.rst index 7bdaf81..0c80b9a 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -63,12 +63,6 @@ the hybrid GCMC method it implements (see :doc:`bibliography`). gcmc_acceptance_convention replica_exchange_ladder_spacing -.. toctree:: - :caption: Tutorials - :maxdepth: 1 - - tutorials/cluster_install_iqtc - .. toctree:: :caption: Examples :maxdepth: 1 diff --git a/docs/tutorials/cluster_install_iqtc.rst b/docs/tutorials/cluster_install_iqtc.rst deleted file mode 100644 index 2d1db53..0000000 --- a/docs/tutorials/cluster_install_iqtc.rst +++ /dev/null @@ -1,343 +0,0 @@ -Cluster install: IQTC portal (mcpy + Alchemi on Ubuntu GPU nodes) -================================================================= - -This recipe sets up an environment that runs ``mcpy`` GCMC simulations with -the NVIDIA Alchemi GPU stack (``AlchemiCalculator`` / -``AlchemiFCalculator``) on the IQTC portal cluster's **iqtc13.q** partition -(RTX 4090 nodes ``merry04`` / ``merry05``). - -It is written for that cluster's actual layout — the login node and the -compute nodes have different operating systems, so a naive ``conda -activate`` in a SLURM script does not work. - -The same approach generalises to any cluster where the login and compute -nodes diverge in OS, drivers, or mounts. - - -Cluster facts you need to know ------------------------------- - -.. list-table:: - :header-rows: 1 - :widths: 18 10 22 8 16 8 10 - - * - Node - - Role - - OS - - glibc - - ``/aplic`` conda - - ``module`` - - Internet - * - ``portal03`` - - login - - Debian 10 - - 2.28 - - yes (``/aplic/anaconda/2020.02``) - - yes - - yes - * - ``merry04``, ``merry05`` - - iqtc13.q GPU - - Ubuntu 24.04 - - 2.39 - - **NO** (path absent) - - **NO** - - **NO** - -Consequences: - -- Build and install on the **login node** (only place with internet + conda). -- Run on the **compute nodes** using the env's ``python`` directly — **do - not** call ``conda activate`` from inside the SLURM script (the path - doesn't exist on Ubuntu nodes and the script will fail silently). -- Pre-download any model checkpoints on the login node; compute nodes can't - reach the internet. -- ``merry04/05`` driver is **570.x → CUDA 12.8**. Use PyTorch built for - ``cu128`` (``cu130`` wheels will see ``torch.cuda.is_available() == - False``). - -``/home`` is NFS-mounted on every node, so a conda env created on -``portal03`` is reachable from ``merry04/05``. - - -One-time setup on ``portal03`` ------------------------------- - -Workspace prep -~~~~~~~~~~~~~~ - -NFS-safe scratch for pip. The default ``/tmp`` is only ~1 GB and ``/home`` -is NFS, which trips pip with "Directory not empty" errors during installs. - -.. code-block:: bash - - mkdir -p /dev/shm/$USER-pip - export TMPDIR=/dev/shm/$USER-pip - export TEMP=$TMPDIR TMP=$TMPDIR - -Create a fresh conda env -~~~~~~~~~~~~~~~~~~~~~~~~ - -Always use a **fresh env**. Mixing ``nvalchemi`` / MACE / cuequivariance -into an existing env almost always breaks because of pinned CUDA-specific -dependencies. - -.. code-block:: bash - - source /aplic/anaconda/2020.02/etc/profile.d/conda.sh - conda create -n alchemi128 python=3.11 pip -y - -Install PyTorch with CUDA 12.8 wheels -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. code-block:: bash - - ENV_PY=$HOME/.conda/envs/alchemi128/bin/python - PYTHONNOUSERSITE=1 $ENV_PY -m pip install --no-cache-dir \ - torch --index-url https://download.pytorch.org/whl/cu128 - -The ``+cu128`` build matches the driver on ``merry04/05``. Don't take -whatever ``pip install torch`` gives you by default — that's ``+cu130`` -and will run on CPU only. - -Install nvalchemi + MACE + mcpy -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. code-block:: bash - - PYTHONNOUSERSITE=1 $ENV_PY -m pip install --no-cache-dir --prefer-binary \ - 'nvalchemi-toolkit[mace]' - - # mcpy from your clone (editable) - PYTHONNOUSERSITE=1 $ENV_PY -m pip install --no-cache-dir --prefer-binary \ - -e $HOME/bin/mcpy - -Resulting key versions (as of this writing): - -.. code-block:: text - - torch 2.11.0+cu128 - nvalchemi-toolkit 0.1.0 - nvalchemi-toolkit-ops 0.3.1 - cuequivariance-torch 0.10.0 - cuequivariance-ops-torch-cu12 0.10.0 # matches your CUDA major - mace-torch 0.3.15 - e3nn 0.4.4 - ase 3.28.0 - numpy 2.4.6 - -Pre-download the MACE-MP checkpoint -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -Compute nodes have no internet, so cache the model on the login node. The -cache directory (``~/.cache/mace/``) is on NFS, so ``merry04/05`` will see -it. - -.. code-block:: bash - - PYTHONNOUSERSITE=1 $ENV_PY -c " - from mace.calculators.foundations_models import download_mace_mp_checkpoint - print(download_mace_mp_checkpoint('medium-mpa-0')) - " - # → /home/$USER/.cache/mace/macempa0mediummodel (≈ 80 MB) - -Verify on a compute node -~~~~~~~~~~~~~~~~~~~~~~~~ - -.. code-block:: bash - - ssh merry04 'PYTHONNOUSERSITE=1 $HOME/.conda/envs/alchemi128/bin/python -c " - import torch - print(\"torch:\", torch.__version__, \"cu:\", torch.version.cuda) - print(\"cuda available:\", torch.cuda.is_available()) - print(\"cuda devices:\", torch.cuda.device_count()) - from mcpy.calculators import AlchemiCalculator, AlchemiFCalculator - print(\"mcpy alchemi calculators: ok\") - "' - -Expected: - -.. code-block:: text - - torch: 2.11.0+cu128 cu: 12.8 - cuda available: True - cuda devices: 3 - mcpy alchemi calculators: ok - -If ``cuda available: False``, the driver is older than CUDA 12.8 — -reinstall torch from the matching ``cu12X`` wheel index (e.g. ``cu124`` for -driver 550). - - -SLURM submit script -------------------- - -Save as ``submit_gcmc_alchemi.sh``: - -.. code-block:: bash - - #!/bin/bash - #SBATCH --job-name=gcmc_alchemi - #SBATCH --partition=iqtc13.q - #SBATCH --gres=gpu:1 - #SBATCH --cpus-per-task=4 - #SBATCH --time=24:00:00 - #SBATCH --output=%x_%j.out - #SBATCH --error=%x_%j.err - - set -eo pipefail - - PYTHON=$HOME/.conda/envs/alchemi128/bin/python - SCRIPT=$HOME/bin/mcpy/examples/03_gcmc_surface_mace.py - RESULTS_DIR=$SLURM_SUBMIT_DIR/results/job_${SLURM_JOB_ID:-local} - - export PYTHONNOUSERSITE=1 - export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-4} - export TMPDIR=/tmp/$USER - mkdir -p "$TMPDIR" "$RESULTS_DIR" - - echo "=== node=$(hostname) $(date) ===" - nvidia-smi -L - "$PYTHON" -c "import torch; print('torch', torch.__version__, 'cu_avail', torch.cuda.is_available())" - - cd "$RESULTS_DIR" - "$PYTHON" "$SCRIPT" - -The script takes no command-line flags: every setting is a constant at the -top of the file (``CHECKPOINT``, ``TEMPERATURE``, ``STEPS``, ``DELTA_MU_O``, -``FMAX``, ``SEED``, ...). Edit your copy before submitting, and note that the -output files land in the working directory, which is why the job changes into -``$RESULTS_DIR`` first. - -Submit with ``sbatch submit_gcmc_alchemi.sh``. - -Key SLURM-script rules: - -- **No** ``conda activate``. Use ``$HOME/.conda/envs/alchemi128/bin/python`` - directly. -- **No** ``module load``. Ubuntu compute nodes don't have environment - modules. -- **Always export** ``PYTHONNOUSERSITE=1`` to ignore - ``~/.local/lib/python3.11/site-packages``, which on many user homes - contains stale packages that shadow the env. -- ``TMPDIR=/tmp/$USER`` — local disk on the compute node, not NFS. -- ``torch.compile`` is on by default and handles GCMC's dynamic atom count - (after the first size change it switches to dynamic shapes; one-time - warmup, then ~1.3-1.4x faster). On nodes with tight host-memory limits the - compile guard-building phase can OOM at startup — if you hit that (see - troubleshooting below), set ``COMPILE_MODEL = False`` in the script. - - -Batched replica exchange on one GPU ------------------------------------ - -The script above runs a single replica. ``mcpy`` also runs **replica exchange -on one GPU**: all replicas live in a single Python process, share one -``AlchemiCalculator``, and have their trial-move energies evaluated in a single -batched forward pass per MC step. This is the GPU path the package is built -around. Unlike the MPI ``ReplicaExchange`` (one rank per replica), there is **no -mpirun** and still only ``--gres=gpu:1``. - -The script is ``examples/04_replica_exchange_batched.py``. It takes no -command-line flags: the replica temperatures are the ``TEMPERATURES`` list at -the top of the file, and its length is the replica count. - -Save as ``submit_re_gcmc_batched.sh``: - -.. code-block:: bash - - #!/bin/bash - #SBATCH --job-name=re_gcmc_batched - #SBATCH --partition=iqtc13.q - #SBATCH --gres=gpu:1 - #SBATCH --cpus-per-task=4 - #SBATCH --time=24:00:00 - #SBATCH --output=%x_%j.out - #SBATCH --error=%x_%j.err - - set -eo pipefail - - PYTHON=$HOME/.conda/envs/alchemi128/bin/python - SCRIPT=$HOME/bin/mcpy/examples/04_replica_exchange_batched.py - RESULTS_DIR=$SLURM_SUBMIT_DIR/results/job_${SLURM_JOB_ID:-local} - - export PYTHONNOUSERSITE=1 - export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-4} - export TMPDIR=/tmp/$USER - mkdir -p "$TMPDIR" "$RESULTS_DIR" - - echo "=== node=$(hostname) $(date) ===" - nvidia-smi -L - "$PYTHON" -c "import torch; print('torch', torch.__version__, 'cu_avail', torch.cuda.is_available())" - - cd "$RESULTS_DIR" - "$PYTHON" "$SCRIPT" - -Submit with ``sbatch submit_re_gcmc_batched.sh``. - -Notes specific to the batched path: - -- **One GPU, one process.** All replicas in ``TEMPERATURES`` share a single - ``AlchemiCalculator`` and are evaluated in one batched forward pass. Keep - ``--gres=gpu:1``; adding ranks or extra GPUs does not help this path. -- **``torch.compile`` stays on by default**, same as the single-replica run: - it handles GCMC's dynamic atom count via dynamic shapes after a one-time - warmup. ``COMPILE_MODEL = False`` is only the escape hatch if the compile - guard-building phase OOMs on a node with tight host-memory limits. -- **Replica count = length of ``TEMPERATURES``.** More replicas means a - larger batch and more GPU memory; if startup OOMs, shorten that list - before changing anything else. -- **Outputs** land in ``$RESULTS_DIR``: per-replica ``gcmc_batched_*.out`` and - ``.xyz`` files, plus a combined ``replica_exchange_batched.log`` with the - exchange history. - - -Common failure modes --------------------- - -.. list-table:: - :header-rows: 1 - :widths: 32 36 32 - - * - Symptom - - Cause - - Fix - * - Job exits 127 with empty ``.out``/``.err`` - - Script tried ``conda activate`` on a node where ``/aplic/anaconda/`` - doesn't exist; ``set -e`` killed it. - - Use the env's ``python`` binary directly (no activation). - * - ``python: /lib64/libc.so.6: version 'GLIBC_2.28' not found`` - - Job landed on an old node (e.g. ``g9noder25o``) with glibc < 2.28. - - Restrict the partition to modern nodes (``iqtc13.q`` is safe). - * - ``torch.cuda.is_available() == False`` despite GPU present, with - ``nvidia-smi`` driver "too old" - - PyTorch built for newer CUDA than the driver supports. - - Reinstall ``torch`` from the matching ``cu12X`` index. - * - ``urllib.error.URLError: Name or service not known`` while loading - the MACE checkpoint - - Compute node has no internet. - - Pre-download the model on the login node (see above). - * - ``cuequivariance_ops_torch_cu11`` errors at runtime, or ``e3nn`` - version too old - - ``~/.local`` user-site shadows the env. - - Always export ``PYTHONNOUSERSITE=1``. - * - ``pip install ... Directory not empty`` - - ``TMPDIR`` is on NFS. - - ``export TMPDIR=/dev/shm/$USER-pip``. - * - ``InternalTorchDynamoError: MemoryError: std::bad_alloc`` or - ``Warp CUDA error 2: out of memory`` at startup - - ``torch.compile`` tries to JIT-compile the model and the guard-building - pass OOMs against the job's host-memory limit. - - Set ``COMPILE_MODEL = False`` in the script (passes ``compile_model=False``), or - request more host memory for the job. - - -Reproducing this exact env later --------------------------------- - -.. code-block:: bash - - source /aplic/anaconda/2020.02/etc/profile.d/conda.sh - conda env remove -n alchemi128 -y - # then re-run the "One-time setup on portal03" section - -The whole install takes ~5 minutes on a good connection and uses ~7 GB.