From 91b5db9f004cda717fab326048fc4915cc374ef8 Mon Sep 17 00:00:00 2001 From: cmp0xff Date: Mon, 29 Jul 2024 22:00:53 +0200 Subject: [PATCH] fix(typing): #69 np.ndarray is not a subclass of typing.Sequence (#70) fix(typing): #69 --- hamilflow/models/brownian_motion.py | 11 +++--- hamilflow/models/free_particle.py | 13 +++---- hamilflow/models/harmonic_oscillator.py | 46 ++++++++++++++----------- hamilflow/models/pendulum.py | 13 +++---- tests/test_models/test_free_particle.py | 3 +- tests/test_models/test_pendulum.py | 30 ++++++++-------- 6 files changed, 62 insertions(+), 54 deletions(-) diff --git a/hamilflow/models/brownian_motion.py b/hamilflow/models/brownian_motion.py index becbec6..9a5a6b3 100644 --- a/hamilflow/models/brownian_motion.py +++ b/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: diff --git a/hamilflow/models/free_particle.py b/hamilflow/models/free_particle.py index 8435358..bf1f21b 100644 --- a/hamilflow/models/free_particle.py +++ b/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). diff --git a/hamilflow/models/harmonic_oscillator.py b/hamilflow/models/harmonic_oscillator.py index 1ffd9ba..b272d50 100644 --- a/hamilflow/models/harmonic_oscillator.py +++ b/hamilflow/models/harmonic_oscillator.py @@ -77,15 +77,15 @@ 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(), @@ -93,7 +93,7 @@ def definition(self) -> dict[str, dict[str, float | int]]: } @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,8 +153,8 @@ 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": @@ -162,7 +162,7 @@ def __init__( 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,7 +356,7 @@ 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( @@ -364,7 +364,9 @@ def definition( 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}) diff --git a/hamilflow/models/pendulum.py b/hamilflow/models/pendulum.py index d0d9a90..b97b064 100644 --- a/hamilflow/models/pendulum.py +++ b/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 diff --git a/tests/test_models/test_free_particle.py b/tests/test_models/test_free_particle.py index 4668eb7..571bf78 100644 --- a/tests/test_models/test_free_particle.py +++ b/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), ) diff --git a/tests/test_models/test_pendulum.py b/tests/test_models/test_pendulum.py index d41ebcf..a077b65 100644 --- a/tests/test_models/test_pendulum.py +++ b/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)