doc(kepler-problem): #82 add tutorial (#98)

* docs(kepler2d): add tutorial for Kepler2D; improve docstring of Kepler2D #82 * docs(kepler-problem): add tutorials for kepler problems: ellipse and hyperbolic 82 * docs(kepler-problem): added more examples for kepler tutorial #82 * docs(kepler-problem): fix typing of kepler problem tutorial #82 * Update docs/tutorials/kepler_problem.py Co-authored-by: cmp0xff <[email protected]> * Update docs/tutorials/kepler_problem.py Co-authored-by: cmp0xff <[email protected]> * docs(kepler-problem): add annimations of kepler orbits #82 * docs(kepler-problem): adjust kepler demo time interval and animation speed #82 * Update docs/tutorials/kepler_problem.py Co-authored-by: cmp0xff <[email protected]> --------- Co-authored-by: cmp0xff <[email protected]>
kausalflow · Sep 25, 2024 · ff67b48 · ff67b48
1 parent 6ca2f1e
commit ff67b48
Show file tree

Hide file tree

Showing 2 changed files with 133 additions and 8 deletions.
diff --git a/docs/tutorials/kepler_problem.py b/docs/tutorials/kepler_problem.py
@@ -13,22 +13,147 @@
 # ---
 
 # %% [markdown]
-# # Motion in a Central Field
+# # Kepler Problem
 #
-# In this tutorial, we generate data for objects moving in a central field.
+# In this tutorial we will generate data for the Kepler problem.
 #
 
 # %% [markdown]
 # ## Usage
 
 # %%
+import numpy as np
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
 
+from hamilflow.models.kepler_problem import Kepler2D
+
+# %% [markdown]
+# To make it easy to use the Kepler system, we implemented an interface `Kepler2D.from_geometry` to specify the configuration using the geometry of the orbit. The geometry of the orbit requires three parameters:
+#
+# - `positive_angular_mom`: whether the angular momentum is positive
+# - `ecc`: the eccentricity of the orbit
+# - `parameter`: the semi-latus rectum of the orbits, for circular orbits, this is the radius
+
+# %%
+system = {
+    "alpha": 1.0,
+    "mass": 1.0,
+}
+
+k2d = Kepler2D.from_geometry(
+    system=system,
+    geometries={
+        "positive_angular_mom": True,
+        "ecc": 0,
+        "parameter": 1.0,
+    },
+)
+
+# %% [markdown]
+# We define the time steps for a Kepler system. In this example, we create 401 time steps, also with some some negative timestamps for a more symmetric trajectory.
 
 # %%
-# 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(-20, 20, 401)
+t = np.arange(-20, 20, 0.125)
+
+# %% [markdown]
+# Orbit time series data is generated by calling the object `k2d`, as the data generation is done in the `__call__` method.
+
+# %%
+df_p_1 = k2d(t=t)
+df_p_1  # noqa: B018
+
+
+# %%
+def visualize_orbit(dataframe: pd.DataFrame) -> go.Figure:
+    """Plot out the trajectory in a polar coordinate.
+
+    :param dataframe: dataframe containing the trajectory data
+    """
+    df_chart = dataframe.assign(
+        phi_degree=dataframe.phi / (2 * np.pi) * 360,
+    )
+    range_r = 0, df_chart.r.max() * 1.1
+    fig = px.scatter_polar(
+        df_chart,
+        r="r",
+        theta="phi_degree",
+        hover_data=["t", "r", "phi"],
+        start_angle=0,
+        animation_frame="t",
+        color_discrete_sequence=["red"],
+        range_r=range_r,
+    )
 
-# t = np.linspace(0, 1, 11)
+    trajectory_traces = list(
+        px.line_polar(
+            df_chart,
+            r="r",
+            theta="phi_degree",
+            hover_data=["t", "r", "phi"],
+            start_angle=0,
+            range_r=range_r,
+        ).select_traces(),
+    )
+
+    fig.update_layout(
+        polar={
+            "angularaxis_thetaunit": "radians",
+        },
+    )
+
+    fig.add_traces(trajectory_traces)
+
+    fig.layout.updatemenus[0].buttons[0].args[1]["frame"]["duration"] = 50
+    fig.layout.updatemenus[0].buttons[0].args[1]["transition"]["duration"] = 50
+
+    return fig
+
+
+# %%
+fig = visualize_orbit(df_p_1)
+fig.show()
+
+# %% [markdown]
+# We also plot out the ellipse, hyperbolic, parabolic orbits.
+
+# %%
+visualize_orbit(
+    Kepler2D.from_geometry(
+        system=system,
+        geometries={
+            "positive_angular_mom": True,
+            "ecc": 0.5,
+            "parameter": 1.0,
+        },
+    )(t=t),
+).show()
+
+# %%
+visualize_orbit(
+    Kepler2D.from_geometry(
+        system=system,
+        geometries={
+            "positive_angular_mom": True,
+            "ecc": 1,
+            "parameter": 1.0,
+        },
+    )(t=t),
+).show()
+
+# %%
+visualize_orbit(
+    Kepler2D.from_geometry(
+        system=system,
+        geometries={
+            "positive_angular_mom": True,
+            "ecc": 2,
+            "parameter": 1.0,
+        },
+    )(t=np.arange(-5, 5, 0.125)),
+).show()
 
 # %% [markdown]
 # ## Formalism

diff --git a/hamilflow/models/kepler_problem/model.py b/hamilflow/models/kepler_problem/model.py
@@ -110,14 +110,14 @@ def from_geometry(
     ) -> "Self":
         r"""Alternative initialiser from system and geometry specifications.
 
-        Given the eccentricity $e$ and the conic section parameter $p$,
+        Given the eccentricity $e$ and the conic section semi-latus rectum $p$,
         $$l = \pm \sqrt{mp}\,,\quad E = (e^2-1) \left|E_\text{min}\right|\,.$$
 
         :param system: the Kepler problem system specification
         :param geometries: geometric specifications
             `positive_angular_mom`: whether the angular momentum is positive
             `ecc`: eccentricity of the conic section
-            `parameter`: parameter of the conic section
+            `parameter`: semi-latus rectum of the conic section
         """
         mass, alpha = system["mass"], system["alpha"]
         positive_angular_mom = bool(geometries["positive_angular_mom"])
@@ -195,7 +195,7 @@ def period(self) -> float:
     # FIXME is it called parameter in English?
     @cached_property
     def parameter(self) -> float:
-        r"""Conic section parameter of the Kepler problem.
+        r"""Conic section semi-latus rectum of the Kepler problem.
 
         $$ p = \frac{l^2}{\alpha m}\,. $$
         """