import numpy as np
import matplotlib.pyplot as plt
def mandelbrot( h,w, maxit=50 ):
    """Returns an image of the Mandelbrot fractal of size (h,w)."""

    #y, x = np.ogrid[ -1.4:1.4:h*1j, -2:0.8:w*1j ]

    y = np.linspace(-1.4, 1.4, h).reshape(h, 1)
    x = np.linspace(-2, 0.8, w).reshape(1, w)
    c = x + y * 1j
    z = c
    divtime = np.full(z.shape, maxit, dtype=int)

    for i in range(maxit):
        z = z * z + c
        diverge = z * np.conj(z) > 4          # who is diverging
        div_now = diverge & (divtime==maxit)  # who is diverging now
        divtime[div_now] = i                  # note when
        z[diverge] = 2                        # avoid diverging too much

    return divtime

plt.imshow(mandelbrot(400,400), cmap='plasma')
plt.show()