use crate::math::{cos, floor, sin, sin_cos, sqrt};
use alloc::vec;
use alloc::vec::Vec;
use crate::gfx::cubemap;
use crate::math::vec3::{cross as cross3, dot as dot3, length};
pub const DEFAULT_IRRADIANCE_PHI_SAMPLES: u32 = 64;
pub const DEFAULT_IRRADIANCE_THETA_SAMPLES: u32 = 16;
fn sample_cube(faces: &[Vec<f32>; 6], face_size: u32, dir: [f32; 3]) -> [f32; 3] {
let ax = dir[0].abs();
let ay = dir[1].abs();
let az = dir[2].abs();
let (face, ma, s, t) = if ax >= ay && ax >= az {
if dir[0] > 0.0 {
(0usize, ax, -dir[2], -dir[1])
} else {
(1, ax, dir[2], -dir[1])
}
} else if ay >= az {
if dir[1] > 0.0 {
(2usize, ay, dir[0], dir[2])
} else {
(3, ay, dir[0], -dir[2])
}
} else if dir[2] > 0.0 {
(4usize, az, dir[0], -dir[1])
} else {
(5, az, -dir[0], -dir[1])
};
let inv = 0.5 / ma.max(1e-20);
let fs = face_size as f32;
let fx = (s * inv + 0.5) * fs - 0.5;
let fy = (t * inv + 0.5) * fs - 0.5;
let x0 = (floor(fx) as i32).clamp(0, face_size as i32 - 1);
let y0 = (floor(fy) as i32).clamp(0, face_size as i32 - 1);
let x1 = (x0 + 1).clamp(0, face_size as i32 - 1);
let y1 = (y0 + 1).clamp(0, face_size as i32 - 1);
let dx = (fx - floor(fx)).clamp(0.0, 1.0);
let dy = (fy - floor(fy)).clamp(0.0, 1.0);
let p = |x: i32, y: i32| -> [f32; 3] {
let off = ((y as usize) * face_size as usize + x as usize) * 4;
let face_data = &faces[face];
[face_data[off], face_data[off + 1], face_data[off + 2]]
};
let p00 = p(x0, y0);
let p10 = p(x1, y0);
let p01 = p(x0, y1);
let p11 = p(x1, y1);
let w00 = (1.0 - dx) * (1.0 - dy);
let w10 = dx * (1.0 - dy);
let w01 = (1.0 - dx) * dy;
let w11 = dx * dy;
[
p00[0] * w00 + p10[0] * w10 + p01[0] * w01 + p11[0] * w11,
p00[1] * w00 + p10[1] * w10 + p01[1] * w01 + p11[1] * w11,
p00[2] * w00 + p10[2] * w10 + p01[2] * w01 + p11[2] * w11,
]
}
fn normalize3(v: [f32; 3]) -> [f32; 3] {
let l = length(v).max(1e-20);
[v[0] / l, v[1] / l, v[2] / l]
}
fn make_tbn(n: [f32; 3]) -> ([f32; 3], [f32; 3]) {
let up = if n[2].abs() < 0.999 {
[0.0, 0.0, 1.0]
} else {
[1.0, 0.0, 0.0]
};
let t = normalize3(cross3(up, n));
let b = cross3(n, t);
(t, b)
}
fn hammersley(i: u32, n: u32) -> [f32; 2] {
let mut bits = i;
bits = bits.rotate_right(16);
bits = ((bits & 0x5555_5555) << 1) | ((bits & 0xAAAA_AAAA) >> 1);
bits = ((bits & 0x3333_3333) << 2) | ((bits & 0xCCCC_CCCC) >> 2);
bits = ((bits & 0x0F0F_0F0F) << 4) | ((bits & 0xF0F0_F0F0) >> 4);
bits = ((bits & 0x00FF_00FF) << 8) | ((bits & 0xFF00_FF00) >> 8);
let radical_inverse = (bits as f32) * 2.328_306_4e-10; [i as f32 / n as f32, radical_inverse]
}
fn importance_sample_ggx(
xi: [f32; 2],
n: [f32; 3],
basis: ([f32; 3], [f32; 3]),
a2m1: f32,
) -> [f32; 3] {
let (t, b) = basis;
let phi = 2.0 * core::f32::consts::PI * xi[0];
let cos_theta = sqrt((1.0 - xi[1]) / (1.0 + a2m1 * xi[1]));
let sin_theta = sqrt((1.0 - cos_theta * cos_theta).max(0.0));
let (sin_phi, cos_phi) = sin_cos(phi);
let h_local = [sin_theta * cos_phi, sin_theta * sin_phi, cos_theta];
normalize3([
t[0] * h_local[0] + b[0] * h_local[1] + n[0] * h_local[2],
t[1] * h_local[0] + b[1] * h_local[1] + n[1] * h_local[2],
t[2] * h_local[0] + b[2] * h_local[1] + n[2] * h_local[2],
])
}
fn clamp_radiance(rgb: [f32; 3], clamp: f32) -> [f32; 3] {
if clamp <= 0.0 {
return rgb;
}
let lum = 0.2126 * rgb[0] + 0.7152 * rgb[1] + 0.0722 * rgb[2];
if lum > clamp {
let s = clamp / lum;
return [rgb[0] * s, rgb[1] * s, rgb[2] * s];
}
rgb
}
pub struct FaceRow<'a> {
face: usize,
y: u32,
texels: &'a mut [f32],
}
pub fn face_rows(faces: &mut [Vec<f32>; 6], face_size: u32) -> Vec<FaceRow<'_>> {
let stride = face_size as usize * 4;
faces
.iter_mut()
.enumerate()
.flat_map(|(face, data)| {
data.chunks_mut(stride)
.enumerate()
.map(move |(y, texels)| FaceRow {
face,
y: y as u32,
texels,
})
})
.collect()
}
enum Kernel {
Irradiance(IrradianceKernel),
Ggx(GgxKernel),
}
struct IrradianceKernel {
phi_samples: u32,
theta_samples: u32,
inv_n_phi: f32,
inv_n_theta: f32,
weight: f32,
}
struct GgxKernel {
samples: u32,
a2m1: f32,
clamp: f32,
}
pub struct CubeBake<'a> {
source: &'a [Vec<f32>; 6],
source_face_size: u32,
output_face_size: u32,
kernel: Kernel,
}
impl<'a> CubeBake<'a> {
pub fn irradiance(
source: &'a [Vec<f32>; 6],
source_face_size: u32,
output_face_size: u32,
phi_samples: u32,
theta_samples: u32,
) -> Self {
let inv_n_phi = 1.0 / phi_samples as f32;
let inv_n_theta = 1.0 / theta_samples as f32;
let weight = core::f32::consts::PI * core::f32::consts::PI * inv_n_phi * inv_n_theta;
Self {
source,
source_face_size,
output_face_size,
kernel: Kernel::Irradiance(IrradianceKernel {
phi_samples,
theta_samples,
inv_n_phi,
inv_n_theta,
weight,
}),
}
}
pub fn ggx(
source: &'a [Vec<f32>; 6],
source_face_size: u32,
output_face_size: u32,
roughness: f32,
samples: u32,
clamp: f32,
) -> Self {
let a = roughness * roughness;
Self {
source,
source_face_size,
output_face_size,
kernel: Kernel::Ggx(GgxKernel {
samples,
a2m1: a * a - 1.0,
clamp,
}),
}
}
pub fn output_face_size(&self) -> u32 {
self.output_face_size
}
pub fn output_faces(&self) -> [Vec<f32>; 6] {
let f = self.output_face_size as usize;
core::array::from_fn(|_| vec![0.0; f * f * 4])
}
pub fn bake<S: super::schedule::RowScheduler>(&self, scheduler: &S) -> [Vec<f32>; 6] {
let mut faces = self.output_faces();
let mut rows = face_rows(&mut faces, self.output_face_size());
scheduler.run(&mut rows, &|row| self.compute_row(row));
faces
}
pub fn compute_row(&self, row: &mut FaceRow<'_>) {
match &self.kernel {
Kernel::Irradiance(k) => self.irradiance_row(row, k),
Kernel::Ggx(k) => self.ggx_row(row, k),
}
}
pub fn compute(&self) -> [Vec<f32>; 6] {
let mut faces = self.output_faces();
for row in &mut face_rows(&mut faces, self.output_face_size) {
self.compute_row(row);
}
faces
}
fn irradiance_row(&self, row: &mut FaceRow<'_>, k: &IrradianceKernel) {
for x in 0..self.output_face_size {
let n = cubemap::texel_dir(row.face, x, row.y, self.output_face_size);
let (tan, bit) = make_tbn(n);
let mut sum = [0.0f32; 3];
for phi_i in 0..k.phi_samples {
let phi = 2.0 * core::f32::consts::PI * (phi_i as f32 + 0.5) * k.inv_n_phi;
let sin_phi = sin(phi);
let cos_phi = cos(phi);
for theta_i in 0..k.theta_samples {
let theta =
0.5 * core::f32::consts::PI * (theta_i as f32 + 0.5) * k.inv_n_theta;
let sin_theta = sin(theta);
let cos_theta = cos(theta);
let l_local = [sin_theta * cos_phi, sin_theta * sin_phi, cos_theta];
let dir = [
tan[0] * l_local[0] + bit[0] * l_local[1] + n[0] * l_local[2],
tan[1] * l_local[0] + bit[1] * l_local[1] + n[1] * l_local[2],
tan[2] * l_local[0] + bit[2] * l_local[1] + n[2] * l_local[2],
];
let env = sample_cube(self.source, self.source_face_size, normalize3(dir));
let w = cos_theta * sin_theta;
sum[0] += env[0] * w;
sum[1] += env[1] * w;
sum[2] += env[2] * w;
}
}
let off = x as usize * 4;
row.texels[off] = sum[0] * k.weight;
row.texels[off + 1] = sum[1] * k.weight;
row.texels[off + 2] = sum[2] * k.weight;
row.texels[off + 3] = 1.0;
}
}
fn ggx_row(&self, row: &mut FaceRow<'_>, k: &GgxKernel) {
for x in 0..self.output_face_size {
let n = cubemap::texel_dir(row.face, x, row.y, self.output_face_size);
let basis = make_tbn(n);
let mut accum = [0.0f32; 3];
let mut total_weight = 0.0f32;
for i in 0..k.samples {
let xi = hammersley(i, k.samples);
let h = importance_sample_ggx(xi, n, basis, k.a2m1);
let ndh = dot3(n, h);
if ndh <= 0.0 {
continue;
}
let l = normalize3([
2.0 * ndh * h[0] - n[0],
2.0 * ndh * h[1] - n[1],
2.0 * ndh * h[2] - n[2],
]);
let ndl = dot3(n, l).max(0.0);
if ndl > 0.0 {
let env =
clamp_radiance(sample_cube(self.source, self.source_face_size, l), k.clamp);
accum[0] += env[0] * ndl;
accum[1] += env[1] * ndl;
accum[2] += env[2] * ndl;
total_weight += ndl;
}
}
let off = x as usize * 4;
if total_weight > 0.0 {
let inv = 1.0 / total_weight;
row.texels[off] = accum[0] * inv;
row.texels[off + 1] = accum[1] * inv;
row.texels[off + 2] = accum[2] * inv;
} else {
let n_sample = sample_cube(self.source, self.source_face_size, n);
row.texels[off] = n_sample[0];
row.texels[off + 1] = n_sample[1];
row.texels[off + 2] = n_sample[2];
}
row.texels[off + 3] = 1.0;
}
}
}
pub fn prefilter_mip0(
source: &[Vec<f32>; 6],
source_face_size: u32,
clamp: f32,
clamp_mip0: bool,
) -> [Vec<f32>; 6] {
let f = source_face_size as usize;
let mut mip0: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; f * f * 4]);
for face in 0..6 {
for i in 0..f * f {
let off = i * 4;
let mut rgb = [
source[face][off],
source[face][off + 1],
source[face][off + 2],
];
if clamp_mip0 {
rgb = clamp_radiance(rgb, clamp);
}
mip0[face][off] = rgb[0];
mip0[face][off + 1] = rgb[1];
mip0[face][off + 2] = rgb[2];
mip0[face][off + 3] = 1.0;
}
}
mip0
}
pub fn prefilter_roughness(mip: u32, mip_count: u32) -> f32 {
mip as f32 / (mip_count - 1) as f32
}
pub fn compute_irradiance(
source: &[Vec<f32>; 6],
source_face_size: u32,
output_face_size: u32,
phi_samples: u32,
theta_samples: u32,
) -> [Vec<f32>; 6] {
CubeBake::irradiance(
source,
source_face_size,
output_face_size,
phi_samples,
theta_samples,
)
.compute()
}
pub fn compute_prefilter(
source: &[Vec<f32>; 6],
source_face_size: u32,
mip_count: u32,
samples_per_texel: u32,
clamp: f32,
clamp_mip0: bool,
) -> Vec<[Vec<f32>; 6]> {
let mut mips: Vec<[Vec<f32>; 6]> = Vec::with_capacity(mip_count as usize);
mips.push(prefilter_mip0(source, source_face_size, clamp, clamp_mip0));
for mip in 1..mip_count {
mips.push(
CubeBake::ggx(
source,
source_face_size,
source_face_size >> mip,
prefilter_roughness(mip, mip_count),
samples_per_texel,
clamp,
)
.compute(),
);
}
mips
}
#[cfg(test)]
mod tests {
use super::super::schedule::{RowScheduler, Serial};
use super::*;
struct Reversed;
impl RowScheduler for Reversed {
fn run<T: Send>(&self, items: &mut [T], compute: &(dyn Fn(&mut T) + Send + Sync)) {
items.iter_mut().rev().for_each(compute);
}
}
#[test]
fn a_bake_does_not_depend_on_the_row_order() {
let source = solid_cube(8, [0.4, 0.6, 0.9]);
let bake = CubeBake::ggx(&source, 8, 8, 0.5, 16, 0.0);
assert_eq!(bake.bake(&Serial), bake.bake(&Reversed));
}
#[test]
fn a_bake_fills_every_output_face() {
let source = solid_cube(8, [1.0, 1.0, 1.0]);
let bake = CubeBake::irradiance(&source, 8, 4, 8, 8);
let faces = bake.bake(&Serial);
for face in &faces {
assert_eq!(face.len(), 4 * 4 * 4);
assert!(face.chunks_exact(4).all(|px| px[0] > 0.0));
}
}
fn solid_cube(face_size: u32, color: [f32; 3]) -> [Vec<f32>; 6] {
let f = face_size as usize;
core::array::from_fn(|_| {
let mut face = Vec::with_capacity(f * f * 4);
for _ in 0..f * f {
face.extend_from_slice(&[color[0], color[1], color[2], 1.0]);
}
face
})
}
fn face_mean(face: &[f32]) -> [f32; 3] {
let n = face.len() / 4;
let mut m = [0.0f32; 3];
for px in face.chunks_exact(4) {
m[0] += px[0];
m[1] += px[1];
m[2] += px[2];
}
[m[0] / n as f32, m[1] / n as f32, m[2] / n as f32]
}
fn face_variance_red(face: &[f32]) -> f32 {
let n = face.len() / 4;
let mean = face.chunks_exact(4).map(|p| p[0]).sum::<f32>() / n as f32;
face.chunks_exact(4)
.map(|p| (p[0] - mean).powi(2))
.sum::<f32>()
/ n as f32
}
fn firefly_cube(face: usize) -> [Vec<f32>; 6] {
let mut s: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; face * face * 4]);
for fd in s.iter_mut() {
for p in fd.chunks_exact_mut(4) {
p[0] = 1.0;
p[1] = 1.0;
p[2] = 1.0;
p[3] = 1.0;
}
}
let off = ((face / 2) * face + face / 2) * 4;
s[4][off] = 2000.0;
s[4][off + 1] = 2000.0;
s[4][off + 2] = 2000.0;
s
}
fn peak(fd: &[f32]) -> f32 {
fd.chunks_exact(4)
.map(|p| p[0].max(p[1]).max(p[2]))
.fold(0.0f32, f32::max)
}
fn gradient_cube(face_size: u32) -> [Vec<f32>; 6] {
let f = face_size as usize;
core::array::from_fn(|face| {
let mut data = vec![0.0f32; f * f * 4];
for y in 0..f {
for x in 0..f {
let off = (y * f + x) * 4;
data[off] = face as f32 + x as f32 / f as f32;
data[off + 1] = y as f32 / f as f32;
data[off + 2] = (x + y) as f32 / (2 * f) as f32;
data[off + 3] = 1.0;
}
}
data
})
}
fn assert_rows_match_whole_image(bake: &CubeBake<'_>) {
let whole = bake.compute();
let mut chunked = bake.output_faces();
let mut rows = face_rows(&mut chunked, bake.output_face_size());
assert_eq!(
rows.len(),
6 * bake.output_face_size() as usize,
"one row per face line"
);
for row in rows.iter_mut().rev() {
bake.compute_row(row);
}
assert_eq!(chunked, whole, "chunked bake diverged from the whole image");
}
#[test]
fn chunked_irradiance_matches_the_whole_image() {
let source = gradient_cube(8);
assert_rows_match_whole_image(&CubeBake::irradiance(&source, 8, 4, 16, 8));
}
#[test]
fn chunked_ggx_matches_the_whole_image() {
let source = gradient_cube(16);
for roughness in [0.25f32, 0.5, 1.0] {
assert_rows_match_whole_image(&CubeBake::ggx(&source, 16, 8, roughness, 32, 0.0));
}
}
#[test]
fn chunked_ggx_matches_the_whole_image_under_the_firefly_cap() {
let source = firefly_cube(16);
assert_rows_match_whole_image(&CubeBake::ggx(&source, 16, 8, 0.5, 64, 8.0));
}
#[test]
fn face_rows_cover_every_texel_exactly_once() {
let mut faces: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; 4 * 4 * 4]);
for row in &mut face_rows(&mut faces, 4) {
assert_eq!(row.texels.len(), 4 * 4, "a row is one line of RGBA texels");
for t in row.texels.iter_mut() {
*t += 1.0;
}
}
assert!(
faces.iter().flatten().all(|&t| t == 1.0),
"every texel written exactly once"
);
}
#[test]
fn hammersley_first_sample_is_zero() {
let s = hammersley(0, 1024);
assert!(s[0].abs() < 1e-6, "x was {}", s[0]);
assert!(s[1].abs() < 1e-6, "y was {}", s[1]);
}
#[test]
fn hammersley_last_sample_is_just_under_one() {
let s = hammersley(1023, 1024);
assert!(s[0] > 0.99 && s[0] < 1.0, "x was {}", s[0]);
}
#[test]
fn importance_sample_ggx_at_xi_zero_returns_n() {
let n = [0.0, 0.0, 1.0];
let a = 0.5f32 * 0.5;
let h = importance_sample_ggx([0.0, 0.0], n, make_tbn(n), a * a - 1.0);
assert!((h[0] - 0.0).abs() < 1e-5);
assert!((h[1] - 0.0).abs() < 1e-5);
assert!((h[2] - 1.0).abs() < 1e-5);
}
#[test]
fn irradiance_solid_color_is_pi_times_color() {
let source = solid_cube(8, [1.0, 0.5, 0.25]);
let irr = compute_irradiance(&source, 8, 4, 64, 16);
let mean = face_mean(&irr[0]);
let expected = [
core::f32::consts::PI * 1.0,
core::f32::consts::PI * 0.5,
core::f32::consts::PI * 0.25,
];
for c in 0..3 {
let delta = (mean[c] - expected[c]).abs() / expected[c];
assert!(
delta < 0.05,
"channel {} mean {} expected {}",
c,
mean[c],
expected[c]
);
}
}
#[test]
fn prefilter_mip_zero_matches_source_with_alpha_one() {
let source = solid_cube(16, [0.7, 0.3, 0.1]);
let mips = compute_prefilter(&source, 16, 3, 16, 0.0, false);
for face in &mips[0] {
for px in 0..16 * 16 {
let off = px * 4;
assert!((face[off] - 0.7).abs() < 1e-6);
assert!((face[off + 1] - 0.3).abs() < 1e-6);
assert!((face[off + 2] - 0.1).abs() < 1e-6);
assert!((face[off + 3] - 1.0).abs() < 1e-6);
}
}
}
#[test]
fn prefilter_solid_color_stays_solid_at_high_roughness() {
let source = solid_cube(16, [0.5, 0.5, 0.5]);
let mips = compute_prefilter(&source, 16, 4, 32, 0.0, false);
let mean = face_mean(&mips[3][0]);
for (c, m) in mean.iter().enumerate() {
assert!((m - 0.5).abs() < 0.02, "channel {} mean {}", c, m);
}
}
#[test]
fn prefilter_roughness_spans_zero_to_one() {
assert_eq!(prefilter_roughness(0, 5), 0.0);
assert_eq!(prefilter_roughness(4, 5), 1.0);
assert_eq!(prefilter_roughness(2, 5), 0.5);
}
#[test]
fn prefilter_blurs_a_red_seam() {
let face = 16usize;
let mut source: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; face * face * 4]);
for face_data in source.iter_mut() {
for p in face_data.chunks_exact_mut(4) {
p[3] = 1.0;
}
}
for y in 0..face {
let off = (y * face + 8) * 4;
source[4][off] = 20.0;
}
let mips = compute_prefilter(&source, face as u32, 3, 256, 0.0, false);
let v0 = face_variance_red(&mips[0][4]);
let v2 = face_variance_red(&mips[2][4]);
assert!(
v2 < v0 * 0.5,
"prefilter did not blur: mip 0 var={}, mip 2 var={}",
v0,
v2
);
}
#[test]
fn clamp_radiance_caps_luminance_and_keeps_hue() {
let dim = [1.0, 0.5, 0.25];
assert_eq!(clamp_radiance(dim, 10.0), dim);
let hot = [100.0, 50.0, 25.0];
assert_eq!(clamp_radiance(hot, 0.0), hot);
let capped = clamp_radiance(hot, 10.0);
let lum = 0.2126 * capped[0] + 0.7152 * capped[1] + 0.0722 * capped[2];
assert!((lum - 10.0).abs() < 1e-3, "luminance {} != cap 10", lum);
assert!((capped[0] / capped[1] - hot[0] / hot[1]).abs() < 1e-4);
assert!((capped[1] / capped[2] - hot[1] / hot[2]).abs() < 1e-4);
}
#[test]
fn prefilter_clamp_suppresses_a_firefly() {
let face = 16usize;
let unclamped = compute_prefilter(&firefly_cube(face), face as u32, 3, 256, 0.0, false);
let clamped = compute_prefilter(&firefly_cube(face), face as u32, 3, 256, 8.0, false);
let p_unclamped = peak(&unclamped[1][4]);
let p_clamped = peak(&clamped[1][4]);
assert!(
p_clamped < p_unclamped * 0.5,
"clamp did not suppress the firefly: unclamped {}, clamped {}",
p_unclamped,
p_clamped
);
assert!(
p_clamped >= 0.9,
"clamp crushed the background: clamped peak {}",
p_clamped
);
}
#[test]
fn prefilter_mip0_clamp_caps_a_mirror_firefly_only_when_requested() {
let face = 16usize;
let unclamped_mip0 = compute_prefilter(&firefly_cube(face), face as u32, 2, 16, 8.0, false);
assert!(
peak(&unclamped_mip0[0][4]) > 1000.0,
"mip 0 should be unclamped when clamp_mip0 is false: peak {}",
peak(&unclamped_mip0[0][4])
);
let clamped_mip0 = compute_prefilter(&firefly_cube(face), face as u32, 2, 16, 8.0, true);
let p = peak(&clamped_mip0[0][4]);
assert!(p <= 8.0 + 1e-3, "mip 0 firefly not capped: peak {p}");
assert!(p >= 0.9, "mip 0 clamp crushed the background: peak {p}");
}
}