diff --git a/docs/references/models/harmonic_oscillator_chain.md b/docs/references/models/harmonic_oscillator_chain.md new file mode 100644 index 0000000..57535b1 --- /dev/null +++ b/docs/references/models/harmonic_oscillator_chain.md @@ -0,0 +1,3 @@ +# Harmonic Oscillator Chain + +::: hamilflow.models.harmonic_oscillator_chain diff --git a/docs/tutorials/harmonic_oscillator_chain.py b/docs/tutorials/harmonic_oscillator_chain.py new file mode 100644 index 0000000..87d794f --- /dev/null +++ b/docs/tutorials/harmonic_oscillator_chain.py @@ -0,0 +1,136 @@ +# --- +# jupyter: +# jupytext: +# formats: ipynb,py:percent +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.16.4 +# kernelspec: +# display_name: .venv +# language: python +# name: python3 +# --- + +# %% [markdown] +# # Harmonic oscillator circle +# +# In this tutorial, we show a few interesting examples of the harmonic oscillator circle. + +# %% +import math + +import numpy as np +from plotly import express as px + +from hamilflow.models.harmonic_oscillator_chain import HarmonicOscillatorsChain + +# %% [markdown] +# ## Basic setups + +# %% +t = np.linspace(0, 3, 257) +omega = 2 * math.pi +pattern_x = r"x\d+" + + +# %% [markdown] +# ## Fundamental right-moving wave + +# %% +ics = [dict(x0=0, v0=0), dict(amp=(1, 0))] + [dict(amp=(0, 0))] * 5 +hoc = HarmonicOscillatorsChain(omega, ics, True) + +df_res = hoc(t) +df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)] +px.imshow(df_x_wide, origin="lower", labels={"y": "t"}) + +# %% [markdown] +# ## Fundamental left-moving wave + +# %% +ics = [dict(x0=0, v0=0), dict(amp=(0, 1))] + [dict(amp=(0, 0))] * 5 +hoc = HarmonicOscillatorsChain(2 * math.pi, ics, True) + +df_res = hoc(t) +df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)] +px.imshow(df_x_wide, origin="lower", labels={"y": "t"}) + +# %% [markdown] +# ## Faster right-moving wave +# Also known as the first harmonic. + +# %% +ics = [dict(x0=0, v0=0), dict(amp=(0, 0)), dict(amp=(1, 0))] + [dict(amp=(0, 0))] * 4 +hoc = HarmonicOscillatorsChain(omega, ics, True) + +df_res = hoc(t) +df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)] +px.imshow(df_x_wide, origin="lower", labels={"y": "t"}) + +# %% [markdown] +# ## Fundamental stationary wave, odd dof + +# %% +ics = [dict(x0=0, v0=0), dict(amp=(1, 1))] + [dict(amp=(0, 0))] * 5 +hoc = HarmonicOscillatorsChain(omega, ics, True) + +df_res = hoc(t) +df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)] +px.imshow(df_x_wide, origin="lower", labels={"y": "t"}) + + +# %% [markdown] +# ## Fundamental stationary wave, even dof +# There are stationary nodes at $i = 3, 9$. + +# %% +ics = [dict(x0=0, v0=0), dict(amp=(1, 1))] + [dict(amp=(0, 0))] * 5 +hoc = HarmonicOscillatorsChain(omega, ics, False) + +df_res = hoc(t) +df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)] +px.imshow(df_x_wide, origin="lower", labels={"y": "t"}) + + +# %% [markdown] +# ## Linearly moving chain and fundamental right-moving wave + +# %% +ics = [dict(x0=0, v0=2), dict(amp=(1, 0))] + [dict(amp=(0, 0))] * 5 +hoc = HarmonicOscillatorsChain(omega, ics, False) + +df_res = hoc(t) +df_x_wide = df_res.loc[:, df_res.columns.str.match(pattern_x)] +px.imshow(df_x_wide, origin="lower", labels={"y": "t"}) + +# %% [markdown] +# ## Mathematical-physical discription +# A one-dimensional circle of $N$ interacting harmonic oscillators can be described by the Lagrangian action +# $$S_L[x_i] = \int_{t_0}^{t_1}\mathbb{d} t \left\\{ \sum_{i=0}^{N-1} \frac{1}{2}m \dot x_i^2 - \frac{1}{2}m\omega^2\left(x_i - x_{i+1}\right)^2 \right\\}\\,,$$ +# where $x_N \coloneqq x_0$. +# +# This system can be solved in terms of _travelling waves_, obtained by discrete Fourier transform. +# + +# We can complexify the system +# $$S_L[x_i] = S_L[x_i, \phi_j] \equiv S_L[X^\ast_i, X_j] = \int_{t_0}^{t_1}\mathbb{d} t \left\\{ \frac{1}{2}m \dot X^\ast_i \delta_{ij} \dot X_j - \frac{1}{2}m X^\ast_i A_{ij} X_j\right\\}\\,,$$ +# where $A_{ij} / \omega^2$ is equal to $(-2)$ if $i=j$, $1$ if $|i-j|=1$ or $|i-j|=N$, and $0$ otherwise; +# $X_i \coloneqq x_i \mathbb{e}^{-\phi_i}$, $X^\ast_i \coloneqq x_i \mathbb{e}^{+\phi_i}$. + +# $A_{ij}$ can be diagonalised by the inverse discrete Fourier transform +# $$X_i = (F^{-1})_{ik} Y_k = \frac{1}{\sqrt{N}}\sum_k \mathbb{e}^{i \frac{2\mathbb{\pi}}{N} k\mathbb{i}} Y_k\\,.$$ + +# Calculating gives +# $$S_L[X^\ast_i, X_j] = S_L[Y^\ast_i, Y_j] = \sum_{k=0}^{N-1} \int_{t_0}^{t_1}\mathbb{d} t \left\\{ \frac{1}{2}m \dot Y^\ast_k \dot Y_k - \frac{1}{2}m \omega^2\cdot4\sin^2\frac{2\mathbb{\pi}k}{N} Y^\ast_k Y_k\right\\}\\,.$$ +# We can arrive at an action for the Fourier modes +# $$S_L[Y, Y^*] = \sum_{k=0}^{N-1} \int_{t_0}^{t_1}\mathbb{d} t \left\\{ \frac{1}{2}m \dot Y_k^* Y_k - \frac{1}{2}m \omega^2\cdot4\sin^2\frac{2\mathbb{\pi}k}{N} Y_k^* Y_k\right\\}\\,.$$ + +# The origional system can then be solved by $N$ complex oscillators. + +# Since the original degrees of freedom are real, the initial conditions of the propagating waves need to satisfy +# $Y_k = Y^*_{-k \mod N}$, see [Wikipedia](https://en.wikipedia.org/wiki/Discrete_Fourier_transform#DFT_of_real_and_purely_imaginary_signals). + +# %% [markdown] +# # End of Notebook diff --git a/hamilflow/models/discrete/__init__.py b/hamilflow/models/discrete/__init__.py deleted file mode 100644 index f5a53f4..0000000 --- a/hamilflow/models/discrete/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"System with discrete degrees of freedom, e.g. particles." diff --git a/hamilflow/models/discrete/d0/__init__.py b/hamilflow/models/discrete/d0/__init__.py deleted file mode 100644 index 695f0ab..0000000 --- a/hamilflow/models/discrete/d0/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"One-particle systems." diff --git a/hamilflow/models/harmonic_oscillator_chain.py b/hamilflow/models/harmonic_oscillator_chain.py new file mode 100644 index 0000000..5adb8f6 --- /dev/null +++ b/hamilflow/models/harmonic_oscillator_chain.py @@ -0,0 +1,139 @@ +from functools import cached_property +from typing import Mapping, Sequence, cast + +import numpy as np +import pandas as pd +from numpy.typing import ArrayLike +from scipy.fft import ifft + +from .free_particle import FreeParticle +from .harmonic_oscillator import ComplexSimpleHarmonicOscillator + + +class HarmonicOscillatorsChain: + r"""Generate time series data for a coupled harmonic oscillator chain + with periodic boundary condition. + + A one-dimensional circle of $N$ interacting harmonic oscillators can be described by the Lagrangian action + $$S_L[x_i] = \int_{t_0}^{t_1}\mathbb{d} t \left\{ \sum_{i=0}^{N-1} \frac{1}{2}m \dot x_i^2 - \frac{1}{2}m\omega^2\left(x_i - x_{i+1}\right)^2 \right\}\,,$$ + where $x_N \coloneqq x_0$. + + This system can be solved in terms of _travelling waves_, obtained by discrete Fourier transform. + + Since the original degrees of freedom are real, the initial conditions of the propagating waves need to satisfy + $Y_k = Y^*_{-k \mod N}$, see [Wikipedia](https://en.wikipedia.org/wiki/Discrete_Fourier_transform#DFT_of_real_and_purely_imaginary_signals). + + :param omega: frequence parameter + :param initial_conditions: a sequence of initial conditions on the Fourier modes. + The first element in the sequence is that of the zero mode, taking a position and a velocity. + Rest of the elements are that of the independent travelling waves, taking two amplitudes and two initial phases. + :param odd_dof: The system will have `2 * len(initial_conditions) + int(odd_dof) - 2` degrees of freedom. + """ + + def __init__( + self, + omega: float, + initial_conditions: Sequence[Mapping[str, float | tuple[float, float]]], + odd_dof: bool, + ) -> None: + self.n_dof = 2 * len(initial_conditions) + odd_dof - 2 + if not odd_dof: + prefix = "For even degrees of freedom, " + if self.n_dof == 0: + raise ValueError(prefix + "at least 1 travelling wave is needed") + amp = cast(tuple[float, float], initial_conditions[-1]["amp"]) + if amp[0] != amp[1]: + msg = "k == N // 2 must have equal positive and negative amplitudes." + raise ValueError(prefix + msg) + self.omega = omega + self.odd_dof = odd_dof + + self.free_mode = FreeParticle(cast(Mapping[str, float], initial_conditions[0])) + + self.independent_csho_modes = [ + self._sho_factory( + k, + cast(tuple[float, float], ic["amp"]), + cast(tuple[float, float] | None, ic.get("phi")), + ) + for k, ic in enumerate(initial_conditions[1:], 1) + ] + + def _sho_factory( + self, + k: int, + amp: tuple[float, float], + phi: tuple[float, float] | None = None, + ) -> ComplexSimpleHarmonicOscillator: + return ComplexSimpleHarmonicOscillator( + dict( + omega=2 * self.omega * np.sin(np.pi * k / self.n_dof), + ), + dict(x0=amp) | (dict(phi=phi) if phi else {}), + ) + + @cached_property + def definition( + self, + ) -> dict[ + str, + float + | dict[str, dict[str, float | list[float]]] + | list[dict[str, dict[str, float | tuple[float, float]]]], + ]: + """model params and initial conditions defined as a dictionary.""" + return dict( + omega=self.omega, + n_dof=self.n_dof, + free_mode=self.free_mode.definition, + independent_csho_modes=[ + rwm.definition for rwm in self.independent_csho_modes + ], + ) + + def _z( + self, t: "Sequence[float] | ArrayLike[float]" + ) -> tuple[np.ndarray, np.ndarray]: + t = np.array(t, copy=False).reshape(-1) + all_travelling_waves = [self.free_mode._x(t).reshape(1, -1)] + + if self.independent_csho_modes: + independent_cshos = np.array( + [o._z(t) for o in self.independent_csho_modes], copy=False + ) + all_travelling_waves.extend( + (independent_cshos, independent_cshos[::-1].conj()) + if self.odd_dof + else ( + independent_cshos[:-1], + independent_cshos[[-1]], + independent_cshos[-2::-1].conj(), + ) + ) + + travelling_waves = np.concatenate(all_travelling_waves) + original_zs = ifft(travelling_waves, axis=0, norm="ortho") + return original_zs, travelling_waves + + def _x( + self, t: "Sequence[float] | ArrayLike[float]" + ) -> tuple[np.ndarray, np.ndarray]: + original_xs, travelling_waves = self._z(t) + + return np.real(original_xs), travelling_waves + + def __call__(self, t: "Sequence[float] | ArrayLike[float]") -> pd.DataFrame: + """Generate time series data for the harmonic oscillator chain. + + Returns float(s) representing the displacement at the given time(s). + + :param t: time. + """ + original_xs, travelling_waves = self._x(t) + data = { + f"{name}{i}": values + for name, xs in zip(("x", "y"), (original_xs, travelling_waves)) + for i, values in enumerate(xs) + } + + return pd.DataFrame(data, index=t) diff --git a/mkdocs.yml b/mkdocs.yml index 634fc71..9e5724e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -84,10 +84,12 @@ nav: - "Complex Harmonic Oscillator": tutorials/complex_harmonic_oscillator.py - "Brownian Motion": tutorials/brownian_motion.py - "Pendulum": tutorials/pendulum.py + - "Harmonic Oscillator Chain": tutorials/harmonic_oscillator_chain.py - References: - "Introduction": references/index.md - "Models": - "Harmonic Oscillator": references/models/harmonic_oscillator.md - "Brownian Motion": references/models/brownian_motion.md - "Pendulum": references/models/pendulum.md + - "Harmonic Oscillator Chain": references/models/harmonic_oscillator_chain.md - "Changelog": changelog.md diff --git a/tests/test_models/test_harmonic_oscillator_chain.py b/tests/test_models/test_harmonic_oscillator_chain.py new file mode 100644 index 0000000..35a4f6c --- /dev/null +++ b/tests/test_models/test_harmonic_oscillator_chain.py @@ -0,0 +1,98 @@ +from itertools import chain, product +from typing import Iterable, Mapping, Sequence + +import numpy as np +import pytest + +from hamilflow.models.harmonic_oscillator_chain import HarmonicOscillatorsChain + +_wave_params: list[tuple[tuple[int, int], ...]] = [ + ((0, 0),), + ((1, 0),), + ((0, 1), (0, 1)), +] +_possible_wave_modes: list[dict[str, tuple[int, int]]] = [ + dict(zip(("amp", "phi"), param, strict=False)) for param in _wave_params +] + + +class TestHarmonicOscillatorChain: + @pytest.fixture(params=(1, 2)) + def omega(self, request: pytest.FixtureRequest) -> int: + return request.param + + @pytest.fixture(params=((0, 0), (0, 1), (1, 0), (1, 1))) + def free_mode(self, request: pytest.FixtureRequest) -> dict[str, int]: + return dict(zip(("x0", "v0"), request.param)) + + @pytest.fixture( + params=chain.from_iterable( + product(_possible_wave_modes, repeat=r) for r in range(3) + ) + ) + def wave_modes( + self, request: pytest.FixtureRequest + ) -> list[dict[str, tuple[int, int]]]: + return request.param + + @pytest.fixture(params=(False, True)) + def odd_dof(self, request: pytest.FixtureRequest) -> bool: + return request.param + + @pytest.fixture() + def legal_wave_modes_and_odd_def( + self, wave_modes: Iterable[Mapping[str, tuple[int, int]]], odd_dof: bool + ) -> tuple[Iterable[Mapping[str, tuple[int, int]]], bool]: + return wave_modes if odd_dof else chain(wave_modes, [dict(amp=(1, 1))]), odd_dof + + @pytest.fixture(params=(0, 1, (0, 1))) + def times(self, request: pytest.FixtureRequest) -> int | tuple[int]: + return request.param + + def test_init( + self, + omega: int, + free_mode: Mapping[str, int], + legal_wave_modes_and_odd_def: tuple[ + Iterable[Mapping[str, tuple[int, int]]], bool + ], + ) -> None: + wave_modes, odd_dof = legal_wave_modes_and_odd_def + assert HarmonicOscillatorsChain(omega, [free_mode, *wave_modes], odd_dof) + + @pytest.fixture() + def hoc_and_zs( + self, + omega: int, + free_mode: Mapping[str, int], + legal_wave_modes_and_odd_def: tuple[ + Iterable[Mapping[str, tuple[int, int]]], bool + ], + times: int | Sequence[int], + ) -> tuple[HarmonicOscillatorsChain, np.ndarray, np.ndarray]: + wave_modes, odd_dof = legal_wave_modes_and_odd_def + hoc = HarmonicOscillatorsChain(omega, [free_mode, *wave_modes], odd_dof) + return (hoc, *hoc._z(times)) + + def test_real( + self, hoc_and_zs: tuple[HarmonicOscillatorsChain, np.ndarray, np.ndarray] + ) -> None: + _, original_zs, _ = hoc_and_zs + assert np.all(original_zs.imag == 0.0) + + def test_dof( + self, hoc_and_zs: tuple[HarmonicOscillatorsChain, np.ndarray, np.ndarray] + ) -> None: + hoc, original_zs, _ = hoc_and_zs + assert original_zs.shape[0] == hoc.n_dof + + @pytest.mark.parametrize("wave_mode", [None, *_possible_wave_modes[1:]]) + def test_raise( + self, + omega: int, + free_mode: Mapping[str, int] | None, + wave_mode: Mapping[str, int], + ) -> None: + ics = [free_mode, *([wave_mode] if wave_mode else [])] + with pytest.raises(ValueError): + HarmonicOscillatorsChain(omega, ics, False)