feat(model): #59 make a class for linear motion (#60)

* feat: #59 free particle

* chore(poetry): #59 🔒

hamilflow/models/discrete/__init__.py

@@ -0,0 +1 @@
+"System with discrete degrees of freedom, e.g. particles."

hamilflow/models/discrete/d0/__init__.py

@@ -0,0 +1 @@
+"One-particle systems."

hamilflow/models/discrete/d0/free_particle.py

@@ -0,0 +1,62 @@
+from functools import cached_property
+from typing import Mapping, Sequence, cast
+
+import numpy as np
+import pandas as pd
+from pydantic import BaseModel, Field, model_validator
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+
+class FreeParticleIC(BaseModel):
+    """The initial condition for a free particle
+
+    :cvar x0: the initial displacement
+    :cvar v0: the initial velocity
+    """
+
+    x0: float | int | Sequence[float | int] = Field()
+    v0: float | int | Sequence[float | int] = Field()
+
+    @model_validator(mode="after")
+    def check_dimensions_match(self) -> Self:
+        assert (
+            len(self.x0) == len(cast(Sequence, self.v0))
+            if isinstance(self.x0, Sequence)
+            else not isinstance(self.v0, Sequence)
+        )
+        return self
+
+
+class FreeParticle:
+    r"""Base class to generate time series data for a free particle.
+
+    :param initial_condition: the initial condition of the free particle.
+    """
+
+    def __init__(
+        self, initial_condition: Mapping[str, float | int | Sequence[float | int]]
+    ) -> None:
+        self.initial_condition = FreeParticleIC.model_validate(initial_condition)
+
+    @cached_property
+    def definition(self) -> dict[str, dict[str, int | float | Sequence[int | 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:
+        return np.outer(t, self.initial_condition.v0) + self.initial_condition.x0
+
+    def __call__(self, t: float | int | Sequence[float | int]) -> 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]))
+
+        return pd.DataFrame(data, columns=columns).assign(t=t).sort_index(axis=1)

pyproject.toml

@@ -19,7 +19,7 @@ scipy = "^1.12.0"
 
 
 [tool.poetry.group.test.dependencies]
-pytest = "^8.1.1"
+pytest = "^8.2.1"
 pytest-cov = "^4.1.0"
 
 

tests/test_models/discrete/d0/test_free_particle.py

@@ -0,0 +1,61 @@
+from typing import Mapping, Sequence
+
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+from pydantic import ValidationError
+
+from hamilflow.models.discrete.d0.free_particle import FreeParticle, FreeParticleIC
+
+
+class TestFreeParticleIC:
+    @pytest.mark.parametrize(("x0", "v0"), [(1, 2), ((1,), (2,)), ((1, 2), (2, 3))])
+    def test_constructor(
+        self, x0: int | Sequence[int], v0: int | Sequence[int]
+    ) -> None:
+        assert FreeParticleIC(x0=x0, v0=v0)
+
+    @pytest.mark.parametrize(("x0", "v0"), [(1, (2,)), ((1,), (2, 3))])
+    def test_raise(self, x0: int | Sequence[int], v0: int | Sequence[int]) -> None:
+        with pytest.raises(ValidationError):
+            FreeParticleIC(x0=x0, v0=v0)
+
+
+class TestFreeParticle:
+    @pytest.mark.parametrize(
+        ("x0", "v0", "expected"),
+        [
+            (1, 2, dict(initial_condition=dict(x0=1, v0=2))),
+            ((1,), (2,), dict(initial_condition=dict(x0=(1,), v0=(2,)))),
+        ],
+    )
+    def test_definition(
+        self,
+        x0: int | Sequence[int],
+        v0: int | Sequence[int],
+        expected: Mapping[str, Mapping[str, int | Sequence[int]]],
+    ) -> None:
+        assert FreeParticle(initial_condition=dict(x0=x0, v0=v0)).definition == expected
+
+    @pytest.mark.parametrize(
+        ("x0", "v0", "t", "expected"),
+        [
+            (1, 2, (3,), pd.DataFrame(dict(t=[3], x1=[7]))),
+            (
+                (1, 2),
+                (2, 3),
+                (3, 4),
+                pd.DataFrame(dict(t=(3, 4), x1=(7, 9), x2=(11, 14))),
+            ),
+        ],
+    )
+    def test_call(
+        self,
+        x0: int | Sequence[int],
+        v0: int | Sequence[int],
+        t: int | Sequence[int],
+        expected: pd.DataFrame,
+    ) -> None:
+        assert_frame_equal(
+            FreeParticle(initial_condition=dict(x0=x0, v0=v0))(t), expected
+        )