Рисуем Мандельброта с помощью Python и Numpy

Что уж говорить: страничка Sunandstuff про фракталы, мелькающая сегодня целый день в моей ленте, и правда красивая (хотя объяснения немного хромают). У меня с этим сюжетом особые отношения: мне в детстве в руки попала книжка «Красота фракталов», в которой, кроме кучи непонятных формул (до сих пор их не понимаю), были ещё очень красивые картинки, а также — самое главное — короткий алгоритм, который позволял эти картинки строить самому. Это казалось мне чудом: программа в 10 строк рисует бесконечно сложные изображения! С тех пор я как одержимый программирую фрактал Мандельброта на всём, что движется. Сначала это был Basic на 486-м компьютере (на одну картинку уходила ровно ночь), потом я выучил Pascal (и с удивлением обнаружил, что он работает раз в сто быстрее бейсика!), потом C, потом даже пытался выучить ассемблер, чтобы рисовать фракталы ещё быстрее (не преуспел, впрочем). Картинка с Мандельбротом, построенном на калькуляторе, украшает мою первую статью в «Компьютерре». Потом я узнал, что по-научному та область, в которой появляется множества Мандельброта и Жюлиа, называется «голоморфной динамикой», и кроме красивых картинок там есть ещё и красивые теоремы, а также до сих пор неразгаданные загадки. Впрочем, своих работ в этой области у меня нет, хоть я и занимаюсь близкими динамическими вещами. Что, однако, не помешает мне рассказать, как строить эти картинки с помощью Python.

Описание алгоритма

Итак, план такой. Пусть $c$ — некоторое комплексное число. Рассмотрим последовательность чисел $z_0, z_1, z_2, \ldots$, которая строится следующим образом:

$$z_0 = 0,\quad z_{k+1}=z_k^2 + c, \quad k = 0, 1, 2, \ldots$$

На каждом шаге мы берём предыдущее число, возводим в квадрат и прибавляем $c$. В зависимости от значения $c$, последовательность чисел $\{z_k\}$ может быть ограниченной или неограниченной. Если она является ограниченной, мы говорим, что $c$ принадлежит множеству Мандельброта $M$.

Поскольку число $c$ комплексное, у него есть вещественная и мнимая части. Каждое комплексное число задаётся точкой декартовой плоскости: по горизонтальной координате будем откладывать вещественную часть, а по вертикальной — мнимую. Таким образом, множество $M$ является множеством на вещественной плоскости.

Наивная реализация: много циклов

Python умеет работать с комплексными числами, так что запрограммировать этот алгоритм совсем легко. Давайте попробуем!

In [1]:
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline 
%config InlineBackend.figure_format='retina'
# библиотеки
In [2]:
# инициализиация

pmin, pmax, qmin, qmax = -2.5, 1.5, -2, 2
# пусть c = p + iq и p меняется в диапазоне от pmin до pmax,
# а q меняется в диапазоне от qmin до qmax

ppoints, qpoints = 200, 200
# число точек по горизонтали и вертикали

max_iterations = 300
# максимальное количество итераций

infinity_border = 10
# если ушли на это расстояние, считаем, что ушли на бесконечность
In [3]:
image = np.zeros((ppoints, qpoints))
# image — это двумерный массив, в котором будет записана наша картинка
# по умолчанию он заполнен нулями

for ip, p in enumerate(np.linspace(pmin, pmax, ppoints)):
    for iq, q in enumerate(np.linspace(qmin, qmax, qpoints)):
        c = p + 1j * q
        # буквой j обозначается мнимая единица: чтобы Python понимал, что речь
        # идёт о комплексном числе, а не о переменной j, мы пишем 1j
        
        z = 0
        for k in range(max_iterations):
            z = z**2 + c
            # Самая Главная Формула
            
            if abs(z) > infinity_border:
                # если z достаточно большое, считаем, что последовательость
                # ушла на бесконечность
                # или уйдёт
                # можно доказать, что infinity_border можно взять равным 4
                
                image[ip,iq] = 1
                # находимся вне M: отметить точку как белую
                break
plt.xticks([])
plt.yticks([])
# выключим метки на осях

plt.imshow(-image.T, cmap='Greys')
# транспонируем картинку, чтобы оси были направлены правильно
# перед image стоит знак минус, чтобы множество Мандельброта рисовалось
# чёрным цветом
Out[3]:
<matplotlib.image.AxesImage at 0x105f697f0>

Картинка похожа, но далеко не такая красивая, какими обычно рисуют фракталы. Сейчас чёрным цветом закрашены точки множества Мандельброта, а белым — все остальные точки. Чтобы получить красивую картинку, нужно закрашивать точки, не входящие в множество Мандельброта, разными цветами, в зависимости от скорости «ухода на бесконечность».

In [4]:
image = np.zeros((ppoints, qpoints))
# image — это двумерный массив, в котором будет записана наша картинка
# по умолчанию он заполнен нулями

for ip, p in enumerate(np.linspace(pmin, pmax, ppoints)):
    for iq, q in enumerate(np.linspace(qmin, qmax, qpoints)):
        c = p + 1j * q
        # буквой j обозначается мнимая единица: чтобы Python понимал, что речь
        # идёт о комплексном числе, а не о переменной j, мы пишем 1j
        
        z = 0
        for k in range(max_iterations):
            z = z**2 + c
            # Самая Главная Формула
            
            if abs(z) > infinity_border:
                image[ip,iq] = k
                break
plt.xticks([])
plt.yticks([])
# выключим метки на осях

plt.imshow(-image.T, cmap='flag')
# транспонируем картинку, чтобы оси были направлены правильно
# перед image стоит знак минус, чтобы множество Мандельброта рисовалось
# чёрным цветом
# параметр cmap задаёт палитру
Out[4]:
<matplotlib.image.AxesImage at 0x10655c198>

Это уже похоже на правду! Однако, наблюдается проблема: код работает довольно медленно. Быстрее, конечно, чем на Basic на 486 компьютере, но всё-таки непростительно медленно.

In [5]:
%%timeit

image = np.zeros((ppoints, qpoints))
# image — это двумерный массив, в котором будет записана наша картинка
# по умолчанию он заполнен нулями

for ip, p in enumerate(np.linspace(pmin, pmax, ppoints)):
    for iq, q in enumerate(np.linspace(qmin, qmax, qpoints)):
        c = p + 1j * q
        # буквой j обозначается мнимая единица: чтобы Python понимал, что речь
        # идёт о комплексном числе, а не о переменной j, мы пишем 1j
        
        z = 0
        for k in range(max_iterations):
            z = z**2 + c
            # Самая Главная Формула
            
            if abs(z) > infinity_border:
                image[ip,iq] = k
                break
1 loop, best of 3: 859 ms per loop

Почти целая секунда! И это только картинка 200 на 200 пикселей! Если я захочу нарисовать картинку 1000 на 1000 пикселей, потребуется в 25 раз больше времени. А если я захочу сделать анимацию? (А я захочу!)

Реализация на базе numpy

Проблема в том, что Python очень медленно работает с циклами. Поэтому хорошо бы от них избавиться. Для этого можно использовать массивы numpy. Вместо того, чтобы делать три вложенных цикла, мы создадим двумерный массив, в котором будут содержаться все значения c, которые мы хотим обработать, и будем вычислять для каждого из них последовательность $\{z_k\}$ параллельно.

In [6]:
%%timeit

image = np.zeros((ppoints, qpoints))
# image — это двумерный массив, в котором будет записана наша картинка
# по умолчанию он заполнен нулями

p, q = np.mgrid[pmin:pmax:(ppoints*1j), qmin:qmax:(qpoints*1j)]
# np.mgrid создаёт сетку значений p и q, ppoints*1j здесь означает
# что мы хотим получить ppoints точек — это такая магия

c = p + 1j*q
z = np.zeros_like(c)
# теперь c и z — это двумерные матрицы

for k in range(max_iterations):
    z = z**2 + c
    # Самая Главная Формула осталась без изменений
    # но в данном случае операции производятся с матрицами
    # и действуют поэлементно
    
    mask = (np.abs(z) > infinity_border) & (image == 0)
    # это означает следующее: мы находим все ячейки в матрице z, 
    # у которых модуль очень большой, и одновременно в соответствующей
    # ячейке в матрице image находится ноль
    
    image[mask] = k
    # заносим все найденные ячейки в image значение k
    # это аналог оператора if из предыдущей версии кода
    
    z[mask] = np.nan
    # те ячейки, про которые мы уже понимаем, что там 
    # z «ушло на бесконечность», мы не будем дальше обрабатывать
    # для этого вносим в них специальное значение np.nan
    # это поможет нам избежать ошибок переполнения
1 loop, best of 3: 216 ms per loop
/usr/local/lib/python3.5/site-packages/ipykernel/__main__.py:276: RuntimeWarning: invalid value encountered in greater

Вот это другое дело — уже можно жить. Посмотрим, что получилось:

In [7]:
plt.xticks([])
plt.yticks([])
plt.imshow(-image.T, cmap='flag')
Out[7]:
<matplotlib.image.AxesImage at 0x106ae9b00>

Теперь давайте сделаем картинку побольше. Чтобы не копировать каждый раз весь код, оформим его в функцию.

In [8]:
def mandelbrot(pmin, pmax, ppoints, qmin, qmax, qpoints, 
               max_iterations=200, infinity_border=10):
    image = np.zeros((ppoints, qpoints))
    p, q = np.mgrid[pmin:pmax:(ppoints*1j), qmin:qmax:(qpoints*1j)]
    c = p + 1j*q
    z = np.zeros_like(c)
    for k in range(max_iterations):
        z = z**2 + c
        mask = (np.abs(z) > infinity_border) & (image == 0)
        image[mask] = k
        z[mask] = np.nan
    return -image.T
In [9]:
plt.figure(figsize=(10, 10))
image = mandelbrot(-2.5, 1.5, 1000, -2, 2, 1000)
plt.xticks([])
plt.yticks([])
plt.imshow(image, cmap='flag', interpolation='none')
/usr/local/lib/python3.5/site-packages/ipykernel/__main__.py:9: RuntimeWarning: invalid value encountered in greater
Out[9]:
<matplotlib.image.AxesImage at 0x10654dc50>

Поиграем с палитрой

In [10]:
from itertools import cycle
import matplotlib.colors as clr
colorpoints = [(1-(1-q)**4, c) for q, c in zip(np.linspace(0, 1, 20), 
                                               cycle(['#ffff88', '#000000', 
                                                      '#ffaa00',]))]
cmap = clr.LinearSegmentedColormap.from_list('mycmap', 
                                             colorpoints, N=2048)

# LinearSegmentedColormap создаёт палитру по заданным точкам и заданным цветам
# можете попробовать выбрать другие цвета

plt.figure(figsize=(10, 10))
plt.xticks([])
plt.yticks([])
plt.imshow(image, cmap=cmap, interpolation='none')
Out[10]:
<matplotlib.image.AxesImage at 0x1070649b0>

Мультики!

Следующая ячейка будет считаться довольно долго.

In [21]:
from matplotlib import animation, rc

rc('animation', html='html5')
# отображать анимацию в виде html5 video

fig = plt.figure(figsize=(10, 10))
max_frames = 200
max_zoom = 300
pmin, pmax, qmin, qmax = -2.5, 1.5, -2, 2

images = []
# кэш картинок

def init():
    return plt.gca()

def animate(i):
    if i > max_frames // 2:
        # фаза zoom out, можно достать картинку из кэша
        
        plt.imshow(images[max_frames//2-i], cmap=cmap)
        return
    
    p_center, q_center = -0.793191078177363, 0.16093721735804
    zoom = (i / max_frames * 2)**3 * max_zoom + 1
    scalefactor = 1 / zoom
    pmin_ = (pmin - p_center) * scalefactor + p_center
    qmin_ = (qmin - q_center) * scalefactor + q_center
    pmax_ = (pmax - p_center) * scalefactor + p_center
    qmax_ = (qmax - q_center) * scalefactor + q_center
    image = mandelbrot(pmin_, pmax_, 500, qmin_, qmax_, 500)
    plt.imshow(image, cmap=cmap)
    images.append(image)
    
    # добавить картинку в кэш
    return plt.gca()

animation.FuncAnimation(fig, animate, init_func=init,
                               frames=max_frames, interval=50)
/usr/local/lib/python3.5/site-packages/ipykernel/__main__.py:9: RuntimeWarning: invalid value encountered in greater
Out[21]:

На сегодня всё! Дальше экспериментируйте сами :)

Комментарии