LOREM-JAX

JAX implementation of LOREM (Learning Long-Range Representations with Equivariant Messages), a machine learning interatomic potential with equivariant long-range message passing.

Built on JAX, Flax, e3x, and jax-pme.

Installation

Requires Python >= 3.11.

pip install .

Usage

ASE calculator

import jax
from ase.build import bulk
from lorem.models.mlip import Lorem
from lorem.calculator import Calculator

model = Lorem(cutoff=5.0)
params = model.init(jax.random.key(42), *model.dummy_inputs())
calc = Calculator.from_model(model, params=params)

atoms = bulk("Ar") * [2, 2, 2]
calc.calculate(atoms)
print(calc.results["energy"], calc.results["forces"].shape)

To load a trained model from a checkpoint:

calc = Calculator.from_checkpoint("path/to/checkpoint")

Training

Training a model involves three steps: preparing the data, configuring the model and training settings, and running the training script.

1. Prepare data

Training data is stored in marathon format. Convert your extended XYZ dataset using a preparation script (see examples/train-mlp/prepare.py for a template):

from marathon.data import datasets, get_splits
from marathon.grain import prepare

# datasets is a Path resolved from the $DATASETS environment variable
prepare(train_atoms, folder=datasets / "my_project/train", ...)
prepare(valid_atoms, folder=datasets / "my_project/valid", ...)

The $DATASETS environment variable sets the root directory where prepared datasets are stored. All dataset paths in settings.yaml are resolved relative to this directory.

2. Configure the experiment

Each experiment lives in its own directory containing two YAML files:

model.yaml defines the model architecture:

model:
  lorem.Lorem:
    cutoff: 5.0
    max_degree: 4
    max_degree_lr: 2
    num_features: 128
    num_spherical_features: 4
    num_message_passing: 1

Use lorem.LoremBEC instead of lorem.Lorem to train a model that additionally predicts Born effective charges.

settings.yaml configures training:

train: "my_project/train"           # path relative to $DATASETS
valid: "my_project/valid"           # path relative to $DATASETS
seed: 23
batcher:
  batch_size: 4
loss_weights: {"energy": 0.5, "forces": 0.5}
optimizer: adam                      # adam or muon
start_learning_rate: 1e-3
min_learning_rate: 1e-6
max_epochs: 2000
valid_every_epoch: 2
decay_style: linear                  # linear, exponential, or warmup_cosine
use_wandb: True

3. Run training

cd my_experiment
DATASETS=/path/to/datasets lorem-train

Training writes checkpoints, logs, and plots to a run/ directory inside the experiment folder. If a run/ directory already exists, training resumes from the latest checkpoint.

See examples/train-mlp/ and examples/train-bec/ for complete examples including data preparation and configuration files.

Model variants

• Lorem -- the standard MLIP model (energy + forces + stress)
• LoremBEC -- predicts Born effective charges in addition to energy/forces

Key hyperparameters

Installing the i-PI driver

After installation of the package, install the i-PI driver via:

lorem-install-ipi-driver

This copies the LOREM driver into the i-PI pes directory. You can rerun lorem-install-ipi-driver anytime (it is idempotent) if you switch environments or reinstall i-PI.

Development

Format and lint:

ruff format . && ruff check --fix .

Run tests:

python -m pytest tests/ -v --override-ini="addopts="

Or use tox:

tox -e lint       # check formatting + linting
tox -e tests      # run unit tests
tox -e examples   # run examples as smoke tests
tox -e format     # auto-format

License

BSD-3-Clause

This project is maintained by @PicoCentauri and @sirmarcel, who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you did not receive an answer after a couple of days.