Files
pyscript/pyscriptjs/examples/fractals.py
Peter W 6f6efa4525 add config file for pre-commit (#235)
* add config file

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add isort

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2022-05-05 15:12:25 -05:00

142 lines
3.8 KiB
Python

from typing import Tuple
import numpy as np
from numpy.polynomial import Polynomial
def mandelbrot(
width: int,
height: int,
*,
x: float = -0.5,
y: float = 0,
zoom: int = 1,
max_iterations: int = 100
) -> np.array:
"""
From https://www.learnpythonwithrune.org/numpy-compute-mandelbrot-set-by-vectorization/.
"""
# To make navigation easier we calculate these values
x_width, y_height = 1.5, 1.5 * height / width
x_from, x_to = x - x_width / zoom, x + x_width / zoom
y_from, y_to = y - y_height / zoom, y + y_height / zoom
# Here the actual algorithm starts
x = np.linspace(x_from, x_to, width).reshape((1, width))
y = np.linspace(y_from, y_to, height).reshape((height, 1))
c = x + 1j * y
# Initialize z to all zero
z = np.zeros(c.shape, dtype=np.complex128)
# To keep track in which iteration the point diverged
div_time = np.zeros(z.shape, dtype=int)
# To keep track on which points did not converge so far
m = np.full(c.shape, True, dtype=bool)
for i in range(max_iterations):
z[m] = z[m] ** 2 + c[m]
diverged = np.greater(
np.abs(z), 2, out=np.full(c.shape, False), where=m
) # Find diverging
div_time[diverged] = i # set the value of the diverged iteration number
m[np.abs(z) > 2] = False # to remember which have diverged
return div_time
def julia(
width: int,
height: int,
*,
c: complex = -0.4 + 0.6j,
x: float = 0,
y: float = 0,
zoom: int = 1,
max_iterations: int = 100
) -> np.array:
"""
From https://www.learnpythonwithrune.org/numpy-calculate-the-julia-set-with-vectorization/.
"""
# To make navigation easier we calculate these values
x_width, y_height = 1.5, 1.5 * height / width
x_from, x_to = x - x_width / zoom, x + x_width / zoom
y_from, y_to = y - y_height / zoom, y + y_height / zoom
# Here the actual algorithm starts
x = np.linspace(x_from, x_to, width).reshape((1, width))
y = np.linspace(y_from, y_to, height).reshape((height, 1))
z = x + 1j * y
# Initialize z to all zero
c = np.full(z.shape, c)
# To keep track in which iteration the point diverged
div_time = np.zeros(z.shape, dtype=int)
# To keep track on which points did not converge so far
m = np.full(c.shape, True, dtype=bool)
for i in range(max_iterations):
z[m] = z[m] ** 2 + c[m]
m[np.abs(z) > 2] = False
div_time[m] = i
return div_time
Range = Tuple[float, float]
def newton(
width: int,
height: int,
*,
p: Polynomial,
a: complex,
xr: Range = (-2.5, 1),
yr: Range = (-1, 1),
max_iterations: int = 100
) -> (np.array, np.array):
""" """
# To make navigation easier we calculate these values
x_from, x_to = xr
y_from, y_to = yr
# Here the actual algorithm starts
x = np.linspace(x_from, x_to, width).reshape((1, width))
y = np.linspace(y_from, y_to, height).reshape((height, 1))
z = x + 1j * y
# Compute the derivative
dp = p.deriv()
# Compute roots
roots = p.roots()
epsilon = 1e-5
# Set the initial conditions
a = np.full(z.shape, a)
# To keep track in which iteration the point diverged
div_time = np.zeros(z.shape, dtype=int)
# To keep track on which points did not converge so far
m = np.full(a.shape, True, dtype=bool)
# To keep track which root each point converged to
r = np.full(a.shape, 0, dtype=int)
for i in range(max_iterations):
z[m] = z[m] - a[m] * p(z[m]) / dp(z[m])
for j, root in enumerate(roots):
converged = (np.abs(z.real - root.real) < epsilon) & (
np.abs(z.imag - root.imag) < epsilon
)
m[converged] = False
r[converged] = j + 1
div_time[m] = i
return div_time, r