Migrate to __future__.annotations

conftest.py

@@ -3,6 +3,8 @@
 Based on https://docs.pytest.org/en/latest/example/simple.html.
 """
 
+from __future__ import annotations
+
 import pytest
 
 

janus_core/__init__.py

@@ -1,5 +1,7 @@
 """Tools for machine learnt interatomic potentials."""
 
+from __future__ import annotations
+
 from importlib.metadata import version
 
 __version__ = version("janus-core")

janus_core/calculations/base.py

@@ -1,6 +1,8 @@
 """Prepare structures for MLIP calculations."""
 
-from typing import Any, Optional
+from __future__ import annotations
+
+from typing import Any
 
 from ase import Atoms
 
@@ -71,22 +73,22 @@ def __init__(
         self,
         *,
         calc_name: str = "base",
-        struct: Optional[MaybeSequence[Atoms]] = None,
-        struct_path: Optional[PathLike] = None,
+        struct: MaybeSequence[Atoms] | None = None,
+        struct_path: PathLike | None = None,
         arch: Architectures = "mace_mp",
         device: Devices = "cpu",
-        model_path: Optional[PathLike] = None,
-        read_kwargs: Optional[ASEReadArgs] = None,
+        model_path: PathLike | None = None,
+        read_kwargs: ASEReadArgs | None = None,
         sequence_allowed: bool = True,
-        calc_kwargs: Optional[dict[str, Any]] = None,
-        set_calc: Optional[bool] = None,
+        calc_kwargs: dict[str, Any] | None = None,
+        set_calc: bool | None = None,
         attach_logger: bool = False,
-        log_kwargs: Optional[dict[str, Any]] = None,
+        log_kwargs: dict[str, Any] | None = None,
         track_carbon: bool = True,
-        tracker_kwargs: Optional[dict[str, Any]] = None,
-        file_prefix: Optional[PathLike] = None,
-        additional_prefix: Optional[str] = None,
-        param_prefix: Optional[str] = None,
+        tracker_kwargs: dict[str, Any] | None = None,
+        file_prefix: PathLike | None = None,
+        additional_prefix: str | None = None,
+        param_prefix: str | None = None,
     ) -> None:
         """
         Read the structure being simulated and attach an MLIP calculator.

janus_core/calculations/descriptors.py

@@ -1,7 +1,9 @@
 """Calculate MLIP descriptors for structures."""
 
+from __future__ import annotations
+
 from collections.abc import Sequence
-from typing import Any, Optional
+from typing import Any
 
 from ase import Atoms
 import numpy as np
@@ -73,23 +75,23 @@ class Descriptors(BaseCalculation):
 
     def __init__(
         self,
-        struct: Optional[MaybeSequence[Atoms]] = None,
-        struct_path: Optional[PathLike] = None,
+        struct: MaybeSequence[Atoms] | None = None,
+        struct_path: PathLike | None = None,
         arch: Architectures = "mace_mp",
         device: Devices = "cpu",
-        model_path: Optional[PathLike] = None,
-        read_kwargs: Optional[ASEReadArgs] = None,
-        calc_kwargs: Optional[dict[str, Any]] = None,
-        set_calc: Optional[bool] = None,
+        model_path: PathLike | None = None,
+        read_kwargs: ASEReadArgs | None = None,
+        calc_kwargs: dict[str, Any] | None = None,
+        set_calc: bool | None = None,
         attach_logger: bool = False,
-        log_kwargs: Optional[dict[str, Any]] = None,
+        log_kwargs: dict[str, Any] | None = None,
         track_carbon: bool = True,
-        tracker_kwargs: Optional[dict[str, Any]] = None,
+        tracker_kwargs: dict[str, Any] | None = None,
         invariants_only: bool = True,
         calc_per_element: bool = False,
         calc_per_atom: bool = False,
         write_results: bool = False,
-        write_kwargs: Optional[ASEWriteArgs] = None,
+        write_kwargs: ASEWriteArgs | None = None,
     ) -> None:
         """
         Initialise class.

janus_core/calculations/eos.py

@@ -1,7 +1,9 @@
 """Equation of State."""
 
+from __future__ import annotations
+
 from copy import copy
-from typing import Any, Optional
+from typing import Any
 
 from ase import Atoms
 from ase.eos import EquationOfState
@@ -105,31 +107,31 @@ class EoS(BaseCalculation):
 
     def __init__(
         self,
-        struct: Optional[Atoms] = None,
-        struct_path: Optional[PathLike] = None,
+        struct: Atoms | None = None,
+        struct_path: PathLike | None = None,
         arch: Architectures = "mace_mp",
         device: Devices = "cpu",
-        model_path: Optional[PathLike] = None,
-        read_kwargs: Optional[ASEReadArgs] = None,
-        calc_kwargs: Optional[dict[str, Any]] = None,
-        set_calc: Optional[bool] = None,
+        model_path: PathLike | None = None,
+        read_kwargs: ASEReadArgs | None = None,
+        calc_kwargs: dict[str, Any] | None = None,
+        set_calc: bool | None = None,
         attach_logger: bool = False,
-        log_kwargs: Optional[dict[str, Any]] = None,
+        log_kwargs: dict[str, Any] | None = None,
         track_carbon: bool = True,
-        tracker_kwargs: Optional[dict[str, Any]] = None,
+        tracker_kwargs: dict[str, Any] | None = None,
         min_volume: float = 0.95,
         max_volume: float = 1.05,
         n_volumes: int = 7,
         eos_type: EoSNames = "birchmurnaghan",
         minimize: bool = True,
         minimize_all: bool = False,
-        minimize_kwargs: Optional[dict[str, Any]] = None,
+        minimize_kwargs: dict[str, Any] | None = None,
         write_results: bool = True,
         write_structures: bool = False,
-        write_kwargs: Optional[OutputKwargs] = None,
+        write_kwargs: OutputKwargs | None = None,
         plot_to_file: bool = False,
-        plot_kwargs: Optional[dict[str, Any]] = None,
-        file_prefix: Optional[PathLike] = None,
+        plot_kwargs: dict[str, Any] | None = None,
+        file_prefix: PathLike | None = None,
     ) -> None:
         """
         Initialise class.

janus_core/calculations/geom_opt.py

@@ -1,6 +1,8 @@
 """Prepare and run geometry optimization."""
 
-from typing import Any, Callable, Optional, Union
+from __future__ import annotations
+
+from typing import Any, Callable
 import warnings
 
 from ase import Atoms, filters, units
@@ -97,30 +99,30 @@ class GeomOpt(BaseCalculation):
 
     def __init__(
         self,
-        struct: Optional[Atoms] = None,
-        struct_path: Optional[PathLike] = None,
+        struct: Atoms | None = None,
+        struct_path: PathLike | None = None,
         arch: Architectures = "mace_mp",
         device: Devices = "cpu",
-        model_path: Optional[PathLike] = None,
-        read_kwargs: Optional[ASEReadArgs] = None,
-        calc_kwargs: Optional[dict[str, Any]] = None,
-        set_calc: Optional[bool] = None,
+        model_path: PathLike | None = None,
+        read_kwargs: ASEReadArgs | None = None,
+        calc_kwargs: dict[str, Any] | None = None,
+        set_calc: bool | None = None,
         attach_logger: bool = False,
-        log_kwargs: Optional[dict[str, Any]] = None,
+        log_kwargs: dict[str, Any] | None = None,
         track_carbon: bool = True,
-        tracker_kwargs: Optional[dict[str, Any]] = None,
+        tracker_kwargs: dict[str, Any] | None = None,
         fmax: float = 0.1,
         steps: int = 1000,
         symmetrize: bool = False,
         symmetry_tolerance: float = 0.001,
         angle_tolerance: float = -1.0,
-        filter_func: Optional[Union[Callable, str]] = FrechetCellFilter,
-        filter_kwargs: Optional[dict[str, Any]] = None,
-        optimizer: Union[Callable, str] = LBFGS,
-        opt_kwargs: Optional[ASEOptArgs] = None,
+        filter_func: Callable | str | None = FrechetCellFilter,
+        filter_kwargs: dict[str, Any] | None = None,
+        optimizer: Callable | str = LBFGS,
+        opt_kwargs: ASEOptArgs | None = None,
         write_results: bool = False,
-        write_kwargs: Optional[OutputKwargs] = None,
-        traj_kwargs: Optional[OutputKwargs] = None,
+        write_kwargs: OutputKwargs | None = None,
+        traj_kwargs: OutputKwargs | None = None,
     ) -> None:
         """
         Initialise GeomOpt class.

janus_core/calculations/md.py

@@ -1,13 +1,15 @@
 """Run molecular dynamics simulations."""
 
+from __future__ import annotations
+
 import datetime
 from functools import partial
 from itertools import combinations_with_replacement
 from math import isclose
 from os.path import getmtime
 from pathlib import Path
 import random
-from typing import Any, Optional, Union
+from typing import Any
 from warnings import warn
 
 from ase import Atoms, units
@@ -175,51 +177,51 @@ class MolecularDynamics(BaseCalculation):
 
     def __init__(
         self,
-        struct: Optional[Atoms] = None,
-        struct_path: Optional[PathLike] = None,
+        struct: Atoms | None = None,
+        struct_path: PathLike | None = None,
         arch: Architectures = "mace_mp",
         device: Devices = "cpu",
-        model_path: Optional[PathLike] = None,
-        read_kwargs: Optional[ASEReadArgs] = None,
-        calc_kwargs: Optional[dict[str, Any]] = None,
-        set_calc: Optional[bool] = None,
+        model_path: PathLike | None = None,
+        read_kwargs: ASEReadArgs | None = None,
+        calc_kwargs: dict[str, Any] | None = None,
+        set_calc: bool | None = None,
         attach_logger: bool = False,
-        log_kwargs: Optional[dict[str, Any]] = None,
+        log_kwargs: dict[str, Any] | None = None,
         track_carbon: bool = True,
-        tracker_kwargs: Optional[dict[str, Any]] = None,
-        ensemble: Optional[Ensembles] = None,
+        tracker_kwargs: dict[str, Any] | None = None,
+        ensemble: Ensembles | None = None,
         steps: int = 0,
         timestep: float = 1.0,
         temp: float = 300,
         equil_steps: int = 0,
         minimize: bool = False,
         minimize_every: int = -1,
-        minimize_kwargs: Optional[dict[str, Any]] = None,
+        minimize_kwargs: dict[str, Any] | None = None,
         rescale_velocities: bool = False,
         remove_rot: bool = False,
         rescale_every: int = 10,
-        file_prefix: Optional[PathLike] = None,
+        file_prefix: PathLike | None = None,
         restart: bool = False,
         restart_auto: bool = True,
-        restart_stem: Optional[PathLike] = None,
+        restart_stem: PathLike | None = None,
         restart_every: int = 1000,
         rotate_restart: bool = False,
         restarts_to_keep: int = 4,
-        final_file: Optional[PathLike] = None,
-        stats_file: Optional[PathLike] = None,
+        final_file: PathLike | None = None,
+        stats_file: PathLike | None = None,
         stats_every: int = 100,
-        traj_file: Optional[PathLike] = None,
+        traj_file: PathLike | None = None,
         traj_append: bool = False,
         traj_start: int = 0,
         traj_every: int = 100,
-        temp_start: Optional[float] = None,
-        temp_end: Optional[float] = None,
-        temp_step: Optional[float] = None,
-        temp_time: Optional[float] = None,
-        write_kwargs: Optional[OutputKwargs] = None,
-        post_process_kwargs: Optional[PostProcessKwargs] = None,
-        correlation_kwargs: Optional[list[CorrelationKwargs]] = None,
-        seed: Optional[int] = None,
+        temp_start: float | None = None,
+        temp_end: float | None = None,
+        temp_step: float | None = None,
+        temp_time: float | None = None,
+        write_kwargs: OutputKwargs | None = None,
+        post_process_kwargs: PostProcessKwargs | None = None,
+        correlation_kwargs: list[CorrelationKwargs] | None = None,
+        seed: int | None = None,
     ) -> None:
         """
         Initialise molecular dynamics simulation configuration.
@@ -493,7 +495,7 @@ def __init__(
                 "filemode": "a",
             }
 
-        self.dyn: Union[Langevin, VelocityVerlet, ASE_NPT]
+        self.dyn: Langevin | VelocityVerlet | ASE_NPT
         self.n_atoms = len(self.struct)
 
         self.offset = 0
@@ -633,7 +635,7 @@ def _optimize_structure(self) -> None:
             optimizer = GeomOpt(self.struct, **self.minimize_kwargs)
             optimizer.run()
 
-    def _set_param_prefix(self, file_prefix: Optional[PathLike] = None) -> str:
+    def _set_param_prefix(self, file_prefix: PathLike | None = None) -> str:
         """
         Set ensemble parameters for output files.
 
@@ -1156,8 +1158,8 @@ def __init__(
         bulk_modulus: float = 2.0,
         pressure: float = 0.0,
         ensemble: Ensembles = "npt",
-        file_prefix: Optional[PathLike] = None,
-        ensemble_kwargs: Optional[dict[str, Any]] = None,
+        file_prefix: PathLike | None = None,
+        ensemble_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ) -> None:
         """
@@ -1211,7 +1213,7 @@ def __init__(
             **ensemble_kwargs,
         )
 
-    def _set_param_prefix(self, file_prefix: Optional[PathLike] = None) -> str:
+    def _set_param_prefix(self, file_prefix: PathLike | None = None) -> str:
         """
         Set ensemble parameters for output files.
 
@@ -1297,7 +1299,7 @@ def __init__(
         *args,
         friction: float = 0.005,
         ensemble: Ensembles = "nvt",
-        ensemble_kwargs: Optional[dict[str, Any]] = None,
+        ensemble_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ) -> None:
         """
@@ -1391,7 +1393,7 @@ def __init__(
         self,
         *args,
         ensemble: Ensembles = "nve",
-        ensemble_kwargs: Optional[dict[str, Any]] = None,
+        ensemble_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ) -> None:
         """
@@ -1442,7 +1444,7 @@ def __init__(
         *args,
         thermostat_time: float = 50.0,
         ensemble: Ensembles = "nvt-nh",
-        ensemble_kwargs: Optional[dict[str, Any]] = None,
+        ensemble_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ) -> None:
         """
@@ -1546,8 +1548,8 @@ def __init__(
         bulk_modulus: float = 2.0,
         pressure: float = 0.0,
         ensemble: Ensembles = "nph",
-        file_prefix: Optional[PathLike] = None,
-        ensemble_kwargs: Optional[dict[str, Any]] = None,
+        file_prefix: PathLike | None = None,
+        ensemble_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ) -> None:
         """