Skip to content

Commit

Permalink
DOC: add a fractal example
Browse files Browse the repository at this point in the history
  • Loading branch information
tacaswell committed Oct 28, 2022
1 parent eafb128 commit fd779c2
Showing 1 changed file with 50 additions and 0 deletions.
50 changes: 50 additions & 0 deletions examples/mandelbrot.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
"""
===================================
(Infinitly) Zoomable Mandelbrot Set
===================================
"""

import matplotlib.pyplot as plt
import numpy as np

from data_prototype.wrappers import ImageWrapper
from data_prototype.containers import FuncContainer

from matplotlib.colors import Normalize

maxiter = 75


def mandelbrot_set(X, Y, maxiter, *, horizon=3, power=2):
C = X + Y[:, None] * 1j
N = np.zeros_like(C, dtype=int)
Z = np.zeros_like(C)
for n in range(maxiter):
I = abs(Z) < horizon
N += I
Z[I] = Z[I] ** power + C[I]
N[N == maxiter] = -1
return Z, N


fc = FuncContainer(
{},
xyfuncs={
"xextent": ((2,), lambda x, y: [x[0], x[-1]]),
"yextent": ((2,), lambda x, y: [y[0], y[-1]]),
"image": (("N", "M"), lambda x, y: mandelbrot_set(x, y, maxiter)[1]),
},
)
cmap = plt.get_cmap()
cmap.set_under("w")
im = ImageWrapper(fc, norm=Normalize(0, maxiter), cmap=cmap)

fig, ax = plt.subplots()
ax.add_artist(im)
ax.set_xlim(-1, 1)
ax.set_ylim(-1, 1)
ax.set_aspect("equal")
fig.colorbar(im)
plt.show()

0 comments on commit fd779c2

Please sign in to comment.