Use compute shader instead of numpy vectorization to significantly improve speed and remove numpy dependency

game/game_of_life.py

@@ -1,27 +1,90 @@
 from utils.constants import ROWS, COLS
-import numpy as np
-                   
-def create_numpy_grid():
-    return np.zeros((ROWS, COLS), dtype=np.uint8)
+from pyglet.gl import glBindBufferBase, GL_SHADER_STORAGE_BUFFER, GL_NEAREST
+import pyglet
 
-def count_neighbors(grid):
-    padded = np.pad(grid, pad_width=1, mode='constant', constant_values=0)
-    
-    neighbors = (
-        padded[0:-2, 0:-2] +  # top-left
-        padded[0:-2, 1:-1] +  # top
-        padded[0:-2, 2:] +    # top-right
-        padded[1:-1, 0:-2] +  # left
-        padded[1:-1, 2:] +    # right
-        padded[2:, 0:-2] +    # bottom-left
-        padded[2:, 1:-1] +    # bottom
-        padded[2:, 2:]        # bottom-right
-    )
-    
-    return neighbors
+shader_source = f"""#version 430 core
 
-def update_generation(cell_grid: np.array):
-    neighbors = count_neighbors(cell_grid)
-    new_grid = ((cell_grid == 1) & ((neighbors == 2) | (neighbors == 3))) | \
-                ((cell_grid == 0) & (neighbors == 3))
-    return new_grid.astype(np.uint8)
+layout(std430, binding = 3) buffer CellGridIn {{
+    int cell_grid_in[{ROWS * COLS}];
+}};
+
+layout(std430, binding = 4) buffer CellGridOut {{
+    int cell_grid_out[{ROWS * COLS}];
+}};
+
+uniform int mouse_row;
+uniform int mouse_col;
+uniform int mouse_interaction;
+uniform int rows;
+uniform int cols;
+uniform bool running;
+
+layout (local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+layout(location = 0, rgba32f) uniform image2D img_output;
+
+void main() {{
+    ivec2 texel_coord = ivec2(gl_GlobalInvocationID.xy);
+
+    int row = texel_coord.y * rows / imageSize(img_output).y;
+    int col = texel_coord.x * cols / imageSize(img_output).x;
+    int current_index = (row * cols) + col;
+    int next = 0;
+    int alive_neighbors = 0;
+    int mouse_interaction_index = (mouse_row * cols) + mouse_col;
+
+    if (mouse_interaction != -1 && current_index == mouse_interaction_index) {{
+        next = mouse_interaction;
+    }}
+    else if (!running) {{
+        next = cell_grid_in[current_index];
+    }}
+    else {{
+        for (int dy = -1; dy <= 1; dy++) {{
+            for (int dx = -1; dx <= 1; dx++) {{
+                if (dx == 0 && dy == 0) continue;
+
+                int nx = texel_coord.x + dx;
+                int ny = texel_coord.y + dy;
+                if (nx >= 0 && nx < cols && ny >= 0 && ny < rows) {{
+                    int neighbor_index = ny * cols + nx;
+                    alive_neighbors += cell_grid_in[neighbor_index];
+                }}
+            }}
+        }}
+
+        if (cell_grid_in[current_index] == 0 && alive_neighbors == 3) {{
+            next = 1;
+        }}
+        else if (cell_grid_in[current_index] == 1 && (alive_neighbors == 3 || alive_neighbors == 2)) {{
+            next = 1;
+        }}  
+    }}
+
+    vec4 value;
+    if (next == 1) {{
+        value = vec4(1.0, 1.0, 1.0, 1.0);
+    }}
+    else {{
+        value = vec4(0.19, 0.31, 0.31, 1.0);
+    }}
+
+    cell_grid_out[current_index] = next;
+
+    imageStore(img_output, texel_coord, value);
+}}
+"""
+
+def create_shader(grid):
+    shader_program = pyglet.graphics.shader.ComputeShaderProgram(shader_source)
+
+    game_of_life_image = pyglet.image.Texture.create(COLS, ROWS, internalformat=pyglet.gl.GL_RGBA32F, min_filter=GL_NEAREST, mag_filter=GL_NEAREST)
+
+    uniform_location = shader_program['img_output']
+    game_of_life_image.bind_image_texture(unit=uniform_location)
+
+    ssbo_in = pyglet.graphics.BufferObject(len(grid) * 4, usage=pyglet.gl.GL_DYNAMIC_COPY)
+    ssbo_out = pyglet.graphics.BufferObject(len(grid) * 4, usage=pyglet.gl.GL_DYNAMIC_COPY)
+    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, ssbo_in.id)
+    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 4, ssbo_out.id)
+
+    return shader_program, game_of_life_image, ssbo_in, ssbo_out