Merge branch 'materialsproject:main' into main

src/atomate2/abinit/flows/core.py

@@ -50,8 +50,7 @@ def make(
         structure: Structure,
         restart_from: str | Path | None = None,
     ) -> Flow:
-        """
-        Create a band structure flow.
+        """Create a band structure flow.
 
         Parameters
         ----------
@@ -109,8 +108,7 @@ def make(
         structure: Structure | None = None,
         restart_from: str | Path | None = None,
     ) -> Flow:
-        """
-        Create a relaxation flow.
+        """Create a relaxation flow.
 
         Parameters
         ----------

src/atomate2/abinit/jobs/base.py

@@ -144,8 +144,7 @@ def make(
         restart_from: str | Path | list[str] | None = None,
         history: JobHistory | None = None,
     ) -> jobflow.Job:
-        """
-        Return an ABINIT jobflow.Job.
+        """Get an ABINIT jobflow.Job.
 
         Parameters
         ----------

src/atomate2/abinit/jobs/core.py

@@ -123,8 +123,7 @@ def make(
         history: JobHistory | None = None,
         mode: str = "uniform",
     ) -> Job:
-        """
-        Run a non-scf ABINIT job.
+        """Run a non-scf ABINIT job.
 
         Parameters
         ----------

src/atomate2/abinit/schemas/calculation.py

@@ -6,10 +6,7 @@
 import os
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Union
-
-if TYPE_CHECKING:
-    pass
+from typing import Optional, Union
 
 from abipy.electrons.gsr import GsrFile
 from abipy.flowtk import events
@@ -18,6 +15,7 @@
 from jobflow.utils import ValueEnum
 from pydantic import BaseModel, Field
 from pymatgen.core import Structure
+from typing_extensions import Self
 
 from atomate2.abinit.utils.common import LOG_FILE_NAME, MPIABORTFILE, get_event_report
 
@@ -110,9 +108,8 @@ class CalculationOutput(BaseModel):
     def from_abinit_gsr(
         cls,
         output: GsrFile,  # Must use auto_load kwarg when passed
-    ) -> CalculationOutput:
-        """
-        Create an Abinit output document from Abinit outputs.
+    ) -> Self:
+        """Create an Abinit output document from Abinit outputs.
 
         Parameters
         ----------
@@ -211,9 +208,8 @@ def from_abinit_files(
         abinit_gsr_file: Path | str = "out_GSR.nc",
         abinit_log_file: Path | str = LOG_FILE_NAME,
         abinit_abort_file: Path | str = MPIABORTFILE,
-    ) -> tuple[Calculation, dict[AbinitObject, dict]]:
-        """
-        Create an Abinit calculation document from a directory and file paths.
+    ) -> tuple[Self, dict[AbinitObject, dict]]:
+        """Create an Abinit calculation document from a directory and file paths.
 
         Parameters
         ----------
@@ -263,15 +259,14 @@ def from_abinit_files(
             msg = f"{cls} exception while parsing event_report:\n{exc}"
             logger.critical(msg)
 
-        return (
-            cls(
-                dir_name=str(dir_name),
-                task_name=task_name,
-                abinit_version=abinit_gsr.abinit_version,
-                has_abinit_completed=has_abinit_completed,
-                completed_at=completed_at,
-                output=output_doc,
-                event_report=report,
-            ),
-            None,  # abinit_objects,
+        instance = cls(
+            dir_name=str(dir_name),
+            task_name=task_name,
+            abinit_version=abinit_gsr.abinit_version,
+            has_abinit_completed=has_abinit_completed,
+            completed_at=completed_at,
+            output=output_doc,
+            event_report=report,
         )
+
+        return instance, None  # abinit_objects,

src/atomate2/abinit/schemas/task.py

@@ -5,22 +5,22 @@
 import logging
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, Optional, Union
 
 from abipy.abio.inputs import AbinitInput
 from abipy.flowtk import events
 from emmet.core.math import Matrix3D, Vector3D
 from emmet.core.structure import StructureMetadata
 from pydantic import BaseModel, Field
 from pymatgen.core import Structure
+from typing_extensions import Self
 
 from atomate2.abinit.files import load_abinit_input
 from atomate2.abinit.schemas.calculation import AbinitObject, Calculation, TaskState
 from atomate2.abinit.utils.common import LOG_FILE_NAME, MPIABORTFILE
 from atomate2.utils.datetime import datetime_str
 from atomate2.utils.path import get_uri, strip_hostname
 
-_T = TypeVar("_T", bound="AbinitTaskDoc")
 logger = logging.getLogger(__name__)
 
 
@@ -42,7 +42,7 @@ class InputDoc(BaseModel):
     )
 
     @classmethod
-    def from_abinit_calc_doc(cls, calc_doc: Calculation) -> InputDoc:
+    def from_abinit_calc_doc(cls, calc_doc: Calculation) -> Self:
         """Create a summary from an abinit CalculationDocument.
 
         Parameters
@@ -111,7 +111,7 @@ class OutputDoc(BaseModel):
     )
 
     @classmethod
-    def from_abinit_calc_doc(cls, calc_doc: Calculation) -> OutputDoc:
+    def from_abinit_calc_doc(cls, calc_doc: Calculation) -> Self:
         """Create a summary from an abinit CalculationDocument.
 
         Parameters
@@ -237,11 +237,11 @@ class AbinitTaskDoc(StructureMetadata):
 
     @classmethod
     def from_directory(
-        cls: type[_T],
+        cls,
         dir_name: Path | str,
         additional_fields: dict[str, Any] = None,
         **abinit_calculation_kwargs,
-    ) -> AbinitTaskDoc:
+    ) -> Self:
         """Create a task document from a directory containing Abinit files.
 
         Parameters

src/atomate2/abinit/sets/base.py

@@ -541,8 +541,7 @@ def get_abinit_input(
         kpoints_settings: dict | KSampling | None = None,
         input_index: int | None = None,
     ) -> AbinitInput:
-        """
-        Generate the AbinitInput for the input set.
+        """Generate the AbinitInput for the input set.
 
         Uses the defined factory function and additional parameters from user
         and subclasses.

src/atomate2/abinit/sets/core.py

@@ -45,8 +45,7 @@ def get_abinit_input(
         kpoints_settings: dict | KSampling | None = None,
         input_index: int | None = None,
     ) -> AbinitInput:
-        """
-        Generate the AbinitInput for the input set.
+        """Generate the AbinitInput for the input set.
 
         Removes some standard variables related to relaxation.
         """
@@ -207,8 +206,7 @@ def get_abinit_input(
         kpoints_settings: dict | KSampling | None = None,
         input_index: int | None = None,
     ) -> AbinitInput:
-        """
-        Generate the AbinitInput for the input set.
+        """Generate the AbinitInput for the input set.
 
         Sets tolmxf and determines the index of the MultiDataset.
         """

src/atomate2/abinit/utils/common.py

@@ -6,24 +6,25 @@
 import os
 from typing import TYPE_CHECKING
 
+from abipy.abio.outputs import AbinitOutputFile
+from abipy.dfpt.ddb import DdbFile
+from abipy.electrons.gsr import GsrFile
+from abipy.flowtk import events
+from abipy.flowtk.utils import Directory, File
+from monty.json import MSONable
+from monty.serialization import MontyDecoder
+
 if TYPE_CHECKING:
     from pathlib import Path
 
     from abipy.abio.inputs import AbinitInput
     from abipy.core.structure import Structure
     from abipy.flowtk.events import EventReport
     from jobflow import Flow, Job
+    from typing_extensions import Self
 
     from atomate2.abinit.utils.history import JobHistory
 
-from abipy.abio.outputs import AbinitOutputFile
-from abipy.dfpt.ddb import DdbFile
-from abipy.electrons.gsr import GsrFile
-from abipy.flowtk import events
-from abipy.flowtk.utils import Directory, File
-from monty.json import MSONable
-from monty.serialization import MontyDecoder
-
 TMPDIR_NAME = "tmpdata"
 OUTDIR_NAME = "outdata"
 INDIR_NAME = "indata"
@@ -163,7 +164,7 @@ def as_dict(self) -> dict:
         return self.to_dict()
 
     @classmethod
-    def from_dict(cls, d: dict) -> AbinitRuntimeError:
+    def from_dict(cls, d: dict) -> Self:
         """Create instance of the error from its dictionary representation."""
         dec = MontyDecoder()
         warnings = (
@@ -245,7 +246,7 @@ def to_dict(self) -> dict:
         return d
 
     @classmethod
-    def from_dict(cls, d: dict) -> UnconvergedError:
+    def from_dict(cls, d: dict) -> Self:
         """Create instance of the error from its dictionary representation."""
         dec = MontyDecoder()
         warnings = (
@@ -319,7 +320,7 @@ def as_dict(self) -> dict:
         }
 
     @classmethod
-    def from_dict(cls, d: dict) -> RestartInfo:
+    def from_dict(cls, d: dict) -> Self:
         """Create instance of the error from its dictionary representation."""
         return cls(
             previous_dir=d["previous_dir"],

src/atomate2/abinit/utils/history.py

@@ -10,15 +10,16 @@
 
 from monty.json import MontyDecoder, MSONable, jsanitize
 
+from atomate2.abinit.utils.common import OUTDIR_NAME
+
 if TYPE_CHECKING:
     from pathlib import Path
 
     from abipy.abio.inputs import AbinitInput
     from abipy.flowtk.events import AbinitEvent
     from abipy.flowtk.utils import Directory
     from jobflow import Flow, Job
-
-from atomate2.abinit.utils.common import OUTDIR_NAME
+    from typing_extensions import Self
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +47,7 @@ def as_dict(self) -> dict:
         }
 
     @classmethod
-    def from_dict(cls, d: dict) -> JobHistory:
+    def from_dict(cls, d: dict) -> Self:
         """Create instance of the history from its dictionary representation."""
         dec = MontyDecoder()
         return cls([dec.process_decoded(i) for i in d["items"]])
@@ -232,7 +233,7 @@ def as_dict(self) -> dict:
         return d
 
     @classmethod
-    def from_dict(cls, d: dict) -> JobEvent:
+    def from_dict(cls, d: dict) -> Self:
         """Create instance of the job event from its dictionary representation."""
         dec = MontyDecoder()
         details = dec.process_decoded(d["details"]) if "details" in d else None

src/atomate2/aims/flows/core.py

@@ -14,6 +14,7 @@
 
 if TYPE_CHECKING:
     from pymatgen.core import Molecule, Structure
+    from typing_extensions import Self
 
     from atomate2.aims.jobs.base import BaseAimsMaker
 
@@ -68,7 +69,7 @@ def from_parameters(
         cls,
         parameters: dict[str, Any],
         species_defaults: list[str] | tuple[str, str] = ("light", "tight"),
-    ) -> DoubleRelaxMaker:
+    ) -> Self:
         """Create the maker from an ASE parameter set.
 
         Creates a DoubleRelaxFlow for the same parameters with two different

src/atomate2/aims/jobs/base.py

@@ -88,8 +88,7 @@ def make(
         structure: Structure | Molecule,
         prev_dir: str | Path | None = None,
     ) -> Response:
-        """
-        Run an FHI-aims calculation.
+        """Run an FHI-aims calculation.
 
         Parameters
         ----------

src/atomate2/aims/schemas/calculation.py

@@ -17,6 +17,7 @@
 from pymatgen.electronic_structure.dos import Dos
 from pymatgen.io.aims.outputs import AimsOutput
 from pymatgen.io.common import VolumetricData
+from typing_extensions import Self
 
 if TYPE_CHECKING:
     from emmet.core.math import Matrix3D, Vector3D
@@ -133,9 +134,8 @@ def from_aims_output(
         cls,
         output: AimsOutput,  # Must use auto_load kwarg when passed
         # store_trajectory: bool = False,
-    ) -> CalculationOutput:
-        """
-        Create an FHI-aims output document from FHI-aims outputs.
+    ) -> Self:
+        """Create an FHI-aims output document from FHI-aims outputs.
 
         Parameters
         ----------
@@ -240,9 +240,8 @@ def from_aims_files(
         store_trajectory: bool = False,
         # store_scf: bool = False,
         store_volumetric_data: Optional[Sequence[str]] = STORE_VOLUMETRIC_DATA,
-    ) -> tuple[Calculation, dict[AimsObject, dict]]:
-        """
-        Create an FHI-aims calculation document from a directory and file paths.
+    ) -> tuple[Self, dict[AimsObject, dict]]:
+        """Create an FHI-aims calculation document from a directory and file paths.
 
         Parameters
         ----------
@@ -319,21 +318,18 @@ def from_aims_files(
             traj = _parse_trajectory(aims_output=aims_output)
             aims_objects[AimsObject.TRAJECTORY] = traj  # type: ignore  # noqa: PGH003
 
-        return (
-            cls(
-                dir_name=str(dir_name),
-                task_name=task_name,
-                aims_version=aims_output.aims_version,
-                has_aims_completed=has_aims_completed,
-                completed_at=completed_at,
-                output=output_doc,
-                output_file_paths={
-                    k.name.lower(): v for k, v in output_file_paths.items()
-                },
-            ),
-            aims_objects,
+        instance = cls(
+            dir_name=str(dir_name),
+            task_name=task_name,
+            aims_version=aims_output.aims_version,
+            has_aims_completed=has_aims_completed,
+            completed_at=completed_at,
+            output=output_doc,
+            output_file_paths={k.name.lower(): v for k, v in output_file_paths.items()},
         )
 
+        return instance, aims_objects
+
 
 def _get_output_file_paths(volumetric_files: list[str]) -> dict[AimsObject, str]:
     """Get the output file paths for FHI-aims output files.