Add support for Torchsim

[orionarcher]

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]

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

[orionarcher]

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]

@orionarcher #1196 ?

[orionarcher]

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]

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]

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]

@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]

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]

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]

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

[orionarcher]

Oop, sorry I missed that, now changed!

[esoteric-ephemera]

This all looks OK to me, I'll have to see why the approx NEB tests are failing now

@JaGeo I agree that it's not great to have this bifurcation of similar MLFF libraries, I'm also not sure how best to avoid that. If you want, we could try to iterate a bit on this and the multi ASE makers you've been working on

[JaGeo]

@esoteric-ephemera Yeah. For me, it is also unclear how to integrate torchsim into other property workflows. I think it would be worth thinking about this a bit more

[JaGeo]

Why is it called "Integrate" and not MD?

[orionarcher]

The TorchSim integrate function also supports some Monte Carlo methods, so it's not strictly limited to MD

[JaGeo]

There is another point. Even if we merge everything in this state, I think we definitely need some new documentation here as well. What to install, how to use it in the best way etc. Some examples.

@orionarcher ?

[JaGeo]

I have made up my mind. If one can show that with this current schema, one can plug TorchSim Flows into larger workflows (e.g., the phonon workflow where a batch mode already exists), I am completely fine with merging. I think this is a basic requirement that such a new schema should fulfill. (plus the documentation, of course)

[orionarcher]

Thanks for the reviews, I understand your concerns about the bifurcation of MLFF libraries.

There is another point. Even if we merge everything in this state, I think we definitely need some new documentation here as well. What to install, how to use it in the best way etc. Some examples.

Happy to add this sort of documentation. I can write a tutorial to walk through the usage.

If one can show that with this current schema, one can plug TorchSim Flows into larger workflows (e.g., the phonon workflow where a batch mode already exists), I am completely fine with merging.

I'll explore this over the next few days! Would you mind pointing me to the batched Phonon Maker? I am looking at phonons.py and the function signatures suggest that batching isn't supported.

class BasePhononMaker(Maker, ABC):

    def make(
        self,
        structure: Structure,
        prev_dir: str | Path | None = None,
        born: list[Matrix3D] | None = None,
        epsilon_static: Matrix3D | None = None,
        total_dft_energy_per_formula_unit: float | None = None,
        supercell_matrix: Matrix3D | None = None,
    ) -> Flow:
    
    ```

[JaGeo]

Hi @orionarcher,
Please have a look here or at the PR that I linked above.

atomate2/src/atomate2/aims/flows/phonons.py

Line 23 in 1822f9d

class PhononMaker(BasePhononMaker):

All other suggestions sound good.

[orionarcher]

I did some digging on the PhononMaker and found the following TaskDoc fields accessed through different Makers.

# bulk_relax_maker
structure = bulk.output.structure
prev_dir = bulk.output.dir_name
optimization_run_job_dir = bulk.output.dir_name
optimization_run_uuid = bulk.output.uuid

# static energy maker
total_dft_energy = static_job.output.output.energy
static_run_job_dir = static_job.output.dir_name
static_run_uuid = static_job.output.uuid
prev_dir = static_job.output.dir_name

# phonon_displacement_maker
outputs["displacement_number"] = list(range(len(displacements)))
outputs["uuids"] = [phonon_job.output.uuid] * len(displacements)
outputs["dirs"] = [phonon_job.output.dir_name] * len(displacements)
outputs["forces"] = phonon_job.output.output.all_forces

# born_maker
epsilon_static = born_job.output.calcs_reversed[0].output.epsilon_static
born = born_job.output.calcs_reversed[0].output.outcar["born"]

Thus the task doc must include:

task_doc.structure
task_doc.output.energy
task_doc.dir_name
task_doc.uuid
task_doc.output.all_forces

# and for born charges
task_doc.calcs_reversed[0].output.epsilon_static
task_doc.calcs_reversed[0].output.outcar["born"]

The first five seem doable and I am happy to add them (and let me know if it seems like I am missing something). TorchSim, however, doesn't currently support Born Charges, so we would be able to support the born_maker argument with TorchSim.

I propose refactoring the TaskDoc to expose those arguments in the same positions as the ForceFieldTaskDoc but leaving out the Born charges for now. Does that, along with the addition of more docs and a tutorial, seem like a reasonable path forward? Once the batched schema is merged in, it should then be quite straightforward to align the schemas.

Another thought is that we could make the structure inputs and outputs Structure | List[Structure] such that TorchSim can be used as a single structure drop-in, if desired.

[JaGeo]

To me, this sounds very good! Thanks!

We do not have any born charges in the force-field based MLIPs as well. This will be fine. We might have them from ML at some point in the future.