TorchSim

TorchSim is a GPU-native atomistic simulation engine built on PyTorch. It provides batched MD, relaxation, and more with significant speedups over ASE. See the TorchSim documentation for full details.

SevenNet provides its own SevenNetModel wrapper for TorchSim, located in sevenn.torchsim.

Installation

TorchSim is an optional dependency of SevenNet (requires Python >= 3.12). Install it via:

pip install sevenn[torchsim]

Usage

Loading a model

SevenNetModel accepts the same model specifiers as SevenNetCalculator: a pretrained model name, a checkpoint path, or a model object.

from sevenn.torchsim import SevenNetModel

model = SevenNetModel(model="7net-omni", modal="mpa")

You can enable accelerators (cuEquivariance, flashTP, or OpenEquivariance) via the corresponding flags. For more information about accelerators, follow here.

model = SevenNetModel(model="7net-omni", modal="mpa", enable_oeq=True)
# or enable_cueq=True or enable_flash=True

The device parameter defaults to 'auto' (CUDA if available, otherwise CPU).

Batched MD

import torch_sim as ts
from ase.build import bulk

atoms = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2))

final_state = ts.integrate(
    system=[atoms] * 10,
    model=model,
    n_steps=100,
    timestep=0.002,
    temperature=300,
    integrator=ts.Integrator.nvt_langevin,
)

Relaxation

relaxed_state = ts.optimize(
    system=[atoms] * 10,
    model=model,
    optimizer=ts.Optimizer.fire,
)