Calculators

mcpy.calculators wraps energy backends for the ensembles. Conceptual background, including the relaxation-inside-energy design, is in Computing energies and relaxing trial moves.

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

BaseCalculator(calculator, steps, fmax)

Adapter for any ASE calculator. get_potential_energy(atoms) relaxes atoms with LBFGS up to steps iterations or fmax, then returns the relaxed energy.

  • calculator: a constructed ASE calculator.

  • steps (int): maximum LBFGS steps.

  • fmax (float): force tolerance in eV/Å.

MACECalculator

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

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

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 Computing energies and relaxing trial moves.

  • energy_only (bool): drop force computation (no autograd graph) for a ~12% memory saving. Energy is unchanged; forces are unavailable.

AlchemiFCalculator

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 Computing energies and relaxing trial moves.

run_md

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.