How to speed up plt.scatter() with colormap?

I want to plot the coordinates (following with x0 and y0)

colors = cm.rainbow(np.linspace(0, 1, len(x0)))
for k in np.arange(len(x0)):
    for i in range(len(x0[k])):
        plt.scatter(x0[k][i], y0[k][i], color=colors[k], s = 2)

This loop working approximately 50000-60000 times.

But it takes time to finish 2 hours. I run this script on my macbook air m1.

I think it very slow to run this script.

Do you have any solutions to speed up this code or another libraries to scatter coordinates ex. ggplot etc.?

This is some information for x0 and y0.

for i in range(10000,10021):
        print(f'{x0[i]} {y0[i]}')

-----------------------------------
Result :
[145.5, 146] [39.5, 39]
[146.5] [39.5]
[147] [39]
[147.5] [39.5]
[148.5] [39.5]
[149.5] [39.5]
[150, 150.5] [39, 39.5]
[151] [39]
[151.5] [39.5]
[152.5] [39.5]
[153] [39]
[153.5, 154] [39.5, 39]
[154.5] [39.5]
[155, 155.5] [39, 39.5]
[156] [39]
[156.5] [39.5]
[157] [39]
[157.5] [39.5]
[158.5] [39.5]
[159] [39]
[159.5] [39.5]

plt.show() should be (Don't have white and black spot):

Solution

Try using pandas to explode() the 2D data/colors into long-form vectors. Then you can replace the loop with a single call to scatter():

import pandas as pd

df = pd.DataFrame({'x0': x0, 'color': cm.rainbow(np.linspace(0, 1, len(x0))).tolist()})
df = df.explode('x0')
df['y0'] = pd.Series(y0).explode()

plt.scatter(df.x0, df.y0, color=df.color, s=2)

The long-form dataframe will look something like this:

      x0    y0                                              color
0  145.5  39.5                               [0.5, 0.0, 1.0, 1.0]
0    146    39                               [0.5, 0.0, 1.0, 1.0]
1  146.5  39.5  [0.40588235294117647, 0.1473016980546375, 0.99...
2    147    39  [0.303921568627451, 0.30315267411304353, 0.988...
3  147.5  39.5  [0.2019607843137255, 0.45124405704532283, 0.97...
4  148.5  39.5  [0.09999999999999998, 0.5877852522924731, 0.95...
5  149.5  39.5  [0.0019607843137254832, 0.7092813076058534, 0....
6    150    39  [0.09607843137254901, 0.8053809193888326, 0.89...
6  150.5  39.5  [0.09607843137254901, 0.8053809193888326, 0.89...
7    151    39  [0.19803921568627447, 0.8896040127307095, 0.85...
...  ...   ...                                                ...