521 lines
16 KiB
Rust
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
|
|
}
|
|
}
|