Python Datashader trying to superimpose images onto each other

Question

I'm rather new to datashader and am doing this as an introductory project and I've hit a snag. I'm trying to plot the numbers from 1 to 10,000 in a spiral like the 3b1b video here, but I want the prime numbers to be in yellow and the others to be in blue, in order for a better overall visualization. I did this by trying to superimpose the primes in yellow on top of a previously generated image of all the numbers in a spiral in blue. However, when I do this I get the error: Unable to allocate 1.49 TiB for an array with shape (800, 800, 800, 800) and data type uint32. My code is roughly as follows (I've left the unrelated parts out but I may have missed an import or list):

import sympy
import pandas as pd
from math import sin, cos
import datashader as ds
import datashader.composite as ops
from datashader.utils import export_image


#Generate a list of primes
primes = [i for i in range(10000) if sympy.isprime(i)]

#generate polar x and y coords
allx = [cos(num) * num * 1/10000 for num in range(10000)]
ally = [sin(num) * num * 1/10000 for num in range(10000)]

xp = [cos(num) * num * 1/10000 for num in range(10000)]
yp = [sin(num) * num * 1/10000 for num in range(10000)]

#Store the lists into distinct dataframes
primedf = pd.DataFrame({'Xp': xp, "Yp": yp})
df = pd.DataFrame({'Allx':allx, 'Ally':ally})

#create canvas to plot things on
cvs = ds.Canvas(plot_width=800, plot_height=800)

#create aggregate points, these are what are used to generate the image
agg = cvs.points(df, "Allx", "Ally")

#generate the first image
img = ds.tf.shade(agg)

#Do the same things for the prime sets
aggprime = cvs.points(primedf, "Xp", "Yp")
imgprime = ds.tf.shade(aggprime, cmap="yellow")

#Now superimpose imgprime onto img, this is causing the error
img * imgprime

I've tried a few other methods, including exporting it to a png and reading with iio/cv2, gnuplot, and I did something that I forgot with datashader.composite.source.ops, and none of them seem to work which leads me to believe I am on a completely wrong approach. If anyone would have any advice on how to do this I would be super grateful.

Answer 1

Not sure if anyone cares but I found an answer, if you export the datashades as images and reload them with cv2, you can add the array and it will superimpose them.To export:

export_image(img, "InsertCoolName", background="black", export_path=".")

And then to reload and superimpose:

backgroundimg = cv2.imread('background.png')
foreground = cv2.imread('foreground.png')
backgroundimg = cv2.cvtColor(backgroundimg, cv2.COLOR_BGR2RGB)
foregroundimg = cv2.cvtColor(foregroundimg, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(8,8))
final = np.array(primeimg + backgroundimg)
plt.axis("off")
plt.imshow(final)
plt.title("Prime Distribution Spiral")
plt.savefig('FinalSpiral.png', bbox_inches='tight', pad_inches=0) 
plt.show()

and the output:

Isn't it cute!