Newtonův fraktál

newton.py

from sys import maxsize

from PIL import Image

from python.less7.madelbrot import rgb


def newt_1(z):
    return z ** 3 - 1


def newt_2(z):
    return z ** 3 - z


def newton_fractal(width, steps, zoom, complex_function=newt_1, name="newton.png"):
    real_denom = abs(zoom[1][0] - zoom[0][0])
    imag_denom = abs(zoom[1][1] - zoom[0][1])

    box_ratio = imag_denom / real_denom

    height = round(width * box_ratio)

    im = Image.new("RGB", (width, height), "white")

    eps = 1e-5

    max_s = 0
    min_s = maxsize

    points = []

    for x in range(width):
        for y in range(height):
            real = x / (width / real_denom) + zoom[0][0]
            imag = y / (height / imag_denom) + zoom[0][1]
            zn = complex(real, imag)
            i = 0
            for i in range(steps):
                if zn ** 2 == 0:
                    break

                new_z = zn - complex_function(zn) / (3 * zn ** 2)
                if abs(zn - new_z) < eps:  # stop when close enough to any root
                    break
                zn = new_z

            min_s = min(min_s, i)
            max_s = max(max_s, i)

            points.append(((x, y), i))

    for p, i in points:
        im.putpixel(p, rgb(min_s, max_s, i))

    im.save(name)

# newton_fractal(1000, 30, ((-2, -2), (2, 2)), newt_1, name="newton_1.png")
# newton_fractal(1000, 30, ((-2, -2), (2, 2)), newt_2, name="newton_2.png")