Add support for Torchsim #1335

orionarcher · 2025-11-23T18:14:03Z

Summary

This PR adds support for TorchSim, namely:

Utilities to convert the live TorchSim objects (e.g. TrajectoryReporter) into configurable schema with equivalent features
Input and output schema for integrate, optimize and static functions
Makers for integrate, optimize and static jobs

NOTE: this PR uses StrEnum's which are not supported in python 3.10, see #1334

Additional dependencies introduced (if any)

torchsim==0.4.1
python>=3.11

Checklist

Work-in-progress pull requests are encouraged, but please put [WIP] in the pull request
title.

Before a pull request can be merged, the following items must be checked:

Code is in the standard Python style.
The easiest way to handle this is to run the following in the correct sequence on
your local machine. Start with running ruff and ruff format on your new code. This will
automatically reformat your code to PEP8 conventions and fix many linting issues.
Doc strings have been added in the Numpy docstring format.
Run ruff on your code.
Type annotations are highly encouraged. Run mypy to
type check your code.
Tests have been added for any new functionality or bug fixes.
All linting and tests pass.

Note that the CI system will run all the above checks. But it will be much more
efficient if you already fix most errors prior to submitting the PR. It is highly
recommended that you use the pre-commit hook provided in the repository. Simply run
pre-commit install and a check will be run prior to allowing commits.

esoteric-ephemera · 2025-12-03T17:08:33Z

src/atomate2/torchsim/core.py

+    if model_type == TSModelType.FAIRCHEMV1:
+        from torch_sim.models.fairchem_legacy import FairChemV1Model
+
+        return FairChemV1Model(model=model_path, **model_kwargs)
+    if model_type == TSModelType.FAIRCHEM:
+        from torch_sim.models.fairchem import FairChemModel
+
+        return FairChemModel(model=model_path, **model_kwargs)
+    if model_type == TSModelType.GRAPHPESWRAPPER:
+        from torch_sim.models.graphpes import GraphPESWrapper
+
+        return GraphPESWrapper(model=model_path, **model_kwargs)
+    if model_type == TSModelType.MACE:
+        from torch_sim.models.mace import MaceModel
+
+        return MaceModel(model=model_path, **model_kwargs)
+    if model_type == TSModelType.MATTERSIM:
+        from torch_sim.models.mattersim import MatterSimModel
+
+        return MatterSimModel(model=model_path, **model_kwargs)
+    if model_type == TSModelType.METATOMIC:
+        from torch_sim.models.metatomic import MetatomicModel
+
+        return MetatomicModel(model=model_path, **model_kwargs)
+    if model_type == TSModelType.NEQUIPFRAMEWORK:
+        from torch_sim.models.nequip_framework import NequIPFrameworkModel
+
+        return NequIPFrameworkModel(model=model_path, **model_kwargs)
+    if model_type == TSModelType.ORB:
+        from torch_sim.models.orb import OrbModel
+
+        return OrbModel(model=model_path, **model_kwargs)


For readability, can this block be refactored into a dict like:

import importlib model_to_import_str = { "FAIRCHEMV1": "torch_sim.models.fairchem_legacy.FairChemV1Model", "FAIRCHEM": "torch_sim.models.fairchem.FairChemModel", ... } model_module, model_class = model_to_import_str[TSModelType[model_type]].rsplit(".",1) return getattr(importlib.import_module(model_module),model_class)(model=model_path, **model_kwargs)

esoteric-ephemera · 2025-12-03T17:12:08Z

src/atomate2/torchsim/schema.py

+    all_properties: list[dict[str, np.ndarray]] = Field(
+        ..., description="List of calculated properties for each structure."
+    )
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)


Can the properties not be cast to a list or other built-in? arbitrary_types_allowed eliminates the benefits of type checking here

esoteric-ephemera · 2025-12-03T17:15:46Z

src/atomate2/torchsim/schema.py

+    )
+
+    calcs_reversed: list[
+        TSIntegrateCalculation | TSOpimizeCalculation | TSStaticCalculation


This is fine to have different calculation types, but I'd prefer to avoid making many schemas for small variations in the model fields - can you use the emmet.core.vasp.task_types to merge down these three into one TorchSimCalculation schema?

Avoiding union types like this lets us better support cloud native data formats (e.g., parquet) and most modern compression tools (again arrow/parquet) eliminate the storage penalty associated with nullable fields

esoteric-ephemera · 2025-12-03T17:17:02Z

pyproject.toml

    "quippy-ase>=0.9.14; python_version < '3.12'",
    "sevenn>=0.9.3",
    "torchdata<=0.7.1",                            # TODO: remove when issue fixed
+    "torch_sim==0.4.1",


Move to its own optional import block since the module is currently distinct from the forcefields

esoteric-ephemera

Some general comments besides the ones on specific lines:

The TS prefixing in makers, enums, etc. is confusing given that TS is probably more familiar as "transition state" than "torchsim" - can you change this to TorchSim?
Are you envisioning this being added as a separate (as its currently implemented) or additive module to the existing forcefields stuff? If the latter, the schemas would have to be merged for the job outputs

reduce redundancy in initialization logic

orionarcher · 2025-12-05T20:46:02Z

The TS prefixing in makers, enums, etc. is confusing given that TS is probably more familiar as "transition state" than "torchsim" - can you change this to TorchSim?

Good call.

Are you envisioning this being added as a separate (as its currently implemented) or additive module to the existing forcefields stuff? If the latter, the schemas would have to be merged for the job outputs

I think this should be a separate module. It would be really messy to integrate it with the existing forcefields stuff. TorchSim generally expects many -> many calculations (list[structure] -> list[structure]) and has different output files and such.

JaGeo · 2025-12-05T20:51:23Z

@orionarcher #1196 ?

orionarcher · 2025-12-05T21:03:11Z

Encouraging! In that case, let me reframe. It looks like it would be possible, but I anticipate it would be a major headache, one I am reluctant to take on.

Though they both run MLIPs, ASE and TorchSim are different software packages with pretty different APIs. I don't see a major reason to have them share schema or logic. The forcefields module is hewn pretty closely to the ASE schema and embraces the paradigms of that package. While in principle I appreciate they are doing basically the same thing (take in strucuture -> repeatedly evaluate MLIP -> generate trajectory + final structure) so many of the norms and expectations of the software are different that I don't think integration would make either interface any better.

JaGeo · 2025-12-05T21:07:23Z

The motivation should always be the following:
If you have a similar schema people can replace their current code using forcefields easily with torchsim. This includes larger workflows.

orionarcher · 2025-12-05T21:21:11Z

I hear you and I am sympathetic to that argument. In this case, I feel there is a tradeoff between the immediate adoptability of the TorchSim interface and it's overall quality. After looking back and forth at the ASE and TorchSim schema for the past 15 minutes, I don't think it's possible to adapt the TorchSim API to fit the ASE schemas without adding complexity, reducing readability, and making the overall API less natural and maintainable.

I would love for users currently using ASE to be able to quickly and reliably switch to TorchSim but it's not clear to me that equating the schemas is the best way to do that. I would be happy to write a transition guide outlining the schema differences and how to transition from ASE -> TorchSim and add it to this PR.

orionarcher added 4 commits November 21, 2025 16:09

untested update of core logic and organization to match quacc

6fd7ccd

ignore uv lock file

ae5e289

write tests and get them passing

3aca26e

lint

a5d313c

esoteric-ephemera reviewed Dec 3, 2025

View reviewed changes

esoteric-ephemera requested changes Dec 3, 2025

View reviewed changes

orionarcher and others added 6 commits December 5, 2025 15:06

TS -> TorchSim

1e266e6

split out torchsim tests in pyproject

c2935b9

refactor to single TorchSimCalculation schema to avoid union types

703d8a8

remove arrays and arbitrary types

188f070

reduce redundancy in initialization logic

add torchsim dep to testing

589fc35

Merge branch 'main' into torchsim

e3b2330

fix import in testing.yml and pyproject.toml

bfbd2a3

skip torchsim tests if not on python 3.12

b4171c5

Add support for Torchsim #1335

Are you sure you want to change the base?

Add support for Torchsim #1335

Uh oh!

Conversation

orionarcher commented Nov 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Additional dependencies introduced (if any)

Checklist

Uh oh!

esoteric-ephemera Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

esoteric-ephemera Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

esoteric-ephemera Dec 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

orionarcher Dec 5, 2025

Choose a reason for hiding this comment

Uh oh!

esoteric-ephemera Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

orionarcher Dec 5, 2025

Choose a reason for hiding this comment

Uh oh!

esoteric-ephemera left a comment

Choose a reason for hiding this comment

Uh oh!

orionarcher commented Dec 5, 2025

Uh oh!

JaGeo commented Dec 5, 2025

Uh oh!

orionarcher commented Dec 5, 2025

Uh oh!

JaGeo commented Dec 5, 2025

Uh oh!

orionarcher commented Dec 5, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

orionarcher commented Nov 23, 2025 •

edited

Loading

esoteric-ephemera Dec 3, 2025 •

edited

Loading