pbrt/src/lights/sampler.rs

521 lines
16 KiB
Rust

use super::{
Bounds3f, Float, Light, LightBounds, LightSampleContext, LightTrait, Normal3f, PI, Point3f,
SampledSpectrum, SampledWavelengths, Vector3f, VectorLike, safe_sqrt, square,
};
use crate::geometry::primitives::OctahedralVector;
use crate::geometry::{DirectionCone, Normal};
use crate::utils::math::sample_discrete;
use std::collections::HashMap;
use std::sync::Arc;
use crate::core::pbrt::{ONE_MINUS_EPSILON, clamp_t, lerp};
use crate::utils::sampling::AliasTable;
use enum_dispatch::enum_dispatch;
#[derive(Clone, Copy, Debug, Default)]
#[repr(C)]
pub struct CompactLightBounds {
pub w: OctahedralVector,
pub phi: Float,
// [0..15] = qCosTheta_o
// [15..30] = qCosTheta_e
// [30..31] = twoSided
// [31..32] = Unused/Padding
packed_info: u32,
pub qb: [[u16; 3]; 2],
}
const _: () = assert!(std::mem::size_of::<CompactLightBounds>() == 24);
impl CompactLightBounds {
pub fn new(lb: &LightBounds, all_b: &Bounds3f) -> Self {
let q_cos_o = Self::quantize_cos(lb.cos_theta_o);
let q_cos_e = Self::quantize_cos(lb.cos_theta_e);
let two_sided = if lb.two_sided { 1 } else { 0 };
// | - twoSided (1) - | - qCosTheta_e (15) - | - qCosTheta_o (15) - |
let packed_info = (q_cos_o & 0x7FFF) | ((q_cos_e & 0x7FFF) << 15) | (two_sided << 30);
let mut qb = [[0u16; 3]; 2];
for i in 0..3 {
qb[0][i] = Self::quantize_bounds(lb.bounds.p_min[i], all_b.p_min[i], all_b.p_max[i])
.floor() as u16;
qb[1][i] = Self::quantize_bounds(lb.bounds.p_max[i], all_b.p_min[i], all_b.p_max[i])
.ceil() as u16;
}
Self {
w: OctahedralVector::new(lb.w.normalize()),
phi: lb.phi,
packed_info,
qb,
}
}
#[inline(always)]
pub fn two_sided(&self) -> bool {
(self.packed_info >> 30) & 1 == 1
}
#[inline(always)]
fn q_cos_theta_o(&self) -> u32 {
self.packed_info & 0x7FFF
}
#[inline(always)]
fn q_cos_theta_e(&self) -> u32 {
(self.packed_info >> 15) & 0x7FFF
}
#[inline]
pub fn cos_theta_o(&self) -> Float {
2.0 * (self.q_cos_theta_o() as Float / 32767.0) - 1.0
}
#[inline]
pub fn cos_theta_e(&self) -> Float {
2.0 * (self.q_cos_theta_e() as Float / 32767.0) - 1.0
}
pub fn importance(&self, p: Point3f, n: Normal3f, all_b: &Bounds3f) -> Float {
let bounds = self.bounds(all_b);
let cos_o = self.cos_theta_o();
let cos_e = self.cos_theta_e();
let pc = bounds.centroid();
let d2 = p.distance_squared(pc).max(bounds.diagonal().norm() * 0.5);
let cos_sub_clamped = |sin_a: Float, cos_a: Float, sin_b: Float, cos_b: Float| {
if cos_a > cos_b {
1.0
} else {
cos_a * cos_b + sin_a * sin_b
}
};
let sin_sub_clamped = |sin_a: Float, cos_a: Float, sin_b: Float, cos_b: Float| {
if cos_a > cos_b {
0.0
} else {
sin_a * cos_b - cos_a * sin_b
}
};
let wi = (p - pc).normalize();
let w_vec = self.w.to_vector();
let mut cos_w = w_vec.dot(wi);
if self.two_sided() {
cos_w = cos_w.abs();
}
let sin_w = safe_sqrt(1.0 - square(cos_w));
let cos_b = DirectionCone::bound_subtended_directions(&bounds, p).cos_theta;
let sin_b = safe_sqrt(1. - square(cos_b));
let sin_o = safe_sqrt(1. - square(cos_o));
let cos_x = cos_sub_clamped(sin_w, cos_w, sin_o, cos_o);
let sin_x = sin_sub_clamped(sin_w, cos_w, sin_o, cos_o);
let cos_p = cos_sub_clamped(sin_x, cos_x, sin_b, cos_b);
if cos_p <= cos_e {
return 0.;
}
let mut importance = self.phi * cos_p / d2;
if n != Normal3f::zero() {
let cos_i = wi.abs_dot(n.into());
let sin_i = safe_sqrt(1. - square(cos_i));
let cos_pi = cos_sub_clamped(sin_i, cos_i, sin_b, cos_b);
importance *= cos_pi;
}
importance
}
pub fn bounds(&self, all_b: &Bounds3f) -> Bounds3f {
let mut p_min = Point3f::default();
let mut p_max = Point3f::default();
for i in 0..3 {
let t_min = self.qb[0][i] as Float / 65535.0;
let t_max = self.qb[1][i] as Float / 65535.0;
p_min[i] = lerp(t_min, all_b.p_min[i], all_b.p_max[i]);
p_max[i] = lerp(t_max, all_b.p_min[i], all_b.p_max[i]);
}
Bounds3f::from_points(p_min, p_max)
}
fn quantize_cos(c: Float) -> u32 {
(32767.0 * ((c + 1.0) * 0.5)).floor() as u32
}
fn quantize_bounds(c: Float, min: Float, max: Float) -> Float {
if min == max {
return 0.0;
}
65535.0 * clamp_t((c - min) / (max - min), 0.0, 1.0)
}
}
#[derive(Debug, Clone)]
pub struct SampledLight {
pub light: Arc<Light>,
pub p: Float,
}
impl SampledLight {
pub fn new(light: Arc<Light>, p: Float) -> Self {
Self { light, p }
}
}
#[enum_dispatch]
pub trait LightSamplerTrait: Send + Sync + std::fmt::Debug {
fn sample_with_context(&self, ctx: &LightSampleContext, u: Float) -> Option<SampledLight>;
fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Arc<Light>) -> Float;
fn sample(&self, u: Float) -> Option<SampledLight>;
fn pmf(&self, light: &Arc<Light>) -> Float;
}
#[derive(Clone, Debug)]
#[enum_dispatch(LightSamplerTrait)]
pub enum LightSampler {
Uniform(UniformLightSampler),
Power(PowerLightSampler),
BVH(BVHLightSampler),
}
#[derive(Clone, Debug)]
pub struct UniformLightSampler {
lights: Vec<Arc<Light>>,
}
impl UniformLightSampler {
pub fn new(lights: &[Arc<Light>]) -> Self {
Self {
lights: lights.to_vec(),
}
}
}
impl LightSamplerTrait for UniformLightSampler {
fn sample_with_context(&self, _ctx: &LightSampleContext, u: Float) -> Option<SampledLight> {
self.sample(u)
}
fn pmf_with_context(&self, _ctx: &LightSampleContext, light: &Arc<Light>) -> Float {
self.pmf(light)
}
fn sample(&self, u: Float) -> Option<SampledLight> {
if self.lights.is_empty() {
return None;
}
let light_index = (u as usize * self.lights.len()).min(self.lights.len() - 1);
Some(SampledLight {
light: self.lights[light_index].clone(),
p: 1. / self.lights.len() as Float,
})
}
fn pmf(&self, _light: &Arc<Light>) -> Float {
if self.lights.is_empty() {
return 0.;
}
1. / self.lights.len() as Float
}
}
#[derive(Clone, Debug)]
pub struct PowerLightSampler {
lights: Vec<Arc<Light>>,
light_to_index: HashMap<usize, usize>,
alias_table: AliasTable,
}
impl PowerLightSampler {
pub fn new(lights: &[Arc<Light>]) -> Self {
if lights.is_empty() {
return Self {
lights: Vec::new(),
light_to_index: HashMap::new(),
alias_table: AliasTable::new(&[]),
};
}
let mut lights_vec = Vec::with_capacity(lights.len());
let mut light_to_index = HashMap::with_capacity(lights.len());
let mut light_power = Vec::with_capacity(lights.len());
let lambda = SampledWavelengths::sample_visible(0.5);
for (i, light) in lights.iter().enumerate() {
lights_vec.push(light.clone());
let ptr = Arc::as_ptr(light) as usize;
light_to_index.insert(ptr, i);
let phi = SampledSpectrum::safe_div(&light.phi(lambda), &lambda.pdf());
light_power.push(phi.average());
}
let alias_table = AliasTable::new(&light_power);
Self {
lights: lights_vec,
light_to_index,
alias_table,
}
}
}
impl LightSamplerTrait for PowerLightSampler {
fn sample_with_context(&self, _ctx: &LightSampleContext, u: Float) -> Option<SampledLight> {
self.sample(u)
}
fn pmf_with_context(&self, _ctx: &LightSampleContext, light: &Arc<Light>) -> Float {
self.pmf(light)
}
fn sample(&self, u: Float) -> Option<SampledLight> {
if self.alias_table.size() == 0 {
return None;
}
let (light_index, pmf, _) = self.alias_table.sample(u);
Some(SampledLight {
light: self.lights[light_index].clone(),
p: pmf,
})
}
fn pmf(&self, light: &Arc<Light>) -> Float {
if self.alias_table.size() == 0 {
return 0.;
}
let ptr = Arc::as_ptr(light) as usize;
if let Some(&index) = self.light_to_index.get(&ptr) {
self.alias_table.pmf(index)
} else {
0.0
}
}
}
#[derive(Clone, Copy, Debug, Default)]
#[repr(C, align(32))]
pub struct LightBVHNode {
pub light_bounds: CompactLightBounds,
// Bit 31 (MSB) : isLeaf (1 bit)
// Bits 0..31 : childOrLightIndex (31 bits)
packed_data: u32,
}
const _: () = assert!(std::mem::size_of::<LightBVHNode>() == 32);
impl LightBVHNode {
/// Mask to isolate the Leaf Flag (Bit 31)
const LEAF_MASK: u32 = 0x8000_0000;
/// Mask to isolate the Index (Bits 0-30)
const INDEX_MASK: u32 = 0x7FFF_FFFF;
pub fn make_leaf(light_index: u32, cb: CompactLightBounds) -> Self {
debug_assert!(
(light_index & Self::LEAF_MASK) == 0,
"Light index too large"
);
Self {
light_bounds: cb,
// Set index and flip the MSB to 1
packed_data: light_index | Self::LEAF_MASK,
}
}
pub fn make_interior(child_index: u32, cb: CompactLightBounds) -> Self {
debug_assert!(
(child_index & Self::LEAF_MASK) == 0,
"Child index too large"
);
Self {
light_bounds: cb,
// Set index, MSB remains 0
packed_data: child_index,
}
}
#[inline(always)]
pub fn is_leaf(&self) -> bool {
(self.packed_data & Self::LEAF_MASK) != 0
}
#[inline(always)]
pub fn light_index(&self) -> u32 {
debug_assert!(self.is_leaf());
self.packed_data & Self::INDEX_MASK
}
#[inline(always)]
pub fn child_index(&self) -> u32 {
debug_assert!(!self.is_leaf());
self.packed_data & Self::INDEX_MASK
}
#[inline(always)]
pub fn child_or_light_index(&self) -> u32 {
self.packed_data & Self::INDEX_MASK
}
pub fn sample(&self, _ctx: &LightSampleContext, _u: Float) -> Option<SampledLight> {
todo!("Implement LightBVHNode::Sample logic")
}
}
#[derive(Clone, Debug)]
pub struct BVHLightSampler {
lights: Vec<Arc<Light>>,
infinite_lights: Vec<Arc<Light>>,
all_light_bounds: Bounds3f,
nodes: Vec<LightBVHNode>,
light_to_bit_trail: HashMap<usize, usize>,
}
impl BVHLightSampler {
fn evaluate_cost(&self, b: &LightBounds, bounds: &Bounds3f, dim: usize) -> Float {
let theta_o = b.cos_theta_o.acos();
let theta_e = b.cos_theta_e.acos();
let theta_w = (theta_o + theta_e).min(PI);
let sin_o = safe_sqrt(1. - square(b.cos_theta_o));
let m_omega = 2. * PI * (1. - b.cos_theta_o)
+ PI / 2.
* (2. * theta_w - (theta_o - 2. * theta_w).cos() - 2. * theta_o * sin_o
+ b.cos_theta_o);
let kr = bounds.diagonal().max_component_value() / bounds.diagonal()[dim];
b.phi * m_omega * kr * b.bounds.surface_area()
}
}
impl LightSamplerTrait for BVHLightSampler {
fn sample_with_context(&self, ctx: &LightSampleContext, mut u: Float) -> Option<SampledLight> {
let empty_nodes = if self.nodes.is_empty() { 0. } else { 1. };
let inf_size = self.infinite_lights.len() as Float;
let light_size = self.lights.len() as Float;
let p_inf = inf_size / (inf_size + empty_nodes);
if u < p_inf {
u /= p_inf;
let ind = (u * light_size).min(light_size - 1.) as usize;
let pmf = p_inf / inf_size;
Some(SampledLight::new(self.infinite_lights[ind].clone(), pmf))
} else {
if self.nodes.is_empty() {
return None;
}
let p = ctx.p();
let n = ctx.ns;
u = ((u - p_inf) / (1. - p_inf)).min(ONE_MINUS_EPSILON);
let mut node_ind = 0;
let mut pmf = 1. - p_inf;
loop {
let node = self.nodes[node_ind];
if !node.is_leaf() {
let children: [LightBVHNode; 2] = [
self.nodes[node_ind + 1],
self.nodes[node.child_or_light_index() as usize],
];
let ci: [Float; 2] = [
children[0]
.light_bounds
.importance(p, n, &self.all_light_bounds),
children[1]
.light_bounds
.importance(p, n, &self.all_light_bounds),
];
if ci[0] == 0. && ci[1] == 0. {
return None;
}
let mut node_pmf: Float = 0.;
let child = sample_discrete(&ci, u, Some(&mut node_pmf), Some(&mut u));
pmf *= node_pmf;
node_ind = if child == 0 {
node_ind + 1
} else {
node.child_or_light_index() as usize
};
} else {
if node_ind > 0
|| node.light_bounds.importance(p, n, &self.all_light_bounds) > 0.
{
return Some(SampledLight::new(
self.lights[node.child_or_light_index() as usize].clone(),
pmf,
));
}
return None;
}
}
}
}
fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Arc<Light>) -> Float {
let ptr = Arc::as_ptr(light) as usize;
let empty_nodes = if self.nodes.is_empty() { 0. } else { 1. };
if self.light_to_bit_trail.contains_key(&ptr) {
return 1. / (self.infinite_lights.len() as Float + empty_nodes);
}
let mut bit_trail = self.light_to_bit_trail[&ptr];
let p = ctx.p();
let n = ctx.ns;
let p_inf = self.infinite_lights.len() as Float
/ (self.infinite_lights.len() as Float + empty_nodes);
let mut pmf = 1. - p_inf;
let mut node_ind = 0;
loop {
let node = self.nodes[node_ind];
if node.is_leaf() {
return pmf;
}
let child0 = self.nodes[node_ind + 1];
let child1 = self.nodes[node.child_or_light_index() as usize];
let ci = [
child0.light_bounds.importance(p, n, &self.all_light_bounds),
child1.light_bounds.importance(p, n, &self.all_light_bounds),
];
pmf *= ci[bit_trail & 1] / (ci[0] + ci[1]);
node_ind = if (bit_trail & 1) != 0 {
node.child_or_light_index() as usize
} else {
node_ind + 1
};
bit_trail >>= 1;
}
}
fn sample(&self, u: Float) -> Option<SampledLight> {
if self.lights.is_empty() {
return None;
}
let light_ind =
(u * self.lights.len() as Float).min(self.lights.len() as Float - 1.) as usize;
Some(SampledLight::new(
self.lights[light_ind].clone(),
1. / self.lights.len() as Float,
))
}
fn pmf(&self, _light: &Arc<Light>) -> Float {
if self.lights.is_empty() {
return 0.;
}
1. / self.lights.len() as Float
}
}