From 08e0e08e02f3232ce7a3606d1fe0d1fc73bd55f9 Mon Sep 17 00:00:00 2001 From: LM Date: Tue, 16 Apr 2024 21:04:01 +0200 Subject: [PATCH] refactor(harmonic-oscillator): separate the class of simple and damped 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 | 29 ++- hamilflow/models/harmonic_oscillator.py | 195 ++++++++++++------ tests/test_models/test_harmonic_oscillator.py | 37 ++-- 3 files changed, 170 insertions(+), 91 deletions(-) diff --git a/docs/tutorials/harmonic_oscillator.py b/docs/tutorials/harmonic_oscillator.py index 36493ea..9bf2db7 100644 --- a/docs/tutorials/harmonic_oscillator.py +++ b/docs/tutorials/harmonic_oscillator.py @@ -22,7 +22,10 @@ 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 @@ -30,11 +33,23 @@ # %% [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')})") ) diff --git a/hamilflow/models/harmonic_oscillator.py b/hamilflow/models/harmonic_oscillator.py index 9997cf5..072d869 100644 --- a/hamilflow/models/harmonic_oscillator.py +++ b/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 diff --git a/tests/test_models/test_harmonic_oscillator.py b/tests/test_models/test_harmonic_oscillator.py index da5b2c5..24dd291 100644 --- a/tests/test_models/test_harmonic_oscillator.py +++ b/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)