TrIP: Transferable Interatomic Potentials

Transferable Interatomic Potential (TrIP) trains lightweight supervised heads on top of pretrained interatomic potential embeddings. Use it when you have materials structures and want to learn scalar labels attached to atoms, bonds, or whole structures.

What it learns

• Per-atom values: one scalar for each atom, for example charge, local energy, oxidation proxy, or any array stored in atoms.arrays[...].
• Per-bond values: one scalar for selected atom pairs, for example interaction strength, bond order, or pair labels stored in a pair-list array.
• Per-structure values: one scalar for each structure, for example formation energy, band gap, adsorption energy, or any scalar stored in atoms.info[...].

The old U and Vlist names still work, but they are just examples of per-atom and per-bond targets.

Install

Install the package with the backend you plan to use:

pip install -e .[mace]
pip install -e .[mattersim]
pip install -e .[orb]

For development with every optional backend:

pip install -e .[all]

Pick a backbone

from trip.models import mace, mattersim, orb

model = mace("small")
# model = mattersim("1m")
# model = orb("orb_v3_conservative_inf_omat")

By default, TrIP trains only the prediction head. Pass full_fine_tuning=True to update the backbone weights too.

Prepare the data

TrIP works with ASE Atoms objects, .xyz/.extxyz files, or a dataset created by trip.dataset.dataset.

from trip.dataset import dataset

data = dataset("materials.extxyz")
data = data.filter_target(
    target="atom_property",
    elements=["Li", "O"],
    require_nonzero=True,
)

train, val, test = data.split(
    frac_val=0.2,
    frac_test=0.1,
    seed=0,
    target="atom_property",
    elements=["Li", "O"],
)

For small datasets, split with target and elements. This keeps validation and test sets from accidentally missing important elements or target ranges.

Train per-atom values

result = model.train(
    train=train,
    val=val,
    task="per-atom",
    target="atom_property",
    interest_elements=["Li", "O"],
    target_mode="residual",
    epochs=200,
    batch_size=64,
    learning_rate=5e-4,
    head_only=True,
)

Important options:

• target: name of the ASE per-atom array to learn.
• interest_elements: atoms included in the loss and reported metrics.
• target_mode="absolute": learn the raw target.
• target_mode="residual": learn the target minus a learned element baseline.

Train per-bond values

result = model.train(
    train=train,
    val=val,
    task="per-bond",
    target="bond_values",
    bonds=[("Li", "O"), ("O", "O")],
    target_mode="residual",
    combine_method="average",
    epochs=50,
    head_only=True,
)

Important options:

• target: name of the pair-list array in each structure.
• bonds: element pairs included in training.
• combine_method: how duplicate pair labels are combined.
• use_log_dist=True: include log distance instead of raw distance.

task="per-edge" and bonds=[...] remain supported aliases for older code.

Train per-structure values

result = model.train(
    train=train,
    val=val,
    task="per-structure",
    target="formation_energy",
    structure_pooling="mean",
    target_mode="absolute",
    epochs=100,
    batch_size=32,
    learning_rate=5e-4,
    head_only=True,
)

Important options:

• target: name of the scalar in atoms.info[...].
• structure_pooling: how atom embeddings become one structure embedding.

Available pooling methods:

Method	Operation	Use when
"mean"	Average all atom embeddings in the structure.	The target is intensive or normalized by structure size, for example formation energy per atom or band gap. This is the default.
"sum"	Sum all atom embeddings in the structure.	The target is extensive and should scale with the number of atoms, for example total energy.

Any other value raises an error. Pooling is applied only for task="per-structure"; per-atom and per-bond tasks use atom or pair features directly.