I'm playing around with recursive fractals in Python and wanted to add some colour. However, the code that I wrote to make the curve display with a gradient of colour makes the program run much slower than without it. I'm not sure why and I'm still pretty new to Python, so I would appreciate any tips to improve my code.
Here's the code. This specific seed creates the koch curve:
import numpy as np
import matplotlib.pyplot as plt
#------------ plot stuff (i prefer dark mode) ------------------#
fig = plt.figure(figsize = (12,12),facecolor = "black")
ax = plt.axes()
ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_facecolor("black")
ax.spines["bottom"].set_color("white")
ax.spines["top"].set_color("white")
ax.spines["left"].set_color("white")
ax.spines["right"].set_color("white")
ax.tick_params(axis = "x", colors = "white")
ax.tick_params(axis = "y", colors = "white")
#------------------------ fractal function ---------------------#
def fractal(order):
def points(order):
if (order==0):
return np.array([-1 0j,1 0j])
else:
p1 = points(order-1)
p2 = p1
p1 = p1[:-1]
p2 = p2[1:]
dp = p2 - p1
#--------------------- seed ---------------------#
#--- this specific seed creates the koch curve ---#
b = 4 # number of segments in the seed
d1 = dp/3
d2 = (1.5 np.sqrt(3)/2*1j)*dp/3
d3 = 2*dp/3
d4 = dp
#d5 = (optional)
#d6 = (optional)
#d7 = (optional)
#d8 = (optional)
new_points = np.empty(len(p1)*(b 1), dtype=np.complex128)
new_points[::b 1] = p1
new_points[1::b 1] = p1 d1
new_points[2::b 1] = p1 d2
new_points[3::b 1] = p1 d3
new_points[4::b 1] = p1 d4
#new_points[5::b 1] = p1 d5 (optional)
#new_points[6::b 1] = p1 d6 (optional)
#new_points[7::b 1] = p1 d7 (optional)
#new_points[8::b 1] = p1 d8 (optional)
return new_points
return_points = points(order)
x, y = return_points.real, return_points.imag
return x, y
#------------------ plotting the function ------------------#
x, y = fractal(order=5)
m = len(x)
c = 255/m
for i in range(0,m-1):
# rgb values using i to create a gradient:
red = 255 (2.5/m)*c*i*(i-m)
green = 0.5*c*i 127
blue = 150
# convert rgb to hex:
col = f'#{round(red):02x}{round(green):02x}{round(blue):02x}'
plt.plot(x[i:i 2],y[i:i 2], color=col)
#when i remove this block and simply write
# plt.plot(x,y)
#it completes a 10th order koch curve in less than 2 seconds
CodePudding user response：
It is slow because you are plotting 2345 different lines.
To improve performance, you should use matplotlib's LineCollection which allows to create a gradient-colored line. Think of it as a collection of segments (conceptually, this is very similar to what you did in your code).
To assign a color to each segment we are going to set the array parameter: in this case, it will be a sequential array starting from 0 going to 1 with 2345 elements.
We also need to create a custom colormap from which the colors will be calculated starting from the array parameter. To achieve that, we are going to sub-class matplotlib's Colormap class and use your formulas to compute colors.
This approach will be much much faster in comparison to yours:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import Colormap
def get_segments(x, y):
points = np.ma.array((x, y)).T.reshape(-1, 1, 2)
return np.ma.concatenate([points[:-1], points[1:]], axis=1)
#------------ plot stuff (i prefer dark mode) ------------------#
fig = plt.figure(figsize = (12,12),facecolor = "black")
ax = plt.axes()
ax.set_xlim(-1.1,1.1)
ax.set_ylim(-1.1,1.1)
ax.set_facecolor("black")
ax.spines["bottom"].set_color("white")
ax.spines["top"].set_color("white")
ax.spines["left"].set_color("white")
ax.spines["right"].set_color("white")
ax.tick_params(axis = "x", colors = "white")
ax.tick_params(axis = "y", colors = "white")
#------------------------ fractal function ---------------------#
def fractal(order):
def points(order):
if (order==0):
return np.array([-1 0j,1 0j])
else:
p1 = points(order-1)
p2 = p1
p1 = p1[:-1]
p2 = p2[1:]
dp = p2 - p1
#--------------------- seed ---------------------#
#--- this specific seed creates the koch curve ---#
b = 4 # number of segments in the seed
d1 = dp/3
d2 = (1.5 np.sqrt(3)/2*1j)*dp/3
d3 = 2*dp/3
d4 = dp
#d5 = (optional)
#d6 = (optional)
#d7 = (optional)
#d8 = (optional)
new_points = np.empty(len(p1)*(b 1), dtype=np.complex128)
new_points[::b 1] = p1
new_points[1::b 1] = p1 d1
new_points[2::b 1] = p1 d2
new_points[3::b 1] = p1 d3
new_points[4::b 1] = p1 d4
#new_points[5::b 1] = p1 d5 (optional)
#new_points[6::b 1] = p1 d6 (optional)
#new_points[7::b 1] = p1 d7 (optional)
#new_points[8::b 1] = p1 d8 (optional)
return new_points
return_points = points(order)
x, y = return_points.real, return_points.imag
return x, y
#------------------ plotting the function ------------------#
x, y = fractal(order=5)
class MyColorMap(Colormap):
def __init__(self):
pass
def __call__(self, X, alpha=None, bytes=False):
# generate custom colors
m = len(X)
# X is a modified version of what we set on the
# `array` keyword argument.
# X will be bound between [0, 1]. Need to scale it up
# in order to use your formula. We want X going
# from 0 to 2345 (or whatever the number of points)
X = X * m
c = 255 / m
red = 255 (2.5/m)*c*X*(X-m)
green = 0.5*c*X 127
blue = 150 * np.ones_like(X)
alpha = 255 * np.ones_like(X)
# colors need to be an array with shape [N x 4],
# where N is the number of points. The values must be bound
# between [0, 1].
colors = np.stack((red, green, blue, alpha)).T / 255
return colors
m = len(x)
cmap = MyColorMap()
segments = get_segments(x, y)
lc = LineCollection(segments, array=np.linspace(0, 1, m), cmap=cmap)
ax.add_collection(lc)
plt.show()