From 5ddc2acc0892dc35c9ad5a7c56943231c9881bb7 Mon Sep 17 00:00:00 2001 From: LM Date: Sun, 18 Aug 2024 21:14:28 +0200 Subject: [PATCH] feat(central-field): add central force field model (#50) * feat(central-field): add theories of central force field #13 * feat(central-field): add solutions to central field problem add solutions to the Kepler problem #13 * feat(central-field): add central field system, ic, and useful quantities #13 * docs(central-field): add 2d central field motion derivations and formulations #13 * feat(central-field): numerical integration of the Kepler equation #13 * feat(central-field): adjust naming of var #13 * feat(central-field): joint solution of r and phi #13 --------- Co-authored-by: cmp0xff --- README.md | 9 ++ docs/references/models/central_field.md | 3 + docs/tutorials/central_force_motion.py | 155 ++++++++++++++++++++++++ docs/tutorials/index.md | 8 ++ hamilflow/models/central_field.py | 152 +++++++++++++++++++++++ mkdocs.yml | 2 + pyproject.toml | 4 + tests/test_models/test_central_field.py | 48 ++++++++ 8 files changed, 381 insertions(+) create mode 100644 docs/references/models/central_field.md create mode 100644 docs/tutorials/central_force_motion.py create mode 100644 hamilflow/models/central_field.py create mode 100644 tests/test_models/test_central_field.py diff --git a/README.md b/README.md index 2c88656..aecdda4 100644 --- a/README.md +++ b/README.md @@ -14,3 +14,12 @@ Dataset of simple physical systems. - Create a new tag - Push to `release/X.Y.Z` - Merge to `main`, no need to squash, maybe don't delete the branch under release + +## Update Tutorials + +The tutorials are auto generated by mkdocs based on percent scripts located in the `docs/tutorials` folder. The scripts can be modified and run in vscode without problems. However, if one decided to use jupyter to create/develop tutorials, jupytext is the tool that helps sync between jupyter notebooks and percent scripts. + +1. To create jupyter notebooks for the first time, run `poetry run jupytext --to ipynb docs/tutorials/*.py`. +2. To sync between notebooks and percent scripts, run `poetry run jupytext --sync docs/tutorials/*`. +3. To create percent script from jupyter notebooks, run `poetry run jupytext --to py:percent docs/tutorials/your_notebook.ipynb`. +4. Jupyter Lab has an extension that automatically syncs between the different versions. Please refer to jupytext documentation for more info. diff --git a/docs/references/models/central_field.md b/docs/references/models/central_field.md new file mode 100644 index 0000000..fb34732 --- /dev/null +++ b/docs/references/models/central_field.md @@ -0,0 +1,3 @@ +# Brownian Motion + +::: hamilflow.models.central_field diff --git a/docs/tutorials/central_force_motion.py b/docs/tutorials/central_force_motion.py new file mode 100644 index 0000000..e1e9e67 --- /dev/null +++ b/docs/tutorials/central_force_motion.py @@ -0,0 +1,155 @@ +# --- +# jupyter: +# jupytext: +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.16.2 +# kernelspec: +# display_name: .venv +# language: python +# name: python3 +# --- + +# %% [markdown] +# # Motion in a Central Field +# +# In this tutorial, we generate data for objects moving in a central field. +# + +# %% [markdown] +# ## Usage + +# %% +import numpy as np +import pandas as pd +import plotly.express as px + +from hamilflow.models.central_field import CentralField2D + +# %% +system = {"alpha": 1.0, "mass": 1.0} +ic = {"r_0": 1.0, "phi_0": 0.0, "drdt_0": 0.0, "dphidt_0": 0.1} + +t = np.linspace(0, 1, 11) + +# %% +cf = CentralField2D( + system=system, + initial_condition=ic, +) + +df_orbit = cf(t) + +df_orbit.head() + + +# %% +df_orbit.plot.line(x="t", y="r") + +# %% +df_orbit.plot.line(x="phi", y="r") + +# %% +fig = px.line_polar(df_orbit, r="r", theta="phi") + +fig.update_polars(angularaxis=dict(thetaunit="radians")) +fig.show() + +# %% [markdown] +# ## Formalism + +# %% [markdown] +# In this section, we briefly discuss the derivations of motion in a central field. Please refer to *Mechanics: Vol 1* by Landau and Lifshitz for more details[^1]. +# +# The Lagrangian for an object in a central field is +# +# $$ +# \mathcal L = \frac{1}{2} m ({\dot r}^2 + r^2 {\dot \phi}^2) - V(r), +# $$ +# +# where $r$ and $\phi$ are the polar coordinates, $m$ is the mass of the object, and $V(r)$ is the potential energy. The equations of motion are +# +# $$ +# \begin{align} +# \frac{\mathrm d r}{\mathrm d t} &= \sqrt{ \frac{2}{m} (E - V(r)) - \frac{L^2}{m^2 r^2} } \\ +# \frac{\mathrm d \phi}{\mathrm d t} &= \frac{L}{m r^2}, +# \end{align} +# $$ +# +# where $E$ is the total energy and $L$ is the angular momentum, +# +# $$ +# \begin{align} +# E =& \frac{1}{2} m \left(\left(\frac{\mathrm d r}{ \mathrm dt} \right)^2 + r^2 \left( \frac{\mathrm d r}{\mathrm dt} \right)^2\right) + V(r) \\ +# L =& m r^2 \frac{d\phi}{dt} +# \end{align} +# $$ +# +# Both $E$ and $L$ are conserved. We obtain the coordinates as a function of time by solving the two equations. +# +# + +# %% [markdown] +# For a inverse-square force, the potential energy is +# +# $$ +# V(r) = - \frac{\alpha}{r}, +# $$ +# +# where $\alpha$ is a constant that specifies the strength of the force. For Newtonian gravity, $\alpha = G m_0$ with $G$ being the gravitational constant. +# +# First of all, we solve $\phi(r)$, +# +# $$ +# \phi = \cos^{-1}\left( \frac{L/r - m\alpha/L}{2 m E + m \alpha^2/L^2} \right). +# $$ +# +# Define +# +# $$ +# p = \frac{L^2}{m\alpha} +# $$ +# +# and +# +# $$ +# e = \sqrt{ 1 + \frac{2 E L^2}{m \alpha^2}}, +# $$ +# +# we rewrite the solution $\phi(r)$ as +# +# $$ +# r = \frac{p}{1 + e \cos{\phi}}. +# $$ +# +# With the above relationship between $r$ and $\phi$, and $\frac{\mathrm d \phi}{\mathrm d} = \frac{L}{mr^2}$, we find that +# +# $$ +# \frac{m\alpha^2}{L^3} \mathrm d t = \frac{1}{(1 + e \cos{\phi})^2} \mathrm d \phi. +# $$ +# +# The integration on the right hand side depends on the domain of $e$. +# +# $$ +# \int\frac{1}{(1 + e \cos{\phi})^2} \mathrm d \phi = \begin{cases} +# \frac{1}{(1 - e^2)^{3/2}} \left( 2 \tan^{-1} \sqrt{\frac{1 - e}{1 + e}} \tan\frac{\phi}{2} - \frac{e\sqrt{1 - e^2}\sin\phi}{1 + e\cos\phi} \right), & \text{if } e<1 \\ +# \frac{1}{2}\tan{\frac{\phi}{2}} + \frac{1}{6}\tan^3{\frac{\phi}{2}}, & \text{if } e=1 \\ +# \frac{1}{(e^2 - 1)^{3/2}}\left( \frac{e\sqrt{e^2-1}\sin\phi}{1 + e\cos\phi} - \ln \left( \frac{\sqrt{1 + e} + \sqrt{e-1}\tan\frac{\phi}{2}}{\sqrt{1 + e} - \sqrt{e-1}\tan\frac{\phi}{2} } \right) \right), & \text{if } e> 1. +# \end{cases} +# $$ +# +# The value of $t(\phi)$ is easily obtained from the above formulae. + +# %% [markdown] +# There exists many numerical methods to solve the Kepler orbits as functions of time, $r(t)$ and $\phi(t)$. For our use case of the solutions, we choose to integrate the equation of motion directly. + +# %% [markdown] +# References: +# +# 1. Landau LD, Lifshitz EM. Mechanics: Vol 1. 3rd ed. Oxford, England: Butterworth-Heinemann; 1982. +# 2. Klioner SA. Basic Celestial Mechanics. arXiv [astro-ph.IM]. 2016. Available: http://arxiv.org/abs/1609.00915 + +# %% [markdown] +# diff --git a/docs/tutorials/index.md b/docs/tutorials/index.md index 508c035..7bb7d06 100644 --- a/docs/tutorials/index.md +++ b/docs/tutorials/index.md @@ -1,3 +1,11 @@ # Tutorials We provide some tutorials to help you get started. + +For easier reading, we following the following notations whenever possible. + +| Notation | Description | +| ------------ | ---------------- | +| $\mathcal L$ | Lagrangian | +| $L$ | Angular Momentum | +| $V$ | Potential | diff --git a/hamilflow/models/central_field.py b/hamilflow/models/central_field.py new file mode 100644 index 0000000..d723dfd --- /dev/null +++ b/hamilflow/models/central_field.py @@ -0,0 +1,152 @@ +from functools import cached_property +from typing import Dict, Optional + +import numpy as np +import numpy.typing as npt +import pandas as pd +import scipy as sp +from pydantic import BaseModel, Field + + +class CentralField2DSystem(BaseModel): + r"""Definition of the central field system + + Potential: + + $$ + V(r) = - \frac{\alpha}{r}. + $$ + + For reference, if an object is orbiting our Sun, the constant $\alpha = G M_{\odot} ~ 1.327×10^20 m^3/s^2$ in SI, + which is also called 1 TCB, or 1 solar mass parameter. For computational stability, we recommend using + TCB as the unit instead of the large SI values. + + !!! note "Units" + + When specifying the parameters of the system, be ware of the consistency of the units. + + :cvar alpha: the proportional constant of the potential energy. + :cvar mass: the mass of the orbiting object + """ + + alpha: float = Field(gt=0, default=1) + mass: float = Field(gt=0, default=1) + + +class CentralField2DIC(BaseModel): + """The initial condition for a Brownian motion + + :cvar r_0: the initial radial coordinate + :cvar phi_0: the initial phase + :cvar drdt_0: the initial radial velocity + :cvar dphidt_0: the initial phase velocity + """ + + r_0: float = Field(gt=0, default=1.0) + phi_0: float = Field(ge=0, default=0.0) + drdt_0: float = 1.0 + dphidt_0: float = 0.0 + + +class CentralField2D: + r"""Central field motion in two dimensional space. + + !!! info "Inverse Sqare Law" + + We only consider central field of the form $-\frac{\alpha}{r}$. + + :param system: the Central field motion system definition + :param initial_condition: the initial condition for the simulation + """ + + def __init__( + self, + system: Dict[str, float], + initial_condition: Optional[Dict[str, float]] = {}, + ): + self.system = CentralField2DSystem.model_validate(system) + self.initial_condition = CentralField2DIC.model_validate(initial_condition) + + @cached_property + def _angular_momentum(self) -> float: + """computes the angular momentum of the motion. Since the angular momentum is + conserved, it doesn't change through time. + """ + return ( + self.system.mass + * self.initial_condition.r_0**2 + * self.initial_condition.dphidt_0 + ) + + @cached_property + def _energy(self) -> float: + """computes the total energy""" + drdt_0 = self.initial_condition.drdt_0 + dphidt_0 = self.initial_condition.dphidt_0 + r_0 = self.initial_condition.r_0 + + potential_energy = self._potential(r_0) + + return ( + 0.5 * self.system.mass * (drdt_0**2 + r_0**2 * dphidt_0**2) + + potential_energy + ) + + def _potential(self, r: npt.ArrayLike) -> npt.ArrayLike: + return -1 * self.system.alpha / r + + def drdt(self, t: npt.ArrayLike, r: npt.ArrayLike) -> npt.ArrayLike: + return np.sqrt( + 2 / self.system.mass * (self._energy - self._potential(r)) + - self._angular_momentum**2 / self.system.mass**2 / r**2 + ) + + def r(self, t: npt.ArrayLike) -> npt.ArrayLike: + t_span = t.min(), t.max() + sol = sp.integrate.solve_ivp( + self.drdt, + t_span=t_span, + y0=[self.initial_condition.r_0], + t_eval=t, + ) + + return sol.y[0] + + def phi(self, t: npt.ArrayLike, r: npt.ArrayLike) -> npt.ArrayLike: + return ( + self.initial_condition.phi_0 + + self._angular_momentum / self.system.mass / r**2 * t + ) + + def state_derivative(self, t: npt.ArrayLike, state: npt.ArrayLike) -> npt.ArrayLike: + """derivative of the state vector (r, phi). + + :param t: time sequence to be solved at + :param state: the state vector (r, phi) + """ + r, _ = state + return np.array( + [ + np.sqrt( + 2 / self.system.mass * (self._energy - self._potential(r)) + - self._angular_momentum**2 / self.system.mass**2 / r**2 + ), + self._angular_momentum / self.system.mass / r**2, + ] + ) + + def state(self, t: npt.ArrayLike) -> npt.ArrayLike: + t_span = t.min(), t.max() + sol = sp.integrate.solve_ivp( + self.state_derivative, + t_span=t_span, + y0=[self.initial_condition.r_0, self.initial_condition.phi_0], + t_eval=t, + ) + + return sol.y + + def __call__(self, t: npt.ArrayLike) -> npt.ArrayLike: + r, phi = self.state(t) + + return pd.DataFrame(dict(t=t, r=r, phi=phi)) diff --git a/mkdocs.yml b/mkdocs.yml index 9e5724e..19fe319 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -84,6 +84,7 @@ nav: - "Complex Harmonic Oscillator": tutorials/complex_harmonic_oscillator.py - "Brownian Motion": tutorials/brownian_motion.py - "Pendulum": tutorials/pendulum.py + - "Central Force Motion": tutorials/central_force_motion.py - "Harmonic Oscillator Chain": tutorials/harmonic_oscillator_chain.py - References: - "Introduction": references/index.md @@ -91,5 +92,6 @@ nav: - "Harmonic Oscillator": references/models/harmonic_oscillator.md - "Brownian Motion": references/models/brownian_motion.md - "Pendulum": references/models/pendulum.md + - "Central Field": references/models/central_field.md - "Harmonic Oscillator Chain": references/models/harmonic_oscillator_chain.md - "Changelog": changelog.md diff --git a/pyproject.toml b/pyproject.toml index 5d7f144..85c215d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,10 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" +[tool.jupytext] +formats = "ipynb,py:percent" + + [[tool.mypy.overrides]] module = ["plotly.*", "scipy.*"] ignore_missing_imports = true diff --git a/tests/test_models/test_central_field.py b/tests/test_models/test_central_field.py new file mode 100644 index 0000000..efee095 --- /dev/null +++ b/tests/test_models/test_central_field.py @@ -0,0 +1,48 @@ +import numpy as np +import pandas as pd +import pytest + +from hamilflow.models.central_field import ( + CentralField2D, + CentralField2DIC, + CentralField2DSystem, +) + + +@pytest.fixture +def central_field_2d_system_params(): + return {"alpha": 1.0, "mass": 1.0} + + +@pytest.fixture +def central_field_2d_ic_params(): + return {"r_0": 1.0, "phi_0": 0.0, "v_r_0": 0.0, "v_phi_0": 1.0} + + +@pytest.fixture +def t(): + return np.linspace(0, 10, 101) + + +def test_central_field_2d_conserved( + central_field_2d_ic_params, central_field_2d_system_params +): + cf = CentralField2D( + system=central_field_2d_system_params, + initial_condition=central_field_2d_ic_params, + ) + + assert cf._energy == 0.0 + assert cf._angular_momentum == 1.0 + + +def test_central_field_2d_orbit( + central_field_2d_ic_params, central_field_2d_system_params, t +): + + cf = CentralField2D( + system=central_field_2d_system_params, + initial_condition=central_field_2d_ic_params, + ) + + cf(t)