542 lines
16 KiB
Rust
542 lines
16 KiB
Rust
use crate::core::geometry::primitives::OctahedralVector;
|
|
use crate::core::geometry::{Bounds3f, Normal3f, Point3f, Vector3f, VectorLike};
|
|
use crate::core::geometry::{DirectionCone, Normal};
|
|
use crate::core::light::Light;
|
|
use crate::core::light::{LightBounds, LightSampleContext};
|
|
use crate::spectra::{SampledSpectrum, SampledWavelengths};
|
|
use crate::utils::math::{clamp, lerp, sample_discrete};
|
|
use crate::utils::math::{safe_sqrt, square};
|
|
use crate::utils::ptr::{Ptr, Slice};
|
|
use crate::utils::sampling::AliasTable;
|
|
use crate::{Float, ONE_MINUS_EPSILON, PI};
|
|
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((c - min) / (max - min), 0.0, 1.0)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct SampledLight {
|
|
pub light: Ptr<Light>,
|
|
pub p: Float,
|
|
}
|
|
|
|
impl SampledLight {
|
|
pub fn new(light: Light, p: Float) -> Self {
|
|
Self { light, p }
|
|
}
|
|
}
|
|
|
|
#[enum_dispatch]
|
|
pub trait LightSamplerTrait {
|
|
fn sample_with_context(&self, ctx: &LightSampleContext, u: Float) -> Option<SampledLight>;
|
|
fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Light) -> Float;
|
|
fn sample(&self, u: Float) -> Option<SampledLight>;
|
|
fn pmf(&self, light: &Light) -> Float;
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
#[enum_dispatch(LightSamplerTrait)]
|
|
pub enum LightSampler {
|
|
Uniform(UniformLightSampler),
|
|
Power(PowerLightSampler),
|
|
BVH(BVHLightSampler),
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct UniformLightSampler {
|
|
lights: *const Light,
|
|
lights_len: u32,
|
|
}
|
|
|
|
impl UniformLightSampler {
|
|
pub fn new(lights: *const Light, lights_len: u32) -> Self {
|
|
Self { lights, lights_len }
|
|
}
|
|
|
|
#[inline(always)]
|
|
fn light(&self, idx: usize) -> Light {
|
|
unsafe { *self.lights.add(idx) }
|
|
}
|
|
}
|
|
|
|
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: &Light) -> Float {
|
|
self.pmf(light)
|
|
}
|
|
fn sample(&self, u: Float) -> Option<SampledLight> {
|
|
if self.lights_len == 0 {
|
|
return None;
|
|
}
|
|
|
|
let light_index = (u as u32 * self.lights_len).min(self.lights_len - 1) as usize;
|
|
Some(SampledLight {
|
|
light: self.light(light_index),
|
|
p: 1. / self.lights_len as Float,
|
|
})
|
|
}
|
|
fn pmf(&self, _light: &Light) -> Float {
|
|
if self.lights_len == 0 {
|
|
return 0.;
|
|
}
|
|
1. / self.lights_len as Float
|
|
}
|
|
}
|
|
|
|
#[repr(C)]
|
|
#[derive(Clone, Copy, Debug)]
|
|
pub struct Alias {
|
|
pub q: Float,
|
|
pub alias: u32,
|
|
}
|
|
|
|
#[repr(C)]
|
|
#[derive(Clone, Debug, Copy)]
|
|
pub struct PowerLightSampler {
|
|
pub lights: Slice<Light>,
|
|
pub lights_len: u32,
|
|
pub alias_table: AliasTable,
|
|
}
|
|
|
|
unsafe impl Send for PowerLightSampler {}
|
|
unsafe impl Sync for PowerLightSampler {}
|
|
|
|
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: &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);
|
|
|
|
let light_ref = &self.lights[light_index as usize];
|
|
Some(SampledLight {
|
|
light: Ptr::from(light_ref),
|
|
p: pmf,
|
|
})
|
|
}
|
|
|
|
fn pmf(&self, light: &Light) -> Float {
|
|
if self.lights_len == 0 {
|
|
return 0.0;
|
|
}
|
|
|
|
let light_ptr = light as *const Light;
|
|
let start = self.lights.as_ptr();
|
|
|
|
let end = unsafe { start.add(self.lights.len as usize) };
|
|
|
|
if light_ptr >= start && light_ptr < end {
|
|
let index = unsafe { light_ptr.offset_from(start) };
|
|
return self.alias_table.pmf(index as u32);
|
|
}
|
|
|
|
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 {
|
|
pub nodes: *const LightBVHNode,
|
|
pub lights: *const Light,
|
|
pub infinite_lights: *const Light,
|
|
pub bit_trails: *const u64,
|
|
pub nodes_len: u32,
|
|
pub lights_len: u32,
|
|
pub infinite_lights_len: u32,
|
|
pub all_light_bounds: Bounds3f,
|
|
}
|
|
|
|
unsafe impl Send for BVHLightSampler {}
|
|
unsafe impl Sync for BVHLightSampler {}
|
|
|
|
impl BVHLightSampler {
|
|
#[inline(always)]
|
|
fn node(&self, idx: usize) -> &LightBVHNode {
|
|
unsafe { &*self.nodes.add(idx) }
|
|
}
|
|
|
|
#[inline(always)]
|
|
fn light(&self, idx: usize) -> Light {
|
|
unsafe { *self.lights.add(idx) }
|
|
}
|
|
|
|
#[inline(always)]
|
|
fn infinite_light(&self, idx: usize) -> Light {
|
|
unsafe { *self.infinite_lights.add(idx) }
|
|
}
|
|
|
|
#[inline(always)]
|
|
fn bit_trail(&self, idx: usize) -> u64 {
|
|
unsafe { *self.bit_trails.add(idx) }
|
|
}
|
|
|
|
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_len == 0 { 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;
|
|
return Some(SampledLight::new(self.infinite_light(ind), pmf));
|
|
}
|
|
|
|
if self.nodes_len == 0 {
|
|
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.node(node_ind);
|
|
if !node.is_leaf() {
|
|
let child0_idx = node_ind + 1;
|
|
let child1_idx = node.child_or_light_index() as usize;
|
|
let child0 = self.node(child0_idx);
|
|
let child1 = self.node(child1_idx);
|
|
|
|
let ci: [Float; 2] = [
|
|
child0.light_bounds.importance(p, n, &self.all_light_bounds),
|
|
child1.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 { child0_idx } else { child1_idx };
|
|
} else {
|
|
if node_ind > 0 || node.light_bounds.importance(p, n, &self.all_light_bounds) > 0. {
|
|
let light_idx = node.child_or_light_index() as usize;
|
|
return Some(SampledLight::new(self.light(light_idx), pmf));
|
|
}
|
|
return None;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Light) -> Float {
|
|
let light_ptr = light as *const Light;
|
|
let empty_nodes = if self.nodes_len == 0 { 0. } else { 1. };
|
|
let n_infinite = self.infinite_lights_len as Float;
|
|
|
|
let inf_start = self.infinite_lights;
|
|
let inf_end = unsafe { self.infinite_lights.add(self.infinite_lights_len as usize) };
|
|
if light_ptr >= inf_start && light_ptr < inf_end {
|
|
return 1.0 / (n_infinite + empty_nodes);
|
|
}
|
|
|
|
let finite_start = self.lights;
|
|
let finite_end = unsafe { self.lights.add(self.lights_len as usize) };
|
|
|
|
if light_ptr < finite_start || light_ptr >= finite_end {
|
|
return 0.0;
|
|
}
|
|
|
|
let light_index = unsafe { light_ptr.offset_from(finite_start) as usize };
|
|
|
|
let mut bit_trail = self.bit_trail(light_index);
|
|
|
|
let p_inf = n_infinite / (n_infinite + empty_nodes);
|
|
let mut pmf = 1.0 - p_inf;
|
|
let mut node_ind = 0;
|
|
let p = ctx.p();
|
|
let n = ctx.ns;
|
|
|
|
loop {
|
|
let node = self.node(node_ind);
|
|
if node.is_leaf() {
|
|
return pmf;
|
|
}
|
|
let child0 = self.node(node_ind + 1);
|
|
let child1 = self.node(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),
|
|
];
|
|
|
|
let sum_importance = ci[0] + ci[1];
|
|
if sum_importance == 0.0 {
|
|
return 0.0;
|
|
}
|
|
|
|
let which_child = (bit_trail & 1) as usize;
|
|
|
|
// Update probability: prob of picking the correct child
|
|
pmf *= ci[which_child] / sum_importance;
|
|
|
|
// Advance
|
|
node_ind = if which_child == 1 {
|
|
node.child_or_light_index() as usize
|
|
} else {
|
|
node_ind + 1
|
|
};
|
|
|
|
bit_trail >>= 1;
|
|
}
|
|
}
|
|
|
|
fn sample(&self, u: Float) -> Option<SampledLight> {
|
|
if self.lights_len == 0 {
|
|
return None;
|
|
}
|
|
|
|
let light_ind = (u * self.lights_len as Float).min(self.lights_len as Float - 1.) as usize;
|
|
|
|
Some(SampledLight::new(
|
|
self.light(light_ind),
|
|
1. / self.lights_len as Float,
|
|
))
|
|
}
|
|
|
|
fn pmf(&self, _light: &Light) -> Float {
|
|
if self.lights_len == 0 {
|
|
return 0.;
|
|
}
|
|
|
|
1. / self.lights_len as Float
|
|
}
|
|
}
|