Python/fractals/koch_snowflake.py
Maxim Smolskiy f2246ce7fd
Enable ruff ICN001 rule (#11329)
* Enable ruff ICN001 rule

* [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>
2024-03-28 18:03:23 +00:00

116 lines
4.3 KiB
Python

"""
Description
The Koch snowflake is a fractal curve and one of the earliest fractals to
have been described. The Koch snowflake can be built up iteratively, in a
sequence of stages. The first stage is an equilateral triangle, and each
successive stage is formed by adding outward bends to each side of the
previous stage, making smaller equilateral triangles.
This can be achieved through the following steps for each line:
1. divide the line segment into three segments of equal length.
2. draw an equilateral triangle that has the middle segment from step 1
as its base and points outward.
3. remove the line segment that is the base of the triangle from step 2.
(description adapted from https://en.wikipedia.org/wiki/Koch_snowflake )
(for a more detailed explanation and an implementation in the
Processing language, see https://natureofcode.com/book/chapter-8-fractals/
#84-the-koch-curve-and-the-arraylist-technique )
Requirements (pip):
- matplotlib
- numpy
"""
from __future__ import annotations
import matplotlib.pyplot as plt
import numpy as np
# initial triangle of Koch snowflake
VECTOR_1 = np.array([0, 0])
VECTOR_2 = np.array([0.5, 0.8660254])
VECTOR_3 = np.array([1, 0])
INITIAL_VECTORS = [VECTOR_1, VECTOR_2, VECTOR_3, VECTOR_1]
# uncomment for simple Koch curve instead of Koch snowflake
# INITIAL_VECTORS = [VECTOR_1, VECTOR_3]
def iterate(initial_vectors: list[np.ndarray], steps: int) -> list[np.ndarray]:
"""
Go through the number of iterations determined by the argument "steps".
Be careful with high values (above 5) since the time to calculate increases
exponentially.
>>> iterate([np.array([0, 0]), np.array([1, 0])], 1)
[array([0, 0]), array([0.33333333, 0. ]), array([0.5 , \
0.28867513]), array([0.66666667, 0. ]), array([1, 0])]
"""
vectors = initial_vectors
for _ in range(steps):
vectors = iteration_step(vectors)
return vectors
def iteration_step(vectors: list[np.ndarray]) -> list[np.ndarray]:
"""
Loops through each pair of adjacent vectors. Each line between two adjacent
vectors is divided into 4 segments by adding 3 additional vectors in-between
the original two vectors. The vector in the middle is constructed through a
60 degree rotation so it is bent outwards.
>>> iteration_step([np.array([0, 0]), np.array([1, 0])])
[array([0, 0]), array([0.33333333, 0. ]), array([0.5 , \
0.28867513]), array([0.66666667, 0. ]), array([1, 0])]
"""
new_vectors = []
for i, start_vector in enumerate(vectors[:-1]):
end_vector = vectors[i + 1]
new_vectors.append(start_vector)
difference_vector = end_vector - start_vector
new_vectors.append(start_vector + difference_vector / 3)
new_vectors.append(
start_vector + difference_vector / 3 + rotate(difference_vector / 3, 60)
)
new_vectors.append(start_vector + difference_vector * 2 / 3)
new_vectors.append(vectors[-1])
return new_vectors
def rotate(vector: np.ndarray, angle_in_degrees: float) -> np.ndarray:
"""
Standard rotation of a 2D vector with a rotation matrix
(see https://en.wikipedia.org/wiki/Rotation_matrix )
>>> rotate(np.array([1, 0]), 60)
array([0.5 , 0.8660254])
>>> rotate(np.array([1, 0]), 90)
array([6.123234e-17, 1.000000e+00])
"""
theta = np.radians(angle_in_degrees)
c, s = np.cos(theta), np.sin(theta)
rotation_matrix = np.array(((c, -s), (s, c)))
return np.dot(rotation_matrix, vector)
def plot(vectors: list[np.ndarray]) -> None:
"""
Utility function to plot the vectors using matplotlib.pyplot
No doctest was implemented since this function does not have a return value
"""
# avoid stretched display of graph
axes = plt.gca()
axes.set_aspect("equal")
# matplotlib.pyplot.plot takes a list of all x-coordinates and a list of all
# y-coordinates as inputs, which are constructed from the vector-list using
# zip()
x_coordinates, y_coordinates = zip(*vectors)
plt.plot(x_coordinates, y_coordinates)
plt.show()
if __name__ == "__main__":
import doctest
doctest.testmod()
processed_vectors = iterate(INITIAL_VECTORS, 5)
plot(processed_vectors)