Merge branch 'main' into feature/cmp0xff/76-kepler-problem

.pre-commit-config.yaml

@@ -9,11 +9,6 @@ repos:
       - id: end-of-file-fixer
       - id: requirements-txt-fixer
       - id: trailing-whitespace
-  - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.8.0
-    hooks:
-      - id: mypy
-        files: ^(hamilflow/|tests/)
   - repo: https://github.com/ambv/black
     rev: 24.4.2
     hooks:
@@ -25,3 +20,28 @@ repos:
       - id: isort
         name: isort (python)
         args: ["--multi-line", "3", "--profile", "black", "--treat-comment-as-code", "# %%", "--float-to-top"]
+  - repo: local
+    # We do not use pre-commit/mirrors-mypy,
+    # as it comes with opinionated defaults
+    # (like --ignore-missing-imports)
+    # and is difficult to configure to run
+    # with the dependencies correctly installed.
+    hooks:
+      - id: mypy
+        name: mypy
+        entry: mypy
+        language: python
+        # language_version: python3.12
+        additional_dependencies:
+          - mypy
+          - pandas-stubs
+          - pydantic
+          - pytest
+        types:
+          - python
+        # use require_serial so that script
+        # is only called once per commit
+        require_serial: true
+        # Print the number of files as a sanity-check
+        verbose: true
+        exclude: ^docs/tutorials

docs/tutorials/harmonic_oscillator_chain.py

@@ -43,7 +43,7 @@
 hoc = HarmonicOscillatorsChain(omega, ics, True)
 
 df_res = hoc(t)
-df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)]
+df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)]  # type: ignore [index]
 px.imshow(df_x_wide, origin="lower", labels={"y": "t"})
 
 # %% [markdown]

hamilflow/models/brownian_motion.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 import scipy as sp
+from numpy import typing as npt
 from pydantic import BaseModel, Field, computed_field, field_validator
 
 from hamilflow.models.utils.typing import TypeTime
@@ -60,12 +61,12 @@ class BrownianMotionIC(BaseModel):
 
     @field_validator("x0")
     @classmethod
-    def check_x0_types(cls, v: float | Sequence[float]) -> np.ndarray:
+    def check_x0_types(cls, v: float | Sequence[float]) -> float | Sequence[float]:
         if not isinstance(v, (float, int, Sequence)):
             # TODO I do not think this raise can be reached
             raise ValueError(f"Value of x0 should be int/float/list of int/float: {v=}")
 
-        return np.array(v, copy=False)
+        return v
 
 
 class BrownianMotion:
@@ -146,7 +147,9 @@ class BrownianMotion:
     def __init__(
         self,
         system: Mapping[str, float],
-        initial_condition: Mapping[str, float] | None = None,
+        initial_condition: (
+            Mapping[str, "Sequence[float] | npt.ArrayLike"] | None
+        ) = None,
     ):
         initial_condition = initial_condition or {}
         self.system = BrownianMotionSystem.model_validate(system)
@@ -155,13 +158,13 @@ def __init__(
     @property
     def dim(self) -> int:
         """Dimension of the Brownian motion"""
-        return self.initial_condition.x0.size
+        return np.array(self.initial_condition.x0, copy=False).size
 
     @property
     def _axis_names(self) -> list[str]:
         return [f"x_{i}" for i in range(self.dim)]
 
-    def _trajectory(self, n_new_steps: int, seed: int) -> np.ndarray:
+    def _trajectory(self, n_new_steps: int, seed: int) -> "npt.NDArray[np.float64]":
         """The trajectory of the particle.
 
         We first compute the delta displacement in each step.

hamilflow/models/free_particle.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 import pandas as pd
-from numpy.typing import ArrayLike
+from numpy import typing as npt
 from pydantic import BaseModel, Field, model_validator
 
 try:
@@ -48,16 +48,19 @@ def definition(self) -> dict[str, dict[str, float | list[float]]]:
         """model params and initial conditions defined as a dictionary."""
         return dict(initial_condition=self.initial_condition.model_dump())
 
-    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray:
-        return np.outer(t, self.initial_condition.v0) + self.initial_condition.x0
+    def _x(self, t: "Sequence[float] |  npt.ArrayLike") -> "npt.NDArray[np.float64]":
+        t = np.array(t, copy=False)
+        v0 = np.array(self.initial_condition.v0, copy=False)
+        x0 = np.array(self.initial_condition.x0, copy=False)
+        return np.outer(t, v0) + x0
 
-    def __call__(self, t: "Sequence[float] | ArrayLike[float]") -> pd.DataFrame:
+    def __call__(self, t: "Sequence[float] |  npt.ArrayLike") -> pd.DataFrame:
         """Generate time series data for the free particle.
 
         :param t: time(s).
         """
 
         data = self._x(t)
-        columns = (f"x{i+1}" for i in range(data.shape[1]))
+        columns = [f"x{i+1}" for i in range(data.shape[1])]
 
         return pd.DataFrame(data, columns=columns).assign(t=t).sort_index(axis=1)

hamilflow/models/harmonic_oscillator.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import pandas as pd
-from numpy.typing import ArrayLike
+from numpy import typing as npt
 from pydantic import BaseModel, Field, computed_field, field_validator
 
 
@@ -93,7 +93,7 @@ def definition(self) -> dict[str, dict[str, float]]:
         }
 
     @abstractmethod
-    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
+    def _x(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike:
         r"""Solution to simple harmonic oscillators."""
         ...
 
@@ -162,7 +162,7 @@ def __init__(
                 f"System is not a Simple Harmonic Oscillator: {self.system}"
             )
 
-    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray:
+    def _x(self, t: "Sequence[float] | npt.ArrayLike") -> np.ndarray:
         r"""Solution to simple harmonic oscillators:
 
         $$
@@ -235,7 +235,7 @@ def __init__(
                 f"This is a simple harmonic oscillator, use `SimpleHarmonicOscillator`."
             )
 
-    def _x_under_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
+    def _x_under_damped(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike:
         r"""Solution to under damped harmonic oscillators:
 
         $$
@@ -249,6 +249,7 @@ def _x_under_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
         \Omega = \omega\sqrt{ 1 - \zeta^2}.
         $$
         """
+        t = np.array(t, copy=False)
         omega_damp = self.system.omega * np.sqrt(1 - self.system.zeta)
         return (
             self.initial_condition.x0 * np.cos(omega_damp * t)
@@ -260,7 +261,7 @@ def _x_under_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
             * np.sin(omega_damp * t)
         ) * np.exp(-self.system.zeta * self.system.omega * t)
 
-    def _x_critical_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
+    def _x_critical_damped(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike:
         r"""Solution to critical damped harmonic oscillators:
 
         $$
@@ -274,11 +275,12 @@ def _x_critical_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLi
         \Omega = \omega\sqrt{ 1 - \zeta^2}.
         $$
         """
+        t = np.array(t, copy=False)
         return self.initial_condition.x0 * np.exp(
             -self.system.zeta * self.system.omega * t
         )
 
-    def _x_over_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
+    def _x_over_damped(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike:
         r"""Solution to over harmonic oscillators:
 
         $$
@@ -292,6 +294,7 @@ def _x_over_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
         \Gamma = \omega\sqrt{ \zeta^2 - 1 }.
         $$
         """
+        t = np.array(t, copy=False)
         gamma_damp = self.system.omega * np.sqrt(self.system.zeta - 1)
 
         return (
@@ -304,7 +307,7 @@ def _x_over_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
             * np.sinh(gamma_damp * t)
         ) * np.exp(-self.system.zeta * self.system.omega * t)
 
-    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
+    def _x(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike:
         r"""Solution to damped harmonic oscillators."""
         t = np.array(t, copy=False)
         if self.system.type == "under_damped":
@@ -364,9 +367,7 @@ def definition(
             initial_condition=self.initial_condition.model_dump(),
         )
 
-    def _z(
-        self, t: "Sequence[float] | ArrayLike[float] | ArrayLike[float]"
-    ) -> ArrayLike:
+    def _z(self, t: "Sequence[float] | npt.ArrayLike") -> npt.ArrayLike:
         r"""Solution to complex simple harmonic oscillators:
 
         $$
@@ -379,16 +380,14 @@ def _z(
         phases = -omega * t - phi[0], omega * t + phi[1]
         return x0[0] * np.exp(1j * phases[0]) + x0[1] * np.exp(1j * phases[1])
 
-    def __call__(
-        self, t: "Sequence[float] | ArrayLike[float] | ArrayLike[float]"
-    ) -> pd.DataFrame:
+    def __call__(self, t: "Sequence[float] | npt.ArrayLike") -> pd.DataFrame:
         """Generate time series data for the harmonic oscillator.
 
         Returns a list of floats representing the displacement at each time.
 
         :param t: time(s).
         """
-        t = t if isinstance(t, (Sequence, np.ndarray)) else [t]
+        t = np.array(t, copy=False)
         data = self._z(t)
 
         return pd.DataFrame({"t": t, "z": data})

hamilflow/models/harmonic_oscillator_chain.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 import pandas as pd
-from numpy.typing import ArrayLike
+from numpy import typing as npt
 from scipy.fft import ifft
 
 from .free_particle import FreeParticle
@@ -92,8 +92,8 @@ def definition(
         )
 
     def _z(
-        self, t: "Sequence[float] | ArrayLike[float]"
-    ) -> tuple[np.ndarray, np.ndarray]:
+        self, t: "Sequence[float] | npt.ArrayLike"
+    ) -> "tuple[npt.NDArray[np.complex64], npt.NDArray[np.complex64]]":
         t = np.array(t, copy=False).reshape(-1)
         all_travelling_waves = [self.free_mode._x(t).reshape(1, -1)]
 
@@ -116,24 +116,27 @@ def _z(
         return original_zs, travelling_waves
 
     def _x(
-        self, t: "Sequence[float] | ArrayLike[float]"
-    ) -> tuple[np.ndarray, np.ndarray]:
+        self, t: "Sequence[float] | npt.ArrayLike"
+    ) -> "tuple[npt.NDArray[np.float64], npt.NDArray[np.complex64]]":
         original_xs, travelling_waves = self._z(t)
 
         return np.real(original_xs), travelling_waves
 
-    def __call__(self, t: "Sequence[float] | ArrayLike[float]") -> pd.DataFrame:
+    def __call__(self, t: "Sequence[float] | npt.ArrayLike") -> pd.DataFrame:
         """Generate time series data for the harmonic oscillator chain.
 
         Returns float(s) representing the displacement at the given time(s).
 
         :param t: time.
         """
+        t = np.array(t, copy=False)
         original_xs, travelling_waves = self._x(t)
-        data = {
-            f"{name}{i}": values
+        data = {  # type: ignore [var-annotated]
+            f"{name}{i}": cast(
+                "npt.NDArray[np.float64] | npt.NDArray[np.complex64]", values
+            )
             for name, xs in zip(("x", "y"), (original_xs, travelling_waves))
-            for i, values in enumerate(xs)
+            for i, values in enumerate(xs)  # type: ignore [arg-type]
         }
 
         return pd.DataFrame(data, index=t)

hamilflow/models/pendulum.py

@@ -1,10 +1,10 @@
 import math
 from functools import cached_property
-from typing import Mapping, Sequence
+from typing import Any, Mapping, Sequence
 
 import numpy as np
 import pandas as pd
-from numpy.typing import ArrayLike
+from numpy import typing as npt
 from pydantic import BaseModel, Field, computed_field
 from scipy.special import ellipj, ellipk
 
@@ -68,7 +68,7 @@ def __init__(
         self.initial_condition = PendulumIC.model_validate(initial_condition)
 
     @cached_property
-    def definition(self) -> dict[str, float]:
+    def definition(self) -> dict[str, dict[str, Any]]:
         """Model params and initial conditions defined as a dictionary."""
         return dict(
             system=self.system.model_dump(),
@@ -105,10 +105,12 @@ def period(self) -> float:
         """
         return 4 * ellipk(self._math_m) / self.omega0
 
-    def _math_u(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
+    def _math_u(
+        self, t: "Sequence[float] | npt.ArrayLike"
+    ) -> "npt.NDArray[np.float64]":
         return self.omega0 * np.array(t, copy=False)
 
-    def u(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
+    def u(self, t: "Sequence[float] | npt.ArrayLike") -> "npt.NDArray[np.float64]":
         r"""The convenient generalised coordinate $u$,
         $\sin u \coloneqq \frac{\sin\frac{\theta}{2}}{\sin\frac{\theta_0}{2}}$.
 
@@ -121,7 +123,7 @@ def u(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
 
         return ph
 
-    def theta(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
+    def theta(self, t: "Sequence[float] | npt.ArrayLike") -> "npt.NDArray[np.float64]":
         r"""Angle $\theta$.
 
         :param t: time

pyproject.toml

@@ -62,3 +62,8 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.jupytext]
 formats = "ipynb,py:percent"
+
+
+[[tool.mypy.overrides]]
+module = ["plotly.*", "scipy.*"]
+ignore_missing_imports = true

tests/models/test_harmonic_oscillator.py

@@ -177,7 +177,7 @@ def test_degenerate_real(
             dict(omega=omega), dict(x0=(x0, x0), phi=(phi, phi))
         )
         sho = SimpleHarmonicOscillator(dict(omega=omega), dict(x0=2 * x0, phi=phi))
-        z = csho._z(times)
+        z = np.array(csho._z(times), copy=False)
         x = sho._x(times)
 
         assert np.all(z.imag == 0.0)

tests/models/test_harmonic_oscillator_chain.py

@@ -1,5 +1,5 @@
 from itertools import chain, product
-from typing import Iterable, Mapping, Sequence
+from typing import Iterable, Mapping, Sequence, cast
 
 import numpy as np
 import pytest
@@ -88,10 +88,7 @@ def test_dof(
 
     @pytest.mark.parametrize("wave_mode", [None, *_possible_wave_modes[1:]])
     def test_raise(
-        self,
-        omega: int,
-        free_mode: Mapping[str, int],
-        wave_mode: Mapping[str, int],
+        self, omega: int, free_mode: Mapping[str, int], wave_mode: Mapping[str, int]
     ) -> None:
         ics = [free_mode, *([wave_mode] if wave_mode else [])]
         with pytest.raises(ValueError):