From c85988c8581aa87f6e8fa8981cc40d2f91f484af Mon Sep 17 00:00:00 2001 From: ali-ramadhan Date: Sun, 30 Jun 2024 08:20:27 -0600 Subject: [PATCH 1/2] Add double pendulum example --- examples/double_pendulum.py | 83 +++++++++++++++++++++++++++++++++++++ poetry.lock | 22 +++++++++- pyproject.toml | 1 + 3 files changed, 105 insertions(+), 1 deletion(-) create mode 100644 examples/double_pendulum.py diff --git a/examples/double_pendulum.py b/examples/double_pendulum.py new file mode 100644 index 0000000..b6c586e --- /dev/null +++ b/examples/double_pendulum.py @@ -0,0 +1,83 @@ +import numpy as np +import matplotlib.pyplot as plt + +from scipy.integrate import solve_ivp +from tqdm import tqdm +from matplotloom import Loom + +g = 9.80665 # standard acceleration of gravity [m/s²] +l1, l2 = 1, 1 # pendulum arms lengths [m] +m1, m2 = 1, 1 # pendulum masses [kg] + +# Calculate dy/dt where y = [θ1, ω1, θ2, ω2]. +def derivatives(t, state): + θ1, ω1, θ2, ω2 = state + dydt = np.zeros_like(state) + + dydt[0] = ω1 + dydt[2] = ω2 + + Δθ = θ2 - θ1 + + denominator1 = (m1 + m2) * l1 - m2 * l1 * np.cos(Δθ)**2 + dydt[1] = (m2 * l1 * ω1**2 * np.sin(Δθ) * np.cos(Δθ) + + m2 * g * np.sin(θ2) * np.cos(Δθ) + + m2 * l2 * ω2**2 * np.sin(Δθ) + - (m1 + m2) * g * np.sin(θ1)) / denominator1 + + denominator2 = (l2 / l1) * denominator1 + dydt[3] = (-m2 * l2 * ω2**2 * np.sin(Δθ) * np.cos(Δθ) + + (m1 + m2) * g * np.sin(θ1) * np.cos(Δθ) + - (m1 + m2) * l1 * ω1**2 * np.sin(Δθ) + - (m1 + m2) * g * np.sin(θ2)) / denominator2 + + return dydt + +t_span = (0, 20) +y0 = [np.pi/2, 0, np.pi/2, 0] +sol = solve_ivp(derivatives, t_span, y0, dense_output=True) + +times = np.linspace(t_span[0], t_span[1], 1000) + +θ1, ω1, θ2, ω2 = sol.sol(times) +x1 = l1 * np.sin(θ1) +y1 = -l1 * np.cos(θ1) +x2 = x1 + l2 * np.sin(θ2) +y2 = y1 - l2 * np.cos(θ2) + +loom = Loom( + "double_pendulum.mp4", + fps = 60, + overwrite = True, + savefig_kwargs = {"bbox_inches": "tight"} +) + +with loom: + for i, t in tqdm(enumerate(times), total=len(times)): + fig, ax = plt.subplots(figsize=(8, 8)) + + ax.plot( + [0, x1[i], x2[i]], + [0, y1[i], y2[i]], + linestyle = "solid", + marker = "o", + color = "black", + linewidth = 3 + ) + + ax.plot( + x2[:i+1], + y2[:i+1], + linestyle = "solid", + linewidth = 2, + color = "red", + alpha = 0.5 + ) + + ax.set_title(f"Double Pendulum: t = {t:.3f}s") + + ax.set_xlim(-2.2, 2.2) + ax.set_ylim(-2.2, 2.2) + ax.set_aspect("equal", adjustable="box") + + loom.save_frame(fig) diff --git a/poetry.lock b/poetry.lock index 4277576..2e0bfed 100644 --- a/poetry.lock +++ b/poetry.lock @@ -982,6 +982,26 @@ files = [ {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, ] +[[package]] +name = "tqdm" +version = "4.66.4" +description = "Fast, Extensible Progress Meter" +optional = false +python-versions = ">=3.7" +files = [ + {file = "tqdm-4.66.4-py3-none-any.whl", hash = "sha256:b75ca56b413b030bc3f00af51fd2c1a1a5eac6a0c1cca83cbb37a5c52abce644"}, + {file = "tqdm-4.66.4.tar.gz", hash = "sha256:e4d936c9de8727928f3be6079590e97d9abfe8d39a590be678eb5919ffc186bb"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "platform_system == \"Windows\""} + +[package.extras] +dev = ["pytest (>=6)", "pytest-cov", "pytest-timeout", "pytest-xdist"] +notebook = ["ipywidgets (>=6)"] +slack = ["slack-sdk"] +telegram = ["requests"] + [[package]] name = "traitlets" version = "5.14.2" @@ -1037,4 +1057,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "efd5694a606c587e4b089d18ad96e2cd15998b5182f4cbaac60e7612bc92bdd0" +content-hash = "ad2b90685d2de92a1e7e4921fc5b0ea3f267847e59af11ba4a054ddeb077fb1b" diff --git a/pyproject.toml b/pyproject.toml index 6019b32..f57f333 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ scipy = [ {version = "^1.13.0", python = "~3.9"}, {version = "^1.14.0", python = "^3.10"} ] +tqdm = "^4.66.4" [build-system] requires = ["poetry-core"] From 9b9ac21ddd8b819969454e0a59fd7f48ff9720d7 Mon Sep 17 00:00:00 2001 From: ali-ramadhan Date: Sun, 30 Jun 2024 08:22:15 -0600 Subject: [PATCH 2/2] Test double pendulum --- tests/test_examples.py | 86 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 85 insertions(+), 1 deletion(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index ff0218c..c8a91bc 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -5,8 +5,10 @@ from pathlib import Path from cmocean import cm +from scipy.integrate import solve_ivp from scipy.special import j0 from joblib import Parallel, delayed +from tqdm import tqdm from matplotloom import Loom def test_sine_wave(): @@ -97,6 +99,87 @@ def create_frame(x, y, t): assert Path("bessel_wave.mp4").is_file() assert Path("bessel_wave.mp4").stat().st_size > 0 +def test_double_pendulum(): + g = 9.80665 # standard acceleration of gravity [m/s²] + l1, l2 = 1, 1 # pendulum arms lengths [m] + m1, m2 = 1, 1 # pendulum masses [kg] + + # Calculate dy/dt where y = [θ1, ω1, θ2, ω2]. + def derivatives(t, state): + θ1, ω1, θ2, ω2 = state + dydt = np.zeros_like(state) + + dydt[0] = ω1 + dydt[2] = ω2 + + Δθ = θ2 - θ1 + + denominator1 = (m1 + m2) * l1 - m2 * l1 * np.cos(Δθ)**2 + dydt[1] = (m2 * l1 * ω1**2 * np.sin(Δθ) * np.cos(Δθ) + + m2 * g * np.sin(θ2) * np.cos(Δθ) + + m2 * l2 * ω2**2 * np.sin(Δθ) + - (m1 + m2) * g * np.sin(θ1)) / denominator1 + + denominator2 = (l2 / l1) * denominator1 + dydt[3] = (-m2 * l2 * ω2**2 * np.sin(Δθ) * np.cos(Δθ) + + (m1 + m2) * g * np.sin(θ1) * np.cos(Δθ) + - (m1 + m2) * l1 * ω1**2 * np.sin(Δθ) + - (m1 + m2) * g * np.sin(θ2)) / denominator2 + + return dydt + + t_span = (0, 1) + y0 = [np.pi/2, 0, np.pi/2, 0] + sol = solve_ivp(derivatives, t_span, y0, dense_output=True) + + times = np.linspace(t_span[0], t_span[1], 10) + + θ1, ω1, θ2, ω2 = sol.sol(times) + x1 = l1 * np.sin(θ1) + y1 = -l1 * np.cos(θ1) + x2 = x1 + l2 * np.sin(θ2) + y2 = y1 - l2 * np.cos(θ2) + + loom = Loom( + "double_pendulum.mp4", + fps = 60, + overwrite = True, + savefig_kwargs = {"bbox_inches": "tight"} + ) + + with loom: + for i, t in tqdm(enumerate(times), total=len(times)): + fig, ax = plt.subplots(figsize=(8, 8)) + + ax.plot( + [0, x1[i], x2[i]], + [0, y1[i], y2[i]], + linestyle = "solid", + marker = "o", + color = "black", + linewidth = 3 + ) + + ax.plot( + x2[:i+1], + y2[:i+1], + linestyle = "solid", + linewidth = 2, + color = "red", + alpha = 0.5 + ) + + ax.set_title(f"Double Pendulum: t = {t:.3f}s") + + ax.set_xlim(-2.2, 2.2) + ax.set_ylim(-2.2, 2.2) + ax.set_aspect("equal", adjustable="box") + + loom.save_frame(fig) + + assert Path("double_pendulum.mp4").is_file() + assert Path("double_pendulum.mp4").stat().st_size > 0 + def test_parallel_sine_wave(): def plot_frame(phase, frame_number, loom): fig, ax = plt.subplots() @@ -117,4 +200,5 @@ def plot_frame(phase, frame_number, loom): for i, phase in enumerate(phases) ) - assert Path("parallel_sine_wave.gif").is_file() \ No newline at end of file + assert Path("parallel_sine_wave.gif").is_file() + assert Path("parallel_sine_wave.gif").stat().st_size > 0