Búsqueda de vecinos, colisiones y SPH

Con la rejilla construida, encontrar los vecinos de una partícula es un bucle triple sobre las celdas adyacentes y un bucle interno sobre el tramo de cada celda. Eso desbloquea las dos aplicaciones que justifican todo el andamiaje: colisiones entre partículas, que convierten un sistema de puntos en algo con volumen, e hidrodinámica de partículas suavizadas, que es la forma más directa de simular un líquido en tiempo real. Las dos comparten estructura, y la segunda tiene una particularidad que resume el bloque entero: no cabe en un solo dispatch.

El bucle de vecinos

struct Params { conteo: u32, radio: f32 };
@group(0) @binding(0) var<uniform> p: Params;
@group(0) @binding(1) var<uniform> g: Rejilla;
@group(0) @binding(2) var<storage, read> inicios: array<u32>;
@group(0) @binding(3) var<storage, read> parts:   array<Particula>;

// Recorre los vecinos de `pos` dentro del radio y acumula con `acc`.
fn vecinos(yo: u32, pos: vec3f) -> vec3f {
  let c0 = coordCelda(pos);
  let r2 = p.radio * p.radio;
  var acumulado = vec3f(0.0);

  for (var dz: i32 = -1; dz <= 1; dz = dz + 1) {
    for (var dy: i32 = -1; dy <= 1; dy = dy + 1) {
      for (var dx: i32 = -1; dx <= 1; dx = dx + 1) {
        let c = c0 + vec3i(dx, dy, dz);
        // Fuera del dominio: no hay celda, se salta.
        if (any(c < vec3i(0)) || any(c >= vec3i(g.dims))) { continue; }

        let ci  = indiceCelda(c);
        let ini = inicios[ci];
        let fin = inicios[ci + 1u];

        for (var k: u32 = ini; k < fin; k = k + 1u) {
          if (k == yo) { continue; }
          let d = parts[k].pos - pos;
          let dd = dot(d, d);
          if (dd >= r2 || dd < 1e-12) { continue; }
          acumulado += contribucion(d, dd, k);
        }
      }
    }
  }
  return acumulado;
}

Cuatro detalles que hay que respetar.

El guardia de borde con any evita calcular un índice de celda inválido. Con el clamp de indiceCelda no reventaría, pero saturar aquí sería peor: las celdas del borde se visitarían varias veces y las partículas de fuera contarían dos y tres veces. Fuera del dominio hay que saltar, no saturar.

La comparación de distancia usa el cuadrado, sin raíz. La raíz solo se calcula si la partícula entra de verdad, y muchas veces ni siquiera entonces.

El descarte de dd muy pequeño evita la división por cero cuando dos partículas están exactamente en el mismo punto, cosa que ocurre al inicializar con posiciones en rejilla o después de un rebote simétrico.

Y el bucle interno recorre parts[k] con k consecutivo porque las partículas están reordenadas por celda. Sin esa reordenación habría que leer parts[ordenado[k]] y el kernel sería varias veces más lento.

Colisiones por penalización

La forma más simple de dar volumen a las partículas es una fuerza de repulsión que aparece cuando se solapan. Con radio R por partícula, dos que estén a distancia menor que 2R se empujan.

const R: f32 = 0.05;
const K_COLISION: f32 = 800.0;
const AMORTIGUACION: f32 = 6.0;

fn fuerzaColision(d: vec3f, dd: f32, vRel: vec3f) -> vec3f {
  let dist = sqrt(dd);
  let solape = 2.0 * R - dist;
  if (solape <= 0.0) { return vec3f(0.0); }

  let n = d / dist;                          // de mi hacia el vecino
  // Resorte: empuja separando, proporcional al solape.
  var f = -n * (K_COLISION * solape);
  // Amortiguacion: disipa la velocidad relativa a lo largo de la normal.
  f -= n * (AMORTIGUACION * dot(vRel, n));
  return f;
}

La amortiguación es lo que impide que el sistema gane energía indefinidamente. Sin ella, dos partículas que chocan rebotan con la misma velocidad, y con errores de integración acaban rebotando más fuerte de lo que llegaron.

El problema de la penalización es la rigidez. Una K_COLISION alta hace que las partículas no se atraviesen y exige un paso de tiempo corto, con el límite de estabilidad que ya vimos. Una K_COLISION baja permite pasos largos y las partículas se hunden unas en otras. La regla práctica es fijar K a partir de la máxima interpenetración que estás dispuesto a tolerar con la velocidad típica del sistema, y después comprobar que el paso de tiempo aguanta.

SPH: densidad, presión y viscosidad

La hidrodinámica de partículas suavizadas representa el fluido como un conjunto de partículas que llevan masa, y estima cualquier campo continuo en un punto como una suma ponderada de las partículas cercanas, con un peso que decae con la distancia. Los núcleos de Müller, Charypar y Gross de 2003 siguen siendo los que se usan en tiempo real.

const H:  f32 = 0.1;                 // radio de suavizado
const H2: f32 = H * H;
const PI: f32 = 3.14159265359;

const MASA:      f32 = 0.02;
const RHO0:      f32 = 1000.0;       // densidad de reposo
const RIGIDEZ:   f32 = 3.0;          // constante de gas
const VISCOSIDAD:f32 = 0.25;

// Nucleo Poly6, para la densidad. Solo depende de la distancia.
const W_POLY6: f32 = 315.0 / (64.0 * PI * H*H*H*H*H*H*H*H*H);
fn poly6(dd: f32) -> f32 {
  let x = H2 - dd;
  if (x <= 0.0) { return 0.0; }
  return W_POLY6 * x * x * x;
}

// Gradiente del nucleo Spiky, para la presion. No se anula en r cero,
// que es exactamente lo que evita que las particulas se apelmacen.
const W_SPIKY: f32 = -45.0 / (PI * H*H*H*H*H*H);
fn gradSpiky(d: vec3f, dist: f32) -> vec3f {
  let x = H - dist;
  if (x <= 0.0) { return vec3f(0.0); }
  return (d / dist) * (W_SPIKY * x * x);
}

// Laplaciano del nucleo de viscosidad.
const W_VISC: f32 = 45.0 / (PI * H*H*H*H*H*H);
fn lapVisc(dist: f32) -> f32 {
  let x = H - dist;
  if (x <= 0.0) { return 0.0; }
  return W_VISC * x;
}

Los tres núcleos comparten el soporte compacto: valen cero más allá de H, que es lo que permite usar la rejilla con celda de lado H.

El primer dispatch calcula la densidad de cada partícula sumando la contribución de sus vecinas, y de ella la presión con una ecuación de estado:

@group(0) @binding(4) var<storage, read_write> densidad: array<f32>;
@group(0) @binding(5) var<storage, read_write> presion:  array<f32>;

@compute @workgroup_size(64)
fn calcularDensidad(@builtin(global_invocation_id) gid: vec3u) {
  let i = gid.x;
  if (i >= p.conteo) { return; }
  let pos = parts[i].pos;

  var rho = 0.0;
  // ... bucle de 27 celdas ...
  //     rho += MASA * poly6(dd);
  // la propia particula tambien contribuye:
  rho += MASA * poly6(0.0);

  densidad[i] = rho;
  presion[i]  = RIGIDEZ * (rho - RHO0);
}

El segundo dispatch calcula las fuerzas, y para eso necesita la presión de las vecinas, que el primer dispatch acaba de escribir:

@compute @workgroup_size(64)
fn calcularFuerzas(@builtin(global_invocation_id) gid: vec3u) {
  let i = gid.x;
  if (i >= p.conteo) { return; }
  let pos = parts[i].pos;
  let vel = parts[i].vel;
  let pi  = presion[i];

  var fPresion   = vec3f(0.0);
  var fViscosa   = vec3f(0.0);

  // ... bucle de 27 celdas, para cada vecina k con d y dd ...
  //   let dist = sqrt(dd);
  //   fPresion += gradSpiky(d, dist)
  //             * (-MASA * (pi + presion[k]) / (2.0 * densidad[k]));
  //   fViscosa += (parts[k].vel - vel)
  //             * (VISCOSIDAD * MASA * lapVisc(dist) / densidad[k]);

  let a = (fPresion + fViscosa) / densidad[i] + vec3f(0.0, -9.81, 0.0);
  aceleraciones[i] = a;
}

La fuerza de presión lleva signo negativo porque el gradiente apunta hacia donde la densidad crece y el fluido empuja hacia donde decrece. Y la simetrización con la media de las dos presiones es lo que hace que la fuerza entre dos partículas sea igual y opuesta, o sea que el momento total se conserve. Sin esa simetrización el fluido se autoacelera.

Por qué dos dispatches y no uno

La tentación de fusionar los dos kernels es fuerte: el bucle de 27 celdas se recorre dos veces y es la parte cara. Es imposible, y por la razón que atraviesa todo este bloque.

El kernel de fuerzas necesita presion[k] de todas sus vecinas. Esas vecinas las procesan otras invocaciones, y muy probablemente otros workgroups. Para que la presión de la vecina esté escrita, el dispatch de densidad tiene que haber terminado entero. Y la única sincronización que garantiza eso es la frontera de dispatch.

Fusionarlos daría un resultado en el que cada partícula lee presiones de vecinas en un estado indeterminado: unas ya calculadas, otras del frame anterior, y la mezcla cambiaría de una ejecución a otra. El fluido no explotaría de forma evidente; simplemente se comportaría de forma ligeramente distinta cada vez, con artefactos que aparecen y desaparecen. Es la clase de bug que cuesta semanas.

La estructura correcta de un paso de SPH es esta:

Doce dispatches por paso de simulación. Con un paso de 1/120 y dos pasos por frame, veinticuatro dispatches, todos encolables en un único compute pass. Para cien mil partículas eso corre cómodamente a sesenta cuadros por segundo en una GPU integrada moderna; para un millón hace falta gama alta o bajar el número de pasos por frame.