feat(model): #36 one dimensional harmonic oscillator chain

docs/references/models/harmonic_oscillator_chain.md

@@ -0,0 +1,3 @@
+# Harmonic Oscillator Chain
+
+::: hamilflow.models.harmonic_oscillator_chain

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

hamilflow/models/harmonic_oscillator_chain.py

@@ -0,0 +1,142 @@
+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).
+    """
+
+    def __init__(
+        self,
+        omega: float,
+        initial_conditions: Sequence[Mapping[str, float | tuple[float, float]]],
+        odd_dof: bool,
+    ) -> None:
+        """Instantiate an oscillator chain.
+
+        :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.
+        """
+        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
+        | int
+        | 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)

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

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)