マンデルブロ集合

mandelbrot.py

import numpy as np
from PIL import Image
from tqdm import tqdm


class Mandelbrot(object):
    def __init__(self, r1, r2, i1, i2, k, l):
        self.r1 = min(r1, r2)
        self.r2 = max(r1, r2)
        self.i1 = min(i1, i2)
        self.i2 = max(i1, i2)
        self.k = k
        self.l = l

    def _initial_matrix(self):
        m1 = np.empty((self.k, self.l), dtype=np.complex256)
        for i in range(self.k):
            m1[i, :].real = (self.r2 - self.r1) * i / self.k + self.r1
        for j in range(self.l):
            m1[:, j].imag = (self.i2 - self.i1) * j / self.l + self.i1
        return m1

    def calc(self, n=1000):
        m = self._initial_matrix()
        c = np.array(m, dtype=m.dtype)
        n_diverge = np.zeros_like(m, dtype=np.int)
        diverged_real = np.zeros_like(m, dtype=np.float)
        diverged_imag = np.zeros_like(m, dtype=np.float)

        for i in tqdm(range(n)):
            m = m * m + c
            absolutes = np.absolute(m)
            diverged = np.logical_and(absolutes > 2, n_diverge == 0)

            diverged_real[diverged] = m[diverged].real
            diverged_imag[diverged] = m[diverged].imag
            n_diverge[diverged] = i

            m[n_diverge > 0] = 0

        return n_diverge, diverged_real, diverged_imag

    def to_image(self, n=1000):
        n_diverge, diverged_real, diverged_imag = self.calc(n)

        max_red = np.max(n_diverge)
        red = n_diverge.astype(np.float64) / np.float64(max_red)

        green = diverged_real * diverged_real
        green /= np.max(green)

        blue = diverged_real * diverged_real
        blue /= np.max(blue)

        # emphasis difference
        red = np.sqrt(red * 100) * 10 / 100
        green = np.sqrt(green * 100) * 10 / 100
        blue = np.sqrt(blue * 100) * 10 / 100

        image_array = np.stack((red, green, blue), axis=2) * 255
        return Image.fromarray(
            image_array.astype(np.uint8).transpose((1, 0, 2)))


def main(args):
    mb = Mandelbrot(
        args.real1,
        args.real2,
        args.imag1,
        args.imag2,
        args.x,
        args.y
    )

    image = mb.to_image(args.iter)
    image.save(args.out)


def get_parser():
    import argparse
    parser = argparse.ArgumentParser("MandelbrotDrawer")

    parser.add_argument("--real1", type=float, default=-2.0)
    parser.add_argument("--imag1", type=float, default=-2.0)
    parser.add_argument("--real2", type=float, default=2.0)
    parser.add_argument("--imag2", type=float, default=2.0)
    parser.add_argument("--iter", type=int, default=500)
    parser.add_argument("--x", type=int, default=512)
    parser.add_argument("--y", type=int, default=512)
    parser.add_argument("--out", default="mb.png")

    return parser


if __name__ == '__main__':
    main(get_parser().parse_args())

参考にしたサイト: http://azisava.sakura.ne.jp/mandelbrot/definition.html