Datashader叠加图像遇内存错误，求质数螺旋可视化解决方案

尝试模仿3Blue1Brown的质数螺旋图，用Datashader绘制1-10000的数字螺旋，质数标黄色、非质数标蓝色，叠加图像时出现错误：
Unable to allocate 1.49 TiB for an array with shape (800, 800, 800, 800) and data type uint32

原代码：

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

错误原因 #

直接用*操作符对两个Datashader生成的xarray图像进行运算时，会触发广播机制，将两个形状为(800,800,4)的RGBA数组错误扩展为(800,800,800,800,4)，导致内存占用暴增触发溢出。

解决方案 #

方法1：合并数据一次渲染（推荐，更高效） #

将所有数据放入同一个DataFrame，标记是否为质数，通过分类着色实现，避免多次渲染和叠加：

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

# 生成所有数字的坐标及质数标记
data = []
for num in range(10000):
    x = cos(num) * num / 10000
    y = sin(num) * num / 10000
    is_prime = sympy.isprime(num)
    data.append({"x": x, "y": y, "is_prime": is_prime})

df = pd.DataFrame(data)

# 创建画布
cvs = ds.Canvas(plot_width=800, plot_height=800)

# 按is_prime字段聚合，区分质数/非质数
agg = cvs.points(df, "x", "y", agg=ds.count_cat("is_prime"))

# 设置颜色：非质数蓝色，质数黄色
color_map = {False: "#1f77b4", True: "#ffdd00"}
img = ds.tf.shade(agg, color_key=color_map)

# 导出图像
export_image(img, "prime_spiral", background="white")

方法2：用Datashader Composite工具正确叠加 #

如果必须分开渲染，使用datashader.composite.over函数进行图层叠加，该函数会正确处理RGBA图层的混合，不会触发错误广播：

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

# 生成质数集合
primes = set(i for i in range(10000) if sympy.isprime(i))

# 生成所有数字坐标
all_coords = [(cos(num)*num/10000, sin(num)*num/10000) for num in range(10000)]
# 分离质数和非质数坐标
non_prime_x, non_prime_y = zip(*[coord for idx, coord in enumerate(all_coords) if idx not in primes])
prime_x, prime_y = zip(*[coord for idx, coord in enumerate(all_coords) if idx in primes])

# 创建数据框
non_prime_df = pd.DataFrame({"x": non_prime_x, "y": non_prime_y})
prime_df = pd.DataFrame({"x": prime_x, "y": prime_y})

# 创建画布
cvs = ds.Canvas(plot_width=800, plot_height=800)

# 分别聚合并着色
agg_non_prime = cvs.points(non_prime_df, "x", "y")
img_non_prime = ds.tf.shade(agg_non_prime, cmap=["#1f77b4"])

agg_prime = cvs.points(prime_df, "x", "y")
img_prime = ds.tf.shade(agg_prime, cmap=["#ffdd00"])

# 使用over函数叠加图层（质数图层放在非质数图层上方）
final_img = ops.over(img_prime, img_non_prime)

# 导出图像
export_image(final_img, "prime_spiral", background="white")

关键说明 #

• 方法1通过count_cat聚合分类数据，一次渲染即可区分颜色，是Datashader推荐的高效方式，避免了图层叠加的额外开销。
• 方法2中ops.over会正确处理两个RGBA图像的像素级混合，不会触发错误的广播运算，解决了内存溢出问题。

内容的提问来源于stack exchange，提问作者Robertmg