pbrt/shared/src/spectra/sampled.rs

use crate::core::pbrt::Float;
use crate::core::spectrum::{SpectrumTrait, StandardSpectra};
use crate::utils::math::{clamp, lerp};
use std::ops::{
    Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
};

pub const CIE_Y_INTEGRAL: Float = 106.856895;

pub const N_SPECTRUM_SAMPLES: usize = 4;
pub const LAMBDA_MIN: i32 = 360;
pub const LAMBDA_MAX: i32 = 830;

#[derive(Debug, Copy, Clone, PartialEq)]
#[repr(C)]
pub struct SampledSpectrum {
    pub values: [Float; N_SPECTRUM_SAMPLES],
}

impl Default for SampledSpectrum {
    fn default() -> Self {
        Self {
            values: [0.0; N_SPECTRUM_SAMPLES],
        }
    }
}

impl SampledSpectrum {
    pub fn new(c: Float) -> Self {
        Self {
            values: [c; N_SPECTRUM_SAMPLES],
        }
    }

    pub fn zero() -> Self {
        Self::default()
    }

    #[inline(always)]
    pub fn from_fn<F>(cb: F) -> Self
    where
        F: FnMut(usize) -> Float,
    {
        Self {
            values: std::array::from_fn(cb),
        }
    }

    pub fn from_array(v: &[Float]) -> Self {
        assert!(N_SPECTRUM_SAMPLES == v.len());
        let mut values = [0.0; N_SPECTRUM_SAMPLES];
        for i in 0..N_SPECTRUM_SAMPLES {
            values[i] = v[i];
        }

        Self { values }
    }

    pub fn is_black(&self) -> bool {
        self.values.iter().all(|&sample| sample == 0.0)
    }

    pub fn has_nans(&self) -> bool {
        self.values.iter().any(|&v| v.is_nan())
    }

    pub fn min_component_value(&self) -> Float {
        self.values.iter().fold(Float::INFINITY, |a, &b| a.min(b))
    }

    pub fn max_component_value(&self) -> Float {
        self.values
            .iter()
            .fold(Float::NEG_INFINITY, |a, &b| a.max(b))
    }

    pub fn average(&self) -> Float {
        self.values.iter().sum::<Float>() / (N_SPECTRUM_SAMPLES as Float)
    }

    pub fn exp(&self) -> SampledSpectrum {
        let values = self.values.map(|v| v.exp());
        let ret = Self { values };
        debug_assert!(!ret.has_nans());
        ret
    }

    pub fn pow_int(&self, mut power: usize) -> Self {
        let mut result = Self::new(1.0);
        let mut base = *self;
        while power > 0 {
            if power % 2 == 1 {
                result *= base;
            }
            base *= base;
            power /= 2;
        }
        result
    }

    pub fn clamp_zero(s: &SampledSpectrum) -> Self {
        let ret = SampledSpectrum::from_fn(|i| s[i].max(0.));
        debug_assert!(!ret.has_nans());
        ret
    }

    pub fn sqrt(&self) -> Self {
        let values = self.values.map(|v| v.sqrt());
        let ret = SampledSpectrum { values };
        debug_assert!(!ret.has_nans());
        ret
    }

    pub fn clamp(s: &SampledSpectrum, low: Float, high: Float) -> Self {
        let ret = SampledSpectrum::from_fn(|i| clamp(s[i], low, high));
        debug_assert!(!ret.has_nans());
        ret
    }

    pub fn safe_div(a: &SampledSpectrum, b: &SampledSpectrum) -> SampledSpectrum {
        SampledSpectrum::from_fn(|i| if b[i] != 0. { a[i] / b[i] } else { 0. })
    }

    pub fn y(&self, lambda: &SampledWavelengths, std: &StandardSpectra) -> Float {
        let ys = std.y.sample(lambda);
        let pdf = lambda.pdf();
        SampledSpectrum::safe_div(&(ys * *self), &pdf).average() / CIE_Y_INTEGRAL
    }
}

impl<'a> IntoIterator for &'a SampledSpectrum {
    type Item = &'a Float;
    type IntoIter = std::slice::Iter<'a, Float>;

    fn into_iter(self) -> Self::IntoIter {
        self.values.iter()
    }
}

impl<'a> IntoIterator for &'a mut SampledSpectrum {
    type Item = &'a mut Float;
    type IntoIter = std::slice::IterMut<'a, Float>;

    fn into_iter(self) -> Self::IntoIter {
        self.values.iter_mut()
    }
}

impl Index<usize> for SampledSpectrum {
    type Output = Float;
    fn index(&self, i: usize) -> &Self::Output {
        &self.values[i]
    }
}

impl IndexMut<usize> for SampledSpectrum {
    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
        &mut self.values[i]
    }
}

impl Add for SampledSpectrum {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] += rhs.values[i];
        }
        ret
    }
}

impl AddAssign for SampledSpectrum {
    fn add_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] += rhs.values[i];
        }
    }
}

impl Sub for SampledSpectrum {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] -= rhs.values[i];
        }
        ret
    }
}

impl SubAssign for SampledSpectrum {
    fn sub_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] -= rhs.values[i];
        }
    }
}

impl Sub<SampledSpectrum> for Float {
    type Output = SampledSpectrum;
    fn sub(self, rhs: SampledSpectrum) -> SampledSpectrum {
        SampledSpectrum::from_fn(|i| self - rhs[i])
    }
}

impl Mul for SampledSpectrum {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] *= rhs.values[i];
        }
        ret
    }
}

impl MulAssign for SampledSpectrum {
    fn mul_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] *= rhs.values[i];
        }
    }
}

impl Mul<Float> for SampledSpectrum {
    type Output = Self;
    fn mul(self, rhs: Float) -> Self {
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] *= rhs;
        }
        ret
    }
}

impl Mul<SampledSpectrum> for Float {
    type Output = SampledSpectrum;
    fn mul(self, rhs: SampledSpectrum) -> SampledSpectrum {
        rhs * self
    }
}

impl MulAssign<Float> for SampledSpectrum {
    fn mul_assign(&mut self, rhs: Float) {
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] *= rhs;
        }
    }
}

impl DivAssign for SampledSpectrum {
    fn div_assign(&mut self, rhs: Self) {
        for i in 0..N_SPECTRUM_SAMPLES {
            debug_assert_ne!(0.0, rhs.values[i]);
            self.values[i] /= rhs.values[i];
        }
    }
}

impl Div for SampledSpectrum {
    type Output = Self;
    fn div(self, rhs: Self) -> Self::Output {
        let mut ret = self;
        ret /= rhs;
        ret
    }
}
impl Div<Float> for SampledSpectrum {
    type Output = Self;

    fn div(self, rhs: Float) -> Self::Output {
        debug_assert_ne!(rhs, 0.0);
        let mut ret = self;
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] /= rhs;
        }
        ret
    }
}

impl DivAssign<Float> for SampledSpectrum {
    fn div_assign(&mut self, rhs: Float) {
        debug_assert_ne!(rhs, 0.0);
        for i in 0..N_SPECTRUM_SAMPLES {
            self.values[i] /= rhs;
        }
    }
}

impl Neg for SampledSpectrum {
    type Output = Self;

    fn neg(self) -> Self::Output {
        let mut ret = SampledSpectrum::new(0.0);
        for i in 0..N_SPECTRUM_SAMPLES {
            ret.values[i] = -self.values[i];
        }
        ret
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct SampledWavelengths {
    pub lambda: [Float; N_SPECTRUM_SAMPLES],
    pub pdf: [Float; N_SPECTRUM_SAMPLES],
}

impl SampledWavelengths {
    pub fn pdf(&self) -> SampledSpectrum {
        SampledSpectrum::from_array(&self.pdf)
    }

    pub fn secondary_terminated(&self) -> bool {
        for i in 1..N_SPECTRUM_SAMPLES {
            if self.pdf[i] != 0.0 {
                return false;
            }
        }
        true
    }

    pub fn terminate_secondary(self) -> Self {
        if self.secondary_terminated() {
            return self;
        }

        let mut new_pdf = [0.0; N_SPECTRUM_SAMPLES];

        new_pdf[0] = self.pdf[0] / (N_SPECTRUM_SAMPLES as Float);

        Self {
            pdf: new_pdf,
            ..self
        }
    }

    pub fn terminate_secondary_inplace(&mut self) {}

    pub fn sample_uniform(u: Float, lambda_min: Float, lambda_max: Float) -> Self {
        let mut lambda = [0.0; N_SPECTRUM_SAMPLES];
        lambda[0] = lerp(u, lambda_min, lambda_min);
        let delta = (lambda_max - lambda_min) / N_SPECTRUM_SAMPLES as Float;
        for i in 1..N_SPECTRUM_SAMPLES {
            lambda[i] = lambda[i - 1] + delta;
            if lambda[i] > lambda_max {
                lambda[i] = lambda_min + (lambda[i] - lambda_max);
            }
        }

        let pdf = [1. / (lambda_max - lambda_min); N_SPECTRUM_SAMPLES];

        Self { lambda, pdf }
    }

    pub fn sample_visible_wavelengths(u: Float) -> Float {
        538.0 - 138.888889 * Float::atanh(0.85691062 - 1.82750197 * u)
    }

    pub fn visible_wavelengths_pdf(lambda: Float) -> Float {
        if !(360.0..830.0).contains(&lambda) {
            return 0.0;
        }
        0.0039398042 / (Float::cosh(0.0072 * (lambda - 538.0))).sqrt()
    }

    pub fn sample_visible(u: Float) -> Self {
        let mut lambda = [0.0; N_SPECTRUM_SAMPLES];
        let mut pdf = [0.0; N_SPECTRUM_SAMPLES];

        for i in 0..N_SPECTRUM_SAMPLES {
            let mut up = u + i as Float / N_SPECTRUM_SAMPLES as Float;
            if up > 1.0 {
                up -= 1.0;
            }
            lambda[i] = Self::sample_visible_wavelengths(up);
            pdf[i] = Self::visible_wavelengths_pdf(lambda[i]);
        }
        Self { lambda, pdf }
    }
}

impl Index<usize> for SampledWavelengths {
    type Output = Float;
    fn index(&self, i: usize) -> &Self::Output {
        &self.lambda[i]
    }
}

impl IndexMut<usize> for SampledWavelengths {
    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
        &mut self.lambda[i]
    }
}