Code supplementing the Mandelbrot Set tutorial

bzaczynski · bzaczynski · commit cf57f26830d5 · 2021-09-19T15:46:18.000+02:00
diff --git a/mandelbrot/01_scatter_plot.py b/mandelbrot/01_scatter_plot.py
@@ -0,0 +1,32 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+np.warnings.filterwarnings("ignore")
+
+
+def complex_matrix(xmin, xmax, ymin, ymax, pixel_density):
+    re = np.linspace(xmin, xmax, int((xmax - xmin) * pixel_density))
+    im = np.linspace(ymin, ymax, int((ymax - ymin) * pixel_density))
+    return re[:, None] + im[None, :] * 1j
+
+
+def is_stable(c, num_iterations):
+    z = 0
+    for _ in range(num_iterations):
+        z = z ** 2 + c
+    return abs(z) < 2
+
+
+def get_members(c, num_iterations):
+    return c[is_stable(c, num_iterations)]
+
+
+if __name__ == "__main__":
+    c = complex_matrix(-2, 0.5, -1.5, 1.5, pixel_density=21)
+    members = get_members(c, num_iterations=20)
+
+    plt.scatter(members.real, members.imag, color="black", marker=",", s=1)
+    plt.gca().set_aspect("equal")
+    plt.axis("off")
+    plt.tight_layout()
+    plt.show()
diff --git a/mandelbrot/02_bw_plot.py b/mandelbrot/02_bw_plot.py
@@ -0,0 +1,31 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+np.warnings.filterwarnings("ignore")
+
+
+def complex_matrix(xmin, xmax, ymin, ymax, pixel_density):
+    re = np.linspace(xmin, xmax, int((xmax - xmin) * pixel_density))
+    im = np.linspace(ymin, ymax, int((ymax - ymin) * pixel_density))
+    return re[:, None] + im[None, :] * 1j
+
+
+def is_stable(c, num_iterations):
+    z = 0
+    for _ in range(num_iterations):
+        z = z ** 2 + c
+    return abs(z) < 2
+
+
+def get_members(c, num_iterations):
+    return c[is_stable(c, num_iterations)]
+
+
+if __name__ == "__main__":
+    c = complex_matrix(-2, 0.5, -1.5, 1.5, pixel_density=1280)
+
+    plt.imshow(is_stable(c, num_iterations=20).T, cmap="binary")
+    plt.gca().set_aspect("equal")
+    plt.axis("off")
+    plt.tight_layout()
+    plt.show()
diff --git a/mandelbrot/03_bw.py b/mandelbrot/03_bw.py
@@ -0,0 +1,20 @@
+from PIL import Image
+
+from mandelbrot_01 import MandelbrotSet
+
+if __name__ == "__main__":
+
+    mandelbrot_set = MandelbrotSet(max_iterations=20)
+
+    width, height = 512, 512
+    scale = 0.0075
+    black_and_white = "1"
+
+    image = Image.new(mode=black_and_white, size=(width, height))
+    for y in range(height):
+        for x in range(width):
+            re = scale * (x - width / 2)
+            im = scale * (height / 2 - y)
+            image.putpixel((x, y), complex(re, im) not in mandelbrot_set)
+
+    image.show()
diff --git a/mandelbrot/04_grayscale.py b/mandelbrot/04_grayscale.py
@@ -0,0 +1,21 @@
+from PIL import Image
+
+from mandelbrot_02 import MandelbrotSet
+
+if __name__ == "__main__":
+
+    mandelbrot_set = MandelbrotSet(max_iterations=20)
+
+    width, height = 512, 512
+    scale = 0.0075
+    grayscale = "L"
+
+    image = Image.new(mode=grayscale, size=(width, height))
+    for y in range(height):
+        for x in range(width):
+            re = scale * (x - width / 2)
+            im = scale * (height / 2 - y)
+            probability = 1 - mandelbrot_set.probability(complex(re, im))
+            image.putpixel((x, y), int(probability * 255))
+
+    image.show()
diff --git a/mandelbrot/05_grayscale_smooth.py b/mandelbrot/05_grayscale_smooth.py
@@ -0,0 +1,22 @@
+from PIL import Image
+
+from mandelbrot_03 import MandelbrotSet
+
+if __name__ == "__main__":
+
+    mandelbrot_set = MandelbrotSet(max_iterations=20)
+
+    width, height = 512, 512
+    scale = 0.0075
+    grayscale = "L"
+
+    image = Image.new(mode=grayscale, size=(width, height))
+    for y in range(height):
+        for x in range(width):
+            re = scale * (x - width / 2)
+            im = scale * (height / 2 - y)
+            c = complex(re, im)
+            probability = 1 - mandelbrot_set.probability(c, smooth=True)
+            image.putpixel((x, y), max(0, min(int(probability * 255), 255)))
+
+    image.show()
diff --git a/mandelbrot/06_viewport_aspect.py b/mandelbrot/06_viewport_aspect.py
@@ -0,0 +1,15 @@
+from PIL import Image
+
+from mandelbrot_03 import MandelbrotSet
+from viewport_01 import Viewport
+
+if __name__ == "__main__":
+
+    mandelbrot_set = MandelbrotSet(max_iterations=20)
+
+    image = Image.new(mode="1", size=(512, 512), color=1)
+    for pixel in Viewport(image, xmin=-2.5, xmax=1):
+        if complex(pixel) in mandelbrot_set:
+            pixel.color = 0
+
+    image.show()
diff --git a/mandelbrot/07_viewport_center.py b/mandelbrot/07_viewport_center.py
@@ -0,0 +1,19 @@
+from PIL import Image
+from PIL import ImageEnhance
+
+from mandelbrot_03 import MandelbrotSet
+from viewport_02 import Viewport
+
+if __name__ == "__main__":
+    print("This might take a while...")
+
+    mandelbrot_set = MandelbrotSet(max_iterations=256)
+
+    image = Image.new(mode="L", size=(512, 512))
+    for pixel in Viewport(image, center=-0.7435 + 0.1314j, width=0.002):
+        c = complex(pixel)
+        probability = 1 - mandelbrot_set.probability(c, smooth=True)
+        pixel.color = max(0, min(int(probability * 255), 255))
+
+    enhancer = ImageEnhance.Brightness(image)
+    enhancer.enhance(1.25).show()
diff --git a/mandelbrot/08_color_palette_matplotlib.py b/mandelbrot/08_color_palette_matplotlib.py
@@ -0,0 +1,29 @@
+import matplotlib.cm
+from PIL import Image
+
+from mandelbrot_03 import MandelbrotSet
+from viewport_02 import Viewport
+
+
+def paint(mandelbrot_set, viewport, palette, smooth):
+    for pixel in viewport:
+        probability = mandelbrot_set.probability(complex(pixel), smooth)
+        index = int(min(probability * len(palette), len(palette) - 1))
+        pixel.color = palette[index % len(palette)]
+
+
+def denormalize(palette):
+    return [tuple(int(channel * 255) for channel in color) for color in palette]
+
+
+if __name__ == "__main__":
+    print("This might take a while...")
+
+    colormap = matplotlib.cm.get_cmap("twilight").colors
+    palette = denormalize(colormap)
+
+    mandelbrot_set = MandelbrotSet(max_iterations=512)
+    image = Image.new(mode="RGB", size=(512, 512))
+    viewport = Viewport(image, center=-0.7435 + 0.1314j, width=0.002)
+    paint(mandelbrot_set, viewport, palette, smooth=True)
+    image.show()
diff --git a/mandelbrot/09_color_palette_custom.py b/mandelbrot/09_color_palette_custom.py
@@ -0,0 +1,30 @@
+from PIL import Image
+
+from mandelbrot_03 import MandelbrotSet
+from viewport_02 import Viewport
+
+
+def paint(mandelbrot_set, viewport, palette, smooth):
+    for pixel in viewport:
+        probability = mandelbrot_set.probability(complex(pixel), smooth)
+        index = int(min(probability * len(palette), len(palette) - 1))
+        pixel.color = palette[index % len(palette)]
+
+
+def denormalize(palette):
+    return [tuple(int(channel * 255) for channel in color) for color in palette]
+
+
+if __name__ == "__main__":
+    print("This might take a while...")
+
+    exterior = [(1, 1, 1)] * 50
+    interior = [(1, 1, 1)] * 5
+    gray_area = [(1 - i / 44,) * 3 for i in range(45)]
+    palette = denormalize(exterior + gray_area + interior)
+
+    mandelbrot_set = MandelbrotSet(max_iterations=20)
+    image = Image.new(mode="RGB", size=(512, 512))
+    viewport = Viewport(image, center=-0.75, width=3.5)
+    paint(mandelbrot_set, viewport, palette, smooth=True)
+    image.show()
diff --git a/mandelbrot/10_color_gradient.py b/mandelbrot/10_color_gradient.py
@@ -0,0 +1,46 @@
+import numpy as np
+from PIL import Image
+from scipy.interpolate import interp1d
+
+from mandelbrot_03 import MandelbrotSet
+from viewport_02 import Viewport
+
+
+def paint(mandelbrot_set, viewport, palette, smooth):
+    for pixel in viewport:
+        probability = mandelbrot_set.probability(complex(pixel), smooth)
+        index = int(min(probability * len(palette), len(palette) - 1))
+        pixel.color = palette[index % len(palette)]
+
+
+def denormalize(palette):
+    return [tuple(int(channel * 255) for channel in color) for color in palette]
+
+
+def make_gradient(colors, interpolation="linear"):
+    X = [i / (len(colors) - 1) for i in range(len(colors))]
+    Y = [[color[i] for color in colors] for i in range(3)]
+    channels = [interp1d(X, y, kind=interpolation) for y in Y]
+    return lambda x: [np.clip(channel(x), 0, 1) for channel in channels]
+
+
+if __name__ == "__main__":
+    print("This might take a while...")
+
+    black = (0, 0, 0)
+    blue = (0, 0, 1)
+    maroon = (0.5, 0, 0)
+    navy = (0, 0, 0.5)
+    red = (1, 0, 0)
+
+    colors = [black, navy, blue, maroon, red, black]
+    gradient = make_gradient(colors, interpolation="cubic")
+
+    num_colors = 256
+    palette = denormalize([gradient(i / num_colors) for i in range(num_colors)])
+
+    mandelbrot_set = MandelbrotSet(max_iterations=25)
+    image = Image.new(mode="RGB", size=(512, 512))
+    viewport = Viewport(image, center=-0.75, width=3.5)
+    paint(mandelbrot_set, viewport, palette, smooth=True)
+    image.show()
diff --git a/mandelbrot/11_color_hsb.py b/mandelbrot/11_color_hsb.py
@@ -0,0 +1,32 @@
+from PIL import Image
+from PIL.ImageColor import getrgb
+
+from mandelbrot_03 import MandelbrotSet
+from viewport_02 import Viewport
+
+
+def hsb(hue_degrees: int, saturation: float, brightness: float):
+    return getrgb(
+        f"hsv({hue_degrees % 360}," f"{saturation * 100}%," f"{brightness * 100}%)"
+    )
+
+
+if __name__ == "__main__":
+    print("This might take a while...")
+
+    mandelbrot_set = MandelbrotSet(max_iterations=25)
+    image = Image.new(mode="RGB", size=(512, 512))
+    for pixel in Viewport(image, center=-0.75, width=3.5):
+        p = mandelbrot_set.probability(complex(pixel), smooth=True)
+        probability = max(0, min(p, 1))
+        pixel.color = (
+            (0, 0, 0)
+            if probability == 1
+            else hsb(
+                hue_degrees=int(probability * 360),
+                saturation=probability,
+                brightness=1,
+            )
+        )
+
+    image.show()
diff --git a/mandelbrot/README.md b/mandelbrot/README.md
@@ -0,0 +1,26 @@
+# Draw the Mandelbrot Set in Python
+
+Code snippets and Python scripts supplementing the [Draw the Mandelbrot Set in Python](https://realpython.com/mandelbrot-set-python) tutorial on [Real Python](https://realpython.com/).
+
+## Installation
+
+This code requires the following libraries:
+
+- Matplotlib
+- NumPy
+- SciPy
+- Pillow
+
+To install them, activate your [virtual environment](https://realpython.com/python-virtual-environments-a-primer/), and type the following command:
+
+```
+$ python -m pip install -r requirements.txt
+```
+
+## Usage
+
+Run a given script whose name starts with a number, for example:
+
+```python
+$ python 01_scatter_plot.py
+```
diff --git a/mandelbrot/mandelbrot_01.py b/mandelbrot/mandelbrot_01.py
@@ -0,0 +1,15 @@
+# mandelbrot_01.py
+
+from typing import NamedTuple
+
+
+class MandelbrotSet(NamedTuple):
+    max_iterations: int
+
+    def __contains__(self, c: complex) -> bool:
+        z = 0
+        for _ in range(self.max_iterations):
+            z = z ** 2 + c
+            if abs(z) >= 2:
+                return False
+        return True
diff --git a/mandelbrot/mandelbrot_02.py b/mandelbrot/mandelbrot_02.py
@@ -0,0 +1,21 @@
+# mandelbrot_02.py
+
+from typing import NamedTuple
+
+
+class MandelbrotSet(NamedTuple):
+    max_iterations: int
+
+    def __contains__(self, c: complex) -> bool:
+        return self.probability(c) == 1
+
+    def probability(self, c: complex) -> float:
+        return self.escape_count(c) / self.max_iterations
+
+    def escape_count(self, c: complex) -> int:
+        z = 0
+        for i in range(self.max_iterations):
+            z = z ** 2 + c
+            if abs(z) >= 2:
+                return i
+        return self.max_iterations
diff --git a/mandelbrot/mandelbrot_03.py b/mandelbrot/mandelbrot_03.py
@@ -0,0 +1,22 @@
+# mandelbrot_03.py
+
+from math import log
+from typing import NamedTuple
+
+
+class MandelbrotSet(NamedTuple):
+    max_iterations: int
+
+    def __contains__(self, c: complex) -> bool:
+        return self.probability(c) == 1
+
+    def probability(self, c: complex, smooth: bool = False) -> float:
+        return self.escape_count(c, smooth) / self.max_iterations
+
+    def escape_count(self, c: complex, smooth: bool = False) -> int:
+        z = 0
+        for i in range(self.max_iterations):
+            z = z ** 2 + c
+            if abs(z) >= 2:
+                return i + 1 - log(log(abs(z))) / log(2) if smooth else i
+        return self.max_iterations
diff --git a/mandelbrot/requirements.txt b/mandelbrot/requirements.txt
diff --git a/mandelbrot/viewport_01.py b/mandelbrot/viewport_01.py
diff --git a/mandelbrot/viewport_02.py b/mandelbrot/viewport_02.py
