## Replicating NumPy in GLSL for the AI Voxel Framework

The AI Voxel framework heavily relies on GPU‑based ODE integration and vector‑field evaluation. In the CPU (PyTorch/NumPy) pipeline, operations like matrix multiplications, broadcasting, reductions, and element‑wise math are expressed with NumPy’s concise syntax. On the GPU, we must **replicate this functionality** inside GLSL compute shaders, where we have no NumPy, no dynamic allocations, and limited recursion.

The key is to map NumPy’s array programming model to GLSL’s **buffer‑oriented, thread‑parallel** model. This is exactly what we need to implement the shader shown in Stage 6, and to extend it with more complex operations (e.g., batch processing, higher‑order integration, or custom losses).

Below I show how to **replicate the most common NumPy operations**—element‑wise, reduction, broadcasting, and linear algebra—within GLSL, using the AI Voxel’s Fourier field as a running example.

---

### 1. Element‑wise Operations

NumPy’s bread‑and‑butter: `+, -, *, /, sin, cos, exp, tanh, etc.`  
In GLSL, these are built‑in for scalar and vector types. For arrays (buffers), you loop over the elements.

| NumPy | GLSL (per‑thread) |
|-------|-------------------|
| `a + b` (same shape) | `a[i] + b[i]` |
| `np.sin(a)` | `sin(a[i])` |
| `a * b` | `a[i] * b[i]` |
| `np.exp(a)` | `exp(a[i])` |

**Example: evaluating the Fourier field** (from Stage 6)

```glsl
void evaluateField(in float h[MAX_D], out float F[MAX_D]) {
    for (int d = 0; d < u_D; d++) F[d] = 0.0;   // zero initialization
    for (int k = 0; k < u_K; k++) {
        float projection = 0.0;
        for (int d = 0; d < u_D; d++)
            projection += omega[k * u_D + d] * h[d];   // dot product

        float sinP = sin(projection);
        float cosP = cos(projection);

        for (int d = 0; d < u_D; d++) {
            F[d] += aCoeff[k * u_D + d] * sinP + bCoeff[k * u_D + d] * cosP;
        }
    }
}
```

Here `sin` and `cos` are native GLSL functions, replicating `np.sin` and `np.cos` for each element.

---

### 2. Reductions (sum, mean, max, min)

NumPy reductions like `np.sum`, `np.mean`, `np.max` aggregate over one or more axes.  
In GLSL we must **explicitly loop** over the reduced dimension. For large reductions, we can use **work‑group shared memory** to parallelize.

| NumPy | GLSL (single thread) |
|-------|----------------------|
| `np.sum(a)` | `float sum = 0.0; for(i) sum += a[i];` |
| `np.max(a)` | `float mx = -1e9; for(i) mx = max(mx, a[i]);` |
| `np.mean(a)` | compute sum then divide by size |

**Example: computing the dot product** (used in field evaluation and flux computation)

```glsl
float dotProduct(in float vec1[MAX_D], in float vec2[MAX_D]) {
    float sum = 0.0;
    for (int d = 0; d < u_D; d++)
        sum += vec1[d] * vec2[d];
    return sum;
}
```

For **parallel reductions** (e.g., across thousands of threads), use shared memory:

```glsl
shared float sharedSum[WORK_GROUP_SIZE];

void main() {
    uint localID = gl_LocalInvocationID.x;
    uint groupID = gl_WorkGroupID.x;

    // Load data into shared memory
    sharedSum[localID] = data[groupID * WORK_GROUP_SIZE + localID];
    barrier();

    // Reduce within work group (tree reduction)
    for (uint stride = WORK_GROUP_SIZE/2; stride > 0; stride >>= 1) {
        if (localID < stride)
            sharedSum[localID] += sharedSum[localID + stride];
        barrier();
    }
    if (localID == 0) {
        result[groupID] = sharedSum[0];   // write out partial sum
    }
}
```

---

### 3. Broadcasting

NumPy’s broadcasting lets you operate arrays of different shapes (e.g., `(n,) + (1,)` → `(n,)`).  
In GLSL, we replicate this by **reusing scalars** or **striding** in loops.

| NumPy | GLSL equivalent |
|-------|-----------------|
| `scalar + array` | `float s = …; for(i) result[i] = s + array[i];` |
| `(n,m) + (n,1)` | for i: for j: result[i][j] = A[i][j] + B[i][0]; |

**Example: applying a bias vector** to every batch element (common in MLP layers)

```glsl
// Bias addition: out[i][j] = in[i][j] + bias[j]
for (int i = 0; i < batchSize; i++) {
    for (int j = 0; j < D; j++) {
        out[i * D + j] = in[i * D + j] + bias[j];
    }
}
```

In practice, many AI Voxel operations are already vectorized by processing one query per thread, so broadcasting is often handled by global uniforms.

---

### 4. Linear Algebra (matmul, transpose, etc.)

NumPy’s `@` or `np.matmul` is central to neural networks. In GLSL, you write the triple‑loop.  
For small matrices (e.g., the boundary MLP), you can unroll.

| NumPy | GLSL (naive) |
|-------|--------------|
| `C = A @ B` (A: M×K, B: K×N) | `for(i) for(j) { sum=0; for(k) sum += A[i*K+k]*B[k*N+j]; C[i*N+j]=sum; }` |

**Example: boundary MLP forward pass** (from Stage 6 shader)

```glsl
float evaluateBoundary(in float h[MAX_D]) {
    float hidden[MAX_D];
    // W1 * h + b1
    for (int j = 0; j < u_Dhidden; j++) {
        float sum = bB1[j];
        for (int d = 0; d < u_D; d++)
            sum += bW1[j * u_D + d] * h[d];
        hidden[j] = tanh(sum);
    }
    // W2 * hidden
    float S = 0.0;
    for (int j = 0; j < u_Dhidden; j++)
        S += bW2[j] * hidden[j];
    return S;
}
```

For **batch matmul** (multiple independent matrices), you add an outer loop over batches.

---

### 5. Advanced Indexing & Slicing

NumPy’s fancy indexing (e.g., `a[[0,2]]`) is not directly supported. In GLSL, you compute indices explicitly based on thread IDs or use a lookup buffer.

- **Gather**: read from arbitrary positions → use SSBOs with integer indices.
- **Scatter**: write to arbitrary positions → use atomic operations (e.g., `atomicAdd`) if needed, but for AI Voxel we often have one‑to‑one writes.

**Example: mapping thread ID to a query point in a grid**

```glsl
uint tid = gl_GlobalInvocationID.x;
// Compute grid coordinates from thread ID (like np.indices)
uint x = tid % width;
uint y = (tid / width) % height;
uint z = tid / (width * height);
// Use as input coordinates
h[0] = (float(x) / width) * 2.0 - 1.0;
h[1] = (float(y) / height) * 2.0 - 1.0;
h[2] = (float(z) / depth) * 2.0 - 1.0;
```

This replicates `np.meshgrid` style generation.

---

### 6. Random Number Generation

NumPy’s `np.random` is used for sampling latent vectors or adding noise.  
In GLSL, you can use a **hash‑based PRNG** per thread (e.g., PCG, XORWOW) seeded with the thread ID.

```glsl
uint seed = gl_GlobalInvocationID.x * 1103515245u + 12345u;
float rand() {
    seed = seed * 1664525u + 1013904223u;
    return float(seed & 0xFFFFu) / 65536.0;
}
```

This replicates `np.random.rand()` per thread, with deterministic reproducibility.

---

### 7. Memory Layout and Data Structures

NumPy uses **strided, contiguous arrays** in memory. In GLSL, we must **flatten** multidimensional arrays into 1D buffers (SSBOs) and manually compute offsets.

- `array[shape]` → flatten to `array[offset]` where `offset = ((i * dim1) + j) * dim2 + k`.
- Use `layout(std430)` to map C‑style row‑major layout.

**Example: storing the Fourier coefficients** (Stage 6):

- `omega`: `[K * D]` floats
- `aCoeff`: `[K * D]` floats
- `bCoeff`: `[K * D]` floats

We access them with `omega[k * u_D + d]`.

---

### 8. Putting It All Together: A NumPy‑like GLSL Utility Header

To make shader code more readable, you can define macros and helper functions that mimic NumPy’s semantics.

```glsl
// Define max dimensions (must match CPU side)
#define MAX_D 128
#define MAX_K 256

// Elementwise operations
#define add(a,b,i) ((a)[i] + (b)[i])
#define mul(a,b,i) ((a)[i] * (b)[i])
#define sinv(a,i) (sin((a)[i]))

// Dot product
float dot_vec(in float a[MAX_D], in float b[MAX_D]) {
    float s = 0.0;
    for (int i = 0; i < u_D; i++) s += a[i] * b[i];
    return s;
}

// Matmul (C = A * B) with A: MxK, B: KxN, stored row‑major
void matmul(in float A[], in float B[], out float C[],
            int M, int K, int N) {
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0;
            for (int k = 0; k < K; k++)
                sum += A[i * K + k] * B[k * N + j];
            C[i * N + j] = sum;
        }
    }
}
```

---

### 9. Advanced: Automatic Differentiation (like Autograd)

NumPy does not support autodiff, but PyTorch does. In GLSL we don’t have autograd, but for **training** we could use finite differences (as shown in Stage 4 for `∇S`) or we can offload training to the CPU and only inference on GPU. For the AI Voxel pipeline, the shader is used **only for inference**; the distillation and boundary extraction happen in Python.

---

### 10. Performance Considerations

- **Memory coalescing**: align buffer accesses so that consecutive threads read consecutive memory addresses.
- **Use `vec4` types** for vectors of size 4 to leverage SIMD.
- **Avoid branching** inside tight loops; prefer predication.
- **Use shared memory** for reductions and small reusable data.

---

### Conclusion

Replicating NumPy in GLSL means **manually writing loops** and **explicitly managing memory** via SSBOs. While you lose the syntactic sugar, you gain massive parallelism. In the AI Voxel framework, we successfully replicate:

- **Vectorized field evaluation** → `evaluateField` (element‑wise + reductions)
- **Matrix multiplications** → MLP boundary evaluation (`evaluateBoundary`)
- **Broadcasting** → bias addition, scaling by global constants
- **Reductions** → flux computation (dot product)
- **Random sampling** → PRNG for initial conditions

The GLSL code in Stage 6 is a complete working example that replicates the NumPy‑style operations needed to run an AI Voxel entirely on the GPU.

For more complex operations (e.g., convolution, attention), the same principles apply: flatten the kernels, loop over the receptive field, and use thread‑local accumulators.

By following this mapping, you can port almost any NumPy‑based algorithm to GLSL compute shaders, as long as the data sizes are bounded and the operations are embarrassingly parallel.