diff --git a/README.md b/README.md index 2fc38e3..83d1912 100644 --- a/README.md +++ b/README.md @@ -345,6 +345,90 @@ with loom: https://github.com/ali-ramadhan/matplotloom/assets/20099589/dad48d56-73ae-4cb2-ba71-d21855c72215 +### Lorenz Attractor + +This example shows how you can use `matplotloom` to create an animation of a Lorenz attractor (inspired by [an example from the Makie library](https://docs.makie.org/stable/)). + +```python +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: + steps_per_frame: 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.steps_per_frame)) + + 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.steps_per_frame]) + 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.steps_per_frame / 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) as loom: + attractor = LorenzPlotter() + attractor.initialize(10000) + Parallel(n_jobs=-1)(delayed(attractor.get_frame)(i, loom) for i in attractor.frames) +``` + +https://github.com/user-attachments/assets/69b02d78-386d-4843-90df-b1a43600cfa5 + ## Looming in parallel By passing `parallel=True` when creating a `Loom`, you can save frames using `loom.save_frame(fig, frame_number)` which allows you to plot and save all your frames in parallel. One easy way to leverage this is by using joblib to parallelize the for loop. For example, here's how you can parallelize the simple sine wave example: diff --git a/docs/index.rst b/docs/index.rst index ee039fa..2d80a39 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -348,6 +348,95 @@ matplotloom works out of the box with anything that is built on top of matplotli +Lorenz Attractor +~~~~~~~~~~~~~~~~ + +This example shows how you can use `matplotloom` to create an animation of a Lorenz attractor (inspired by `an example from the Makie library `_). + +.. code-block:: python + + 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: + steps_per_frame: 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.steps_per_frame)) + + 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.steps_per_frame]) + 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.steps_per_frame / 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) as loom: + attractor = LorenzPlotter() + attractor.initialize(10000) + Parallel(n_jobs=-1)(delayed(attractor.get_frame)(i, loom) for i in attractor.frames) + +.. raw:: html + + + Looming in parallel ------------------- diff --git a/examples/lorenz.mp4 b/examples/lorenz.mp4 new file mode 100644 index 0000000..10b875d Binary files /dev/null and b/examples/lorenz.mp4 differ diff --git a/examples/lorenz.py b/examples/lorenz.py new file mode 100644 index 0000000..67c3bc1 --- /dev/null +++ b/examples/lorenz.py @@ -0,0 +1,76 @@ +# 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: + steps_per_frame: 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.steps_per_frame)) + + 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.steps_per_frame]) + 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.steps_per_frame / 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) as loom: + attractor = LorenzPlotter() + attractor.initialize(10000) + Parallel(n_jobs=-1)(delayed(attractor.get_frame)(i, loom) for i in attractor.frames)