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 <[email protected]>

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.

docs/references/models/central_field.md

@@ -0,0 +1,3 @@
+# Brownian Motion
+
+::: hamilflow.models.central_field

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]
+#

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        |

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))

mkdocs.yml

@@ -84,12 +84,14 @@ 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
     - "Models":
       - "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

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