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/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 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/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/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/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/installation.rst b/docs/installation.rst index f73c608..81c1d82 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 @@ -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/04_replica_exchange_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/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..0fda577 100644 --- a/docs/reference/calculators.rst +++ b/docs/reference/calculators.rst @@ -8,128 +8,24 @@ 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 -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._AlchemiBase + :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. +.. autoclass:: mcpy.calculators.alchemi.AlchemiCalculator + :members: + :show-inheritance: -- ``calculator``: a constructed ASE calculator. -- ``steps`` (int): maximum LBFGS steps. -- ``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 ------------------ - -.. 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 | MACECalculator): model file path, or a built - ``MACECalculator`` 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.AlchemiFCalculator + :members: + :show-inheritance: 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 81d56bc..5af0226 100644 --- a/docs/reference/ensembles.rst +++ b/docs/reference/ensembles.rst @@ -6,161 +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, units_type, calculator, user_tag=None, - 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. -- ``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. -- ``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, units_type='metal', - 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: - -- ``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 a7e3f4f..3cc1835 100644 --- a/docs/reference/moves.rst +++ b/docs/reference/moves.rst @@ -11,144 +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. - -Methods: ``do_trial_move(atoms)``, ``acceptance_counter()``, ``get_volume()``, -``interval_ratios()``, ``total_ratios()``, ``reset_counters()``. - - -BaseMove --------- - -.. code-block:: python - - BaseMove(cell, species, seed) - -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)``. - - -InsertionMove -------------- - -.. code-block:: python - - 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 +Molecular exchange +------------------ - BrownianMove(temperature, calculator, steps, d_t, seed) +.. autoclass:: mcpy.moves.MoleculeInsertionMove + :members: + :show-inheritance: -Runs ``steps`` of Velocity-Verlet MD at ``temperature`` (timestep ``d_t`` in fs) -as the trial move, from a Maxwell-Boltzmann velocity draw. +.. autoclass:: mcpy.moves.MoleculeDeletionMove + :members: + :show-inheritance: +.. autoclass:: mcpy.moves.MoleculeDisplacementMove + :members: + :show-inheritance: -AlchemiBrownianMove -------------------- +.. automodule:: mcpy.moves.molecule_utils + :members: -.. code-block:: python +Displacement and reordering +--------------------------- - AlchemiBrownianMove(calculator, temperature, friction=0.01, steps=100, - dt=2.0, seed=0) +.. autoclass:: mcpy.moves.DisplacementMove + :members: + :show-inheritance: -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.PermutationMove + :members: + :show-inheritance: +.. autoclass:: mcpy.moves.ShakeMove + :members: + :show-inheritance: -Experimental moves ------------------- +.. autoclass:: mcpy.moves.BrownianMove + :members: + :show-inheritance: -``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. +.. autoclass:: mcpy.moves.AlchemiBrownianMove + :members: + :show-inheritance: diff --git a/docs/reference/utils.rst b/docs/reference/utils.rst index 091a148..93152e9 100644 --- a/docs/reference/utils.rst +++ b/docs/reference/utils.rst @@ -2,121 +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, warm_up=0) - -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)``. - - -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/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. 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. diff --git a/docs/tutorials/cluster_install_iqtc.rst b/docs/tutorials/cluster_install_iqtc.rst deleted file mode 100644 index 69f992c..0000000 --- a/docs/tutorials/cluster_install_iqtc.rst +++ /dev/null @@ -1,354 +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/gcmc_nano_alchemi.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())" - - "$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 - -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), pass ``--no-compile`` to the simulation 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 CPU ``re_gcmc.py`` (one MPI 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 -with ``--temperatures``; the number of values 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/re_gcmc_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())" - - "$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 - -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. -- **``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 - 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. -- **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. - - Pass ``--no-compile`` to the script (sets ``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. diff --git a/examples/00_hello_gcmc.py b/examples/00_hello_gcmc.py new file mode 100644 index 0000000..f00c924 --- /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/03_gcmc_surface_mace.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/03_gcmc_surface_mace.py b/examples/03_gcmc_surface_mace.py new file mode 100644 index 0000000..9263a7f --- /dev/null +++ b/examples/03_gcmc_surface_mace.py @@ -0,0 +1,159 @@ +"""GCMC on an Ag(111) surface with a machine-learned potential (GPU). + +A production run, in contrast to examples 00-02: a real potential, a real +surface, and both species exchanged with their reservoirs. + + * 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. + +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. + +Requirements: pip install 'nvalchemi-toolkit[mace]', and a CUDA GPU. + + python examples/03_gcmc_surface_mace.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 +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 + +# --- 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 +# 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=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 main(): + configure_logging() + + # 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() + + # 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()) + 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, + ) + + # 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'] + 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 = SLAB_SIZE + tag = f'Ag{nx * ny * nlayers}_111_dmuO_{DELTA_MU_O}_fire_fmax{FMAX}' + + gcmc = GrandCanonicalEnsemble( + atoms=atoms, + cells=cells, + calculator=calculator, + mu={'Ag': mu_ag, 'O': mu_o}, + species=species, + units_type='metal', + temperature=TEMPERATURE, + move_selector=move_selector, + random_seed=seed_ens, + outfile=f'gcmc_{tag}.out', + traj_file=f'gcmc_{tag}.xyz', + outfile_write_interval=WRITE_INTERVAL, + trajectory_write_interval=WRITE_INTERVAL, + ) + gcmc.run(STEPS) + + +if __name__ == '__main__': + main() diff --git a/examples/04_replica_exchange_batched.py b/examples/04_replica_exchange_batched.py new file mode 100644 index 0000000..15b1a5e --- /dev/null +++ b/examples/04_replica_exchange_batched.py @@ -0,0 +1,156 @@ +"""Replica-exchange GCMC on a single GPU, with batched energy evaluation. + +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. + +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. + +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. + +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 +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 + +# --- 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(): + configure_logging() + + 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. 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: this is what makes the batched + # forward pass possible in the first place. + calculator = AlchemiFCalculator( + checkpoint=CHECKPOINT, + device=DEVICE, + enable_cueq=ENABLE_CUEQ, + compile_model=COMPILE_MODEL, + ) + + # 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)) + + # 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]) + + move_selector = MoveSelector( + [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_{DELTA_MU_O}_rank{rank}' + return GrandCanonicalEnsemble( + atoms=atoms, + cells=[cell_ag, cell_o], + calculator=calculator, + mu=mus, + species=['Ag', 'O'], + units_type='metal', + temperature=T, + move_selector=move_selector, + 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, + ) + + replica_exchange = BatchedReplicaExchange( + gcmc_factory, + calculator=calculator, + temperatures=TEMPERATURES, + gcmc_steps=GCMC_STEPS, + exchange_interval=EXCHANGE_INTERVAL, + write_out_interval=WRITE_INTERVAL, + seed=master_seed, + outfile='replica_exchange_batched.log', + ) + replica_exchange.run() + + +if __name__ == '__main__': + main() 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/README.md b/examples/README.md new file mode 100644 index 0000000..f35a2b1 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,25 @@ +# Examples + +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 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 +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. 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_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.py b/examples/re_gcmc_batched.py deleted file mode 100644 index 7dc4ae5..0000000 --- a/examples/re_gcmc_batched.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Batched replica-exchange GCMC on a single GPU. - -Single-process variant of ``re_gcmc.py``. 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 - -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 - -import numpy as np -from ase.build import fcc111 -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 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__, - 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() - - -def main(): - args = parse_args() - 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]] - 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] - base_atoms.set_constraint(FixAtoms(indices=bottom_layer)) - - # ONE model, shared across replicas. Batched eval reuses this calculator. - calculator = AlchemiFCalculator( - checkpoint=args.checkpoint, - device=args.device, - enable_cueq=not args.no_cueq, - compile_model=not args.no_compile, - ) - - 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} - - def gcmc_factory(T, rank): - 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}) - - 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])], - ) - - 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_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'), - ) - - 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_batched.log'), - ) - pt.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() diff --git a/mcpy/calculators/__init__.py b/mcpy/calculators/__init__.py index ec3830f..1052496 100644 --- a/mcpy/calculators/__init__.py +++ b/mcpy/calculators/__init__.py @@ -3,16 +3,14 @@ __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 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 b11f7c3..0000000 --- a/mcpy/calculators/_alchemi_common.py +++ /dev/null @@ -1,249 +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: 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: - 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() - # 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." - ) - - -def _build_nl(batch: Batch, nl_hook: NeighborListHook) -> None: - ctx = DynamicsContext(batch=batch, step_count=0) - nl_hook(ctx, DynamicsStage.BEFORE_COMPUTE) - - -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()) - if fixed: - relaxed = relaxed.copy() - relaxed[list(fixed)] = atoms.positions[list(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) - - # 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) - 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: build NL and compute initial forces before the MD loop. - _build_nl(batch, nl_hook) - opt.compute(batch) - if fixed: - # FreezeAtomsHook zeros frozen forces only at AFTER_POST_UPDATE; the - # bootstrap compute() runs outside the hook loop, so step 0 would - # otherwise kick and displace the frozen rows once. - batch.forces[fixed] = 0.0 - - opt.run(batch) - _write_back_positions(atoms, batch) diff --git a/mcpy/calculators/alchemi.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 3150f4c..0000000 --- a/mcpy/calculators/alchemi_calculator.py +++ /dev/null @@ -1,187 +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, - _make_multi_batch, - _per_graph_energies, - _run_langevin_md, -) - - -class AlchemiCalculator: - """ - Energy-only Alchemi calculator (no geometry relaxation). - - Equivalent to MACECalculator but uses 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_multi_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 deleted file mode 100644 index ff011bf..0000000 --- a/mcpy/calculators/alchemi_f_calculator.py +++ /dev/null @@ -1,344 +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 ( - _build_nl, - _fixed_indices, - _freeze_hook_for, - _load_model, - _make_batch, - _make_multi_batch, - _per_graph_energies, - _run_langevin_md, - _write_back_positions, -) - - -logger = logging.getLogger(__name__) - -_ALCHEMI_OPTIMIZERS = {'fire': AlchemiFIRE, 'fire2': AlchemiFIRE2} - - -class AlchemiFCalculator: - """ - 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. - - 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. - """ - 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()) - - 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) - - # 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) - - _build_nl(batch, nl_hook) - opt.compute(batch) - if fixed: - batch.forces[fixed] = 0.0 # see get_potential_energy bootstrap - - return self._run_compacted(opt, batch, atoms_list) - - 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)) - pos = batch.positions.detach().cpu().numpy() - 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) - - 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, - ) 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/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 1b0bf96..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 @@ -23,7 +22,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) @@ -76,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) @@ -106,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 @@ -119,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): """ @@ -154,9 +131,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..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 @@ -57,7 +56,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 @@ -83,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): """ @@ -121,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 a7af832..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 @@ -32,7 +31,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 @@ -62,28 +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] - - def get_species(self): - """ - Get the species present in the custom cell. - """ - return list(self.species_radii.keys()) + return np.einsum('ij,ij->i', diff, diff) <= self.radius ** 2 def _sample_sphere_points(self, n): """ @@ -121,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/ensembles/base_ensemble.py b/mcpy/ensembles/base_ensemble.py index 70cc7ad..9c68dd7 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 @@ -179,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) @@ -264,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 28a1ca4..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. @@ -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 @@ -235,71 +230,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) - - viable = [i for i in active if i in trial_meta] - if not viable: + proposal = self.replicas[i]._propose() + if proposal is not None: + pending[i] = proposal + + 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 @@ -427,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) @@ -532,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 5cd3a2b..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,13 +23,12 @@ class CanonicalEnsemble(BaseEnsemble): def __init__(self, atoms, calculator, + optimizer, + move_selector, cells=None, - units_type: str = 'metal', random_seed=None, - optimizer=None, fmax=0.1, temperature=300, - move_selector=None, constraints=None, traj_file: str = 'trajectory.xyz', traj_mode: str = 'w', @@ -42,7 +40,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, @@ -59,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 @@ -90,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. @@ -112,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() @@ -137,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() @@ -161,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) @@ -177,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 024b5f9..ef79b46 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, @@ -140,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. @@ -163,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): @@ -241,71 +239,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.""" @@ -328,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/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/moves/move_selector.py b/mcpy/moves/move_selector.py index b4db54e..47fa07a 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 @@ -58,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 @@ -71,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) @@ -113,14 +111,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 +136,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/mcpy/utils/set_unit_constant.py b/mcpy/utils/set_unit_constant.py index dc684e1..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 @@ -14,14 +15,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 @@ -62,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} 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_ag111_gcmc_golden.py b/tests/test_ag111_gcmc_golden.py new file mode 100644 index 0000000..1fae919 --- /dev/null +++ b/tests/test_ag111_gcmc_golden.py @@ -0,0 +1,341 @@ +"""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/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. +""" +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}') diff --git a/tests/test_audit_regressions.py b/tests/test_audit_regressions.py index abf670d..d472f73 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) # -------------------------------------------------------------------------- @@ -216,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 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() 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] 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" # --------------------------------------------------------------------------- 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):