[FIXED] Why is this way of plotting colourful lines so slow?

Issue

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 

Solution

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()

Answered By - Davide_sd