[Bug]: TorchSim backend returns NaN on some structures

[akwarii]

Contact Details

gael.huynh@uclouvain.be

Bug Description

Some structures seems to return NaN when using the TorchSim backend while using Potential.predict_properties works perfectly fine. Note that this error only occurs for some structures (see example below where only 1 structure presents this issue).

MatterSim Version

1.2.5

Python Version

3.12.12

Reproduction Steps

1. Install dependencies pip install mattersim==1.2.5
2. Run the code below
3. The second structure should return NaN values when using the TorchSim backend

import numpy as np
import requests
import torch
from mattersim.forcefield.potential import Potential
from pymatgen.core import Structure


def fetch_alexandria_entry(id_string):
    url = f"https://alexandria.icams.rub.de/pbe/v1/structures/{id_string}"

    response = requests.get(url)
    response.raise_for_status()
    data = response.json()["data"]
    attrs = data["attributes"]

    return Structure(
        lattice=attrs["lattice_vectors"],
        species=attrs["species_at_sites"],
        coords=attrs["cartesian_site_positions"],
        coords_are_cartesian=True,
    )


def direct_inference(model, structures):
    from mattersim.datasets.utils.build import build_dataloader
    from pymatgen.io.ase import AseAtomsAdaptor

    adapter = AseAtomsAdaptor()
    atoms = [adapter.get_atoms(struct) for struct in structures]

    dataloader = build_dataloader(atoms, only_inference=True)
    predictions = model.predict_properties(
        dataloader,
        include_forces=True,
        include_stresses=True,
    )

    return predictions


def torchsim_inference(model, structures):
    import torch_sim as ts
    from mattersim.torchsim import get_torchsim_wrapper

    device = "cuda" if torch.cuda.is_available() else "cpu"
    wrapper = get_torchsim_wrapper(model, device=device)

    state = ts.static(structures, model=wrapper)
    return state


def main():
    alexandria_ids = ["agm003221802", "agm006228157"]
    structures = [fetch_alexandria_entry(id) for id in alexandria_ids]

    model = Potential.from_checkpoint(
        load_path="mattersim-v1.0.0-1M", load_training_state=False
    )

    results = torchsim_inference(model, structures)
    for i in range(len(structures)):
        for key, prop in results[i].items():
            if torch.isnan(prop).any():
                print(f"NaN value found in {key} for structure {i} (TorchSim)")

    results = direct_inference(model, structures)
    for i in range(len(structures)):
        for key, prop in zip(["energy", "forces", "stress"], results[i]):
            if np.isnan(prop).any():
                print(f"NaN value found in {key} for structure {i} (ASE)")


if __name__ == "__main__":
    main()

Expected Behavior

I would expect the TorchSim backend to return the same values as Potential.predict_properties and to not return NaN values on valid structures.

Actual Behavior

Some structures will return NaN properties when using TorchSim.

Error Logs

Code of Conduct

• I agree to follow the project's Code of Conduct