use super::{
super::{
NUM_CHANNELS,
absorption::channel_absorption_alphas_assigned,
camera::{FacetFinish, Ray},
color::{integrate_channels_to_xyz, integrate_channels_to_xyz_families},
environment::{EnvironmentSource, sample_environment_for_nee},
intersect_stone::intersect_stone_soa,
refraction::{RayMaterialContext, RayWavelengthCache},
sampling::{
DISTANCE_SAMPLE_STREAM, NEE_ENV_DIR_U_STREAM, NEE_ENV_DIR_V_STREAM, PHASE_DIR_U_STREAM,
PHASE_DIR_V_STREAM, hash_u32,
},
transport::apply_russian_roulette,
zoned::{ZoneAlphas, ZonedCache, local_extinction, optical_depths},
},
HG_NEE_MAX_CROSSINGS, NeeContext, balance_heuristic, frosted_orthonormal_basis,
};
use crate::optics::{
materials::GemMaterial, polarization::StokesVector, raytracer::camera::HitRecord,
};
use glam::Vec3;
#[expect(
clippy::too_many_arguments,
reason = "argument list mirrors transport_physics.wgsl's own maybe_scatter_or_extinguish \
one-for-one; bundling into a struct would break that direct correspondence"
)]
pub(crate) fn maybe_scatter_or_extinguish(
alphas: &[f32; NUM_CHANNELS],
sigma_s_model: f32,
g: f32,
hero_idx: usize,
ray_dir: Vec3,
hit_t: f32,
path_scale: f32,
rng_seed: u32,
bounce: u32,
stokes: &mut [StokesVector; NUM_CHANNELS],
path_pdf: &mut [f32; NUM_CHANNELS],
) -> Option<(f32, Vec3)> {
let sigma_s = sigma_s_model / path_scale;
let sigma_t_hero = alphas[hero_idx] + sigma_s;
let dist_rand =
(hash_u32(rng_seed ^ hash_u32(bounce ^ DISTANCE_SAMPLE_STREAM)) as f32) / 4_294_967_295.0;
let one_minus_u = (1.0 - dist_rand).max(1e-7);
let t_free = -(one_minus_u.ln()) / sigma_t_hero;
let hit_t_scaled = hit_t * path_scale;
if t_free < hit_t_scaled {
let pdf_hero = sigma_t_hero * one_minus_u;
let mut tr_args = [0f32; NUM_CHANNELS];
for k in 0..NUM_CHANNELS {
tr_args[k] = -(alphas[k] + sigma_s) * t_free;
}
let tr = crate::simd::exp_f32x8(tr_args);
for k in 0..NUM_CHANNELS {
let sigma_t_k = alphas[k] + sigma_s;
let tr_k = tr[k];
let weight = tr_k * sigma_s / pdf_hero;
stokes[k] = StokesVector::unpolarized(stokes[k].intensity() * weight);
path_pdf[k] *= sigma_t_k * tr_k;
}
let u1 =
(hash_u32(rng_seed ^ hash_u32(bounce ^ PHASE_DIR_U_STREAM)) as f32) / 4_294_967_295.0;
let u2 =
(hash_u32(rng_seed ^ hash_u32(bounce ^ PHASE_DIR_V_STREAM)) as f32) / 4_294_967_295.0;
let new_dir = sample_henyey_greenstein_direction(u1, u2, g, ray_dir);
let t_free_model = t_free / path_scale;
Some((t_free_model, new_dir))
} else {
let mut survive_args = [0f32; NUM_CHANNELS];
for k in 0..NUM_CHANNELS {
survive_args[k] = -(alphas[k] + sigma_s) * hit_t_scaled;
}
let survive = crate::simd::exp_f32x8(survive_args);
let survive_hero = survive[hero_idx];
for k in 0..NUM_CHANNELS {
let survive_k = survive[k];
stokes[k] = stokes[k].scale(survive_k / survive_hero.max(1e-30));
path_pdf[k] *= survive_k;
}
None
}
}
pub(in super::super) struct ZonedScatterParams {
pub(in super::super) sigma_s_model: f32,
pub(in super::super) g: f32,
pub(in super::super) hero_idx: usize,
pub(in super::super) ray: Ray,
pub(in super::super) hit_t: f32,
pub(in super::super) path_scale: f32,
pub(in super::super) rng_seed: u32,
pub(in super::super) bounce: u32,
}
pub(in super::super) fn maybe_scatter_or_extinguish_zoned(
zoned: &ZonedCache,
zone_alphas: &ZoneAlphas,
params: &ZonedScatterParams,
stokes: &mut [StokesVector; NUM_CHANNELS],
path_pdf: &mut [f32; NUM_CHANNELS],
) -> Option<(f32, Vec3)> {
let ZonedScatterParams {
sigma_s_model,
g,
hero_idx,
ray,
hit_t,
path_scale,
rng_seed,
bounce,
} = *params;
let sigma_s = sigma_s_model / path_scale;
let depth_to = |t: f32| -> [f32; NUM_CHANNELS] {
let lengths = zoned.lengths_mm(ray.origin, ray.dir, 0.0, t, path_scale);
optical_depths(zone_alphas, &lengths, sigma_s)
};
let dist_rand =
(hash_u32(rng_seed ^ hash_u32(bounce ^ DISTANCE_SAMPLE_STREAM)) as f32) / 4_294_967_295.0;
let one_minus_u = (1.0 - dist_rand).max(1e-7);
let target = -(one_minus_u.ln());
let depth_full = depth_to(hit_t);
if target < depth_full[hero_idx] {
let (mut lo, mut hi) = (0.0f32, hit_t);
for _ in 0..48 {
let mid = f32::midpoint(lo, hi);
if depth_to(mid)[hero_idx] < target {
lo = mid;
} else {
hi = mid;
}
if hi - lo <= f32::EPSILON * hi {
break;
}
}
let t_free = f32::midpoint(lo, hi);
let sigma_t_local =
local_extinction(zoned, zone_alphas, ray, t_free, hit_t, path_scale, sigma_s);
let pdf_hero = sigma_t_local[hero_idx] * one_minus_u;
let depth_free = depth_to(t_free);
let mut tr_args = [0f32; NUM_CHANNELS];
for (a, d) in tr_args.iter_mut().zip(&depth_free) {
*a = -d;
}
let tr = crate::simd::exp_f32x8(tr_args);
for k in 0..NUM_CHANNELS {
let weight = tr[k] * sigma_s / pdf_hero;
stokes[k] = StokesVector::unpolarized(stokes[k].intensity() * weight);
path_pdf[k] *= sigma_t_local[k] * tr[k];
}
let u1 =
(hash_u32(rng_seed ^ hash_u32(bounce ^ PHASE_DIR_U_STREAM)) as f32) / 4_294_967_295.0;
let u2 =
(hash_u32(rng_seed ^ hash_u32(bounce ^ PHASE_DIR_V_STREAM)) as f32) / 4_294_967_295.0;
let new_dir = sample_henyey_greenstein_direction(u1, u2, g, ray.dir);
Some((t_free, new_dir))
} else {
let mut survive_args = [0f32; NUM_CHANNELS];
for (a, d) in survive_args.iter_mut().zip(&depth_full) {
*a = -d;
}
let survive = crate::simd::exp_f32x8(survive_args);
let survive_hero = survive[hero_idx];
for k in 0..NUM_CHANNELS {
stokes[k] = stokes[k].scale(survive[k] / survive_hero.max(1e-30));
path_pdf[k] *= survive[k];
}
None
}
}
#[derive(Clone, Copy)]
struct ScatterSegment {
ray: Ray,
hit_t: f32,
rng_seed: u32,
bounce: u32,
}
fn scatter_or_survive(
zoned: Option<(&ZonedCache, &ZoneAlphas)>,
alphas: &[f32; NUM_CHANNELS],
material: &GemMaterial,
hero_idx: usize,
segment: ScatterSegment,
stokes: &mut [StokesVector; NUM_CHANNELS],
path_pdf: &mut [f32; NUM_CHANNELS],
) -> Option<(f32, Vec3)> {
let ScatterSegment {
ray,
hit_t,
rng_seed,
bounce,
} = segment;
let Some((zoned_cache, zone_alphas)) = zoned else {
return maybe_scatter_or_extinguish(
alphas,
material.scattering_sigma_s,
material.scattering_g,
hero_idx,
ray.dir,
hit_t,
material.absorption_path_scale,
rng_seed,
bounce,
stokes,
path_pdf,
);
};
maybe_scatter_or_extinguish_zoned(
zoned_cache,
zone_alphas,
&ZonedScatterParams {
sigma_s_model: material.scattering_sigma_s,
g: material.scattering_g,
hero_idx,
ray,
hit_t,
path_scale: material.absorption_path_scale,
rng_seed,
bounce,
},
stokes,
path_pdf,
)
}
pub(in super::super) enum ScatterStepOutcome {
NotApplicable,
ReachedBoundary,
ScatteredAndSurvived(Option<(f32, Vec3)>),
ScatteredAndTerminated,
}
#[expect(
clippy::too_many_arguments,
reason = "bundles the fixed-for-the-trace contexts (mat_ctx, cache, nee), this \
bounce's own state, the RNG stream identity, and the per-ray state a \
scatter event can mutate -- the same shape dispatch_bounce's reason \
explains in transport.rs; `lambdas`/`nee_xyz` are the HDR-map NEE \
contribution's inputs/output, threaded through rather than bundled \
into `nee` since they vary in a way `NeeContext` (fixed for the whole \
trace) deliberately does not; `enable_exit_splitting`/`compat` are this \
moment's own family-integration inputs (see `nee_xyz`'s own doc comment); \
`facet_finishes` (the per-facet surface finish) is threaded the same way \
rather than added to `NeeContext` so the GPU Tier 2 harness's \
`NeeContext` literals stay independent of it"
)]
pub(in super::super) fn try_scatter_step(
mat_ctx: &RayMaterialContext,
cache: &RayWavelengthCache,
material: &GemMaterial,
is_extraordinary: bool,
current_ray: &mut Ray,
current_k: &mut Vec3,
hit_t: f32,
rng_seed: u32,
bounce: u32,
stokes: &mut [StokesVector; NUM_CHANNELS],
path_pdf: &mut [f32; NUM_CHANNELS],
split_radiance: &mut [f32; NUM_CHANNELS],
nee: NeeContext<'_>,
lambdas: &[f32; NUM_CHANNELS],
nee_xyz: &mut Vec3,
enable_exit_splitting: bool,
compat: [u8; NUM_CHANNELS],
facet_finishes: &[FacetFinish],
) -> ScatterStepOutcome {
if material.scattering_sigma_s <= 0.0 {
return ScatterStepOutcome::NotApplicable;
}
let alphas = channel_absorption_alphas_assigned(mat_ctx, cache, *current_k, is_extraordinary);
let zoned_alphas = cache
.zoned
.as_ref()
.map(|zoned| zoned.zone_alphas(mat_ctx, cache, *current_k, is_extraordinary));
let nee = if zoned_alphas.is_some() {
NeeContext {
enabled: false,
..nee
}
} else {
nee
};
let scatter = scatter_or_survive(
cache.zoned.as_ref().zip(zoned_alphas.as_ref()),
&alphas,
material,
mat_ctx.hero_idx,
ScatterSegment {
ray: *current_ray,
hit_t,
rng_seed,
bounce,
},
stokes,
path_pdf,
);
let Some((t_free, new_dir)) = scatter else {
return ScatterStepOutcome::ReachedBoundary;
};
let old_dir = current_ray.dir;
let scatter_point = current_ray.origin + t_free * old_dir;
let mut nee_deposit = [0.0f32; NUM_CHANNELS];
nee_contribution_hg_scatter(
nee,
lambdas,
cache.n_o_ch[mat_ctx.hero_idx],
scatter_point,
old_dir,
material.scattering_g,
rng_seed,
bounce,
stokes,
&mut nee_deposit,
&alphas,
material.scattering_sigma_s,
material.absorption_path_scale,
facet_finishes,
);
*nee_xyz += if enable_exit_splitting {
integrate_channels_to_xyz_families(
&nee_deposit,
lambdas,
path_pdf,
mat_ctx.hero_idx,
compat,
)
} else {
integrate_channels_to_xyz(&nee_deposit, lambdas, path_pdf, mat_ctx.hero_idx)
};
current_ray.origin = scatter_point;
current_ray.dir = new_dir;
*current_k = new_dir;
let phase_pdf_for_mis =
(nee.enabled && matches!(nee.environment, EnvironmentSource::HdrMap(_))).then(|| {
(
henyey_greenstein_phase(new_dir.dot(old_dir), material.scattering_g),
new_dir,
)
});
if bounce > 4 && !apply_russian_roulette(bounce, rng_seed, stokes, split_radiance) {
return ScatterStepOutcome::ScatteredAndTerminated;
}
ScatterStepOutcome::ScatteredAndSurvived(phase_pdf_for_mis)
}
#[expect(
clippy::too_many_arguments,
reason = "mirrors try_scatter_step's own argument shape -- context, this bounce's own \
geometric/RNG state, and the per-channel accumulators an NEE contribution \
reads/writes; see that function's identical justification"
)]
pub(crate) fn nee_contribution_hg_scatter(
nee: NeeContext<'_>,
lambdas: &[f32; NUM_CHANNELS],
n_inside_hero: f32,
scatter_point: Vec3,
scatter_dir_in: Vec3,
g: f32,
rng_seed: u32,
bounce: u32,
stokes: &[StokesVector; NUM_CHANNELS],
radiance: &mut [f32; NUM_CHANNELS],
alphas: &[f32; NUM_CHANNELS],
sigma_s: f32,
absorption_path_scale: f32,
facet_finishes: &[FacetFinish],
) {
if !nee.enabled || !matches!(nee.environment, EnvironmentSource::HdrMap(_)) {
return;
}
let u0 =
(hash_u32(rng_seed ^ hash_u32(bounce ^ NEE_ENV_DIR_U_STREAM)) as f32) / 4_294_967_295.0;
let u1 =
(hash_u32(rng_seed ^ hash_u32(bounce ^ NEE_ENV_DIR_V_STREAM)) as f32) / 4_294_967_295.0;
let Some(sample) = sample_environment_for_nee(nee.environment, u0, u1) else {
return;
};
let probe = Ray {
origin: scatter_point + sample.dir * 1e-4,
dir: sample.dir,
};
let Some(first) = intersect_stone_soa(probe, nee.plane_soa, nee.plane_soa.len(), nee.tools)
else {
return;
};
let Some(exit) = (if nee.tools.is_empty() {
single_exit(first, sample.dir, n_inside_hero, facet_finishes)
} else {
walk_exits(nee, probe, first, n_inside_hero, facet_finishes)
}) else {
return;
};
let ExitProbe {
normal: exit_normal,
t_material: hit_t,
cos_i,
cos_t,
t_unpol,
} = exit;
let phase_cos = sample.dir.dot(scatter_dir_in);
let phase_val = henyey_greenstein_phase(phase_cos, g);
let mis_weight = balance_heuristic(sample.pdf, phase_val);
if mis_weight <= 0.0 {
return;
}
let refracted_dir = (n_inside_hero * sample.dir
- n_inside_hero.mul_add(cos_i, -cos_t) * exit_normal)
.normalize_or_zero();
let EnvironmentSource::HdrMap(env_map) = nee.environment else {
return;
};
let env_rgb = env_map.radiance_rgb(refracted_dir);
let hit_t_scaled = hit_t * absorption_path_scale;
let sigma_s_abs = sigma_s / absorption_path_scale;
let mut trans_args = [0f32; NUM_CHANNELS];
for k in 0..NUM_CHANNELS {
trans_args[k] = -(alphas[k] + sigma_s_abs) * hit_t_scaled;
}
let transmittance = crate::simd::exp_f32x8(trans_args);
let common = t_unpol * phase_val * mis_weight / sample.pdf;
for k in 0..NUM_CHANNELS {
let env_k = crate::renderer::env_map::rgb_to_spectral_radiance(env_rgb, lambdas[k]);
radiance[k] =
(stokes[k].intensity() * transmittance[k] * common * env_k).mul_add(1.0, radiance[k]);
}
}
struct ExitProbe {
normal: Vec3,
t_material: f32,
cos_i: f32,
cos_t: f32,
t_unpol: f32,
}
fn exit_transmittance(n_inside: f32, cos_i: f32) -> Option<(f32, f32)> {
let sin2_t = (n_inside * n_inside * cos_i.mul_add(-cos_i, 1.0)).min(1.0);
if sin2_t >= 1.0 {
return None;
}
let cos_t = (1.0 - sin2_t).max(0.0).sqrt();
let r_s = n_inside.mul_add(cos_i, -cos_t) / n_inside.mul_add(cos_i, cos_t);
let r_p = n_inside.mul_add(-cos_t, cos_i) / n_inside.mul_add(cos_t, cos_i);
let r_unpol = (0.5 * r_p.mul_add(r_p, r_s * r_s)).clamp(0.0, 1.0);
Some((1.0 - r_unpol, cos_t))
}
fn hit_is_frosted(facet_finishes: &[FacetFinish], hit: HitRecord) -> bool {
facet_finishes
.get(hit.facet_idx)
.copied()
.unwrap_or_default()
== FacetFinish::Frosted
}
fn single_exit(
hit: HitRecord,
dir: Vec3,
n_inside: f32,
facet_finishes: &[FacetFinish],
) -> Option<ExitProbe> {
if hit_is_frosted(facet_finishes, hit) {
return None;
}
let cos_i = dir.dot(hit.normal).clamp(0.0, 1.0);
let (t_unpol, cos_t) = exit_transmittance(n_inside, cos_i)?;
Some(ExitProbe {
normal: hit.normal,
t_material: hit.t,
cos_i,
cos_t,
t_unpol,
})
}
fn walk_exits(
nee: NeeContext<'_>,
probe: Ray,
first: HitRecord,
n_inside: f32,
facet_finishes: &[FacetFinish],
) -> Option<ExitProbe> {
let dir = probe.dir;
let mut ray = probe;
let mut hit = first;
let mut crossings = 1;
let mut chain = 1.0f32;
let mut material_len = 0.0f32;
let mut after_entry = false;
let mut last: Option<(Vec3, f32, f32)>;
loop {
if hit_is_frosted(facet_finishes, hit) {
return None;
}
let cos = dir.dot(hit.normal);
if cos > 0.0 {
material_len += if after_entry { hit.t + 1e-4 } else { hit.t };
let cos_i = cos.clamp(0.0, 1.0);
let (t, cos_t) = exit_transmittance(n_inside, cos_i)?;
chain *= t;
last = Some((hit.normal, cos_i, cos_t));
after_entry = false;
} else {
let (t, _) = exit_transmittance(n_inside, (-cos).clamp(0.0, 1.0))?;
chain *= t;
last = None;
after_entry = true;
}
ray = Ray {
origin: ray.origin + hit.t * dir + dir * 1e-4,
dir,
};
let Some(next) = intersect_stone_soa(ray, nee.plane_soa, nee.plane_soa.len(), nee.tools)
else {
break;
};
crossings += 1;
if crossings > HG_NEE_MAX_CROSSINGS {
return None;
}
hit = next;
}
let (normal, cos_i, cos_t) = last?;
Some(ExitProbe {
normal,
t_material: material_len,
cos_i,
cos_t,
t_unpol: chain,
})
}
#[must_use]
pub(crate) fn henyey_greenstein_phase(cos_theta: f32, g: f32) -> f32 {
let g2 = g * g;
let base = (2.0 * g).mul_add(-cos_theta, 1.0 + g2).max(1e-6);
let denom = base * base.sqrt();
(1.0 - g2) / (4.0 * std::f32::consts::PI * denom)
}
#[must_use]
pub(crate) fn sample_henyey_greenstein_direction(u1: f32, u2: f32, g: f32, forward: Vec3) -> Vec3 {
let cos_theta = if g.abs() < 1e-3 {
2.0f32.mul_add(-u1, 1.0)
} else {
let one_minus_g2 = g.mul_add(-g, 1.0);
let denom = (2.0 * g).mul_add(u1, 1.0 - g);
let sq = one_minus_g2 / denom;
sq.mul_add(-sq, g.mul_add(g, 1.0)) / (2.0 * g)
};
let cos_theta = cos_theta.clamp(-1.0, 1.0);
let sin_theta = f32::mul_add(cos_theta, -cos_theta, 1.0).max(0.0).sqrt();
let phi = 2.0 * std::f32::consts::PI * u2;
let (sin_p, cos_p) = phi.sin_cos();
let (t, b) = frosted_orthonormal_basis(forward);
(t * (sin_theta * cos_p) + b * (sin_theta * sin_p) + forward * cos_theta).normalize_or_zero()
}
#[cfg(test)]
mod walk_tests {
use super::*;
use crate::{
geometry::{plane::GpuFacetPlane, tool::ToolPrimitive},
optics::raytracer::{build_plane_soa, environment::LightingPreset},
};
#[test]
fn through_groove_probe_matches_analytic_segments_and_interfaces() {
let planes: Vec<GpuFacetPlane> = [
Vec3::X,
Vec3::NEG_X,
Vec3::Y,
Vec3::NEG_Y,
Vec3::Z,
Vec3::NEG_Z,
]
.into_iter()
.map(|n| GpuFacetPlane::new(n, -0.5))
.collect();
let soa = build_plane_soa(&planes);
let tools = [ToolPrimitive::cylinder(Vec3::ZERO, Vec3::Z, 0.1, 2.0)];
let env = LightingPreset::Daylight.studio(1.0, 0.4, 0.35);
let nee = NeeContext {
environment: env,
plane_soa: &soa,
tools: &tools,
enabled: true,
};
let n = 1.76f32;
let start = Vec3::new(-0.4, 0.05, 0.0);
let probe = Ray {
origin: start + Vec3::X * 1e-4,
dir: Vec3::X,
};
let first = intersect_stone_soa(probe, &soa, planes.len(), &tools).expect("groove wall");
let exit = walk_exits(nee, probe, first, n, &[]).expect("probe escapes");
let half_chord = 0.05f32.mul_add(-0.05, 0.1f32 * 0.1).sqrt();
let seg_a = (-half_chord) - (-0.4);
let seg_b = 0.5 - half_chord;
assert!(
(exit.t_material - (seg_a + seg_b)).abs() < 1e-3,
"material path {} vs {}",
exit.t_material,
seg_a + seg_b
);
let cos_wall = (0.05f32 / 0.1).mul_add(-(0.05f32 / 0.1), 1.0).sqrt();
let t_wall = exit_transmittance(n, cos_wall).expect("no TIR").0;
let t_face = exit_transmittance(n, 1.0).expect("no TIR").0;
let analytic = t_wall * t_wall * t_face;
assert!(
(exit.t_unpol - analytic).abs() < 1e-4,
"chain {} vs {}",
exit.t_unpol,
analytic
);
assert!(exit.t_material > first.t + 0.4);
}
}