.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/advanced/04_mace_nvt.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_advanced_04_mace_nvt.py: NVT MD with MACE (with LJ Fallback) ===================================== MACE (Message-passing Atomic Cluster Expansion) is an equivariant message-passing neural network potential with competitive accuracy on bulk and molecular systems within its training distribution. The MACE-MP family of foundation models is pre-trained on the Materials Project and provides reasonable zero-shot predictions across a broad range of compositions, though accuracy should be validated for the specific system of interest before production use — particularly for chemistries or structures not represented in the Materials Project training set. This example shows how to plug a MACE model into an NVT simulation using :class:`~nvalchemi.models.mace.MACEWrapper`. If MACE is not installed (or no checkpoint path is provided), the example falls back to the Lennard-Jones potential so that it can run in CI without any ML-potential dependency. Key concepts demonstrated -------------------------- * Loading a MACE model via ``MACEWrapper.from_checkpoint``. * Reading ``model.model_card.neighbor_config`` to wire a :class:`~nvalchemi.dynamics.hooks.NeighborListHook` automatically — the same code works for LJ (MATRIX format) and MACE (COO format). * Model-agnostic temperature and energy observation. Setting up MACE --------------- Install the optional MACE dependency:: pip install 'nvalchemi-toolkit[mace]' Then point the ``MACE_MODEL_PATH`` environment variable to a MACE-MP checkpoint, e.g. the medium model:: export MACE_MODEL_PATH=medium-0b2 # named checkpoint (auto-downloads) # or a local path: export MACE_MODEL_PATH=/path/to/mace_mp_medium.pt When ``MACE_MODEL_PATH`` is unset the example uses the LJ potential. .. GENERATED FROM PYTHON SOURCE LINES 56-75 .. code-block:: Python from __future__ import annotations import logging import math import os import torch from nvalchemi.data import AtomicData, Batch from nvalchemi.dynamics import NVTLangevin from nvalchemi.dynamics.hooks import NeighborListHook # KB_EV and kinetic_energy_per_graph are internal helpers used by the built-in # integrators. A stable public re-export may be added in a future release. from nvalchemi.dynamics.hooks._utils import KB_EV, kinetic_energy_per_graph logging.basicConfig(level=logging.INFO) .. GENERATED FROM PYTHON SOURCE LINES 76-85 Model selection: MACE or LJ fallback -------------------------------------- The ``MACE_MODEL_PATH`` environment variable selects the model. If it is unset (or MACE is not installed), we fall back to the LJ potential so the example works in CI without any MACE dependency. Both code paths produce a ``model`` object that satisfies the :class:`~nvalchemi.models.base.BaseModelMixin` interface, so all simulation code below is model-agnostic. .. GENERATED FROM PYTHON SOURCE LINES 85-113 .. code-block:: Python MACE_MODEL_PATH = os.environ.get("MACE_MODEL_PATH", "") USE_MACE = False if MACE_MODEL_PATH: try: from nvalchemi.models.mace import MACEWrapper model = MACEWrapper.from_checkpoint( checkpoint_path=MACE_MODEL_PATH, device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), ) print(f"Using MACE model from {MACE_MODEL_PATH}") USE_MACE = True except Exception as exc: print(f"Failed to load MACE ({exc}). Falling back to LJ.") if not USE_MACE: from nvalchemi.models.lj import LennardJonesModelWrapper model = LennardJonesModelWrapper( # type: ignore[assignment] epsilon=0.0104, # eV — argon ε sigma=3.40, # Å — argon σ cutoff=8.5, # Å max_neighbors=32, ) print("Using LJ model (set MACE_MODEL_PATH=/path/to/model.pt to use MACE)") .. rst-class:: sphx-glr-script-out .. code-block:: none Using LJ model (set MACE_MODEL_PATH=/path/to/model.pt to use MACE) .. GENERATED FROM PYTHON SOURCE LINES 114-122 Neighbor list: automatic wiring via ModelCard ---------------------------------------------- :attr:`~nvalchemi.models.base.ModelCard.neighbor_config` encodes the cutoff, list format (COO or MATRIX), and whether to use a half-list. :class:`~nvalchemi.dynamics.hooks.NeighborListHook` reads this automatically. If ``neighbor_config`` is ``None`` (e.g. a demo model that does its own neighbour search), no hook is needed. .. GENERATED FROM PYTHON SOURCE LINES 122-133 .. code-block:: Python neighbor_hook = None if model.model_card.neighbor_config is not None: neighbor_hook = NeighborListHook(model.model_card.neighbor_config) print( f"Wired NeighborListHook: format={model.model_card.neighbor_config.format.name}, " f"cutoff={model.model_card.neighbor_config.cutoff:.2f} Å" ) else: print("Model does not require a NeighborListHook.") .. rst-class:: sphx-glr-script-out .. code-block:: none Wired NeighborListHook: format=MATRIX, cutoff=8.50 Å .. GENERATED FROM PYTHON SOURCE LINES 134-139 Building the system -------------------- When MACE is active we use a small water cluster (3 molecules, 9 atoms). With LJ we use an 8-atom argon cluster. Both result in the same simulation code below. .. GENERATED FROM PYTHON SOURCE LINES 139-215 .. code-block:: Python _R_MIN_AR = 2 ** (1 / 6) * 3.40 # ≈ 3.82 Å def _make_argon_cluster(n_per_side: int = 2, seed: int = 0) -> AtomicData: """8-atom argon cluster on a cubic lattice.""" n = n_per_side**3 spacing = _R_MIN_AR * 1.05 coords = torch.arange(n_per_side, dtype=torch.float32) * spacing gx, gy, gz = torch.meshgrid(coords, coords, coords, indexing="ij") positions = torch.stack([gx.flatten(), gy.flatten(), gz.flatten()], dim=-1) torch.manual_seed(seed) positions = positions + 0.05 * torch.randn_like(positions) # Maxwell-Boltzmann at 300 K: v_std = sqrt(kB * T / m), m_Ar = 39.948 amu _v_std = math.sqrt(KB_EV * 300.0 / 39.948) return AtomicData( positions=positions, atomic_numbers=torch.full((n,), 18, dtype=torch.long), # Ar Z=18 forces=torch.zeros(n, 3), energies=torch.zeros(1, 1), velocities=_v_std * torch.randn(n, 3), ) def _make_water_cluster(n_molecules: int = 3, seed: int = 0) -> AtomicData: """Small water cluster (H₂O): n_molecules × {O, H, H} atoms. Geometry: O at the origin of each molecule; H atoms at ±104.5° / 0.96 Å. Molecules are spaced 3.5 Å apart along the x-axis. """ torch.manual_seed(seed) positions_list = [] atomic_numbers_list = [] o_h_bond = 0.96 # Å half_angle = math.radians(104.5 / 2) for i in range(n_molecules): ox = float(i) * 3.5 # Oxygen o_pos = torch.tensor([ox, 0.0, 0.0]) # Two hydrogens in the xz-plane h1_pos = o_pos + o_h_bond * torch.tensor( [math.sin(half_angle), 0.0, math.cos(half_angle)] ) h2_pos = o_pos + o_h_bond * torch.tensor( [-math.sin(half_angle), 0.0, math.cos(half_angle)] ) positions_list.extend([o_pos, h1_pos, h2_pos]) atomic_numbers_list.extend([8, 1, 1]) # O=8, H=1 positions = torch.stack(positions_list) + 0.02 * torch.randn(len(positions_list), 3) n_atoms = len(atomic_numbers_list) # Maxwell-Boltzmann at 300 K per species: O m=15.999 amu, H m=1.008 amu _masses = torch.tensor([15.999 if z == 8 else 1.008 for z in atomic_numbers_list]) _v_stds = (KB_EV * 300.0 / _masses).sqrt().unsqueeze(-1) # (n_atoms, 1) return AtomicData( positions=positions.float(), atomic_numbers=torch.tensor(atomic_numbers_list, dtype=torch.long), forces=torch.zeros(n_atoms, 3), energies=torch.zeros(1, 1), velocities=(_v_stds * torch.randn(n_atoms, 3)).float(), ) if USE_MACE: data = _make_water_cluster(n_molecules=3) system_label = "water cluster (9 atoms, 3 H₂O)" else: data = _make_argon_cluster(n_per_side=2, seed=0) system_label = "argon cluster (8 atoms)" sim_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") batch = Batch.from_data_list([data], device=sim_device) print(f"\nSystem: {system_label} → {batch.num_nodes} atoms on {sim_device}") .. rst-class:: sphx-glr-script-out .. code-block:: none System: argon cluster (8 atoms) → 8 atoms on cuda .. GENERATED FROM PYTHON SOURCE LINES 216-220 NVT MD simulation ------------------ 100 steps of NVTLangevin at 300 K. The neighbor hook (if any) fires at BEFORE_COMPUTE to rebuild the neighbor list before each forward pass. .. GENERATED FROM PYTHON SOURCE LINES 220-237 .. code-block:: Python nvt = NVTLangevin( model=model, dt=0.1, # fs temperature=300.0, friction=0.5, n_steps=100, random_seed=99, ) if neighbor_hook is not None: nvt.register_hook(neighbor_hook) print("\nRunning 100 NVT steps …") batch = nvt.run(batch) print(f"Run complete: {nvt.step_count} steps") .. rst-class:: sphx-glr-script-out .. code-block:: none Running 100 NVT steps … Run complete: 100 steps .. GENERATED FROM PYTHON SOURCE LINES 238-242 Model-agnostic observation --------------------------- Temperature and mean energy are computed the same way regardless of whether MACE or LJ provided the forces. .. GENERATED FROM PYTHON SOURCE LINES 242-262 .. code-block:: Python # Kinetic energy per graph (eV) → temperature. # 2 KE = 3 N k_B T → T = 2 KE / (3 N k_B) ke_per_graph = kinetic_energy_per_graph( velocities=batch.velocities, masses=batch.atomic_masses, batch_idx=batch.batch, num_graphs=batch.num_graphs, ) # [B, 1] n_atoms_per_graph = batch.num_nodes_per_graph.float() # [B] T_inst = (2.0 * ke_per_graph.squeeze(-1)) / (3.0 * n_atoms_per_graph * KB_EV) mean_E = batch.energies.squeeze(-1).mean().item() mean_T = T_inst.mean().item() print("\nFinal observables (model-agnostic):") print(f" Mean energy : {mean_E:+.4f} eV") print(f" Inst. temp. : {mean_T:.1f} K (target = 300 K)") print(f" Model used : {'MACE' if USE_MACE else 'Lennard-Jones (fallback)'}") .. rst-class:: sphx-glr-script-out .. code-block:: none Final observables (model-agnostic): Mean energy : -0.1421 eV Inst. temp. : 207.8 K (target = 300 K) Model used : Lennard-Jones (fallback) .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.460 seconds) .. _sphx_glr_download_examples_advanced_04_mace_nvt.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 04_mace_nvt.ipynb <04_mace_nvt.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 04_mace_nvt.py <04_mace_nvt.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 04_mace_nvt.zip <04_mace_nvt.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_