Batched workflows for TorchSim (#1505)

* Support batching for the elastic workflow

* Update model construction due to API changes

* Remove unsupported models

* Remove legacy models from TorchSimModelType

* Add tests for torchsim model wrappers

* Revert mattersim changes

* skip test if torchsim is not installed

* Minimalistic elastic workflow for torchsim

* Add tests

* WIP: adapt common elastic jobs

* Fix Optimizer being casted to its string value when passed to a
jobflow.job

* squeeze the stress tensor

* Support TorchSim symmetry constraints

* Convert the stress to kbar

* Unit conversion without ASE

* Batched phonon workflow

* Remove ase and torchsim dependencies from common phonon jobs

* Use lighter model

* Fix tests

* Disable torchsim in test-non-ase jobs

* Fix phonon maker type check and improve readability

* Include cell forces to the convergence check

* Add values check on different phonon derived properties

* Improve socket keyword documentation

* Add elastic workflow tutorial

* Loosen test tolerance

---------

Co-authored-by: J. George <JaGeo@users.noreply.github.com>

Author: akwarii

.github/workflows/testing.yml

@@ -168,7 +168,7 @@ jobs:
         # However this `splitting-algorithm` means that tests cannot depend sensitively on the order they're executed in.
         run: |
           micromamba activate a2
-          pytest --durations=5 -n auto --splits 3 --group ${{ matrix.split }} --durations-path tests/.pytest-split-durations --splitting-algorithm least_duration --ignore=tests/ase --ignore=tests/openff_md --ignore=tests/openmm_md --ignore=tests/forcefields --cov=atomate2 --cov-report=xml
+          pytest --durations=5 -n auto --splits 3 --group ${{ matrix.split }} --durations-path tests/.pytest-split-durations --splitting-algorithm least_duration --ignore=tests/ase --ignore=tests/openff_md --ignore=tests/openmm_md --ignore=tests/forcefields --ignore=tests/torchsim --cov=atomate2 --cov-report=xml
 
 
       - uses: codecov/codecov-action@v1

pyproject.toml

@@ -63,7 +63,7 @@ ase-ext = ["tblite>=0.3.0; platform_system=='Linux'"]
 forcefields-demo = ["chgnet>=0.3.8","atomate2[ase]"]
 
 torchsim = [
-    "torch-sim-atomistic==0.6.0; python_version >= '3.12'"
+    "torch-sim-atomistic[symmetry]==0.6.0; python_version >= '3.12'"
 ]
 jdftx = ["pymatgen==2026.5.4"]
 approxneb = ["pymatgen-analysis-diffusion>=2024.7.15"]

src/atomate2/common/flows/elastic.py

@@ -24,6 +24,7 @@
 
     from atomate2.aims.jobs.base import BaseAimsMaker
     from atomate2.forcefields.jobs import ForceFieldRelaxMaker
+    from atomate2.torchsim import TorchSimOptimizeMaker
     from atomate2.vasp.jobs.base import BaseVaspMaker
 
 
@@ -72,26 +73,38 @@ class BaseElasticMaker(Maker, ABC):
         Keyword arguments passed to :obj:`fit_elastic_tensor`.
     task_document_kwargs : dict
         Additional keyword args passed to :obj:`.ElasticDocument.from_stresses()`.
+    socket : bool
+        If True, uses the socket-io interface to run all deformations in a single
+        job, reducing overhead. In the specific case of TorchSim, this enables batching
+        of all structure relaxations.
+        Note: socket=True is not supported for BaseVaspMaker.
     """
 
     name: str = "elastic"
     order: int = 2
     sym_reduce: bool = True
     symprec: float = SETTINGS.SYMPREC
-    bulk_relax_maker: BaseAimsMaker | BaseVaspMaker | ForceFieldRelaxMaker | None = None
-    elastic_relax_maker: BaseAimsMaker | BaseVaspMaker | ForceFieldRelaxMaker = (
-        None  # constant volume optimization
-    )
+    bulk_relax_maker: (
+        BaseAimsMaker
+        | BaseVaspMaker
+        | ForceFieldRelaxMaker
+        | TorchSimOptimizeMaker
+        | None
+    ) = None
+    elastic_relax_maker: (
+        BaseAimsMaker | BaseVaspMaker | ForceFieldRelaxMaker | TorchSimOptimizeMaker
+    ) = None  # constant volume optimization
     max_failed_deformations: int | float | None = None
     generate_elastic_deformations_kwargs: dict = field(default_factory=dict)
     fit_elastic_tensor_kwargs: dict = field(default_factory=dict)
     task_document_kwargs: dict = field(default_factory=dict)
+    socket: bool = False
 
     def make(
         self,
         structure: Structure,
         prev_dir: str | Path | None = None,
-        equilibrium_stress: Matrix3D = None,
+        equilibrium_stress: Matrix3D | None = None,
         conventional: bool = False,
     ) -> Flow:
         """
@@ -135,15 +148,16 @@ def make(
             **self.generate_elastic_deformations_kwargs,
         )
 
-        vasp_deformation_calcs = run_elastic_deformations(
+        deformation_calcs = run_elastic_deformations(
             structure,
             deformations.output,
             elastic_relax_maker=self.elastic_relax_maker,
             prev_dir=prev_dir,
+            socket=self.socket,
         )
         fit_tensor = fit_elastic_tensor(
             structure,
-            vasp_deformation_calcs.output,
+            deformation_calcs.output,
             equilibrium_stress=equilibrium_stress,
             order=self.order,
             symprec=self.symprec if self.sym_reduce else None,
@@ -156,7 +170,7 @@ def make(
         # allow some of the deformations to fail
         fit_tensor.config.on_missing_references = OnMissing.NONE
 
-        jobs += [deformations, vasp_deformation_calcs, fit_tensor]
+        jobs += [deformations, deformation_calcs, fit_tensor]
 
         return Flow(
             jobs=jobs,

src/atomate2/common/flows/phonons.py

@@ -27,6 +27,7 @@
 
     from atomate2.aims.jobs.base import BaseAimsMaker
     from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker
+    from atomate2.torchsim.core import TorchSimOptimizeMaker, TorchSimStaticMaker
     from atomate2.vasp.jobs.base import BaseVaspMaker
 
 SUPPORTED_CODES = frozenset(("vasp", "aims", "forcefields", "ase", "torchsim"))
@@ -99,13 +100,16 @@ class BasePhononMaker(Maker, ABC):
         A maker to perform a tight relaxation on the bulk.
         Set to ``None`` to skip the
         bulk relaxation
-    static_energy_maker: .ForceFieldRelaxMaker, .BaseAimsMaker, .BaseVaspMaker, or None
+    static_energy_maker: .ForceFieldRelaxMaker, .BaseAimsMaker, .BaseVaspMaker,
+        .TorchSimStaticMaker, or None
         A maker to perform the computation of the DFT energy on the bulk.
         Set to ``None`` to skip the
         static energy computation
-    born_maker: .ForceFieldStaticMaker, .BaseAsimsMaker, .BaseVaspMaker, or None
+    born_maker: .ForceFieldStaticMaker, .BaseAsimsMaker, .BaseVaspMaker,
+        .TorchSimStaticMaker, or None
         Maker to compute the BORN charges.
-    phonon_displacement_maker: .ForceFieldStaticMaker, .BaseAimsMaker, .BaseVaspMaker
+    phonon_displacement_maker: .ForceFieldStaticMaker, .BaseAimsMaker, .BaseVaspMaker,
+        .TorchSimStaticMaker
         Maker used to compute the forces for a supercell.
     generate_frequencies_eigenvectors_kwargs : dict
         Keyword arguments passed to :obj:`generate_frequencies_eigenvectors`.
@@ -132,7 +136,10 @@ class BasePhononMaker(Maker, ABC):
     store_force_constants: bool
         if True, force constants will be stored
     socket: bool
-        If True, use the socket/batch for the calculation
+        If True, uses the socket-io interface to run all displacements in a single
+        job, reducing overhead. In the specific case of TorchSim, this enables batching
+        of all static structure evaluations.
+        Note: socket=True is not supported for BaseVaspMaker.
     """
 
     name: str = "phonon"
@@ -145,14 +152,26 @@ class BasePhononMaker(Maker, ABC):
     allow_orthorhombic: bool = False
     get_supercell_size_kwargs: dict = field(default_factory=dict)
     use_symmetrized_structure: Literal["primitive", "conventional"] | None = None
-    bulk_relax_maker: ForceFieldRelaxMaker | BaseVaspMaker | BaseAimsMaker | None = None
-    static_energy_maker: ForceFieldRelaxMaker | BaseVaspMaker | BaseAimsMaker | None = (
-        None
-    )
-    born_maker: ForceFieldStaticMaker | BaseVaspMaker | None = None
-    phonon_displacement_maker: ForceFieldStaticMaker | BaseVaspMaker | BaseAimsMaker = (
+    bulk_relax_maker: (
+        ForceFieldRelaxMaker
+        | BaseVaspMaker
+        | BaseAimsMaker
+        | TorchSimOptimizeMaker
+        | None
+    ) = None
+    static_energy_maker: (
+        ForceFieldRelaxMaker
+        | BaseVaspMaker
+        | BaseAimsMaker
+        | TorchSimStaticMaker
+        | None
+    ) = None
+    born_maker: ForceFieldStaticMaker | BaseVaspMaker | TorchSimStaticMaker | None = (
         None
     )
+    phonon_displacement_maker: (
+        ForceFieldStaticMaker | BaseVaspMaker | BaseAimsMaker | TorchSimStaticMaker
+    ) = None
     create_thermal_displacements: bool = True
     generate_frequencies_eigenvectors_kwargs: dict = field(
         default_factory=lambda: {

src/atomate2/common/jobs/elastic.py

@@ -18,6 +18,8 @@
 from atomate2 import SETTINGS
 from atomate2.common.analysis.elastic import get_default_strain_states
 from atomate2.common.schemas.elastic import ElasticDocument
+from atomate2.common.utils import check_class_name
+from atomate2.vasp.jobs.base import BaseVaspMaker
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -26,7 +28,7 @@
     from pymatgen.core.structure import Structure
 
     from atomate2.forcefields.jobs import ForceFieldRelaxMaker
-    from atomate2.vasp.jobs.base import BaseVaspMaker
+    from atomate2.torchsim import TorchSimOptimizeMaker
 
 
 logger = logging.getLogger(__name__)
@@ -104,14 +106,19 @@ def run_elastic_deformations(
     structure: Structure,
     deformations: list[Deformation],
     prev_dir: str | Path | None = None,
-    prev_dir_argname: str = None,
-    elastic_relax_maker: BaseVaspMaker | ForceFieldRelaxMaker = None,
+    prev_dir_argname: str | None = None,
+    elastic_relax_maker: BaseVaspMaker
+    | ForceFieldRelaxMaker
+    | TorchSimOptimizeMaker = None,
+    socket: bool = False,
 ) -> Response:
     """
     Run elastic deformations.
 
-    Note, this job will replace itself with N relaxation calculations, where N is
-    the number of deformations.
+    Note, this job will replace itself with N relaxation calculations,
+    or a single batched calculation using the socket interface to run all
+    deformations simultaneously. This results in lower overhead as well as
+    parallel relaxation of the deformations for TorchSim.
 
     Parameters
     ----------
@@ -121,46 +128,81 @@ def run_elastic_deformations(
         The deformations to apply.
     prev_dir : str or Path or None
         A previous directory to use for copying outputs.
-    prev_dir_argname: str
-        argument name for the prev_dir variable
-    elastic_relax_maker : .BaseVaspMaker or .ForceFieldRelaxMaker
-        A VaspMaker or a ForceFieldMaker to use to generate the elastic relaxation jobs.
+    prev_dir_argname: str or None
+        Argument name for the prev_dir variable.
+    elastic_relax_maker : .BaseVaspMaker, .ForceFieldRelaxMaker, or
+        .TorchSimOptimizeMaker
+        A VaspMaker, ForceFieldMaker, or TorchSimMaker to use to generate the elastic
+        relaxation jobs.
+    socket : bool
+        If True, uses the socket-io interface to run all deformations in a single
+        job, reducing overhead. In the specific case of TorchSim, this enables batching
+        of all structure relaxations.
+        Note: socket=True is not supported for BaseVaspMaker.
     """
-    relaxations = []
+    num_deformations = len(deformations)
+    elastic_jobs = []
     outputs = []
-    for idx, deformation in enumerate(deformations):
-        # deform the structure
+
+    if socket and isinstance(elastic_relax_maker, BaseVaspMaker):
+        raise ValueError("socket=True is not supported for BaseVaspMaker.")
+
+    deformed_structures = []
+    for deformation in deformations:
         dst = DeformStructureTransformation(deformation=deformation)
         ts = TransformedStructure(structure, transformations=[dst])
-        deformed_structure = ts.final_structure
+        deformed_structures.append(ts.final_structure)
 
         with contextlib.suppress(Exception):
-            # write details of the transformation to the transformations.json file
-            # this file will automatically get added to the task document and allow
-            # the elastic builder to reconstruct the elastic document; note the ":" is
-            # automatically converted to a "." in the filename.
+            # Write details of the transformation to the transformations.json file
             elastic_relax_maker.write_additional_data["transformations:json"] = ts
 
-        elastic_job_kwargs = {}
-        if prev_dir is not None and prev_dir_argname is not None:
-            elastic_job_kwargs[prev_dir_argname] = prev_dir
-        # create the job
-        relax_job = elastic_relax_maker.make(deformed_structure, **elastic_job_kwargs)
-        relax_job.append_name(f" {idx + 1}/{len(deformations)}")
-        relaxations.append(relax_job)
+    elastic_job_kwargs = {}
+    if prev_dir is not None and prev_dir_argname is not None:
+        elastic_job_kwargs[prev_dir_argname] = prev_dir
 
-        # extract the outputs we want
-        output = {
-            "stress": relax_job.output.output.stress,
-            "deformation": deformation,
-            "uuid": relax_job.output.uuid,
-            "job_dir": relax_job.output.dir_name,
-        }
+    if socket:
+        batched_job = elastic_relax_maker.make(
+            deformed_structures, **elastic_job_kwargs
+        )
+        batched_job.append_name(" batched_socket")
+        elastic_jobs.append(batched_job)
+
+        for idx, deformation in enumerate(deformations):
+            if check_class_name(elastic_relax_maker, "TorchSimOptimizeMaker"):
+                output = {
+                    "stress": batched_job.output.output.stress[idx],
+                    "deformation": deformation,
+                    "uuid": batched_job.output.uuid,
+                    "job_dir": batched_job.output.dir_name,
+                }
+            else:
+                output = {
+                    "stress": batched_job.output[idx].output.stress,
+                    "deformation": deformation,
+                    "uuid": batched_job.output[idx].uuid,
+                    "job_dir": batched_job.output[idx].dir_name,
+                }
+            outputs.append(output)
+
+    else:
+        for idx, deformed_structure in enumerate(deformed_structures):
+            relax_job = elastic_relax_maker.make(
+                deformed_structure, **elastic_job_kwargs
+            )
+            relax_job.append_name(f" {idx + 1}/{num_deformations}")
+            elastic_jobs.append(relax_job)
 
-        outputs.append(output)
+            output = {
+                "stress": relax_job.output.output.stress,
+                "deformation": deformations[idx],
+                "uuid": relax_job.output.uuid,
+                "job_dir": relax_job.output.dir_name,
+            }
+            outputs.append(output)
 
-    relax_flow = Flow(relaxations, outputs)
-    return Response(replace=relax_flow)
+    elastic_flow = Flow(elastic_jobs, outputs)
+    return Response(replace=elastic_flow)
 
 
 @job(output_schema=ElasticDocument)
@@ -221,7 +263,7 @@ def fit_elastic_tensor(
             failed_uuids.append(data["uuid"])
             continue
 
-        stresses.append(Stress(stress_sign_factor * np.array(data["stress"])))
+        stresses.append(Stress(stress_sign_factor * np.squeeze(data["stress"])))
         deformations.append(Deformation(data["deformation"]))
         uuids.append(data["uuid"])
         job_dirs.append(data["job_dir"])