Bug fixes in force field flows + housekeeping (#1286)

* minor bug fixes + pytest xdist for on-runner parallelism

* bump precommit and remove deprecated pragma

* ensure aims test is independent of run order

Author: esoteric-ephemera

.github/workflows/testing.yml

@@ -87,7 +87,7 @@ jobs:
         # However this `splitting-algorithm` means that tests cannot depend sensitively on the order they're executed in.
         run: |
           micromamba activate a2
-          pytest --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 --cov=atomate2 --cov-report=xml
+          pytest -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 --cov=atomate2 --cov-report=xml
 
 
       - uses: codecov/codecov-action@v1
@@ -151,7 +151,7 @@ jobs:
 
         run: |
           micromamba activate a2
-          pytest tests/{openff_md,openmm_md}
+          pytest -n auto tests/{openff_md,openmm_md}
 
   test-notebooks-and-ase:
     # prevent this action from running on forks
@@ -210,14 +210,14 @@ jobs:
           MP_API_KEY: ${{ secrets.MP_API_KEY }}
         run: |
           micromamba activate a2
-          pytest --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb --ignore=./tutorials/force_fields
+          pytest -n auto --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb --ignore=./tutorials/force_fields
 
       - name: Test ASE
         env:
           MP_API_KEY: ${{ secrets.MP_API_KEY }}
         run: |
           micromamba activate a2
-          pytest --splits 1 --group 1 --cov=atomate2 --cov-report=xml tests/ase
+          pytest -n auto --splits 1 --group 1 --cov=atomate2 --cov-report=xml tests/ase
 
       - uses: codecov/codecov-action@v1
         if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2'
@@ -288,7 +288,7 @@ jobs:
         # However this `splitting-algorithm` means that tests cannot depend sensitively on the order they're executed in.
         run: |
           micromamba activate a2
-          pytest --nbmake ./tutorials/force_fields
+          pytest -n auto --nbmake ./tutorials/force_fields
 
 
       - uses: codecov/codecov-action@v1

.pre-commit-config.yaml

@@ -3,20 +3,22 @@ default_language_version:
 exclude: ^(.github/|tests/test_data/abinit/)
 repos:
 - repo: https://github.com/charliermarsh/ruff-pre-commit
-  rev: v0.9.3
+  rev: v0.12.10
   hooks:
   - id: ruff
     args: [--fix]
     exclude: tutorials/grueneisen_workflow.ipynb
   - id: ruff-format
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v5.0.0
+  rev: v6.0.0
   hooks:
   - id: check-yaml
-  - id: fix-encoding-pragma
-    args: [--remove]
   - id: end-of-file-fixer
   - id: trailing-whitespace
+- repo: https://github.com/asottile/pyupgrade
+  rev: v3.20.0
+  hooks:
+    - id: pyupgrade
 - repo: https://github.com/asottile/blacken-docs
   rev: 1.19.1
   hooks:
@@ -31,15 +33,15 @@ repos:
   - id: rst-directive-colons
   - id: rst-inline-touching-normal
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.14.1
+  rev: v1.17.1
   hooks:
   - id: mypy
     files: ^src/
     additional_dependencies:
     - tokenize-rt==4.1.0
     - types-paramiko
 - repo: https://github.com/codespell-project/codespell
-  rev: v2.4.0
+  rev: v2.4.1
   hooks:
   - id: codespell
     stages: [pre-commit, commit-msg]

pyproject.toml

@@ -91,6 +91,7 @@ tests = [
     "pytest-cov==6.2.1",
     "pytest-mock==3.14.1",
     "pytest-split==0.10.0",
+    "pytest-xdist==3.8.0",
     "pytest==8.4.1",
 ]
 strict = [
@@ -105,7 +106,7 @@ strict = [
     "jobflow==0.2.0",
     "lobsterpy==0.5.7",
     "monty==2025.3.3",
-    "mp-api==0.45.7",
+    "mp-api==0.45.8",
     "numpy",
     "phonopy==2.30.1",
     "pydantic-settings==2.10.1",
@@ -130,7 +131,7 @@ strict-forcefields = [
     "mace-torch==0.3.14",
     "matgl==1.3.0",
     "quippy-ase==0.9.14; python_version < '3.12'",
-    "sevenn==0.10.3",
+    "sevenn==0.10.4",
     "torch==2.2.0",
     "torchdata==0.7.1",                            # TODO: remove when issue fixed
 ]
@@ -209,6 +210,7 @@ ignore = [
     "ISC001",
     "PD011",   # pandas-use-of-dot-values
     "PERF203", # try-except-in-loop
+    "PLC0415", # imports not at top level are OK
     "PLR0911", # too many returns
     "PLR0912", # too many branches
     "PLR0913", # too many arguments

src/atomate2/abinit/schemas/calculation.py

@@ -6,7 +6,6 @@
 import os
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Optional, Union
 
 from abipy.electrons.gsr import GsrFile
 from abipy.flowtk import events
@@ -74,31 +73,31 @@ class CalculationOutput(BaseModel):
         None, description="The final DFT energy per atom for the calculation"
     )
 
-    structure: Union[Structure] = Field(
+    structure: Structure = Field(
         None, description="The final structure from the calculation"
     )
 
     efermi: float = Field(
         None, description="The Fermi level from the calculation in eV"
     )
 
-    forces: Optional[list[Vector3D]] = Field(
+    forces: list[Vector3D] | None = Field(
         None, description="Forces acting on each atom"
     )
-    stress: Optional[Matrix3D] = Field(None, description="The stress on the cell")
-    is_metal: Optional[bool] = Field(None, description="Whether the system is metallic")
-    bandgap: Optional[float] = Field(
+    stress: Matrix3D | None = Field(None, description="The stress on the cell")
+    is_metal: bool | None = Field(None, description="Whether the system is metallic")
+    bandgap: float | None = Field(
         None, description="The band gap from the calculation in eV"
     )
-    direct_bandgap: Optional[float] = Field(
+    direct_bandgap: float | None = Field(
         None, description="The direct band gap from the calculation in eV"
     )
-    cbm: Optional[float] = Field(
+    cbm: float | None = Field(
         None,
         description="The conduction band minimum, or LUMO for molecules, in eV "
         "(if system is not metallic)",
     )
-    vbm: Optional[float] = Field(
+    vbm: float | None = Field(
         None,
         description="The valence band maximum, or HOMO for molecules, in eV "
         "(if system is not metallic)",
@@ -194,7 +193,7 @@ class Calculation(BaseModel):
     event_report: events.EventReport = Field(
         None, description="Event report of this abinit job."
     )
-    output_file_paths: Optional[dict[str, str]] = Field(
+    output_file_paths: dict[str, str] | None = Field(
         None,
         description="Paths (relative to dir_name) of the Abinit output files "
         "associated with this calculation",
@@ -257,7 +256,7 @@ def from_abinit_files(
             if report.run_completed:
                 has_abinit_completed = TaskState.SUCCESS
 
-        except (ValueError, RuntimeError, Exception) as exc:
+        except (ValueError, RuntimeError, Exception) as exc:  # noqa: BLE001
             msg = f"{cls} exception while parsing event_report:\n{exc}"
             logger.critical(msg)
 

src/atomate2/abinit/schemas/task.py

@@ -5,7 +5,7 @@
 import logging
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 from abipy.abio.inputs import AbinitInput
 from abipy.flowtk import events
@@ -33,7 +33,7 @@ class InputDoc(BaseModel):
         The final pymatgen Structure of the final system
     """
 
-    structure: Union[Structure] = Field(None, description="The input structure object")
+    structure: Structure = Field(None, description="The input structure object")
     abinit_input: AbinitInput = Field(
         None, description="AbinitInput used to perform calculation."
     )
@@ -88,8 +88,8 @@ class OutputDoc(BaseModel):
         Stress on the unit cell from the last calculation
     """
 
-    structure: Union[Structure] = Field(None, description="The output structure object")
-    trajectory: Optional[Sequence[Union[Structure]]] = Field(
+    structure: Structure = Field(None, description="The output structure object")
+    trajectory: Sequence[Structure] | None = Field(
         None, description="The trajectory of output structures"
     )
     energy: float = Field(
@@ -98,15 +98,15 @@ class OutputDoc(BaseModel):
     energy_per_atom: float = Field(
         None, description="The final DFT energy per atom for the last calculation"
     )
-    bandgap: Optional[float] = Field(
+    bandgap: float | None = Field(
         None, description="The DFT bandgap for the last calculation"
     )
-    cbm: Optional[float] = Field(None, description="CBM for this calculation")
-    vbm: Optional[float] = Field(None, description="VBM for this calculation")
-    forces: Optional[list[Vector3D]] = Field(
+    cbm: float | None = Field(None, description="CBM for this calculation")
+    vbm: float | None = Field(None, description="VBM for this calculation")
+    forces: list[Vector3D] | None = Field(
         None, description="Forces on atoms from the last calculation"
     )
-    stress: Optional[Matrix3D] = Field(
+    stress: Matrix3D | None = Field(
         None, description="Stress on the unit cell from the last calculation"
     )
 
@@ -179,49 +179,45 @@ class AbinitTaskDoc(StructureMetadata):
         Additional json loaded from the calculation directory
     """
 
-    dir_name: Optional[str] = Field(
-        None, description="The directory for this Abinit task"
-    )
-    last_updated: Optional[str] = Field(
+    dir_name: str | None = Field(None, description="The directory for this Abinit task")
+    last_updated: str | None = Field(
         default_factory=datetime_str,
         description="Timestamp for when this task document was last updated",
     )
-    completed_at: Optional[str] = Field(
+    completed_at: str | None = Field(
         None, description="Timestamp for when this task was completed"
     )
-    input: Optional[InputDoc] = Field(
+    input: InputDoc | None = Field(
         None, description="The input to the first calculation"
     )
-    output: Optional[OutputDoc] = Field(
+    output: OutputDoc | None = Field(
         None, description="The output of the final calculation"
     )
-    structure: Union[Structure] = Field(
-        None, description="Final output atoms from the task"
-    )
-    state: Optional[TaskState] = Field(None, description="State of this task")
-    event_report: Optional[events.EventReport] = Field(
+    structure: Structure = Field(None, description="Final output atoms from the task")
+    state: TaskState | None = Field(None, description="State of this task")
+    event_report: events.EventReport | None = Field(
         None, description="Event report of this abinit job."
     )
-    included_objects: Optional[list[AbinitObject]] = Field(
+    included_objects: list[AbinitObject] | None = Field(
         None, description="List of Abinit objects included with this task document"
     )
-    abinit_objects: Optional[dict[AbinitObject, Any]] = Field(
+    abinit_objects: dict[AbinitObject, Any] | None = Field(
         None, description="Abinit objects associated with this task"
     )
-    task_label: Optional[str] = Field(None, description="A description of the task")
-    tags: Optional[list[str]] = Field(
+    task_label: str | None = Field(None, description="A description of the task")
+    tags: list[str] | None = Field(
         None, description="Metadata tags for this task document"
     )
-    author: Optional[str] = Field(
+    author: str | None = Field(
         None, description="Author extracted from transformations"
     )
-    icsd_id: Optional[str] = Field(
+    icsd_id: str | None = Field(
         None, description="International crystal structure database id of the structure"
     )
-    calcs_reversed: Optional[list[Calculation]] = Field(
+    calcs_reversed: list[Calculation] | None = Field(
         None, description="The inputs and outputs for all Abinit runs in this task."
     )
-    transformations: Optional[dict[str, Any]] = Field(
+    transformations: dict[str, Any] | None = Field(
         None,
         description="Information on the structural transformations, parsed from a "
         "transformations.json file",
@@ -231,7 +227,7 @@ class AbinitTaskDoc(StructureMetadata):
         description="Information on the custodian settings used to run this "
         "calculation, parsed from a custodian.json file",
     )
-    additional_json: Optional[dict[str, Any]] = Field(
+    additional_json: dict[str, Any] | None = Field(
         None, description="Additional json loaded from the calculation directory"
     )
 

src/atomate2/abinit/utils/common.py

@@ -352,7 +352,7 @@ def get_final_structure(dir_name: Path | str) -> Structure:
         try:
             gsr_file = GsrFile(gsr_path)
         except Exception:
-            logging.exception("Exception occurred")
+            logging.exception("Exception occurred")  # noqa: LOG015
         else:
             return gsr_file.structure
 
@@ -362,7 +362,7 @@ def get_final_structure(dir_name: Path | str) -> Structure:
         try:
             ddb_file = DdbFile(ddb_path)
         except Exception:
-            logging.exception("Exception occurred")
+            logging.exception("Exception occurred")  # noqa: LOG015
         else:
             return ddb_file.structure
 
@@ -371,7 +371,7 @@ def get_final_structure(dir_name: Path | str) -> Structure:
         try:
             ab_out = AbinitOutputFile.from_file(out_path.path)
         except Exception:
-            logging.exception("Exception occurred")
+            logging.exception("Exception occurred")  # noqa: LOG015
         else:
             return ab_out.final_structure
 
@@ -425,7 +425,7 @@ def get_event_report(ofile: File, mpiabort_file: File) -> EventReport | None:
                     report.append(last_abort_event)
                 else:
                     report.append(last_abort_event)
-    except (ValueError, RuntimeError, Exception) as exc:
+    except (ValueError, RuntimeError, Exception) as exc:  # noqa: BLE001
         # Return a report with an error entry with info on the exception.
         logger.critical(f"{ofile}: Exception while parsing ABINIT events:\n {exc!s}")
         return parser.report_exception(ofile.path, exc)