fix(typing): #69 np.ndarray is not a subclass of typing.Sequence (#70)

fix(typing): #69

hamilflow/models/brownian_motion.py

@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import Mapping
+from typing import Mapping, Sequence
 
 import numpy as np
 import pandas as pd
@@ -54,15 +54,16 @@ class BrownianMotionIC(BaseModel):
         the dimension of the model too.
     """
 
-    x0: float | int | list[float | int] = Field(default=1.0)
+    x0: float | Sequence[float] = Field(default=1.0)
 
     @field_validator("x0")
     @classmethod
-    def check_x0_types(cls, v: float | int | list[float]) -> np.ndarray:
-        if not isinstance(v, (float, int, list)):
+    def check_x0_types(cls, v: float | Sequence[float]) -> np.ndarray:
+        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.asarray(v)
+        return np.array(v, copy=False)
 
 
 class BrownianMotion:

hamilflow/models/free_particle.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pandas as pd
+from numpy.typing import ArrayLike
 from pydantic import BaseModel, Field, model_validator
 
 try:
@@ -18,8 +19,8 @@ class FreeParticleIC(BaseModel):
     :cvar v0: the initial velocity
     """
 
-    x0: float | int | Sequence[float | int] = Field()
-    v0: float | int | Sequence[float | int] = Field()
+    x0: float | Sequence[float] = Field()
+    v0: float | Sequence[float] = Field()
 
     @model_validator(mode="after")
     def check_dimensions_match(self) -> Self:
@@ -38,19 +39,19 @@ class FreeParticle:
     """
 
     def __init__(
-        self, initial_condition: Mapping[str, float | int | Sequence[float | int]]
+        self, initial_condition: Mapping[str, float | Sequence[float]]
     ) -> None:
         self.initial_condition = FreeParticleIC.model_validate(initial_condition)
 
     @cached_property
-    def definition(self) -> dict[str, dict[str, int | float | Sequence[int | float]]]:
+    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: float | int | Sequence[float | int]) -> np.ndarray:
+    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray:
         return np.outer(t, self.initial_condition.v0) + self.initial_condition.x0
 
-    def __call__(self, t: float | int | Sequence[float | int]) -> pd.DataFrame:
+    def __call__(self, t: "Sequence[float] | ArrayLike[float]") -> pd.DataFrame:
         """Generate time series data for the free particle.
 
         :param t: time(s).

hamilflow/models/harmonic_oscillator.py

@@ -77,23 +77,23 @@ class HarmonicOscillatorBase(ABC):
 
     def __init__(
         self,
-        system: Mapping[str, float | int],
-        initial_condition: Mapping[str, float | int] | None = None,
+        system: Mapping[str, float],
+        initial_condition: Mapping[str, float] | None = None,
     ) -> None:
         initial_condition = initial_condition or {}
         self.system = HarmonicOscillatorSystem.model_validate(system)
         self.initial_condition = HarmonicOscillatorIC.model_validate(initial_condition)
 
     @cached_property
-    def definition(self) -> dict[str, dict[str, float | int]]:
+    def definition(self) -> dict[str, dict[str, float]]:
         """model params and initial conditions defined as a dictionary."""
         return {
             "system": self.system.model_dump(),
             "initial_condition": self.initial_condition.model_dump(),
         }
 
     @abstractmethod
-    def _x(self, t: float | int | Sequence[float | int]) -> ArrayLike:
+    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
         r"""Solution to simple harmonic oscillators."""
         ...
 
@@ -153,16 +153,16 @@ class SimpleHarmonicOscillator(HarmonicOscillatorBase):
 
     def __init__(
         self,
-        system: Mapping[str, float | int],
-        initial_condition: Mapping[str, float | int] | None = None,
+        system: Mapping[str, float],
+        initial_condition: Mapping[str, float] | None = None,
     ) -> None:
         super().__init__(system, initial_condition)
         if self.system.type != "simple":
             raise ValueError(
                 f"System is not a Simple Harmonic Oscillator: {self.system}"
             )
 
-    def _x(self, t: float | int | Sequence[float | int]) -> np.ndarray:
+    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray:
         r"""Solution to simple harmonic oscillators:
 
         $$
@@ -225,8 +225,8 @@ class DampedHarmonicOscillator(HarmonicOscillatorBase):
 
     def __init__(
         self,
-        system: Mapping[str, float | int],
-        initial_condition: Mapping[str, float | int] | None = None,
+        system: Mapping[str, float],
+        initial_condition: Mapping[str, float] | None = None,
     ) -> None:
         super().__init__(system, initial_condition)
         if self.system.type == "simple":
@@ -235,7 +235,7 @@ def __init__(
                 f"This is a simple harmonic oscillator, use `SimpleHarmonicOscillator`."
             )
 
-    def _x_under_damped(self, t: float | int | Sequence[float | int]) -> ArrayLike:
+    def _x_under_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
         r"""Solution to under damped harmonic oscillators:
 
         $$
@@ -260,7 +260,7 @@ def _x_under_damped(self, t: float | int | Sequence[float | int]) -> ArrayLike:
             * np.sin(omega_damp * t)
         ) * np.exp(-self.system.zeta * self.system.omega * t)
 
-    def _x_critical_damped(self, t: float | int | Sequence[float | int]) -> ArrayLike:
+    def _x_critical_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
         r"""Solution to critical damped harmonic oscillators:
 
         $$
@@ -278,7 +278,7 @@ def _x_critical_damped(self, t: float | int | Sequence[float | int]) -> ArrayLik
             -self.system.zeta * self.system.omega * t
         )
 
-    def _x_over_damped(self, t: float | int | Sequence[float | int]) -> ArrayLike:
+    def _x_over_damped(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
         r"""Solution to over harmonic oscillators:
 
         $$
@@ -304,7 +304,7 @@ def _x_over_damped(self, t: float | int | Sequence[float | int]) -> ArrayLike:
             * np.sinh(gamma_damp * t)
         ) * np.exp(-self.system.zeta * self.system.omega * t)
 
-    def _x(self, t: float | int | Sequence[float | int]) -> ArrayLike:
+    def _x(self, t: "Sequence[float] | ArrayLike[float]") -> ArrayLike:
         r"""Solution to damped harmonic oscillators."""
         t = np.array(t, copy=False)
         if self.system.type == "under_damped":
@@ -328,8 +328,8 @@ class ComplexSimpleHarmonicOscillatorIC(BaseModel):
     :cvar phi: initial phases
     """
 
-    x0: tuple[float | int, float | int] = Field()
-    phi: tuple[float | int, float | int] = Field(default=(0, 0))
+    x0: tuple[float, float] = Field()
+    phi: tuple[float, float] = Field(default=(0, 0))
 
 
 class ComplexSimpleHarmonicOscillator:
@@ -341,8 +341,8 @@ class ComplexSimpleHarmonicOscillator:
 
     def __init__(
         self,
-        system: Mapping[str, float | int],
-        initial_condition: Mapping[str, tuple[float | int, float | int]],
+        system: Mapping[str, float],
+        initial_condition: Mapping[str, tuple[float, float]],
     ) -> None:
         self.system = HarmonicOscillatorSystem.model_validate(system)
         self.initial_condition = ComplexSimpleHarmonicOscillatorIC.model_validate(
@@ -356,15 +356,17 @@ def __init__(
     @cached_property
     def definition(
         self,
-    ) -> dict[str, dict[str, float | int | tuple[float | int, float | int]]]:
+    ) -> dict[str, dict[str, float | tuple[float, float]]]:
         """model params and initial conditions defined as a dictionary."""
 
         return dict(
             system=self.system.model_dump(),
             initial_condition=self.initial_condition.model_dump(),
         )
 
-    def _z(self, t: float | int | Sequence[float | int]) -> ArrayLike:
+    def _z(
+        self, t: "Sequence[float] | ArrayLike[float] | ArrayLike[float]"
+    ) -> ArrayLike:
         r"""Solution to complex simple harmonic oscillators:
 
         $$
@@ -377,14 +379,16 @@ def _z(self, t: float | int | Sequence[float | int]) -> ArrayLike:
         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: float | int | Sequence[float | int]) -> pd.DataFrame:
+    def __call__(
+        self, t: "Sequence[float] | ArrayLike[float] | ArrayLike[float]"
+    ) -> 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 not isinstance(t, Sequence) else t
+        t = t if isinstance(t, (Sequence, np.ndarray)) else [t]
         data = self._z(t)
 
         return pd.DataFrame({"t": t, "z": data})

hamilflow/models/pendulum.py

@@ -1,5 +1,6 @@
 import math
 from functools import cached_property
+from typing import Mapping, Sequence
 
 import numpy as np
 import pandas as pd
@@ -54,8 +55,8 @@ class Pendulum:
 
     def __init__(
         self,
-        system: int | float | dict[str, int | float],
-        initial_condition: int | float | dict[str, int | float],
+        system: float | Mapping[str, float],
+        initial_condition: float | Mapping[str, float],
     ) -> None:
         if isinstance(system, (float, int)):
             system = {"omega0": system}
@@ -102,10 +103,10 @@ def period(self) -> float:
         """
         return 4 * ellipk(self._math_m) / self.omega0
 
-    def _math_u(self, t: ArrayLike) -> np.ndarray[float]:
-        return self.omega0 * np.asarray(t)
+    def _math_u(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
+        return self.omega0 * np.array(t, copy=False)
 
-    def u(self, t: ArrayLike) -> np.ndarray[float]:
+    def u(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
         r"""The convenient generalised coordinate $u$,
         $\sin u \coloneqq \frac{\sin\frac{\theta}{2}}{\sin\frac{\theta_0}{2}}$.
 
@@ -118,7 +119,7 @@ def u(self, t: ArrayLike) -> np.ndarray[float]:
 
         return ph
 
-    def theta(self, t: ArrayLike) -> np.ndarray[float]:
+    def theta(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
         r"""Angle $\theta$.
 
         :param t: time

tests/test_models/test_free_particle.py

@@ -60,5 +60,6 @@ def test_call(
         expected: pd.DataFrame,
     ) -> None:
         assert_frame_equal(
-            FreeParticle(initial_condition=dict(x0=x0, v0=v0))(t), expected
+            FreeParticle(initial_condition=dict(x0=x0, v0=v0))(t).astype(float),
+            expected.astype(float),
         )

tests/test_models/test_pendulum.py

@@ -1,26 +1,27 @@
 import math
+from typing import Sequence
 
 import numpy as np
 import pytest
-from _pytest.fixtures import SubRequest
 from numpy.testing import assert_array_almost_equal
+from numpy.typing import ArrayLike
 
 from hamilflow.models.pendulum import Pendulum, PendulumIC, PendulumSystem
 
 
 @pytest.fixture(params=[0.3, 0.6, 1.5])
-def omega0(request: SubRequest):
-    yield request.param
+def omega0(request: pytest.FixtureRequest) -> float:
+    return request.param
 
 
 @pytest.fixture(params=[-1.5, -0.3, 0.2, 0.4, +1.4])
-def theta0(request):
-    yield request.param
+def theta0(request: pytest.FixtureRequest) -> float:
+    return request.param
 
 
-@pytest.fixture(params=[[-5.5, 0.0, 0.5, 1.0]])
-def times(request):
-    yield request.param
+@pytest.fixture(params=(-1.2, [-5.5, 0.0, 0.5, 1.0], np.linspace(0, 1, 7)))
+def times(request: pytest.FixtureRequest) -> float:
+    return request.param
 
 
 class TestPendulumSystem:
@@ -50,23 +51,22 @@ def test_period_static(self, omega0: float, theta0: float, period: float) -> Non
         p = Pendulum(omega0, theta0)
         assert pytest.approx(p.period) == period
 
-    def test_transf(self, omega0: float, theta0: float, times: list[float]) -> None:
+    def test_transf(
+        self, omega0: float, theta0: float, times: "Sequence[float] | ArrayLike[float]"
+    ) -> None:
         p = Pendulum(omega0, theta0)
         arr_times = np.asarray(times)
 
         sin_u = np.sin(p.u(arr_times))
         theta_terms = np.sin(p.theta(arr_times) / 2) / p._k
         assert_array_almost_equal(theta_terms, sin_u)
-        # assert_array_almost_equal_nulp(theta_terms, sin_u, 32)
 
     def test_period_dynamic_theta(
-        self, omega0: float, theta0: float, times: list[float]
+        self, omega0: float, theta0: float, times: "Sequence[float] | ArrayLike[float]"
     ) -> None:
         p = Pendulum(omega0, theta0)
-        arr_times = np.asarray(times)
-        arr_times_1 = arr_times + p.period
+        arr_times_1 = np.array(times, copy=False) + p.period
 
-        theta, theta_1 = p.theta(arr_times), p.theta(arr_times_1)
+        theta, theta_1 = p.theta(times), p.theta(arr_times_1)
 
         assert_array_almost_equal(theta, theta_1)
-        # assert_array_almost_equal_nulp(theta, theta_1, 80)