Julia set plotting example (5k res, so kinda slow)

julia.py

"""
Copyright (C) 2018 Torin Cooper-Bennun
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.
    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
    
simple Julia set plotting
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm


def iterate_quad_poly(z_0, c, n_max):
    """
    iterates
        z_{k+1} = z_k^2 + c
    while |z_k| < 2, returning the value of k (k_esc) if this condition is breached; iterates a maximum of *n_max*
    times, such that 0 < k_esc <= *n_max*.

    initial parameters *z_0* and *c* are complex numbers.
    """
    z = z_0
    for k in range(n_max):
        if abs(z) >= 2:
            return k
        z = z*z + c
    return n_max


def julia_calc(c=0.1+0.12j, z_bottomleft=-2.5-2.5j, z_topright=2.5+2.5j, img_width=600):
    aspect = (z_topright - z_bottomleft).real / (z_topright - z_bottomleft).imag
    img_height = int(img_width / aspect)

    escape_times = np.zeros((img_height, img_width))
    n_max = 100

    for i_r, R in enumerate(np.linspace(z_bottomleft.real, z_topright.real, img_width)):
        for i_i, I in enumerate(np.linspace(z_bottomleft.imag, z_topright.imag, img_height)):
            print(f"{(i_i + i_r*img_height)/(img_width*img_height)*100:.2f}% complete     ", end="\r")
            escape_times[i_i, i_r] = iterate_quad_poly(R + I*1j, c, n_max)
    print()

    return escape_times


def julia_plot(escape_times, n_max, fname=None):
    img_height, img_width = escape_times.shape

    dpi = 100
    fig, ax = plt.subplots(figsize=(img_width/dpi, img_height/dpi), dpi=dpi)

    ax.tick_params(which="both", left=False, bottom=False, labelleft=False, labelbottom=False)
    ax.set_aspect(1)

    cmap = cm.get_cmap("seismic")
    cmap.set_over("k")

    print("Plotting...")
    ax.imshow(escape_times, cmap=cmap, vmax=n_max-1)
    ax.set_position([0, 0, 1, 1])

    if fname:
        print("Saving...")
        fig.savefig(fname)


if __name__ == "__main__":
    aspect = 1920/1080
    height = 2.4
    width = height*aspect
    z_bl = -width/2 - height/2 *1j
    z_tr = -z_bl

    print("Generating data...")
    data = julia_calc(c=-0.74543+0.11301j, z_bottomleft=z_bl, z_topright=z_tr, img_width=5120)
    np.savetxt("julia_test_big.txt", data)

    julia_plot(data, 100, "julia_5k.png")