From fde01d7a55f44b18a2261913e2cb460d2343b338 Mon Sep 17 00:00:00 2001 From: "Nathaniel D. Hoffman" Date: Sat, 14 Jun 2025 12:58:26 -0400 Subject: [PATCH 1/3] feat: Add Lorenz attractor example --- examples/lorenz.py | 74 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 examples/lorenz.py diff --git a/examples/lorenz.py b/examples/lorenz.py new file mode 100644 index 0000000..2211db0 --- /dev/null +++ b/examples/lorenz.py @@ -0,0 +1,74 @@ +# Inspired by +from dataclasses import dataclass, field + +import matplotlib.pyplot as plt +import numpy as np +from joblib import Parallel, delayed +from matplotlib.colors import Normalize +from matplotloom import Loom +from mpl_toolkits.mplot3d.art3d import Line3DCollection + + +@dataclass +class Lorenz: + dt: float = 0.01 + sigma: float = 10.0 + rho: float = 28.0 + beta: float = 8.0 / 3.0 + x: float = 1.0 + y: float = 1.0 + z: float = 1.0 + + def step(self): + dx = self.sigma * (self.y - self.x) + dy = self.x * (self.rho - self.z) - self.y + dz = self.x * self.y - self.beta * self.z + self.x += dx * self.dt + self.y += dy * self.dt + self.z += dz * self.dt + + @property + def position(self) -> tuple[float, float, float]: + return self.x, self.y, self.z + + +@dataclass +class LorenzPlotter: + plot_speed: int = 20 + attractor = Lorenz() + points: list[tuple[float, float, float]] = field(default_factory=list) + + def initialize(self, steps: int): + self.points = [self.attractor.position] + for _ in range(steps): + self.attractor.step() + self.points.append(self.attractor.position) + + @property + def frames(self) -> list[int]: + return list(range(1, len(self.points) // self.plot_speed)) + + def get_frame(self, i: int, loom: Loom): + fig, ax = plt.subplots(figsize=(12, 8), subplot_kw={'projection': '3d'}) + points = np.array(self.points[: i * self.plot_speed]) + xs, ys, zs = points.T + segments = np.array([points[:-1], points[1:]]).transpose(1, 0, 2) + norm = Normalize(vmin=0, vmax=len(xs)) + colors = plt.get_cmap('inferno')(norm(np.arange(len(xs) - 1))) + lc = Line3DCollection(segments, colors=colors, linewidth=0.5) + ax.add_collection3d(lc) + ax.set_xlim(-30, 30) + ax.set_ylim(-30, 30) + ax.set_zlim(0, 50) + ax.view_init( + azim=(np.pi * 1.7 + 0.8 * np.sin(2.0 * np.pi * i * self.plot_speed / len(self.frames) / 10)) * 180.0 / np.pi + ) + ax.set_axis_off() + ax.grid(visible=False) + loom.save_frame(fig, i - 1) + + +with Loom('lorenz.mp4', fps=60, parallel=True, overwrite=True) as loom: + attractor = LorenzPlotter() + attractor.initialize(10000) + Parallel(n_jobs=-1)(delayed(attractor.get_frame)(i, loom) for i in attractor.frames) From e2955cab4e6d48e419a34c0a2503732bbf6bba22 Mon Sep 17 00:00:00 2001 From: "Nathaniel D. Hoffman" Date: Sat, 14 Jun 2025 13:07:25 -0400 Subject: [PATCH 2/3] Add option to keep figures open in save_frame --- matplotloom/loom.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/matplotloom/loom.py b/matplotloom/loom.py index 08caf41..669b2fb 100644 --- a/matplotloom/loom.py +++ b/matplotloom/loom.py @@ -158,7 +158,8 @@ def __exit__( def save_frame( self, fig: Figure, - frame_number: Optional[int] = None + frame_number: Optional[int] = None, + close_fig: bool = True, ) -> None: """ Save a single frame of the animation. @@ -188,7 +189,8 @@ def save_frame( print(f"Saving frame {frame_number} to {frame_filepath}") fig.savefig(frame_filepath, **self.savefig_kwargs) - plt.close(fig) + if close_fig: + plt.close(fig) def save_video(self) -> None: """ From 72ef4085380a7f5dacb6ce37617cb78c4e3369e3 Mon Sep 17 00:00:00 2001 From: "Nathaniel D. Hoffman" Date: Sat, 16 Aug 2025 18:26:39 -0400 Subject: [PATCH 3/3] fix: remove examples/lorenz.py --- examples/lorenz.py | 74 ---------------------------------------------- 1 file changed, 74 deletions(-) delete mode 100644 examples/lorenz.py diff --git a/examples/lorenz.py b/examples/lorenz.py deleted file mode 100644 index 2211db0..0000000 --- a/examples/lorenz.py +++ /dev/null @@ -1,74 +0,0 @@ -# Inspired by -from dataclasses import dataclass, field - -import matplotlib.pyplot as plt -import numpy as np -from joblib import Parallel, delayed -from matplotlib.colors import Normalize -from matplotloom import Loom -from mpl_toolkits.mplot3d.art3d import Line3DCollection - - -@dataclass -class Lorenz: - dt: float = 0.01 - sigma: float = 10.0 - rho: float = 28.0 - beta: float = 8.0 / 3.0 - x: float = 1.0 - y: float = 1.0 - z: float = 1.0 - - def step(self): - dx = self.sigma * (self.y - self.x) - dy = self.x * (self.rho - self.z) - self.y - dz = self.x * self.y - self.beta * self.z - self.x += dx * self.dt - self.y += dy * self.dt - self.z += dz * self.dt - - @property - def position(self) -> tuple[float, float, float]: - return self.x, self.y, self.z - - -@dataclass -class LorenzPlotter: - plot_speed: int = 20 - attractor = Lorenz() - points: list[tuple[float, float, float]] = field(default_factory=list) - - def initialize(self, steps: int): - self.points = [self.attractor.position] - for _ in range(steps): - self.attractor.step() - self.points.append(self.attractor.position) - - @property - def frames(self) -> list[int]: - return list(range(1, len(self.points) // self.plot_speed)) - - def get_frame(self, i: int, loom: Loom): - fig, ax = plt.subplots(figsize=(12, 8), subplot_kw={'projection': '3d'}) - points = np.array(self.points[: i * self.plot_speed]) - xs, ys, zs = points.T - segments = np.array([points[:-1], points[1:]]).transpose(1, 0, 2) - norm = Normalize(vmin=0, vmax=len(xs)) - colors = plt.get_cmap('inferno')(norm(np.arange(len(xs) - 1))) - lc = Line3DCollection(segments, colors=colors, linewidth=0.5) - ax.add_collection3d(lc) - ax.set_xlim(-30, 30) - ax.set_ylim(-30, 30) - ax.set_zlim(0, 50) - ax.view_init( - azim=(np.pi * 1.7 + 0.8 * np.sin(2.0 * np.pi * i * self.plot_speed / len(self.frames) / 10)) * 180.0 / np.pi - ) - ax.set_axis_off() - ax.grid(visible=False) - loom.save_frame(fig, i - 1) - - -with Loom('lorenz.mp4', fps=60, parallel=True, overwrite=True) as loom: - attractor = LorenzPlotter() - attractor.initialize(10000) - Parallel(n_jobs=-1)(delayed(attractor.get_frame)(i, loom) for i in attractor.frames)