SimPoly: Simulation of Polymers with Machine Learning Force Fields Derived from First Principles

Companion repository to the paper. It provides the implementation of the Vivace machine learning force field, tools for working with the PolyPack and PolyDiss datasets, and analysis code to reproduce key results.

Installation

Recommended: conda/mamba environment

The pinned, tested environment lives in environment.yml and matches the docker image:

mamba env create -f environment.yml
mamba activate simpoly

src/ is already on sys.path for pytest (see pyproject.toml), so no editable install is needed to run the tests.

LAMMPS is not part of the conda environment and must be built separately against the env's PyTorch/Kokkos toolchain. See docker/Dockerfile for the exact CMake invocation (packages, ABI flags, Kokkos arch, MPI wrapper) and docker/README.md for build args and notes on the reproducible CUDA 13 image bundling the conda env, LAMMPS, and Kokkos.

Advanced: pip install

Editable install with optional extras for the Vivace ML stack (vivace), CUDA-13 GPU wheels (cuda13), and dev tooling (dev):

pip install -e '.[vivace,cuda13,dev]'

The conda environment pins matching versions; pip resolution may differ.

Data & model artifacts

• Datasets: microsoft/simpoly on Hugging Face
• Vivace checkpoints:
  • Hugging Face
  • vivace_v0.1.pt, the ASE / Python loader (vivace.deploy.load_model, vivace.calculator.MLFFCalculator)
  • vivace_v0.1.mliap.pt, the LAMMPS loader (pair_style mliap unified via vivace.mliap)

PolyData

See the poly_data notebook for an introduction to the PolyPack and PolyDiss datasets.

PolyArena

Reference experimental values

See the poly_arena_experiment notebook for the reference room-temperature densities and glass transition temperatures.

Generating LAMMPS inputs

Generate the starting configuration and input files with run.py (use --help for all options).

The main-text protocol is 21step_then_cooling: a Polymatic-style 21-step equilibration at --temp-k, then an NPT cooling scan from --temp-k down to --cool-temp-end-k in --cool-temp-step-k steps, all in one LAMMPS run.

In the paper we pick these per polymer as exp_Tg + 160 K, exp_Tg − 160 K, and 20 K respectively (17 stages). exp_Tg is the tg column of experiment.load_data().

Example — Polystyrene (exp_Tg = 373 K):

python src/simpoly/poly_arena/simulation/run.py \
    --directory ps_tg/ --poly-id PS \
    --temp-k 533 --cool-temp-end-k 213 --cool-temp-step-k 20 \
    --protocol 21step_then_cooling \
    --model-type mlff --mlff-path checkpoints/vivace_v0.1.mliap.pt

Running the simulations

The generated 21step_then_cooling.in is a standard LAMMPS input that loads Vivace via pair_style mliap unified. See the LAMMPS quick start below for the input template and the required lmp invocation with multi-GPU / Kokkos flags. A single run produces the full density-vs-temperature trajectory used for Tg analysis.

Vivace

Requirements: NVIDIA GPU with CUDA. cuequivariance kernels are mandatory; there is no CPU fallback.

ASE quick start

from ase.io import read
from simpoly.vivace.calculator import MLFFCalculator

# pp_*.lmps atom-type → atomic number (type 1 = C, type 2 = H).
Z_OF_TYPE = {1: 6, 2: 1}

atoms = read("tests/vivace/data/pp_274.lmps", format="lammps-data", Z_of_type=Z_OF_TYPE)
atoms.calc = MLFFCalculator(model_path="checkpoints/vivace_v0.1.pt")

energy = float(atoms.get_potential_energy())          # eV
forces = atoms.get_forces()                            # eV / Å, (N, 3)
stress_voigt = atoms.get_stress(voigt=True)            # eV / Å^3, length-6 Voigt

LAMMPS quick start

Adapted from tests/vivace/test_lammps_mliap.py. Vivace loads as a LAMMPS mliap unified pair style; the docker image provides the lmp binary.

Minimal in.lmp:

units           metal
atom_style      atomic
boundary        p p p
newton          on

read_data       tests/vivace/data/pp_274.lmps

mass 1 12.011  # C
mass 2  1.008  # H

pair_style      mliap unified checkpoints/vivace_v0.1.mliap.pt 0
pair_coeff      * * C H

neighbor        2.0 bin
neigh_modify    delay 0 every 1 check yes

thermo          1
thermo_style    custom step pe pxx pyy pzz pxy pxz pyz
run             0

Run it on a single GPU via Kokkos:

lmp -k on g 1 -sf kk -pk kokkos newton on neigh half -in in.lmp -log log.lammps

For a multi-GPU production run (e.g. 4 GPUs over MPI):

mpirun -np 4 lmp -k on g 4 -sf kk -pk kokkos neigh half -in 21step_then_cooling.in

Running the tests

pytest tests/                  # full suite
pytest tests/ -m "not gpu"     # skip GPU-only tests
pytest tests/vivace/ -m gpu    # GPU-only end-to-end checks
# or run lammps test in the docker
docker run --gpus all --rm -u $(id -u):$(id -g) \
     -e HOME=/tmp \
     -v "$(pwd):/simpoly" -w /simpoly \
     -e PYTHONPATH=/simpoly/src \
     simpoly:latest \
     pytest /simpoly/tests/vivace/test_lammps_mliap.py -xvvvs