使用 NVIDIA ALCHEMI 工具包为化学和材料科学构建自定义原子模拟工作流

# Install Atomic Simulation Environment (ASE, used in the examples below)
uv pip install ase

# Using pip
pip install nvalchemi-toolkit

# Using uv
uv venv --seed --python 3.12
uv pip install nvalchemi-toolkit

# Install from source
git clone https://github.com/NVIDIA/nvalchemi-toolkit.git 
cd nvalchemi-toolkit
uv sync --all-extras

# Add nvalchemi as a project dependency
uv add nvalchemi-toolkit

from ase import Atoms
from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics import ConvergenceHook
from nvalchemi.dynamics.optimizers import FIRE2
from nvalchemi.dynamics.integrator import VelocityVerlet

# setup some batch of atomic structures
atomic_data = [AtomicData.from_atoms(Atoms(...), device="cuda") for _ in range(16)]
batch = Batch.from_data_list(atomic_data)

# setup your MLIP and dynamics classes
mlip = ...
# optimizer convergence depends on the force norm and max values
conv_criteria = ConvergenceHook(
    criteria=[
        {"key": "forces", "threshold": 0.05, "reduce_op": "norm"},
        {"key": "forces", "threshold": 0.1, "reduce_op": "max"}
    ]
)
optimizer = FIRE2(
    mlip,
    convergence_hook=conv_criteria,
    n_steps=200
)
velverlet = VelocityVerlet(mlip, n_steps=1000)

fused = optimizer + velverlet
# context manager handles compilation and CUDA stream
with fused:
    # runs 200 steps of optimization and 1000 steps of MD
    fused.run(batch)

pipeline = optimizer | velverlet
# equivalent to manual allocation with explicit producer/consumer
# optimizer.next_rank = 1, velverlet.prior_rank = 0
# DistributedPipeline({0: optimizer, 1: velverlet})
with pipeline:
    pipeline.run(batch)

from torch import distributed as dist
from torch.utils.data.distributed import DistributedSampler
from nvalchemi.data.datapipes import Dataset, DataLoader

# set up distributed; torchrun --nproc-per-node 8 --nnodes 2 ...
dist.initialize_process_group()
# set up data and distributed sampler
dataset = Dataset(...)
data_sampler = DistributedSampler(
   dataset,
   num_replicas=dist.get_world_size(),
   rank=dist.get_rank()
)
loader = DataLoader(
   dataset,
   batch_size=128,
   sampler=sampler,
   use_stream=True
)
# configure your pipeline; 8 ranks do optimization, 8 do langevin dynamics
optimizers = [FIRE2(mlip, ..., next_rank=index + 8) for index in range(8)]
dynamics = [Langevin(mlip, ..., prior_rank=index) for index in range(8)]
pipeline = DistributedPipeline(
    {index: stage for index, stage in enumerate(optimizers + dynamics)}
)
with pipeline:
    for batch in loader:
        pipeline.run(batch)

from enum import Enum

import torch

from nvalchemi.data import Batch
from nvalchemi.dynamics.base import BaseDynamics, DynamicsStage
from nvalchemi.hooks import Hook, HookContext


class MySimulatedAnnealer(Hook):
       def __init__(
           self,
           t_start: float,
           t_end: float,
           cooldown_steps: int,
           frequency: int,
           stage: DynamicsStage
       ) -> None:
           # this hook will fire off every `frequency` MD steps,
           # bringing the temperature from `t_start` to `t_end`
           self.frequency = frequency
           self.t_start = t_start
           self.t_end = t_end
           self.cooldown_steps = cooldown_steps
           self.stage = DynamicsStage.BEFORE_STEP
           self.decay = (t_end / t_start) ** (1.0 / cooldown_steps)
           self._current_temp = t_start

       def __call__(self, ctx: HookContext, stage: Enum) -> None:
           # access the calling dynamics class through `HookContext`
           dynamics = ctx.workflow
           dynamics.target_temperature = max(
               dynamics.target_temperature * self.decay,
               self.t_end
           )


class VelocityVerlet(BaseDynamics)
	__needs_keys__: {"energies", "forces", "masses", "velocities"}
	__provides_keys__: {"positions"}

       def __init__(
           self,
           model: BaseModelMixin,
           n_steps: int,
           dt: float = 1.0,  # timestep
           target_temperature: float = 300.0,  # initial temperature
           tau: float = 10.0,  # coupling constant
           hooks: list[Hook] | None = None,
           convergence_hook: ConvergenceHook | dict | None = None,
           **kwargs,
       ):
           super().__init__(model=model, n_steps=n_steps, hooks=hooks, convergence_hook=convergence_hook)
           self.dt = dt
           self.target_temperature = target_temperature
           self.tau = tau
           self._prev_accelerations = None
       
       def pre_update(self, batch: Batch) -> None:
           # perform the first half of velocity Verlet
           with torch.no_grad():
                accelerations = batch.forces / batch.masses
                self._prev_accelerations = accelerations.clone()
                batch.positions.add_(
                    batch.velocities * dt + 0.5 * accelerations * dt**2.0
                )

       def post_update(self, batch: Batch) -> None:
           # perform second half of velocity Verlet, with thermostat
           # temperature update
           with torch.no_grad():
               new_accelerations = batch.forces / batch.masses
               batch.velocities.add_(0.5 * (self._prev_accelerations + new_accelerations) * self.dt)
               ke_per_atom = 0.5 * batch.masses * (batch.velocities**2).sum(dim=-1, keepdim=True)
               # get the total kinetic energy per system
               total_ke = scatter_add_(...)
               current_temp = 2.0 * total_ke / (batch.num_atoms * 3.0)
               ratio = self.target_temperature / current_temp
               lam = torch.sqrt(
                  torch.tensor(1.0 + (self.dt / self.tau) * (ratio - 1.0))
               ).clamp(min=0.8, max=1.2)  # clamp for stability
               batch.velocities.mul_(lam)

# configure the new dynamics class
my_velverlet = VelocityVerlet(
    ...,
    hooks=[
        MySimulatedAnnealer(t_start=900.0, t_end=300.0, cooldown_steps=10, frequency=100, stage=DynamicsStage.BEFORE_STEP)
    ],
)

from beartype import beartype
from super_mlip import BestMLIPModel

from nvalchemi._typing import ModelOutputs
from nvalchemi.models.base import BaseModelMixin, ModelConfig, NeighborConfig


class BestMLIPWrapper(nn.Module, BaseModelMixin):
    def __init__(self, model: BestMLIPModel, **kwargs):
          super().__init__(**kwargs)
          # ModelConfig declares model capabilities (which are frozen)
          # and runtime control (mutable) for the rest of the framework
          self.model_config = ModelConfig(
	                 outputs=frozenset({"energy", "forces", "hessians"}),
	                 # this is actually the default value
	                 required_inputs=frozenset({"positions", "atomic_numbers"})
	                 autograd_outputs=frozenset({"forces"}),
	                 neighbor_config=NeighborConfig(cutoff=5.0, format="coo")
)

    def adapt_input(self, data: Batch, **kwargs) -> dict[str, Any]:
        # adapts the nvalchemi data structure to what is
        # expected by the model
        model_inputs = super().adapt_input(data, **kwargs)
        # dict structure expected by BestMLIPModel
        model_inputs["atom_numbers"] = data.atomic_numbers
        model_inputs["coords"] = data.positions
        return model_inputs

   def adapt_output(self, model_output: any, data: Batch) -> ModelOutputs:
       # adapt the model outputs from the model's forward pass to
       # format expected by nvalchemi
       output = super().adapt_output(model_output, data)
       energies = model_output["energies"]
       output["energies"] = energies
       # check model config for expected outputs
       if "forces" in self.model_config.active_outputs:
           output["forces"] = model_output["forces"]
       return output

    # beartype decorator is optional, but will runtime type check arguments
    @beartype
    def forward(self, data: Batch, **kwargs) -> ModelOutputs:
        model_inputs = self.adapt_input(data, **kwargs)
        # calls BestMLIPModel's forward definition based on MRO
        model_outputs = super().forward(**model_inputs)
        return self.adapt_output(model_outputs, data)

from nvalchemi import AtomicData, Batch
from nvalchemi import data
from ase.build import slab
atoms = slab(...)

# Create AtomicData object from ase.Atoms object
data = AtomicData.from_atoms(atoms, device="cuda")
>>> data ...

data.node_properties
data.system_properties

# Create a Batch object from a list of AtomicData
batch = Batch.from_data_list([data, data, data])
batch.num_graphs
batch.get_data(0)
# get first three samples
batch[:2]
batch[mask]
batch["energies"] -> ...

batch.from_atoms([ase.Atoms,...])

# Create a dataset from ase.Atoms
writer = data.AtomicDataZarrWriter("atom_dataset.zarr")
# writer will amortize overhead by writing batches of data;
# this is equivalent to writing individual samples but efficiently
writer.write(batch)

# Read the data from zarr
reader = data.AtomicDataZarrReader("atom_dataset.zarr")
# Dataset treats device natively; individual samples # are placed on GPU and it accelerates preprocessing transforms;
# num_workers sets the number of threads used for async prefetching
dataset = data.Dataset(reader, device = "cuda", num_workers=4)
dataloader = data.DataLoader(dataset, batch_size=16)
for batch in dataloader:
	# do something with batch

标签