koch_snowflake.py

"""
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__ importannotations

importmatplotlib.pyplotasplt
importnumpyasnp

# 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]


defiterate(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_inrange(steps):
vectors=iteration_step(vectors)
returnvectors


defiteration_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= []
fori, start_vectorinenumerate(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])
returnnew_vectors


defrotate(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)))
returnnp.dot(rotation_matrix, vector)


defplot(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__":
importdoctest

doctest.testmod()

processed_vectors=iterate(INITIAL_VECTORS, 5)
plot(processed_vectors)