|
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> |
This is probably indicative of bad design but right now the source code for example is replicated four times in the repo. Using the new Lorenz attractor:
examples/directory: https://github.com/ali-ramadhan/matplotloom/blob/main/examples/lorenz.pymatplotloom/docs/index.rst
Lines 351 to 438 in e073617
matplotloom/tests/test_examples.py
Lines 259 to 325 in 7fbde72
The one in the tests is a bit different to speed it up, but otherwise they're all the same. There should be a way to simplify this and make it so we don't forget to add different copies of the same example!