feat(dynamical-system): add first order dynamical system

#103

Author: emptymalei

hamilflow/models/dynamical_systems/__init__.py

@@ -0,0 +1 @@
+"""Collections of dynamical systems."""

hamilflow/models/dynamical_systems/base/__init__.py

@@ -0,0 +1,6 @@
+"""Base classes for first order dynamical systems.
+
+In this module, we define the base classes for first order dynamical systems.
+For example, we define the base class for first order dynamical systems called
+`FirstOrderDynamicsBase`.
+"""

hamilflow/models/dynamical_systems/base/first_order_dynamics.py

@@ -0,0 +1,138 @@
+"""Base classes for first order dynamical systems.
+
+In this module, we introduce the base classes for
+first order dynamical systems, and the corresponding
+abstract classes for system and initial condition.
+"""
+
+from abc import ABC, abstractmethod
+from functools import cached_property
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from numpy import typing as npt
+from pydantic import BaseModel, Field
+
+from hamilflow.models.utils.typing import TypeTime
+
+
+class FirstOrderSystem(BaseModel, ABC):
+    """The base params for a first order dynamical system.
+
+    Create your own system by inheriting from this class
+    and adding your own params.
+
+    ```python
+    class MyCustomSystem(FirstOrderSystem):
+        # your fields here
+        omega: float
+    ```
+    """
+
+
+class FirstOrderIC(BaseModel, ABC):
+    """The base initial condition for a first order dynamical system.
+
+    :cvar x0: the initial state of the system
+    :cvar t0: the initial time (default to 0.0)
+    """
+
+    x0: float | list[float] = Field(...)
+    t0: float = Field(default=0.0)
+
+
+class FirstOrderDynamicsBase(ABC):
+    """Base class to generate time series data for a first order dynamical system.
+
+    :param system: all the params that defines the system.
+    :param initial_condition: the initial condition of the system.
+    """
+
+    def __init__(
+        self,
+        system: FirstOrderSystem,
+        initial_condition: FirstOrderIC,
+    ) -> None:
+
+        self.system = system
+        self.ic = initial_condition
+
+    @cached_property
+    def definition(self) -> dict[str, dict[str, Any]]:
+        """Model params and initial conditions defined as a dictionary.
+
+        :return: dictionary containing system and initial condition parameters.
+        """
+        return {
+            "system": self.system.model_dump(),
+            "initial_condition": self.ic.model_dump(),
+        }
+
+    @abstractmethod
+    def derivatives(
+        self,
+        state: npt.NDArray[np.float64],
+        t: float,
+    ) -> npt.NDArray[np.float64]:
+        """Return the derivative of the state with respect to time.
+
+        :param state: the current state of the system.
+        :param t: the current time.
+        :return: the derivative of the state with respect to time.
+        """
+
+    def step(
+        self,
+        state: npt.NDArray[np.float64],
+        t: float,
+        dt: float,
+    ) -> tuple[npt.NDArray[np.float64], float]:
+        """Advances the system by one time step (dt) using
+        4th-Order Runge-Kutta (RK4) integration.
+
+        :param state: the current state of the system.
+        :param t: the current time.
+        :param dt: the time step size.
+        :return: the new state of the system and the new time.
+        """
+        k1 = self.derivatives(state, t)
+        k2 = self.derivatives(state + 0.5 * dt * k1, t + 0.5 * dt)
+        k3 = self.derivatives(state + 0.5 * dt * k2, t + 0.5 * dt)
+        k4 = self.derivatives(state + dt * k3, t + dt)
+
+        new_state = state + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
+        new_t = t + dt
+
+        return new_state, new_t
+
+    def __call__(self, t: TypeTime) -> pd.DataFrame:
+        """Generate time series data for the dynamical system.
+
+        :param t: sequence of time steps.
+        :return: values of the variables including time `t` and state `x`.
+        """
+        t_arr = np.asarray(t)
+
+        current_state = np.asarray(self.ic.x0, dtype=np.float64)
+        current_t = self.ic.t0
+
+        effective_t = [current_t]
+        history = [current_state]
+        for target_t in t_arr:
+            dt = target_t - current_t
+            if dt > 0:
+                current_state, current_t = self.step(current_state, current_t, dt)
+                effective_t.append(current_t)
+                history.append(current_state)
+
+        data = np.asarray(history)
+
+        if data.ndim == 1:
+            columns = ["x"]
+        else:
+            columns = [f"x{i+1}" for i in range(data.shape[1])]
+
+        return (
+            pd.DataFrame(data, columns=columns).assign(t=effective_t).sort_index(axis=1)
+        )

pyproject.toml

@@ -66,13 +66,16 @@ target-version = "py310"
 select = ["ALL"]
 ignore = [
     "D107",  # undocumented-public-init: we document the class instead
+    "D203",  # one-blank-line-before-class: we use D211 instead
+    "D213",  # multi-line-summary-second-line: we use D212 instead
     "E501",  # line-too-long: we use black instead
     "TID252",  # relative-imports
     "PD901",  # pandas-df-variable-name
     "FBT",
     "FIX",
     "TD",
     "RET",
+    "D205",
 ]
 
 [tool.ruff.lint.per-file-ignores]
@@ -84,6 +87,15 @@ ignore = [
 "tests/**/*.py" = [
     "S101",  # assert: Fine in tests
     "SLF001",  # private-member-access: find in tests
+    "D104",  # missing-module-docstring
+    "INP001",  # implicit-namespace-package: tutorials are not packages
+    "D101",  # missing-class-docstring
+    "D102",  # missing-method-docstring
+    "D103",  # missing-function-docstring
+    "D100",  # missing-module-docstring
+    "ARG002",  # unused-method-argument
+    "ANN201",  # missing-return-type-annotation
+    "PLR2004",  # magic-value
 ]
 
 [tool.ruff.lint.isort]

tests/models/test_dynamical_systems/base/__init__.py

@@ -0,0 +1,96 @@
+import numpy as np
+from numpy import typing as npt
+from pydantic import Field
+
+from hamilflow.models.dynamical_systems.base.first_order_dynamics import (
+    FirstOrderDynamicsBase,
+    FirstOrderIC,
+    FirstOrderSystem,
+)
+
+
+class ExponentialDecaySystem(FirstOrderSystem):
+    k: float = Field(default=1.0)
+
+
+class ExpontialDecayIC(FirstOrderIC):
+    x0: float | list[float] = Field(...)
+    t0: float = Field(default=0.0)
+
+
+class ExponentialDecayModel(FirstOrderDynamicsBase):
+    def derivatives(
+        self,
+        state: npt.NDArray[np.float64],
+        t: float,
+    ) -> npt.NDArray[np.float64]:
+        # dy/dt = -k * y
+        return -self.system.k * state  # type: ignore[attr-defined]
+
+
+def test_first_order_dynamics_instantiation():
+    model = ExponentialDecayModel(
+        system=ExponentialDecaySystem(k=2.0),
+        initial_condition=ExpontialDecayIC(x0=5.0, t0=0.0),
+    )
+    assert model.system.k == 2.0
+    assert model.ic.x0 == 5.0
+    assert model.ic.t0 == 0.0
+
+    definition = model.definition
+    assert definition["system"]["k"] == 2.0
+    assert definition["initial_condition"]["x0"] == 5.0
+
+
+def test_first_order_dynamics_step():
+    model = ExponentialDecayModel(
+        system=ExponentialDecaySystem(k=1.0),
+        initial_condition=ExpontialDecayIC(x0=1.0),
+    )
+    state = np.array([1.0], dtype=np.float64)
+    dt = 0.1
+    t = 0.0
+
+    new_state, new_t = model.step(state, t, dt)
+    assert new_t == 0.1
+
+    # For dx/dt = -x, RK4 step starting at x=1.0 with dt=0.1
+    # RK4 match to 4th order Taylor expansion
+    expected = 1.0 - 0.1 + (0.1**2) / 2.0 - (0.1**3) / 6.0 + (0.1**4) / 24.0
+    np.testing.assert_allclose(new_state, expected, rtol=1e-5)
+
+
+def test_first_order_dynamics_call_scalar_x0():
+    model = ExponentialDecayModel(
+        system=ExponentialDecaySystem(k=1.0),
+        initial_condition=ExpontialDecayIC(x0=3.0, t0=0.0),
+    )
+
+    t_arr = np.linspace(0.0, 1.0, 101)
+    df = model(t_arr)
+
+    assert list(df.columns) == ["t", "x"]
+    np.testing.assert_allclose(df["t"].values, t_arr)
+
+    # Check values against analytical: x = 3 * e^{-t}
+    expected_x = 3.0 * np.exp(-1.0 * t_arr)
+    np.testing.assert_allclose(df["x"].values, expected_x, rtol=1e-4)
+
+
+def test_first_order_dynamics_call_array_x0():
+    model = ExponentialDecayModel(
+        system=ExponentialDecaySystem(k=2.0),
+        initial_condition=ExpontialDecayIC(x0=[2.0, 5.0], t0=0.0),
+    )
+
+    t_arr = np.linspace(0.0, 0.5, 101)
+    df = model(t_arr)
+
+    assert list(df.columns) == ["t", "x1", "x2"]
+
+    # Check values against analytical: x = x0 * e^{-2t}
+    expected_x1 = 2.0 * np.exp(-2.0 * t_arr)
+    expected_x2 = 5.0 * np.exp(-2.0 * t_arr)
+
+    np.testing.assert_allclose(df["x1"].values, expected_x1, rtol=1e-4)
+    np.testing.assert_allclose(df["x2"].values, expected_x2, rtol=1e-4)