feat: Add Lorenz attractor example

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:

docs/index.rst

@@ -348,6 +348,95 @@ matplotloom works out of the box with anything that is built on top of matplotli
      <source src="https://github.com/ali-ramadhan/matplotloom/assets/20099589/dad48d56-73ae-4cb2-ba71-d21855c72215" type="video/mp4">
    </video>
 
+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/>`_).
+
+.. 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
+
+   <video style="width: 100%; height: auto;" controls>
+     <source src="https://github.com/user-attachments/assets/69b02d78-386d-4843-90df-b1a43600cfa5" type="video/mp4">
+   </video>
+
 Looming in parallel
 -------------------
 

examples/lorenz.py

@@ -0,0 +1,76 @@
+# Inspired by <https://docs.makie.org/stable/>
+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)