nordisk/facett

facett-map3d-flight/src/scatter_lut.wgsl

// **GPU transmittance-LUT precompute** (FLIGHT-SIM #7 / roadmap §2-GPU) — the
// compute-shader twin of `facett-map3d-flight::atmos::TransmittanceLut::precompute`
// / `optical_depth_to_top`. One invocation per `(altitude, μ)` texel: ray-march the
// curved planet's Rayleigh + Mie + ozone extinction to the top of the atmosphere and
// store `exp(-optical_depth)` as RGB transmittance.
//
// The CPU model in `atmos.rs` is the reference **oracle**; this shader MUST match it
// texel-for-texel (the GPU-parity test asserts `max_abs_error < tol`). All maths in
// kilometres, mirroring the CPU constants + loop order exactly.

const R_GROUND: f32 = 6360.0;                 // planet surface radius, km
const R_TOP: f32 = 6420.0;                    // top-of-atmosphere radius, km
const H_RAYLEIGH: f32 = 8.0;                  // Rayleigh scale height, km
const H_MIE: f32 = 1.2;                       // Mie scale height, km
const BETA_RAYLEIGH: vec3<f32> = vec3<f32>(5.802e-3, 13.558e-3, 33.1e-3);
const BETA_MIE: f32 = 3.996e-3;
const BETA_MIE_EXT: f32 = 3.996e-3 / 0.9;     // Mie extinction = scattering / albedo
const BETA_OZONE: vec3<f32> = vec3<f32>(0.650e-3, 1.881e-3, 0.085e-3);

// h_size, mu_size, steps, _pad — matches the Rust `LutParams`.
struct Params { dims: vec4<u32> };
@group(0) @binding(0) var<uniform> P: Params;
@group(0) @binding(1) var<storage, read_write> out_lut: array<vec4<f32>>;

// Distance from radius `r`, view-zenith-cosine `mu`, to the top of the atmosphere.
fn dist_to_top(r: f32, mu: f32) -> f32 {
    let disc = r * r * (mu * mu - 1.0) + R_TOP * R_TOP;
    return max(-r * mu + sqrt(max(disc, 0.0)), 0.0);
}

// Distance to the planet surface along the ray, or -1.0 if the ray misses it (so the
// path runs to the top of the atmosphere instead). Mirrors `dist_to_ground`'s `None`.
fn dist_to_ground(r: f32, mu: f32) -> f32 {
    let disc = r * r * (mu * mu - 1.0) + R_GROUND * R_GROUND;
    if (disc < 0.0) { return -1.0; }
    let t = -r * mu - sqrt(disc);
    if (t > 0.0) { return t; }
    return -1.0;
}

@compute @workgroup_size(64)
fn cs_transmittance(@builtin(global_invocation_id) gid: vec3<u32>) {
    let h_size = P.dims.x;
    let mu_size = P.dims.y;
    let steps = P.dims.z;
    let total = h_size * mu_size;
    let idx = gid.x;
    if (idx >= total) { return; }

    let hi = idx / mu_size;                    // altitude row
    let mi = idx % mu_size;                     // μ column

    let max_h = R_TOP - R_GROUND;
    let h = max_h * f32(hi) / f32(h_size - 1u);
    let r = R_GROUND + h;
    let mu = -1.0 + 2.0 * f32(mi) / f32(mu_size - 1u);

    // March only as far as the ground if the ray dives into the planet.
    let tg = dist_to_ground(r, mu);
    var t_max: f32;
    if (tg >= 0.0) { t_max = tg; } else { t_max = dist_to_top(r, mu); }
    let dt = t_max / f32(steps);

    // P(t) = (0,r,0) + dir·t with dir = (sin, mu, 0); height = |P(t)| - R_GROUND.
    let sin_v = sqrt(max(1.0 - mu * mu, 0.0));
    var od = vec3<f32>(0.0);
    for (var i: u32 = 0u; i < steps; i = i + 1u) {
        let t = (f32(i) + 0.5) * dt;
        let px = sin_v * t;
        let py = r + mu * t;
        let radius = sqrt(px * px + py * py);
        let hh = max(radius - R_GROUND, 0.0);
        let dr = exp(-hh / H_RAYLEIGH);
        let dm = exp(-hh / H_MIE);
        let doz = clamp(1.0 - abs(hh - 25.0) / 15.0, 0.0, 1.0);
        od = od + (BETA_RAYLEIGH * dr + vec3<f32>(BETA_MIE_EXT) * dm + BETA_OZONE * doz) * dt;
    }

    out_lut[idx] = vec4<f32>(exp(-od), 1.0);
}