refactor(harmonic-oscillator): sparate the class of simple and damped…

… oscillator

#34

hamilflow/models/harmonic_oscillator.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import Dict, Literal, Optional, Union
 
@@ -63,6 +64,300 @@ class HarmonicOscillatorIC(BaseModel):
     v0: float = 0.0
 
 
+class HarmonicOscillatorBase(ABC):
+    r"""Base class to generate time series data
+    for a [harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
+    """
+
+    def __init__(
+        self,
+        system: Dict[str, float],
+        initial_condition: Optional[Dict[str, float]] = {},
+    ):
+        self.system = HarmonicOscillatorSystem.model_validate(system)
+        self.initial_condition = HarmonicOscillatorIC.model_validate(initial_condition)
+
+    @cached_property
+    def definition(self) -> 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: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        r"""Solution to simple harmonic oscillators."""
+        ...
+
+    def __call__(self, n_periods: int, n_samples_per_period: int) -> pd.DataFrame:
+        """Generate time series data for the harmonic oscillator.
+
+        Returns a list of floats representing the displacement at each time step.
+
+        :param n_periods: Number of periods to generate.
+        :param n_samples_per_period: Number of samples per period.
+        """
+        time_delta = self.system.period / n_samples_per_period
+        time_steps = np.arange(0, n_periods * n_samples_per_period) * time_delta
+
+        data = self._x(time_steps)
+
+        return pd.DataFrame({"t": time_steps, "x": data})
+
+
+class SimpleHarmonicOscillator(HarmonicOscillatorBase):
+    r"""Generate time series data for a [simple harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
+
+    The equation for a general un-driven harmonic oscillator is[^wiki_ho][^libretext_ho]
+
+    $$
+    \frac{\mathrm d x^2}{\mathrm d t^2} + 2\zeta \omega \frac{\mathrm d x}{\mathrm dt} + \omega^2 x = 0,
+    $$
+
+    where $x$ is the displacement, $\omega$ is the angular frequency of an undamped oscillator ($\zeta=0$),
+    and $\zeta$ is the damping ratio.
+
+    [^wiki_ho]: Contributors to Wikimedia projects. Harmonic oscillator.
+                In: Wikipedia [Internet]. 18 Feb 2024 [cited 20 Feb 2024].
+                Available: https://en.wikipedia.org/wiki/Harmonic_oscillator#Damped_harmonic_oscillator
+
+    [^libretext_ho]: Libretexts. 5.3: General Solution for the Damped Harmonic Oscillator. Libretexts. 13 Apr 2021.
+                    Available: https://t.ly/cWTIo. Accessed 20 Feb 2024.
+
+
+    The solution to the above harmonic oscillator is
+
+    $$
+    x(t) = \left( x_0 \cos(\Omega t) + \frac{\zeta \omega x_0 + v_0}{\Omega} \sin(\Omega t) \right)
+        e^{-\zeta \omega t},
+    $$
+
+    where
+
+    $$
+    \Omega = \omega\sqrt{ 1 - \zeta^2}.
+    $$
+
+
+    !!! example "A Simple Harmonic Oscillator ($\zeta=0$)"
+
+        In a one dimensional world, a mass $m$, driven by a force $F=-kx$, is described as
+
+        $$
+        \begin{align}
+        F &= - k x \\
+        F &= m a
+        \end{align}
+        $$
+
+        The mass behaves like a simple harmonic oscillator.
+
+        In general, the solution to a simple harmonic oscillator is
+
+        $$
+        x(t) = A \cos(\omega t + \phi),
+        $$
+
+        where $\omega$ is the angular frequency, $\phi$ is the initial phase, and $A$ is the amplitude.
+
+
+    To use this generator,
+
+    ```python
+    params = {"omega": omega}
+
+    ho = HarmonicOscillator(params=params)
+
+    df = ho(n_periods=1, n_samples_per_period=10)
+    ```
+
+    `df` will be a pandas dataframe with two columns: `t` and `x`.
+
+    :param system: all the params that defines the harmonic oscillator.
+    :param initial_condition: the initial condition of the harmonic oscillator.
+    """
+
+    def __init__(
+        self,
+        system: Dict[str, float],
+        initial_condition: Optional[Dict[str, float]] = {},
+    ):
+        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: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        r"""Solution to simple harmonic oscillators:
+
+        $$
+        x(t) = x_0 \cos(\omega t).
+        $$
+        """
+        return self.initial_condition.x0 * np.cos(self.system.omega * t)
+
+
+class DampedHarmonicOscillator(HarmonicOscillatorBase):
+    r"""Generate time series data for a [simple harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
+
+    The equation for a general un-driven harmonic oscillator is[^wiki_ho][^libretext_ho]
+
+    $$
+    \frac{\mathrm d x^2}{\mathrm d t^2} + 2\zeta \omega \frac{\mathrm d x}{\mathrm dt} + \omega^2 x = 0,
+    $$
+
+    where $x$ is the displacement, $\omega$ is the angular frequency of an undamped oscillator ($\zeta=0$),
+    and $\zeta$ is the damping ratio.
+
+    [^wiki_ho]: Contributors to Wikimedia projects. Harmonic oscillator.
+                In: Wikipedia [Internet]. 18 Feb 2024 [cited 20 Feb 2024].
+                Available: https://en.wikipedia.org/wiki/Harmonic_oscillator#Damped_harmonic_oscillator
+
+    [^libretext_ho]: Libretexts. 5.3: General Solution for the Damped Harmonic Oscillator. Libretexts. 13 Apr 2021.
+                    Available: https://t.ly/cWTIo. Accessed 20 Feb 2024.
+
+
+    The solution to the above harmonic oscillator is
+
+    $$
+    x(t) = \left( x_0 \cos(\Omega t) + \frac{\zeta \omega x_0 + v_0}{\Omega} \sin(\Omega t) \right)
+        e^{-\zeta \omega t},
+    $$
+
+    where
+
+    $$
+    \Omega = \omega\sqrt{ 1 - \zeta^2}.
+    $$
+
+
+    !!! example "A Simple Harmonic Oscillator ($\zeta=0$)"
+
+        In a one dimensional world, a mass $m$, driven by a force $F=-kx$, is described as
+
+        $$
+        \begin{align}
+        F &= - k x \\
+        F &= m a
+        \end{align}
+        $$
+
+        The mass behaves like a simple harmonic oscillator.
+
+        In general, the solution to a simple harmonic oscillator is
+
+        $$
+        x(t) = A \cos(\omega t + \phi),
+        $$
+
+        where $\omega$ is the angular frequency, $\phi$ is the initial phase, and $A$ is the amplitude.
+
+
+    To use this generator,
+
+    ```python
+    params = {"omega": omega}
+
+    ho = HarmonicOscillator(params=params)
+
+    df = ho(n_periods=1, n_samples_per_period=10)
+    ```
+
+    `df` will be a pandas dataframe with two columns: `t` and `x`.
+
+    :param system: all the params that defines the harmonic oscillator.
+    :param initial_condition: the initial condition of the harmonic oscillator.
+    """
+
+    def _x_under_damped(self, t: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        r"""Solution to under damped harmonic oscillators:
+
+        $$
+        x(t) = \left( x_0 \cos(\Omega t) + \frac{\zeta \omega x_0 + v_0}{\Omega} \sin(\Omega t) \right)
+        e^{-\zeta \omega t},
+        $$
+
+        where
+
+        $$
+        \Omega = \omega\sqrt{ 1 - \zeta^2}.
+        $$
+        """
+        omega_damp = self.system.omega * np.sqrt(1 - self.system.zeta)
+        return (
+            self.initial_condition.x0 * np.cos(omega_damp * t)
+            + (
+                self.system.zeta * self.system.omega * self.initial_condition.x0
+                + self.initial_condition.v0
+            )
+            / omega_damp
+            * np.sin(omega_damp * t)
+        ) * np.exp(-self.system.zeta * self.system.omega * t)
+
+    def _x_critical_damped(
+        self, t: Union[float, np.ndarray]
+    ) -> Union[float, np.ndarray]:
+        r"""Solution to critical damped harmonic oscillators:
+
+        $$
+        x(t) = \left( x_0 \cos(\Omega t) + \frac{\zeta \omega x_0 + v_0}{\Omega} \sin(\Omega t) \right)
+        e^{-\zeta \omega t},
+        $$
+
+        where
+
+        $$
+        \Omega = \omega\sqrt{ 1 - \zeta^2}.
+        $$
+        """
+        return self.initial_condition.x0 * np.exp(
+            -self.system.zeta * self.system.omega * t
+        )
+
+    def _x_over_damped(self, t: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        r"""Solution to over harmonic oscillators:
+
+        $$
+        x(t) = \left( x_0 \cosh(\Gamma t) + \frac{\zeta \omega x_0 + v_0}{\Gamma} \sinh(\Gamma t) \right)
+        e^{-\zeta \omega t},
+        $$
+
+        where
+
+        $$
+        \Gamma = \omega\sqrt{ \zeta^2 - 1 }.
+        $$
+        """
+        gamma_damp = self.system.omega * np.sqrt(self.system.zeta - 1)
+
+        return (
+            self.initial_condition.x0 * np.cosh(gamma_damp * t)
+            + (
+                self.system.zeta * self.system.omega * self.initial_condition.x0
+                + self.initial_condition.v0
+            )
+            / gamma_damp
+            * np.sinh(gamma_damp * t)
+        ) * np.exp(-self.system.zeta * self.system.omega * t)
+
+    def _x(self, t: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        r"""Solution to simple harmonic oscillators:"""
+        if self.system.type == "under_damped":
+            x = self._x_under_damped(t)
+        elif self.system.type == "over_damped":
+            x = self._x_over_damped(t)
+        elif self.system.type == "critical_damped":
+            x = self._x_critical_damped(t)
+        else:
+            raise ValueError(
+                "System type is not damped harmonic oscillator: {self.system.type}"
+            )
+
+        return x
+
+
 class HarmonicOscillator:
     r"""Generate time series data for a [harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
 

tests/test_models/test_harmonic_oscillator.py

@@ -2,18 +2,25 @@
 import pytest
 
 from hamilflow.models.harmonic_oscillator import (
-    HarmonicOscillator,
+    DampedHarmonicOscillator,
     HarmonicOscillatorSystem,
+    SimpleHarmonicOscillator,
 )
 
 
 @pytest.mark.parametrize("zeta", [(-0.5), (-2.0)])
-def test_harmonic_oscillator_system_zeta(zeta):
+def test_harmonic_oscillator_system_damping_zeta(zeta):
     with pytest.raises(ValueError):
         HarmonicOscillatorSystem(omega=1, zeta=zeta)
 
     with pytest.raises(ValueError):
-        HarmonicOscillator(system={"omega": 1, "zeta": zeta})
+        SimpleHarmonicOscillator(system={"omega": 1, "zeta": zeta})
+
+
+@pytest.mark.parametrize("zeta", [(0.5), (1.0)])
+def test_simple_harmonic_oscillator_instantiation(zeta):
+    with pytest.raises(ValueError):
+        SimpleHarmonicOscillator(system={"omega": 1, "zeta": zeta})
 
 
 @pytest.mark.parametrize(
@@ -37,7 +44,7 @@ def test_harmonic_oscillator_system_zeta(zeta):
     ],
 )
 def test_simple_harmonic_oscillator(omega, expected):
-    ho = HarmonicOscillator(system={"omega": omega})
+    ho = SimpleHarmonicOscillator(system={"omega": omega})
 
     df = ho(n_periods=1, n_samples_per_period=10)
 
@@ -47,22 +54,6 @@ def test_simple_harmonic_oscillator(omega, expected):
 @pytest.mark.parametrize(
     "omega,zeta,expected",
     [
-        (
-            0.5,
-            0,
-            [
-                {"t": 0.0, "x": 1.0},
-                {"t": 1.2566370614359172, "x": 0.8090169943749475},
-                {"t": 2.5132741228718345, "x": 0.30901699437494745},
-                {"t": 3.7699111843077517, "x": -0.30901699437494734},
-                {"t": 5.026548245743669, "x": -0.8090169943749473},
-                {"t": 6.283185307179586, "x": -1.0},
-                {"t": 7.5398223686155035, "x": -0.8090169943749475},
-                {"t": 8.79645943005142, "x": -0.30901699437494756},
-                {"t": 10.053096491487338, "x": 0.30901699437494723},
-                {"t": 11.309733552923255, "x": 0.8090169943749473},
-            ],
-        ),
         (
             0.5,
             0.5,
@@ -82,7 +73,7 @@ def test_simple_harmonic_oscillator(omega, expected):
     ],
 )
 def test_underdamped_harmonic_oscillator(omega, zeta, expected):
-    ho = HarmonicOscillator(system={"omega": omega, "zeta": zeta})
+    ho = DampedHarmonicOscillator(system={"omega": omega, "zeta": zeta})
 
     df = ho(n_periods=1, n_samples_per_period=10)
 
@@ -111,7 +102,7 @@ def test_underdamped_harmonic_oscillator(omega, zeta, expected):
     ],
 )
 def test_overdamped_harmonic_oscillator(omega, zeta, expected):
-    ho = HarmonicOscillator(system={"omega": omega, "zeta": zeta})
+    ho = DampedHarmonicOscillator(system={"omega": omega, "zeta": zeta})
 
     df = ho(n_periods=1, n_samples_per_period=10)
 
@@ -140,7 +131,7 @@ def test_overdamped_harmonic_oscillator(omega, zeta, expected):
     ],
 )
 def test_criticaldamped_harmonic_oscillator(omega, zeta, expected):
-    ho = HarmonicOscillator(system={"omega": omega, "zeta": zeta})
+    ho = DampedHarmonicOscillator(system={"omega": omega, "zeta": zeta})
 
     df = ho(n_periods=1, n_samples_per_period=10)