refactor(harmonic-oscillator): separate the class of simple and dampe…

…d HO (#48)

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

#34

* docs(harmonic-oscillator): reduce the docstrings of harmonic oscillators

* fix(harmonic-oscillator): remove unused code in harmonic oscillator

#34

* refactor(harmonic-oscillator): change type hint for type float or ndarray

#34

* fix(harmonic-oscillator): use ArrayLike instead of union of float and ndarray

#34

* fix(harmonic-oscillator): add initial phase to harmonic oscillator (simple)

#34

* docs(harmonic-oscillator): adjust tutorials to include ations or equation of motion

#49

docs/tutorials/harmonic_oscillator.py

@@ -22,19 +22,34 @@
 import pandas as pd
 import plotly.express as px
 
-from hamilflow.models.harmonic_oscillator import HarmonicOscillator
+from hamilflow.models.harmonic_oscillator import (
+    DampedHarmonicOscillator,
+    SimpleHarmonicOscillator,
+)
 
 # %%
 n_periods = 3
 n_samples_per_period = 200
 
 # %% [markdown]
 # ## Simple Harmonic Oscillator
+#
+# For an simple harmonic oscillator, the action of a simple harmonic oscillator is
+#
+# $$S_L[x] = \int_{t_0}^{t_1} \mathbb{d}t \left\{\frac{1}{2} m \dot x^2 - \frac{1}{2} m \omega^2 x^2 \right\}\,,$$
+#
+# where the least action principle leads to the following equation of motion,
+#
+# $$
+# \ddot x + \omega^2 x = 0\,.
+# $$
+#
+# A simple harmonic oscillator is a periodic motion.
 
 # %%
 sho_omega = 0.5
 
-sho = HarmonicOscillator(system={"omega": sho_omega})
+sho = SimpleHarmonicOscillator(system={"omega": sho_omega})
 
 # %%
 df_sho = sho(n_periods=n_periods, n_samples_per_period=n_samples_per_period)
@@ -54,6 +69,14 @@
 
 # %% [markdown]
 # ## Damped Harmonic Oscillator
+#
+# A damped harmonic oscillator is a simple harmonic oscillator with damping force that is proportional to its velocity,
+#
+# $$
+# \ddot x + \omega^2 x = - 2\xi\omega \dot x\,.
+# $$
+#
+# In this section, we demonstrate three scenarios of a damped harmonic oscillator.
 
 # %%
 dho_systems = {
@@ -70,7 +93,7 @@
 for s_name, s in dho_systems.items():
 
     dfs_dho.append(
-        HarmonicOscillator(system=s)(
+        DampedHarmonicOscillator(system=s)(
             n_periods=n_periods, n_samples_per_period=n_samples_per_period
         ).assign(system=rf"{s_name} (omega = {s.get('omega')}, zeta = {s.get('zeta')})")
     )

hamilflow/models/harmonic_oscillator.py

@@ -1,8 +1,10 @@
+from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import Dict, Literal, Optional, Union
 
 import numpy as np
 import pandas as pd
+from numpy.typing import ArrayLike
 from pydantic import BaseModel, computed_field, field_validator
 
 
@@ -57,14 +59,122 @@ class HarmonicOscillatorIC(BaseModel):
 
     :cvar x0: the initial displacement
     :cvar v0: the initial velocity
+    :cvar phi: initial phase
     """
 
     x0: float = 1.0
     v0: float = 0.0
+    phi: float = 0.0
 
 
-class HarmonicOscillator:
-    r"""Generate time series data for a [harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
+class HarmonicOscillatorBase(ABC):
+    r"""Base class to generate time series data
+    for a [harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
+
+    :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]] = {},
+    ):
+        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: ArrayLike) -> ArrayLike:
+        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).
+
+
+    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 = SimpleHarmonicOscillator(params=params)
+
+    df = ho(n_periods=1, n_samples_per_period=10)
+    ```
+
+    `df` will be a pandas dataframe with two columns: `t` and `x`.
+    """
+
+    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: ArrayLike) -> ArrayLike:
+        r"""Solution to simple harmonic oscillators:
+
+        $$
+        x(t) = x_0 \cos(\omega t + \phi).
+        $$
+        """
+        return self.initial_condition.x0 * np.cos(
+            self.system.omega * t + self.initial_condition.phi
+        )
+
+
+class DampedHarmonicOscillator(HarmonicOscillatorBase):
+    r"""Generate time series data for a [damped harmonic oscillator](https://en.wikipedia.org/wiki/Harmonic_oscillator).
 
     The equation for a general un-driven harmonic oscillator is[^wiki_ho][^libretext_ho]
 
@@ -96,35 +206,12 @@ class HarmonicOscillator:
     \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}
+    params = {"omega": omega, "zeta"=0.2}
 
-    ho = HarmonicOscillator(params=params)
+    ho = DampedHarmonicOscillator(params=params)
 
     df = ho(n_periods=1, n_samples_per_period=10)
     ```
@@ -140,25 +227,12 @@ def __init__(
         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(),
-        }
-
-    def _x_simple(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)
+        super().__init__(system, initial_condition)
+        if self.system.type == "simple":
+            raise ValueError(
+                f"System is not a Damped Harmonic Oscillator: {self.system}\n"
+                f"This is a simple harmonic oscillator, use `SimpleHarmonicOscillator`."
+            )
 
     def _x_under_damped(self, t: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
         r"""Solution to under damped harmonic oscillators:
@@ -231,26 +305,17 @@ def _x_over_damped(self, t: Union[float, np.ndarray]) -> Union[float, np.ndarray
             * np.sinh(gamma_damp * t)
         ) * np.exp(-self.system.zeta * self.system.omega * t)
 
-    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
-
-        if self.system.type == "simple":
-            data = self._x_simple(time_steps)
-        elif self.system.type == "under_damped":
-            data = self._x_under_damped(time_steps)
+    def _x(self, t: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
+        r"""Solution to damped harmonic oscillators."""
+        if self.system.type == "under_damped":
+            x = self._x_under_damped(t)
         elif self.system.type == "over_damped":
-            data = self._x_over_damped(time_steps)
+            x = self._x_over_damped(t)
         elif self.system.type == "critical_damped":
-            data = self._x_critical_damped(time_steps)
+            x = self._x_critical_damped(t)
         else:
-            raise ValueError(f"system type is not defined: {self.system.type}")
+            raise ValueError(
+                "System type is not damped harmonic oscillator: {self.system.type}"
+            )
 
-        return pd.DataFrame({"t": time_steps, "x": data})
+        return x

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)