如何使用Shader在Pygame中高效绘制大量圆形？

核心问题分析 #

你的问题本质是CPU循环绘制大量图形的开销过高——每次调用pygame.draw.circle都会触发CPU到GPU的小批量数据传输，加上Python循环本身的性能损耗，圆形数量上升后必然出现卡顿。你已经用CuPy完成了GPU端的物理计算，最优解是把绘制逻辑也迁移到GPU，通过批量提交数据的方式减少CPU开销。

可行解决方案 #

方案1：CPU端批量优化（快速见效） #

使用pygame.surfarray直接操作像素数组，将圆形绘制转为批量像素运算，避免循环调用draw.circle。但这种方式仍在CPU执行，性能上限有限。

方案2：OpenGL+Shader GPU加速（推荐） #

结合Pygame与OpenGL Shader，直接在GPU端批量绘制圆形，完全发挥GPU并行计算的优势，可轻松支撑十万级以上的圆形绘制。

具体实现代码（方案2） #

替换原脚本的绘制逻辑与初始化部分，以下是完整优化代码：

import pygame, sys, time
from pygame.locals import*
import cupy as cp
import numpy as np
from OpenGL.GL import *
from OpenGL.GL.shaders import compileProgram, compileShader
import ctypes

pygame.init()
R = (800, 600)
# 初始化带OpenGL支持的Pygame窗口
pygame.display.set_mode(R, DOUBLEBUF|OPENGL)

# 设置OpenGL正交投影，匹配Pygame坐标系统（原点左上角）
glViewport(0, 0, R[0], R[1])
glMatrixMode(GL_PROJECTION)
glLoadIdentity()
glOrtho(0, R[0], R[1], 0, -1, 1)
glMatrixMode(GL_MODELVIEW)
glLoadIdentity()

ballsamount = 10000  # 可轻松扩展至十万级
balls = cp.random.random((4, ballsamount)).astype(cp.float32)

# 初始化核函数（保持原逻辑不变）
initballs = cp.RawKernel('''
      extern "C" __global__
      void init(float* ball){
            int tid = blockDim.x * blockIdx.x + threadIdx.x;
            if (tid%4 == 0){
                  ball[tid] = ball[tid]*800.;
            }
            else if (tid%4 == 1){
                  ball[tid] = ball[tid]*600.;
            }
            else if (tid%4 == 2){
                  ball[tid] = 1.;
            }
            else if (tid%4 == 3){
                  ball[tid] = 0.;
            }
      }
''', 'init')

initballs((ballsamount,), (4,), (balls))

# 更新核函数（增加边界检测）
updateballs = cp.RawKernel('''
      extern "C" __global__
      void update(float* ball){
            int tid = blockDim.x * blockIdx.x + threadIdx.x;
            if ((tid/2)%2 == 0){
                  ball[tid] = ball[tid] + ball[tid+2];
                  ball[tid+1] = ball[tid+1] + ball[tid+3];
                  // 边界检测，防止球飞出窗口
                  if (ball[tid] < 0) ball[tid] = 0;
                  if (ball[tid] > 800) ball[tid] = 800;
                  if (ball[tid+1] < 0) ball[tid+1] = 0;
                  if (ball[tid+1] > 600) ball[tid+1] = 600;
            }
      }
''', 'update')

# 编写OpenGL Shader
vertex_shader = """
#version 330 core
layout (location = 0) in vec2 aPos;
uniform vec2 screenSize;

void main() {
    // 转换Pygame坐标至OpenGL标准化设备坐标
    gl_Position = vec4(aPos.x / screenSize.x * 2.0 - 1.0, 
                       1.0 - aPos.y / screenSize.y * 2.0, 
                       0.0, 1.0);
    gl_PointSize = 10.0; // 圆形直径，对应原半径5
}
"""

fragment_shader = """
#version 330 core
out vec4 FragColor;

void main() {
    // 绘制圆形点（裁剪方形点的边缘）
    vec2 center = gl_PointCoord - vec2(0.5);
    float dist = length(center);
    if (dist > 0.5) discard;
    FragColor = vec4(1.0, 1.0, 1.0, 1.0);
}
"""

# 编译Shader程序
shader_program = compileProgram(
    compileShader(vertex_shader, GL_VERTEX_SHADER),
    compileShader(fragment_shader, GL_FRAGMENT_SHADER)
)

# 创建顶点缓冲区对象（VBO）
vbo = glGenBuffers(1)

def draw():
    glClear(GL_COLOR_BUFFER_BIT)
    glUseProgram(shader_program)
    
    # 设置屏幕大小Uniform变量
    screen_loc = glGetUniformLocation(shader_program, "screenSize")
    glUniform2f(screen_loc, R[0], R[1])
    
    # 提取CuPy中的位置数据，转换为Numpy数组后传入GPU
    positions = cp.vstack((balls[0], balls[1])).T.get()
    glBindBuffer(GL_ARRAY_BUFFER, vbo)
    glBufferData(GL_ARRAY_BUFFER, positions.nbytes, positions, GL_DYNAMIC_DRAW)
    
    # 配置顶点属性指针
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 2*4, ctypes.c_void_p(0))
    
    # 批量绘制所有圆形
    glDrawArrays(GL_POINTS, 0, ballsamount)
    
    pygame.display.flip()

# 主循环（保持原逻辑不变）
while True:
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()
        if event.type == KEYDOWN:
            if event.key == K_ESCAPE:
                pygame.quit()
                sys.exit()
    updateballs((ballsamount,), (4,), (balls))
    cp.cuda.Stream.null.synchronize()
    starttime = time.time()
    draw()
    print(f"绘制耗时: {time.time()-starttime:.6f}秒")

关键优化说明 #

1. Shader并行绘制：片元着色器通过计算像素到点中心的距离裁剪出圆形，GPU并行处理所有像素，效率远超CPU循环绘制。
2. 批量数据传输：一次性将所有圆形的位置数据传入GPU，避免循环传输单条数据的开销。
3. 坐标系统适配：通过顶点着色器将Pygame的屏幕坐标转换为OpenGL标准化设备坐标，确保位置显示正确。
4. 边界检测：在GPU端的更新核函数中增加边界检测，防止圆形飞出窗口。

进阶优化建议 #

若需支持超大规模圆形（如100万+），可使用CUDA与OpenGL内存互操作，直接在GPU内部传递数据，避免CuPy到Numpy的内存拷贝，进一步提升性能。

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