apply updated ruff rules

Author: johannes-moegerle

docs/examples/comparisons/arc_fixed.py

@@ -1,18 +1,20 @@
 # ruff: noqa: INP001
+from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-import arc
 import numpy as np
 import scipy.integrate
 
 if TYPE_CHECKING:
+    import arc
+
     from ryd_numerov.units import NDArray
 
 
 def radialWavefunction(  # noqa: N802
     atom: arc.AlkaliAtom, n: int, l: int, j: float, step: float = 1e-2, use_fixed_arc: bool = False
-) -> tuple["NDArray", "NDArray"]:
+) -> tuple[NDArray, NDArray]:
     n, l, j = int(n), int(l), float(j)
     hartree_energy = 27.211 if not use_fixed_arc else 27.211_386_245_988
     energy = atom.getEnergy(int(n), int(l), float(j)) / hartree_energy

src/ryd_numerov/elements/base_element.py

@@ -1,17 +1,20 @@
+from __future__ import annotations
+
 import inspect
 import logging
 import re
 from abc import ABC
 from fractions import Fraction
 from functools import cache, cached_property
-from pathlib import Path
-from typing import TYPE_CHECKING, ClassVar, Optional, Union, overload
+from typing import TYPE_CHECKING, ClassVar, overload
 
 import numpy as np
 
 from ryd_numerov.units import rydberg_constant, ureg
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from ryd_numerov.model.model import PotentialType
     from ryd_numerov.units import PintFloat
 
@@ -39,21 +42,19 @@ class BaseElement(ABC):
 
     _core_electron_configuration: ClassVar[str]
     """Electron configuration of the core electrons, e.g. 4p6 for Rb or 5s for Sr."""
-    _ionization_energy: tuple[float, Optional[float], str]
+    _ionization_energy: tuple[float, float | None, str]
     """Ionization energy with uncertainty and unit: (value, uncertainty, unit)."""
 
     # Parameters for the extended Rydberg Ritz formula, see calc_n_star
-    _quantum_defects: ClassVar[Optional[dict[tuple[int, float, float], tuple[float, float, float, float, float]]]] = (
-        None
-    )
+    _quantum_defects: ClassVar[dict[tuple[int, float, float], tuple[float, float, float, float, float]] | None] = None
     """Dictionary containing the quantum defects for each (l, j_tot, s_tot) combination, i.e.
     _quantum_defects[(l,j_tot,s_tot)] = (d0, d2, d4, d6, d8)
     """
 
-    _corrected_rydberg_constant: tuple[float, Optional[float], str]
+    _corrected_rydberg_constant: tuple[float, float | None, str]
     r"""Corrected Rydberg constant stored as (value, uncertainty, unit)"""
 
-    potential_type_default: Optional["PotentialType"] = None
+    potential_type_default: PotentialType | None = None
     """Default potential type to use for this element. If None, the potential type must be specified explicitly.
     In general, it looks like marinescu_1993 is better for alkali atoms, and fei_2009 is better for alkaline earth atoms
     """
@@ -78,7 +79,7 @@ class BaseElement(ABC):
     defined in: Y. Fei et al., Chin. Phys. B 18, 4349 (2009), https://iopscience.iop.org/article/10.1088/1674-1056/18/10/025
     """
 
-    _nist_energy_levels_file: Optional[Path] = None
+    _nist_energy_levels_file: Path | None = None
     """Path to the NIST energy levels file for this element.
     The file should be directly downloaded from https://physics.nist.gov/PhysRefData/ASD/levels_form.html
     in the 'Tab-delimited' format and in units of Hartree.
@@ -170,7 +171,7 @@ def _setup_nist_energy_levels(self, file: Path) -> None:  # noqa: C901, PLR0912
 
     @classmethod
     @cache
-    def from_species(cls, species: str, use_nist_data: bool = True) -> "BaseElement":
+    def from_species(cls, species: str, use_nist_data: bool = True) -> BaseElement:
         """Create an instance of the element class from the species string.
 
         This method searches through all subclasses of BaseElement until it finds one with a matching species attribute.
@@ -196,7 +197,7 @@ def from_species(cls, species: str, use_nist_data: bool = True) -> "BaseElement"
         )
 
     @classmethod
-    def _get_concrete_subclasses(cls) -> list[type["BaseElement"]]:
+    def _get_concrete_subclasses(cls) -> list[type[BaseElement]]:
         subclasses = []
         for subclass in cls.__subclasses__():
             if not inspect.isabstract(subclass) and hasattr(subclass, "species"):
@@ -239,12 +240,12 @@ def is_allowed_shell(self, n: int, l: int, s_tot: float) -> bool:
         return (n, l) in self._additional_allowed_shells
 
     @overload
-    def get_ionization_energy(self, unit: None = None) -> "PintFloat": ...
+    def get_ionization_energy(self, unit: None = None) -> PintFloat: ...
 
     @overload
     def get_ionization_energy(self, unit: str) -> float: ...
 
-    def get_ionization_energy(self, unit: Optional[str] = "hartree") -> Union["PintFloat", float]:
+    def get_ionization_energy(self, unit: str | None = "hartree") -> PintFloat | float:
         """Return the ionization energy in the desired unit.
 
         Args:
@@ -263,12 +264,12 @@ def get_ionization_energy(self, unit: Optional[str] = "hartree") -> Union["PintF
         return ionization_energy.to(unit, "spectroscopy").magnitude
 
     @overload
-    def get_corrected_rydberg_constant(self, unit: None = None) -> "PintFloat": ...
+    def get_corrected_rydberg_constant(self, unit: None = None) -> PintFloat: ...
 
     @overload
     def get_corrected_rydberg_constant(self, unit: str) -> float: ...
 
-    def get_corrected_rydberg_constant(self, unit: Optional[str] = "hartree") -> Union["PintFloat", float]:
+    def get_corrected_rydberg_constant(self, unit: str | None = "hartree") -> PintFloat | float:
         r"""Return the corrected Rydberg constant in the desired unit.
 
         The corrected Rydberg constant is defined as
@@ -317,14 +318,14 @@ def reduced_mass_factor(self) -> float:
         )
 
     @overload
-    def calc_energy(self, n: int, l: int, j_tot: float, s_tot: float, unit: None = None) -> "PintFloat": ...
+    def calc_energy(self, n: int, l: int, j_tot: float, s_tot: float, unit: None = None) -> PintFloat: ...
 
     @overload
     def calc_energy(self, n: int, l: int, j_tot: float, s_tot: float, unit: str) -> float: ...
 
     def calc_energy(
-        self, n: int, l: int, j_tot: float, s_tot: float, unit: Optional[str] = "hartree"
-    ) -> Union["PintFloat", float]:
+        self, n: int, l: int, j_tot: float, s_tot: float, unit: str | None = "hartree"
+    ) -> PintFloat | float:
         r"""Calculate the energy of a Rydberg state with for the given n, l, j_tot and s_tot.
 
         I.e. either look up the energy for low lying states in the nist data,
@@ -355,7 +356,7 @@ def calc_energy(
         if j_tot % 1 != (l + s_tot) % 1:
             raise ValueError(f"Invalid quantum numbers: ({l=}, {j_tot=}, {s_tot=})")
 
-        energy_au: Optional[float] = None
+        energy_au: float | None = None
         if n <= self._nist_n_max and self.use_nist_data:
             if (n, l, j_tot, s_tot) in self._nist_energy_levels:
                 energy_au = self._nist_energy_levels[(n, l, j_tot, s_tot)]

src/ryd_numerov/elements/strontium.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from pathlib import Path
-from typing import ClassVar, Optional
+from typing import ClassVar
 
 from ryd_numerov.elements.base_element import BaseElement
 from ryd_numerov.units import electron_mass, rydberg_constant
@@ -15,7 +17,7 @@ class _StrontiumAbstract(BaseElement):
     _nist_energy_levels_file = Path(__file__).parent / "nist_energy_levels" / "strontium.txt"
 
     # https://webbook.nist.gov/cgi/inchi?ID=C7440246&Mask=20
-    _ionization_energy: tuple[float, Optional[float], str] = (5.694_84, 0.000_02, "eV")
+    _ionization_energy: tuple[float, float | None, str] = (5.694_84, 0.000_02, "eV")
 
     potential_type_default = "model_potential_fei_2009"
 

src/ryd_numerov/model/model.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import logging
-from typing import TYPE_CHECKING, Literal, Optional, TypeVar, get_args
+from typing import TYPE_CHECKING, Literal, TypeVar, get_args
 
 import numpy as np
 
@@ -20,9 +22,9 @@ class Model:
 
     def __init__(
         self,
-        element: "BaseElement",
+        element: BaseElement,
         l: int,
-        potential_type: Optional[PotentialType] = None,
+        potential_type: PotentialType | None = None,
     ) -> None:
         r"""Initialize the model.
 
@@ -43,7 +45,7 @@ def __init__(
             raise ValueError(f"Invalid potential type {potential_type}. Must be one of {get_args(PotentialType)}.")
         self.potential_type = potential_type
 
-    def calc_potential_coulomb(self, x: "XType") -> "XType":
+    def calc_potential_coulomb(self, x: XType) -> XType:
         r"""Calculate the Coulomb potential V_Col(x) in atomic units.
 
         The Coulomb potential is given as
@@ -62,7 +64,7 @@ def calc_potential_coulomb(self, x: "XType") -> "XType":
         """
         return -1 / x
 
-    def calc_model_potential_marinescu_1993(self, x: "XType") -> "XType":
+    def calc_model_potential_marinescu_1993(self, x: XType) -> XType:
         r"""Calculate the model potential by Marinescu et al. (1994) in atomic units.
 
         The model potential from
@@ -114,7 +116,7 @@ def calc_model_potential_marinescu_1993(self, x: "XType") -> "XType":
 
         return v_c + v_p
 
-    def calc_model_potential_fei_2009(self, x: "XType") -> "XType":
+    def calc_model_potential_fei_2009(self, x: XType) -> XType:
         r"""Calculate the model potential by Fei et al. (2009) in atomic units.
 
         The four parameter potential from Y. Fei et al., Chin. Phys. B 18, 4349 (2009), https://iopscience.iop.org/article/10.1088/1674-1056/18/10/025
@@ -138,7 +140,7 @@ def calc_model_potential_fei_2009(self, x: "XType") -> "XType":
             denom: XType = 1 - alpha + alpha * np.exp(beta * x**delta + gamma * x ** (2.0 * delta))
             return -1 / x - (self.element.Z - 1) / (x * denom)
 
-    def calc_effective_potential_centrifugal(self, x: "XType") -> "XType":
+    def calc_effective_potential_centrifugal(self, x: XType) -> XType:
         r"""Calculate the effective centrifugal potential V_l(x) in atomic units.
 
         The effective centrifugal potential is given as
@@ -158,7 +160,7 @@ def calc_effective_potential_centrifugal(self, x: "XType") -> "XType":
         x2 = x * x
         return (1 / self.element.reduced_mass_factor) * self.l * (self.l + 1) / (2 * x2)
 
-    def calc_effective_potential_sqrt(self, x: "XType") -> "XType":
+    def calc_effective_potential_sqrt(self, x: XType) -> XType:
         r"""Calculate the effective potential V_sqrt(x) from the sqrt transformation in atomic units.
 
         The sqrt transformation potential arises from the transformation from the wavefunction u(x) to w(z),
@@ -179,7 +181,7 @@ def calc_effective_potential_sqrt(self, x: "XType") -> "XType":
         x2 = x * x
         return (1 / self.element.reduced_mass_factor) * (3 / 32) / x2
 
-    def calc_total_effective_potential(self, x: "XType") -> "XType":
+    def calc_total_effective_potential(self, x: XType) -> XType:
         r"""Calculate the total effective potential V_eff(x) in atomic units.
 
         The total effective potential includes all physical and effective potentials:

src/ryd_numerov/radial/grid.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -65,7 +67,7 @@ def z_max(self) -> float:
         return self.z_list[-1]  # type: ignore [no-any-return]  # FIXME: numpy indexing
 
     @property
-    def z_list(self) -> "NDArray":
+    def z_list(self) -> NDArray:
         """The grid in the scaled dimensionless coordinate z = sqrt{x}.
 
         In this coordinate the grid points are chosen equidistant,
@@ -84,11 +86,11 @@ def x_max(self) -> float:
         return self.z_max**2
 
     @property
-    def x_list(self) -> "NDArray":
+    def x_list(self) -> NDArray:
         """The grid in the dimensionless coordinate x = r/a_0."""
         return self.z_list**2
 
-    def set_grid_range(self, step_start: Optional[int] = None, step_stop: Optional[int] = None) -> None:
+    def set_grid_range(self, step_start: int | None = None, step_stop: int | None = None) -> None:
         """Restrict the grid to the range [step_start, step_stop]."""
         if step_start is None:
             step_start = 0

src/ryd_numerov/radial/numerov.py

@@ -1,9 +1,12 @@
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Callable, Union
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable
 
 from numba import njit
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from ryd_numerov.units import NDArray
 
 
@@ -13,7 +16,7 @@ def _run_numerov_integration_python(
     dx: float,
     y0: float,
     y1: float,
-    g_list: Union[Sequence[float], "NDArray"],
+    g_list: Sequence[float] | NDArray,
     x_min: float,
     verbose: bool = False,
 ) -> list[float]:
@@ -58,7 +61,7 @@ def run_numerov_integration(
     dx: float,
     y0: float,
     y1: float,
-    g_list: Union[Sequence[float], "NDArray"],
+    g_list: Sequence[float] | NDArray,
     x_min: float,
     verbose: bool = False,
 ) -> list[float]:

src/ryd_numerov/radial/radial_matrix_element.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import TYPE_CHECKING, Literal
 
@@ -14,8 +16,8 @@
 
 
 def calc_radial_matrix_element(
-    state1: "RydbergStateBase",
-    state2: "RydbergStateBase",
+    state1: RydbergStateBase,
+    state2: RydbergStateBase,
     k_radial: int = 0,
     integration_method: INTEGRATION_METHODS = "sum",
 ) -> float:
@@ -52,10 +54,10 @@ def calc_radial_matrix_element(
 
 
 def _calc_radial_matrix_element_from_w_z(
-    z1: "NDArray",
-    w1: "NDArray",
-    z2: "NDArray",
-    w2: "NDArray",
+    z1: NDArray,
+    w1: NDArray,
+    z2: NDArray,
+    w2: NDArray,
     k_radial: int = 0,
     integration_method: INTEGRATION_METHODS = "sum",
 ) -> float:
@@ -120,7 +122,7 @@ def _calc_radial_matrix_element_from_w_z(
     return _integrate(integrand, dz, integration_method)
 
 
-def _multiply_by_powers(result: "NDArray", base: "NDArray", exponent: int) -> "NDArray":
+def _multiply_by_powers(result: NDArray, base: NDArray, exponent: int) -> NDArray:
     """Calculate result * base**(exponent) in an optimized way."""
     base_powers = {0: base}
     for i in range(exponent):
@@ -133,7 +135,7 @@ def _multiply_by_powers(result: "NDArray", base: "NDArray", exponent: int) -> "N
     return result
 
 
-def _sanity_check_integration(z1: "NDArray", z2: "NDArray") -> None:
+def _sanity_check_integration(z1: NDArray, z2: NDArray) -> None:
     tol = 1e-10
     assert len(z1) == len(z2), f"Length mismatch: {len(z1)=} != {len(z2)=}"
     assert z1[0] - z2[0] < tol, f"First point mismatch: {z1[0]=} != {z2[0]=}"
@@ -142,7 +144,7 @@ def _sanity_check_integration(z1: "NDArray", z2: "NDArray") -> None:
     assert z1[-1] - z2[-1] < tol, f"Last point mismatch: {z1[-1]=} != {z2[-1]=}"
 
 
-def _integrate(integrand: "NDArray", dz: float, method: INTEGRATION_METHODS) -> float:
+def _integrate(integrand: NDArray, dz: float, method: INTEGRATION_METHODS) -> float:
     """Integrate the given integrand using the specified method."""
     if method == "sum":
         value = np.sum(integrand) * dz