B4 Simulate silicon with MACE

• System: cubic diamond Si, 2 x 2 x 2 cells, 64 atoms.

• Potential: pinned MACE model for energy and forces.

• Run: NVE velocity Verlet, 0.1 fs step, 10 warmup and 200 timed steps.

• Initial velocities: 300 K. NVE has no thermostat, so temperature is not held at 300 K.

The ALCHEMI example builds AtomicData per system, assigns seeded velocities and collects the systems into a Batch (one replica: still one physical simulation):

alchemi_si.py

def make_batch(device, replicas, cells=2):
    atoms = bulk("Si", "diamond", a=5.43, cubic=True).repeat((cells, cells, cells))
    count = len(atoms)
    systems = []
    for index in range(replicas):
        item = AtomicData.from_atoms(atoms, device=device, dtype=torch.float32)
        item.add_node_property("forces", torch.zeros((count, 3), device=device))
        item.add_system_property("energy", torch.zeros((1, 1), device=device))
        generator = torch.Generator(device=device).manual_seed(20260924 + index)
        velocity_scale = math.sqrt(KB_EV * 300.0 / 28.085)
        velocities = torch.randn((count, 3), device=device, generator=generator) * velocity_scale
        velocities -= velocities.mean(dim=0, keepdim=True)
        item.add_node_property("velocities", velocities)
        systems.append(item)
    return Batch.from_data_list(systems, device=device)

Explicit integrator and neighbour-list hook:

• --integrator nve: NVE.

• --integrator langevin: NVT Langevin, temperature and friction set in code.

• Different physics: never share one performance or trajectory claim.

alchemi_si.py

def run_steps(model, batch, steps, seed, integrator="langevin"):
    if integrator == "nve":
        dynamics = NVE(model=model, dt=0.1, n_steps=steps)
    else:
        dynamics = NVTLangevin(
            model=model, dt=0.1, temperature=300.0, friction=0.5,
            n_steps=steps, random_seed=seed,
        )
    dynamics.register_hook(NeighborListHook(
        model.model_config.neighbor_config,
        stage=DynamicsStage.BEFORE_COMPUTE,
    ))
    result = dynamics.run(batch)
    if dynamics.step_count != steps:
        raise RuntimeError("ALCHEMI did not finish the requested steps")
    return result

Seeded velocities and timing code: full ALCHEMI example.

The cell below runs the single trajectory (Arrhenius environment, one GPU allocated). The helper mounts the model and this script read-only into the SIF, so it runs the code above. Terminal equivalent, from the repository root:

bash scripts/run-alchemi.sh --replicas 1 --cells 2 \
  --integrator nve --warmup 10 --steps 200

Output is structured; the notebook tabulates a few fields. Neither submits a job: both need an existing GPU allocation and the setup-page environment variables.

Show code cell source

Hide code cell source

import json
import subprocess
from time import perf_counter
from pathlib import Path
from IPython.display import Markdown, display

lesson = Path.cwd().parent if Path.cwd().name == "episodes" else Path.cwd()
started = perf_counter()
completed = subprocess.run(
    ["bash", str(lesson / "scripts/run-alchemi.sh"), "--replicas", "1", "--cells", "2",
     "--integrator", "nve", "--warmup", "10", "--steps", "200"],
    check=True, capture_output=True, text=True,
)
command_wall = perf_counter() - started
one = json.loads(completed.stdout.strip().splitlines()[-1])
display(Markdown(
    "| Atoms | Timed steps | Whole command (s) | MD steps only (s) | Peak Torch memory (GiB) |\n"
    "| ---: | ---: | ---: | ---: | ---: |\n"
    f"| {one['atoms_per_replica']} | {one['measured_steps']} | "
    f"{command_wall:.3f} | {one['measured_seconds']:.3f} | "
    f"{one['peak_torch_allocated_bytes']/2**30:.2f} |"
))

Note

• Whole command: SIF startup, model loading, warmup and MD.

• MD steps only: excludes those setup costs.

• Neither includes queue wait.

• One short timing does not predict sustained production speed or validate the model against reference science.

• PyTorch peak allocation is not total GPU memory use.