.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/distributed/06_mace_npt_distributed.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_distributed_06_mace_npt_distributed.py: MACE NPT: domain-decomposed constant-pressure MD with a barostat ================================================================= The constant-pressure sibling of example 03. Load a MACE checkpoint, run a short :class:`~nvalchemi.dynamics.NPT` trajectory across multiple ranks under :class:`~nvalchemi.distributed.DomainParallel`, and record the trajectory — including the **evolving cell** — to an xyz file from rank 0. Why NPT-under-DD needs more than NVT: a barostat and a thermostat both couple to *global* thermodynamic quantities — the total kinetic energy, the total degrees of freedom, and the full-system pressure tensor — none of which any single rank can see from its owned subdomain alone. The framework handles this transparently: :class:`~nvalchemi.dynamics.NPT` declares ``__dd_thermo_kind__ = "npt"``, and on ``partition()`` :class:`~nvalchemi.distributed.DomainParallel` installs a dynamics coordinator that * all-reduces the per-rank kinetic energy and kinetic-pressure tensor into mesh-global values (the consolidated virial is already global), * replaces the integrator's per-shard degrees-of-freedom with the global count (and rescales the Nosé–Hoover chain masses accordingly), and * keeps the replicated barostat state and the cell **byte-identical** across ranks by broadcasting them from rank 0 each step. So the wrapper and the integrator stay ensemble-correct; the only user change from example 03 is asking the model for ``stress`` (the barostat needs the virial) and swapping ``NVTLangevin`` for ``NPT``. System: alpha-quartz SiO2 supercell, isotropic barostat at zero external pressure — the box relaxes toward its equilibrium volume while the thermostat holds temperature. .. note:: Run with:: torchrun --nproc_per_node=2 examples/distributed/06_mace_npt_distributed.py For multi-GPU MACE+cuEquivariance, set the env var below to avoid a JIT-compilation race across ranks:: CUEQUIVARIANCE_OPS_PARALLEL_COMPILE=0 \ torchrun --nproc_per_node=N \ examples/distributed/06_mace_npt_distributed.py Output xyz file at ``./mace_npt_trajectory.xyz`` (rank 0 only); each frame carries the cell at that step, so the volume relaxation is visible in OVITO/VMD. .. GENERATED FROM PYTHON SOURCE LINES 66-370 .. code-block:: Python from __future__ import annotations import argparse import os import sys import warnings from pathlib import Path import torch from loguru import logger from nvalchemi.data import AtomicData, Batch from nvalchemi.distributed import DomainConfig, DomainParallel, HookScope from nvalchemi.dynamics import NPT, HostMemory from nvalchemi.dynamics.base import DynamicsStage from nvalchemi.dynamics.hooks import SnapshotHook from nvalchemi.hooks import NeighborListHook # Skip the heavy distributed launch during the Sphinx-Gallery docs build (it has # no torchrun environment), mirroring examples 01-03. _DOCS_BUILD = os.environ.get("NVALCHEMI_SPHINX_BUILD") == "1" _DISTRIBUTED_ENV = "RANK" in os.environ and "WORLD_SIZE" in os.environ # Reuse the SiO2 supercell builder from the benchmark suite (see example 03). # Sphinx-Gallery executes examples without defining ``__file__`` and skips the # launcher path below, so only resolve the benchmark import for a real run. if not _DOCS_BUILD: sys.path.insert( 0, str(Path(__file__).resolve().parents[2] / "benchmark" / "distributed") ) from _benchmark_common import build_sio2_supercell # noqa: E402 # ---------------------------------------------------------------------- # System construction (rank 0 — DomainParallel scatters from there) # ---------------------------------------------------------------------- def build_initial_batch( repeats: tuple[int, int, int], dtype: torch.dtype, device: torch.device ) -> Batch: pos, numbers, masses, cell, velocities = build_sio2_supercell( repeats=repeats, dtype=dtype, seed=0 ) data = AtomicData( positions=pos.to(device), atomic_numbers=numbers.to(device), atomic_masses=masses.to(device), cell=cell.to(device).unsqueeze(0), pbc=torch.tensor([[True, True, True]], device=device), ) data.add_node_property("velocities", velocities.to(device)) # NPT reads a stress tensor each step; pre-allocate so the field exists # before the first forward (the model overwrites it). data["stress"] = torch.zeros(1, 3, 3, dtype=dtype, device=device) return Batch.from_data_list([data], device=device) def cell_volume(batch: Batch) -> float: """Volume (ų) of a single-system batch's (replicated) cell.""" cell = batch.cell if cell.dim() == 3: cell = cell[0] return float(torch.linalg.det(cell).abs()) # ---------------------------------------------------------------------- # Trajectory persistence (rank 0 only) # ---------------------------------------------------------------------- def write_trajectory_xyz(sink: HostMemory, path: Path) -> int: """Decode the :class:`HostMemory` sink into per-frame :class:`ase.Atoms` (with the step's cell) and write an extxyz trajectory. Returns the number of frames written. """ from ase import Atoms from ase.io import write as ase_write trajectory_batch = sink.read() n_frames = trajectory_batch.num_graphs if path.exists(): path.unlink() for frame in range(n_frames): single = trajectory_batch.index_select(torch.tensor([frame])) cell = single.cell if cell.dim() == 3: cell = cell.squeeze(0) atoms = Atoms( numbers=single.atomic_numbers.detach().cpu().numpy(), positions=single.positions.detach().cpu().numpy(), cell=cell.detach().cpu().numpy(), pbc=True, ) atoms.info["frame"] = frame ase_write(str(path), atoms, format="extxyz", append=True) return n_frames # ---------------------------------------------------------------------- # Main # ---------------------------------------------------------------------- def main() -> None: parser = argparse.ArgumentParser( description="MACE NPT (barostat) under DomainParallel." ) parser.add_argument( "--checkpoint", default="medium-0b2", help="MACE foundation model checkpoint name. " "Default fetches MACE-MP-0b2 from HuggingFace.", ) parser.add_argument( "--repeats", type=int, nargs=3, default=[3, 3, 3], help="SiO2 unit-cell repeats along (a, b, c). 3x3x3 → 243 atoms.", ) parser.add_argument("--n-steps", type=int, default=200) parser.add_argument("--temperature-k", type=float, default=300.0) parser.add_argument( "--pressure", type=float, default=0.0, help="Target external pressure in eV/ų (1 bar ≈ 6.32e-7). " "Default 0 → relax toward the equilibrium volume.", ) parser.add_argument( "--dt-fs", type=float, default=1.0, help="MD timestep in femtoseconds." ) parser.add_argument( "--barostat-time-fs", type=float, default=1000.0, help="Barostat coupling time τ_P (fs). Larger = gentler cell motion.", ) parser.add_argument( "--thermostat-time-fs", type=float, default=100.0, help="Thermostat coupling time τ_T (fs).", ) parser.add_argument( "--snapshot-every", type=int, default=10, help="Persist a frame to the trajectory sink every N steps.", ) parser.add_argument( "--output-xyz", type=Path, default=Path("mace_npt_trajectory.xyz"), help="xyz file path (rank 0 only).", ) parser.add_argument( "--dtype", default="float32", choices=["float32", "float64"], help="Model / simulation dtype.", ) args = parser.parse_args() dtype = torch.float64 if args.dtype == "float64" else torch.float32 # Docs build / no torchrun: no process group to join, so skip the launch. if _DOCS_BUILD or not _DISTRIBUTED_ENV: logger.info( "Not running under torchrun — skipping the distributed run. " "Launch with: torchrun --nproc_per_node=N " "examples/distributed/06_mace_npt_distributed.py" ) return # ----- Distributed bootstrap via DistributedManager (see example 03) ----- from nvalchemi.distributed import DistributedManager DistributedManager.initialize() try: dm = DistributedManager() rank, world_size, device = dm.rank, dm.world_size, torch.device(dm.device) mesh = dm.initialize_mesh(mesh_shape=(world_size,), mesh_dim_names=("domain",)) if rank == 0: logger.info( "MACE NPT distributed: world_size={ws} device={dev} " "checkpoint={ckpt} repeats={r} n_steps={n} T={T}K " "P={P} eV/ų dt={dt}fs", ws=world_size, dev=device, ckpt=args.checkpoint, r=tuple(args.repeats), n=args.n_steps, T=args.temperature_k, P=args.pressure, dt=args.dt_fs, ) # ----- Load MACE wrapper ----- with warnings.catch_warnings(): warnings.simplefilter("ignore") from nvalchemi.models.mace import MACEWrapper wrapper = MACEWrapper.from_checkpoint( args.checkpoint, dtype=dtype, device=device ).eval() # The barostat needs the virial → ask MACE for stress (example 03 only # needed energy + forces). Everything else is unchanged. wrapper.set_config("active_outputs", {"energy", "forces", "stress"}) if rank == 0: logger.info("MACE wrapper ready: cutoff={c} Å", c=wrapper.cutoff) # ----- Domain config ----- domain_cfg = DomainConfig(cutoff=float(wrapper.cutoff), skin=0.5, mesh=mesh) # ----- Hooks ----- nl_hook = NeighborListHook( wrapper.model_config.neighbor_config, skin=0.5, stage=DynamicsStage.BEFORE_COMPUTE, ) n_frames_expected = (args.n_steps // args.snapshot_every) + 1 trajectory_sink = HostMemory(capacity=n_frames_expected) snapshot_hook = SnapshotHook( sink=trajectory_sink, frequency=args.snapshot_every, ) # RANK_ZERO scope: gather the full system (with the current cell) onto # rank 0 so each frame is the whole box, not one rank's shard. snapshot_hook.scope = HookScope.RANK_ZERO # ----- Inner integrator: NPT ----- # NPT couples a Nosé–Hoover thermostat chain to the particle velocities # and a barostat to the cell. Under DomainParallel the dynamics # coordinator globalises the kinetic energy / DOF / pressure tensor and # broadcasts the replicated barostat state + cell each step (see the # module docstring), so this stays ensemble-correct across ranks. integrator = NPT( model=wrapper, dt=args.dt_fs, temperature=args.temperature_k, pressure=args.pressure, barostat_time=args.barostat_time_fs, thermostat_time=args.thermostat_time_fs, pressure_coupling="isotropic", chain_length=3, hooks=[nl_hook], n_steps=args.n_steps, ) # ----- DomainParallel wrapping (SnapshotHook on the outer, as in 03) ----- # ``with dynamics:`` makes teardown exception-safe; the process-group # lifecycle stays at launcher scope (``DistributedManager.cleanup()`` below). with DomainParallel( dynamics=integrator, config=domain_cfg, n_steps=args.n_steps, hooks=[snapshot_hook], ) as dynamics: # ----- Build the initial batch on rank 0 and partition ----- initial_batch = ( build_initial_batch(tuple(args.repeats), dtype=dtype, device=device) if rank == 0 else None ) v0 = cell_volume(initial_batch) if rank == 0 else None owned_batch = dynamics.partition(initial_batch) if rank == 0: logger.info( "Partitioned: n_owned (rank 0) = {n} of {tot} atoms; V0 = {v:.2f} ų", n=int(owned_batch.positions.shape[0]), tot=int(initial_batch.positions.shape[0]), v=v0, ) # ----- Run the trajectory ----- final_batch = dynamics.run(owned_batch) # ----- Report volume relaxation + persist trajectory ----- # The cell is replicated (broadcast from rank 0 each step), so any rank's # final cell is the global cell; report it from rank 0. if rank == 0: v1 = cell_volume(final_batch) logger.info( "Volume: {v0:.2f} → {v1:.2f} ų ({pct:+.2f}%) over {n} steps", v0=v0, v1=v1, pct=100.0 * (v1 - v0) / v0, n=args.n_steps, ) n_frames = write_trajectory_xyz(trajectory_sink, args.output_xyz) logger.info( "Done. Wrote {f} xyz frames to {p}.", f=n_frames, p=args.output_xyz ) finally: # Process-group teardown stays at launcher scope. DistributedManager.cleanup() if __name__ == "__main__": main() .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.002 seconds) .. _sphx_glr_download_examples_distributed_06_mace_npt_distributed.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 06_mace_npt_distributed.ipynb <06_mace_npt_distributed.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 06_mace_npt_distributed.py <06_mace_npt_distributed.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 06_mace_npt_distributed.zip <06_mace_npt_distributed.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_