draw mandelbrot set fractal in python

To draw the Mandelbrot set fractal in Python, we can use the following code:

main.py
import numpy as np
import matplotlib.pyplot as plt

def mandelbrot_set(xmin, xmax, ymin, ymax, Nx, Ny, maxiter=200):
    """
    Return image of the Mandelbrot set on the complex plane.
    """
    x, y = np.linspace(xmin, xmax, Nx), np.linspace(ymin, ymax, Ny)
    c = x[:,np.newaxis] + 1j*y[np.newaxis,:]
    z = c.copy()
    image = np.zeros((Nx, Ny))
    for i in range(maxiter):
        z = z**2 + c
        mask = np.abs(z) < 2
        image += mask
        z[~mask] = 0
    return image.T

# Draw the Mandelbrot set
xmin, xmax, ymin, ymax = -2, 1, -1.5, 1.5
Nx, Ny = 1000, 1000
maxiter = 200
image = mandelbrot_set(xmin, xmax, ymin, ymax, Nx, Ny, maxiter)

# Visualize the Mandelbrot set
plt.figure(figsize=(10,10))
plt.imshow(image, cmap=plt.cm.prism, extent=(xmin, xmax, ymin, ymax))
plt.xlabel("Re(c)")
plt.ylabel("Im(c)")
plt.show()
844 chars
31 lines

Explanation:

1. We start by importing numpy and matplotlib.pyplot. We will use numpy to create arrays of complex numbers and matplotlib to visualize the fractal.

2. We define a function mandelbrot_set that takes in the limits of the x and y axes, the size of the grid, and the maximum number of iterations. It returns an array of integers that represents the image of the Mandelbrot set on the complex plane.

3. We create a meshgrid of complex numbers c on the x and y axes using numpy's linspace function.

4. We initialize z as a copy of c.

5. We create an empty image array.

6. We loop over the maximum number of iterations and compute the Mandelbrot set for each pixel by iteratively squaring z and adding c. We keep track of the pixels that have not diverged using a boolean mask.

7. We add the mask to the image array.

8. We set the divergent pixels to zero to avoid computing them in the next iteration.

9. We transpose the image array and return it.

10. We define the limits and size of the plot, and call the mandelbrot_set function to compute the image.

11. We plot the image using imshow and specify the colormap and the extent of the plot.

12. We label the x and y axes and show the plot.

This will generate a colorful image of the Mandelbrot set fractal.