Rewriting attractors

Some context

It was around 2019 when I first started working on attractors. I was in my second year of college, and I had a lot of free time to work on side projects, thanks to Covid. I was also learning about chaos theory and strange attractors, and I thought it would be a good idea to combine my interest in programming with my newfound interest in chaos theory. So I created attractors, and it was a really fun project to work on. I didn't worry too much about the code aspect of it per se, rather just wanted to implement the systems and the ODE solvers I read about in the books. As long as it worked, I was happy. And it worked really well. I was able to generate some really impressive visualizations!

The gallery of attractors created by that package can be found here (they have been slightly post-processed in Illustrator), and the last v1 version of the package itself can be found here, just in case you feel the need to see the older version.

Why a rewrite?

It has been a few years since I wrote it, and I really had no reason to go back to it. attractors was not a dependency or a serious project that I had to maintain. It was essentially, for all intents and purposes of mine, code complete. But there was this lingering thought in my head that I could do better. I could write better code, I could write better documentation, I could write better tests, I could write better examples, I could write better everything. Useful? Maybe not. But it is an exercise in self-improvement, and I think that's what matters. I wanted to see how much I have grown and learnt as a developer in the past few years owing to my newfound work experience, and what better way to test that than to rewrite a project that I wrote when I was a beginner. Add to this the fact that AI has improved to such an extent that I can rewrite the entire package in a fraction of the time it took me to write the first version. I dont have to manually migrate each and every system, I can just ask Claude to do that for me, and helped me it did.

First, to decide the direction of the rewrite, I tested out a couple stuff especially to speed up the ODE solving, which used to be just a standard python loop process. Initally, I went straight to experimenting with Cython, and it was probably the fastest even compared to what I eventually settled with, but I felt it wasn't that easy to work with, and detracts from my other goal that I set for this rewrite: simplicity. I wanted to make the package as simple as possible, and Cython was not that. So I went with Numba, which was a breeze to work with, and the performance was not that far off from Cython. Numba did suffer from an overhead initially and Cython was eventually much faster, but I felt the simplicity of Numba was worth the tradeoff.

With that being decided, I felt that I was ready to migrate the entire project, and for that I needed a new vision, architecture so to say.

The new vision

There was some initial back and forth between my conflicting minds on what structure to go with. I had a few options in mind. Initially, I considered using the same structure as the previous version, but with better code. This one had the least amount of migrating work, but I felt that adding new systems would be a pain, as it was hardcoded insided a class, each system being an instance method (I know, I know, I was a beginner back then).

I also thought about creating a new class for each system and attractor, inherited from a base class, but numba immediately threw a fit at me, saying that it doesn't support class methods (valid). So I had to scrap that idea.

So I had to look at factory patterns, builder patterns, etc. And I settled with a registry pattern, which looked something like this:

@system.register('lorenz')
def lorenz(state, params):
    ...

@solver.register('rk4')
def rk4(system_func, state, params):
    ...

This helped me keep the code functional (which numba prefers), and also allowed me to add new systems easily. Obviously, I needed to write a registry class, and registry happens at runtime, but I felt that was a small price to pay for the flexibility it offered. And using such a system naturally gave in to a project structure that was more modular, and I could easily add new systems and solvers without much hassle, and looked something like this:

attractors/
├── solvers/
│   ├── registry.py
│   ├── rk4.py
│   ├── ...
├── systems/
│   ├── registry.py
│   ├── lorenz.py
│   ├── ...

So the only thing I need now is to create the registries, and an integrator which would take care of the ODE solving, and I was good to go. It took me some time to get there, because I was too used to an object-oriented approach, and working purely with functions was a bit of a challenge. Dealing with numba was especially challenging, making sure i jit-compiled the functions properly, and that I didn't use any unsupported features. But I got there eventually, and it was time to test out the first attractor.

Migrating and validating

After some minor restructuring to make sure the registries do indeed register, and other tinkering, this is what I came up with:

from attractors import SolverRegistry, StaticPlotter, SystemRegistry, ThemeManager, integrate_system

system = SystemRegistry.get("lorenz")
solver = SolverRegistry.get("rk4")

trajectory, time = integrate_system(system, solver, steps=1000000, dt=0.001)

It looked clean, much better than the v1 interface and it worked! Quite fast in fact, that I thought this was a good time to start migrating the rest of the systems and solvers. This was a menial task as I know the exact way of converting all scripts, so I instead sought the help of my favorite assistance, Sonnet 3.6. I came up with a prompt to first migrate then validate each and every system and solver from the old single class files. I dont have the exact migration prompt handy, but it was similar to the validation prompt which I do have here:

# Attractor Migration Verification

For each attractor system, verify the following matches between old (base.py + params.json) and new implementation:
## Equations
- [ ] Match exact equation structure
- [ ] Correct variables
- [ ] Correct parameter usage
- [ ] Data type (double to float64)

## Parameters
- [ ] Parameter names match
- [ ] Number of parameters match
- [ ] Default parameter values match
- [ ] Parameter ordering matches

## Configuration

- [ ] Initial coordinates match
- [ ] Plot limits match (xlim, ylim, zlim)
- [ ] Reference citation matches
- [ ] System name matches

## Example Output Format:

System: [Name]
- Equations: ✓ Match | ⚠️ Mismatch ([details])
- Parameters:
- Names: [old] -> [new]
- Values: [old] -> [new]
- Count: [old] -> [new]
- Initial Coords: [old] -> [new]
- Plot Limits: [old] -> [new]
- Reference: [old] -> [new]
- Status: ✓ Full Match | ⚠️ Needs Review

It worked pretty well, barring some couple of mistakes and overlooked details, but 90% of the migration was smooth (it does make you aware about the fact that you cannot trust AI blindly, atleast not yet). I was able to migrate all the systems and solvers in a couple of days, and I was ready to benchmark the new version with the old one.

Benchmarking

These were the benchmark scripts I used (and apologies in advance if the figure is unreadable, you might have to open them in a new tab):

1. Speedup across all systems

The results were what I expected when I dry-ran numba for the lorenz system before deciding to migrate.

Roughly, the performance of the new version was as follows:

• Average speedup: 50.51x
• Maximum speedup: 83.56x (Three Cell Cnn)
• Minimum speedup: 39.64x (Burke Shaw)

2. Systems and solvers performance

As one would expect, the relation follows a T ∝ $N^a$, where T is execution time and N is number of points. Both new and old versions show linear scaling in log space, indicating polynomial time complexity. There is a growing gap between implementations, exponential over the number of points used, something that is pretty obvious I would say. There is a notable crossover point at low N, where the old version actually performs better. This is likely due to initialization overhead of numba, but this is an edge case.

Overall, it should be fairly obvious that the new version is much faster than the old one!

Plotting and Theming

Moving the themes was an interesting challenge, but since I had followed a sort of registry system, I created a similar one for the themes, and like v1, use a prexisting themes.json to load them. Full credit to Windows Terminal Themes for the themes. So now, either existing themes can be obtained (and modifed) or new themes can be added easily, as shown below:

from attractors import ThemeManager

theme = ThemeManager.get("ayu")

# or

theme = Theme(
    name="monochrome",
    background="#FFFFFF",
    colors=["#616161", "#7a7a7a", "#2e2e2e", "#1c1c1c"],
    foreground="#000000",
) # Black foreground for contrast with white background

Plotting has some interesting changes as well. The core concept has been kept same from v1, that is to have a static plotter and an animator, both using matplotlib backends, but the code should be much better now, and gives more control to the user to directly control a lot of the matplotlib arguments. There are also some new interesting additions as well:

• compression_method: allowing to reduce the number of points plotted while maintaining the visual integrity of the attractor.
• color_mapping: allowing to map a color to a specific axis, time, velocity or even a custom ColorMapper class.

Given a trajectory from the previous example, a sample plot can be as follows:

import matplotlib.pyplot as plt
from attractors import StaticPlotter

plotter = StaticPlotter(
    system,
    theme,
    fig_kwargs={"figsize": (16, 9)},
    color_by="time",
).visualize(
    trajectory,
    compression_method="curvature",
    compression=0.8,
    line_kwargs={"linewidth": 1.5, "antialiased": True, "alpha": 0.4},
)

# Optionally set the axis to be invisible
plotter.ax.get_xaxis().set_visible(False)
plotter.ax.get_yaxis().set_visible(False)

plt.tight_layout(pad=0)
plt.show()

Similarly, a sample animation can be as follows:

import matplotlib.pyplot as plt
from attractors import AnimatedPlotter

def rotate_view(ax) -> None:
    ax.view_init(elev=30, azim=(ax.azim + 10) % 360)


# use the AnimatedPlotter class to create an animation
plotter = AnimatedPlotter(
    system,
    theme,
    fig_kwargs={"figsize": (10, 8)},
).visualize(
    trajectory,
    line_kwargs={"linewidth": 2, "antialiased": True, "alpha": 0.9},
    rotate_view=rotate_view,
    anim_kwargs={"blit": True},
)
plt.show()

In my opinion, the new plotting system is much more flexible and powerful than the previous one.

Other changes

Testing remains the same with pytest framework, but just made faster thanks to Claude. The package is also typed with mypy, which was initially a hassle because of numba, so I had to exclude numba specific lines from the type checking. But otherwise, it was a breeze to work with. The documentation is also much better now, thanks to MkDocs and Material for MkDocs (as I'm more comfortable with markdown than reStructuredText).

Final thoughts

I am really happy with the way attractors-v2 turned out. Initially there was a lot of inertia to start the project, but once I got going, it was a breeze. Try it out, and let me know what you think. I am always open to feedback and suggestions. I hope you enjoy using it as much as I enjoyed writing it.

Visit the project here: attractors

in chaos, emerges beauty