Turn ForceFieldRelaxMaker._calculator method into calculator prop…

…erty (#839)

* remove contextlib.redirect_stdout(io.StringIO()) from ForceFieldMDMaker

* turn ForceFieldRelaxMaker._calculator method into calculator property

* bump matgl>=1.1.1

* ruff opt into docstring-code-format = true

docs/dev/workflow_tutorial.md

@@ -2,14 +2,14 @@
 
 ## Anatomy of an `atomate2` computational workflow (i.e., what do I need to write?)
 
-Every `atomate2` workflow is an instance of jobflow's `Flow ` class, which is a collection of Job and/or other `Flow` objects. So your end goal is to produce a `Flow `.
+Every `atomate2` workflow is an instance of jobflow's `Flow` class, which is a collection of Job and/or other `Flow` objects. So your end goal is to produce a `Flow`.
 
-In the context of computational materials science, `Flow ` objects are most easily created by a `Maker`, which contains a factory method make() that produces a `Flow `, given certain inputs. Typically, the input to `Maker`.make() includes atomic coordinate information in the form of a `pymatgen` `Structure` or `Molecule` object. So the basic signature looks like this:
+In the context of computational materials science, `Flow` objects are most easily created by a `Maker`, which contains a factory method make() that produces a `Flow`, given certain inputs. Typically, the input to `Maker`.make() includes atomic coordinate information in the form of a `pymatgen` `Structure` or `Molecule` object. So the basic signature looks like this:
 
 ```py
 class ExampleMaker(Maker):
     def make(self, coordinates: Structure) -> Flow:
-        # take the input coordinates and return a `Flow `
+        # take the input coordinates and return a `Flow`
         return Flow(...)
 ```
 
@@ -49,15 +49,16 @@ Finally, most `atomate2` workflows return structured output in the form of "Task
 **TODO - extend code block above to illustrate TaskDoc usage**
 
 In summary, a new `atomate2` workflow consists of the following components:
- - A `Maker` that actually generates the workflow
- - One or more `Job` and/or `Flow ` classes that define the discrete steps in the workflow
- - (optionally) an `InputGenerator` that produces a `pymatgen` `InputSet` for writing calculation input files
- - (optionally) a `TaskDocument` that defines a schema for storing the output data
+
+- A `Maker` that actually generates the workflow
+- One or more `Job` and/or `Flow` classes that define the discrete steps in the workflow
+- (optionally) an `InputGenerator` that produces a `pymatgen` `InputSet` for writing calculation input files
+- (optionally) a `TaskDocument` that defines a schema for storing the output data
 
 ## Where do I put my code?
 
 Because of the distributed design of the MP Software Ecosystem, writing a complete new workflow may involve making contributions to more than one GitHub repository. The following guidelines should help you understand where to put your contribution.
 
- - All workflow code (`Job`, `Flow `, `Maker`) belongs in `atomate2`
- - `InputSet` and `InputGenerator` code belongs in `pymatgen`. However, if you need to create these classes from scratch (i.e., you are working with a code that is not already supported in`pymatgen`), then it is recommended to include them in `atomate2` at first to facilitate rapid iteration. Once mature, they can be moved to `pymatgen` or to a `pymatgen` [addon package](https://pymatgen.org/addons).
- - `TaskDocument` schemas should generally be developed in `atomate2` alongside the workflow code. We recommend that you first check emmet to see if there is an existing schema that matches what you need. If so, you can import it. If not, check [`cclib`](https://cclib.github.io/). `cclib` output can be imported via [`atomate2.common.schemas.TaskDocument`](https://github.com/materialsproject/atomate2/blob/main/src/atomate2/common/schemas/cclib.py). If neither code has what you need, then new schemas should be developed within `atomate2` (or `cclib`).
+- All workflow code (`Job`, `Flow`, `Maker`) belongs in `atomate2`
+- `InputSet` and `InputGenerator` code belongs in `pymatgen`. However, if you need to create these classes from scratch (i.e., you are working with a code that is not already supported in`pymatgen`), then it is recommended to include them in `atomate2` at first to facilitate rapid iteration. Once mature, they can be moved to `pymatgen` or to a `pymatgen` [addon package](https://pymatgen.org/addons).
+- `TaskDocument` schemas should generally be developed in `atomate2` alongside the workflow code. We recommend that you first check emmet to see if there is an existing schema that matches what you need. If so, you can import it. If not, check [`cclib`](https://cclib.github.io/). `cclib` output can be imported via [`atomate2.common.schemas.TaskDocument`](https://github.com/materialsproject/atomate2/blob/main/src/atomate2/common/schemas/cclib.py). If neither code has what you need, then new schemas should be developed within `atomate2` (or `cclib`).

pyproject.toml

@@ -51,11 +51,11 @@ defects = [
 ]
 forcefields = [
     "ase>=3.22.1",
-    "torch<=2.2.1", # incompatibility with dgl if newer versions are used
     "chgnet>=0.2.2",
     "mace-torch>=0.3.3",
-    "matgl>=1.0.0",
+    "matgl>=1.1.1",
     "quippy-ase>=0.9.14",
+    "torch<=2.2.1",       # incompatibility with dgl if newer versions are used
 ]
 docs = [
     "FireWorks==2.0.3",
@@ -70,13 +70,17 @@ docs = [
     "sphinx_design==0.5.0",
 ]
 dev = ["pre-commit>=2.12.1"]
-tests = ["FireWorks==2.0.3", "pytest-cov==5.0.0", "pytest==8.0.2", "pytest-mock==3.14.0"]
+tests = [
+    "FireWorks==2.0.3",
+    "pytest-cov==5.0.0",
+    "pytest-mock==3.14.0",
+    "pytest==8.0.2",
+]
 strict = [
     "PyYAML==6.0.1",
     # must use >= for ase to not uninstall main branch install in CI
     # last known working commit: https://gitlab.com/ase/ase@2bab58f4e
     "ase>=3.22.1",
-    "torch==2.2.1",
     "cclib==1.8.1",
     "chgnet==0.3.5",
     "click==8.1.7",
@@ -87,7 +91,7 @@ strict = [
     "jobflow==0.1.17",
     "lobsterpy==0.3.8",
     "mace-torch>=0.3.3",
-    "matgl==1.0.0",
+    "matgl==1.1.1",
     "monty==2024.3.31",
     "mp-api==0.41.2",
     "numpy",
@@ -99,6 +103,7 @@ strict = [
     "python-ulid==2.5.0",
     "quippy-ase==0.9.14",
     "seekpath==2.1.0",
+    "torch==2.2.1",
     "typing-extensions==4.11.0",
 ]
 
@@ -192,6 +197,9 @@ lint.pydocstyle.convention = "numpy"
 lint.isort.known-first-party = ["atomate2"]
 lint.isort.split-on-trailing-comma = false
 
+[tool.ruff.format]
+docstring-code-format = true
+
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401"]
 "**/tests/*" = ["ANN", "ARG001", "D", "INP001", "S101"]

src/atomate2/forcefields/jobs.py

@@ -70,8 +70,7 @@ class ForceFieldRelaxMaker(Maker):
 
     Should be subclassed to use a specific force field. By default,
     the code attempts to use the `self.force_field_name` attr to look
-    up a predefined forcefield. To overwrite this behavior,
-    redefine `self._calculator`.
+    up a predefined forcefield. To overwrite this behavior, redefine `self.calculator`.
 
     Parameters
     ----------
@@ -133,7 +132,7 @@ def make(
 
         with revert_default_dtype():
             relaxer = Relaxer(
-                self._calculator(),
+                self.calculator,
                 relax_cell=self.relax_cell,
                 fix_symmetry=self.fix_symmetry,
                 symprec=self.symprec,
@@ -153,7 +152,8 @@ def make(
             **self.task_document_kwargs,
         )
 
-    def _calculator(self) -> Calculator:
+    @property
+    def calculator(self) -> Calculator:
         """ASE calculator, can be overwritten by user."""
         return ase_calculator(self.force_field_name, **self.calculator_kwargs)
 

src/atomate2/forcefields/md.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-import contextlib
-import io
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
@@ -271,8 +269,8 @@ def make(
         if md_func is NPT:
             # Note that until md_func is instantiated, isinstance(md_func,NPT) is False
             # ASE NPT implementation requires upper triangular cell
-            u, _ = schur(atoms.get_cell(complete=True), output="complex")
-            atoms.set_cell(u.real, scale_atoms=True)
+            schur_decomp, _ = schur(atoms.get_cell(complete=True), output="complex")
+            atoms.set_cell(schur_decomp.real, scale_atoms=True)
 
         if initial_velocities:
             atoms.set_velocities(initial_velocities)
@@ -288,32 +286,29 @@ def make(
                 ZeroRotation(atoms)
 
         with revert_default_dtype():
-            atoms.calc = self._calculator()
+            atoms.calc = self.calculator
 
-            with contextlib.redirect_stdout(io.StringIO()):
-                md_observer = TrajectoryObserver(atoms, store_md_outputs=True)
+            md_observer = TrajectoryObserver(atoms, store_md_outputs=True)
 
-                md_runner = md_func(
-                    atoms=atoms,
-                    timestep=self.time_step * units.fs,
-                    **self.ase_md_kwargs,
-                )
+            md_runner = md_func(
+                atoms=atoms, timestep=self.time_step * units.fs, **self.ase_md_kwargs
+            )
 
-                md_runner.attach(md_observer, interval=self.traj_interval)
+            md_runner.attach(md_observer, interval=self.traj_interval)
 
-                def _callback(dyn: MolecularDynamics = md_runner) -> None:
-                    if self.ensemble == "nve":
-                        return
-                    dyn.set_temperature(temperature_K=self.t_schedule[dyn.nsteps])
-                    if self.ensemble == "nvt":
-                        return
-                    dyn.set_stress(self.p_schedule[dyn.nsteps] * 1e3 * units.bar)
+            def _callback(dyn: MolecularDynamics = md_runner) -> None:
+                if self.ensemble == "nve":
+                    return
+                dyn.set_temperature(temperature_K=self.t_schedule[dyn.nsteps])
+                if self.ensemble == "nvt":
+                    return
+                dyn.set_stress(self.p_schedule[dyn.nsteps] * 1e3 * units.bar)
 
-                md_runner.attach(_callback, interval=1)
-                md_runner.run(steps=self.n_steps)
+            md_runner.attach(_callback, interval=1)
+            md_runner.run(steps=self.n_steps)
 
-                if self.traj_file is not None:
-                    md_observer.save(filename=self.traj_file, fmt=self.traj_file_fmt)
+            if self.traj_file is not None:
+                md_observer.save(filename=self.traj_file, fmt=self.traj_file_fmt)
 
             structure = AseAtomsAdaptor.get_structure(atoms)
 
@@ -332,7 +327,8 @@ def _callback(dyn: MolecularDynamics = md_runner) -> None:
             **self.task_document_kwargs,
         )
 
-    def _calculator(self) -> Calculator:
+    @property
+    def calculator(self) -> Calculator:
         """ASE calculator, can be overwritten by user."""
         return ase_calculator(self.force_field_name, **self.calculator_kwargs)
 

src/atomate2/forcefields/utils.py

@@ -359,8 +359,7 @@ def relax(
         if self.fix_symmetry:
             atoms.set_constraint(FixSymmetry(atoms, symprec=self.symprec))
         atoms.set_calculator(self.calculator)
-        stream = sys.stdout if verbose else io.StringIO()
-        with contextlib.redirect_stdout(stream):
+        with contextlib.redirect_stdout(sys.stdout if verbose else io.StringIO()):
             obs = TrajectoryObserver(atoms)
             if self.relax_cell:
                 atoms = cell_filter(atoms)
@@ -440,8 +439,8 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
             calculator = NequIPCalculator.from_deployed_model(**kwargs)
 
     elif isinstance(calculator_meta, dict):
-        _calculator = MontyDecoder().decode(json.dumps(calculator_meta))
-        calculator = _calculator(**kwargs)
+        calc_cls = MontyDecoder().decode(json.dumps(calculator_meta))
+        calculator = calc_cls(**kwargs)
 
     return calculator
 

tests/forcefields/test_md.py

@@ -160,7 +160,7 @@ def test_nve_and_dynamics_obj(si_structure: Structure, test_dir: Path):
     # This test serves two purposes:
     # 1. Test the NVE calculator
     # 2. Test specifying the `dynamics` kwarg of the `MDMaker` as a str
-    #    vs. as an ase.md.md.MolecularDyanmics object
+    #    vs. as an ase.md.md.MolecularDynamics object
 
     output = {}
     for k in ["from_str", "from_dyn"]: