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.5.0
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.

src/atomate2/torchsim/core.py

src/atomate2/torchsim/schema.py

pyproject.toml

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.

pyproject.toml

orionarcher · 2025-12-08T19:15:40Z

@JaGeo @esoteric-ephemera do you have any idea why the abinit tests might be failing here? This seems unrelated to any changes I made wrt TorchSim?

==================================== ERRORS ====================================
_______________ ERROR collecting tests/abinit/jobs/test_base.py ________________
tests/abinit/jobs/test_base.py:6: in <module>
    from abipy.abio.inputs import AbinitInput
../../../micromamba/envs/a2/lib/python3.12/site-packages/abipy/abio/inputs.py:36: in <module>
    from abipy.flowtk import PseudoTable, Pseudo, AbinitTask, AnaddbTask, ParalHintsParser, NetcdfReader
../../../micromamba/envs/a2/lib/python3.12/site-packages/abipy/flowtk/__init__.py:12: in <module>
    from .qadapters import show_qparams, all_qtypes
../../../micromamba/envs/a2/lib/python3.12/site-packages/abipy/flowtk/qadapters.py:311: in <module>
    _EXCL_NODES_FILE = _ExcludeNodesFile()
                       ^^^^^^^^^^^^^^^^^^^
../../../micromamba/envs/a2/lib/python3.12/site-packages/abipy/flowtk/qadapters.py:289: in __init__
    if not os.path.exists(self.DIRPATH): os.makedirs(self.DIRPATH)
                                         ^^^^^^^^^^^^^^^^^^^^^^^^^
<frozen os>:225: in makedirs
    ???
E   FileExistsError: [Errno 17] File exists: '/home/runner/.abinit/abipy'
================================ tests coverage ================================
_______________ coverage: platform linux, python 3.12.12-final-0 _______________

Coverage XML written to file coverage.xml
=========================== short test summary info ============================
ERROR tests/abinit/jobs/test_base.py - FileExistsError: [Errno 17] File exists: '/home/runner/.abinit/abipy'
============= 133 passed, 12 skipped, 1 error in 370.68s (0:06:10) =============

esoteric-ephemera · 2025-12-08T20:00:46Z

I think you can safely ignore it, can happen with parallel runners and there should be a fix for this in the next release of abipy

orionarcher · 2025-12-08T20:16:34Z

Other than the unrelated failing tests, I think this is ready to merge? The tests are passing locally and I think all the suggestions are addressed.

esoteric-ephemera · 2025-12-10T21:28:24Z

Bumped a few comments for your consideration; totally OK if you disagree with the suggestions therein

JaGeo · 2026-01-14T20:51:11Z

Thanks. That looks already great!

One comment: have you tested integrating one of thrse jobs into the phonon maker or a similar one? Could be a test, for example, to ensure schema structure.

@buildbot-princeton

…rial (#3110) ## Summary of Changes I made some changes to the [`atomate2` interface](materialsproject/atomate2#1335) and I want to make sure they are reflected here! Namely, we automatically calculated energy, forces, and optionally stress and put them in the output dict without needing to deal with the trajectory reporter. I also added a section in the docs for torchsim, let me know if it's in the right place / makes sense. I did verify that all the snippets would run. ### Requirements - [x] My PR is focused on a [single feature addition or bugfix](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/getting-started/best-practices-for-pull-requests#write-small-prs). - [x] My PR has relevant, comprehensive [unit tests](https://quantum-accelerators.github.io/quacc/dev/contributing.html#unit-tests). - [x] My PR is on a [custom branch](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-and-deleting-branches-within-your-repository) (i.e. is _not_ named `main`). Note: If you are an external contributor, you will see a comment from [@buildbot-princeton](https://github.com/buildbot-princeton). This is solely for the maintainers.

orionarcher · 2026-01-15T16:34:38Z

End to end compatibility testing will have to wait until #1196 is merged. TorchSim only has a batched mode so it won't work with the current PhononMaker. That said, I've done some basic compatibility testing and I think it will be an easy adaption once #1196 is in.

I am happy to loop back around and write a test to verify integration once that's in?

JaGeo · 2026-01-15T16:57:51Z

@orionarcher I am not sure I agree that #1196 needs to be in.

My queries here are mainly to avoid breaking changes in the future. The more we think ahead towards other workflows, the less we have to adapt the schema in the future.

@esoteric-ephemera is it for you sufficently similar to the Ase and ForcefieldMakers?

…ched operation

orionarcher · 2026-01-15T18:33:54Z

Ok fair enough. Upon further consideration, it's possible to get compatibility by making the TorchSim functions support Structure | list[Structure] so I went ahead and modified them to that API. I also added some tests for phonon compatibility and it looks to work well!

I appreciate all the feedback and design thoughts on the PR.

JaGeo · 2026-01-15T18:55:33Z

tests/torchsim/test_core.py

+
+    # Create the phonon flow
+    flow = phonon_maker.make(si_structure)
+


Is there a reason not to try running it?

Done! Good catch.

JaGeo · 2026-01-15T18:56:49Z

@orionarcher thanks! That already looks great! I made one (last) comment.

orionarcher · 2026-01-15T19:09:30Z

Good catch, the test was not testing properly. I made the change and it not works with the PhononMaker!

JaGeo · 2026-01-15T19:39:01Z

@orionarcher thanks! We now need to fix the tests.

They are failing as the installations need to much cache. We might need to do some cache cleaning after certain workflows or come up with a different idea.

orionarcher · 2026-01-15T21:27:34Z

Yeah, looks like TorchSim is the straw that broke the camels back. I could split the tests into a separate action but I don't know if that's the most sustainable route.

What would you propose as a solution there? Would you like me to take a swing at it or would you like to try a fix?

JaGeo · 2026-01-15T21:32:10Z

I don't have a good solution. Separate action sounds good for me for now. I won't be able to do much testing here myself for a couple of weeks.

@esoteric-ephemera do you have any opinion here? Also with regard to usage of action time?

orionarcher · 2026-01-15T22:44:36Z

I split the TorchSim stuff out into separate tests. It's dependent on torch>=2.7.0 which conflicts with the strict dependency of torch==2.2.0. Memory stuff aside, it's won't work alongside the current ff tests.

I think it should all be good to go now?

esoteric-ephemera · 2026-01-15T23:57:18Z

Separate workflow action / separate test sounds good to me. Also for the best given the torch version pin needed for some of the MLFFs, as you said @orionarcher

Will take another look through and ping you if I see anything, thanks for your work on this so far!

JaGeo · 2026-01-16T06:48:47Z

Fine from my side now. Once @esoteric-ephemera is happy, we can merge :)

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

src/atomate2/torchsim/core.py Outdated Show resolved Hide resolved

esoteric-ephemera reviewed Dec 3, 2025

View reviewed changes

src/atomate2/torchsim/schema.py Outdated Show resolved Hide resolved

esoteric-ephemera reviewed Dec 3, 2025

View reviewed changes

src/atomate2/torchsim/schema.py Outdated Show resolved Hide resolved

esoteric-ephemera reviewed Dec 3, 2025

View reviewed changes

pyproject.toml Outdated Show resolved Hide resolved

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

orionarcher and others added 2 commits December 5, 2025 16:40

skip torchsim tests if not on python 3.12

b4171c5

Merge branch 'main' into torchsim

ac7c932

esoteric-ephemera reviewed Dec 8, 2025

View reviewed changes

pyproject.toml Outdated Show resolved Hide resolved

orionarcher and others added 2 commits December 8, 2025 13:50

try locking py >= 3.12 in install

daeb490

Merge branch 'main' into torchsim

340c532

add tests for energy, force, stress output

c2de281

make torchsim makers also accept single structures in addition to bat…

8b2b63d

…ched operation

JaGeo reviewed Jan 15, 2026

View reviewed changes

run phonon workflow and test and change nomenclature for compatability

1706f89

update tutorial

b1c1f38

orionarcher added 11 commits January 15, 2026 16:51

try using cpu pytorch without cuda to fix OOM error?

e7e5193

fix typo in testing.yml

8896338

try running test in one liner

2465b1a

try pinning torch to fix CHGNet tests

2249169

try standalone torchsim tests

f79e57a

small

84ba3dc

correct test path

34c8e0e

add ase

b996fe6

add mace

bc3e36a

add phonons

ed267f4

add unit cell filter

8770c9e

delete random depdency addition

d1f2328


		# Create the phonon flow
		flow = phonon_maker.make(si_structure)

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!

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

orionarcher commented Dec 8, 2025

Uh oh!

esoteric-ephemera commented Dec 8, 2025

Uh oh!

orionarcher commented Dec 8, 2025

Uh oh!

esoteric-ephemera commented Dec 10, 2025

Uh oh!

JaGeo commented Jan 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

orionarcher commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JaGeo commented Jan 15, 2026

Uh oh!

orionarcher commented Jan 15, 2026

Uh oh!

JaGeo Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

orionarcher Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

JaGeo commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

orionarcher commented Jan 15, 2026

Uh oh!

JaGeo commented Jan 15, 2026

Uh oh!

orionarcher commented Jan 15, 2026

Uh oh!

JaGeo commented Jan 15, 2026

Uh oh!

orionarcher commented Jan 15, 2026

Uh oh!

esoteric-ephemera commented Jan 15, 2026

Uh oh!

JaGeo commented Jan 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

orionarcher commented Nov 23, 2025 •

edited

Loading

JaGeo commented Jan 14, 2026 •

edited

Loading

orionarcher commented Jan 15, 2026 •

edited

Loading

JaGeo commented Jan 15, 2026 •

edited

Loading