pbrt_r3/core/sampling/
distribution.rs

1use crate::core::base::*;
2
3#[derive(Debug, Default, Clone)]
4pub struct Distribution1D {
5    pub func: Vec<Float>,
6    pub cdf: Vec<Float>,
7    pub func_int: Float,
8}
9
10impl Distribution1D {
11    pub fn new(f: &[Float]) -> Self {
12        let n = f.len();
13        let func = Vec::from(f);
14        let mut cdf = vec![0.0; n + 1];
15        cdf[0] = 0.0;
16        for i in 1..(n + 1) {
17            cdf[i] = cdf[i - 1] + func[i - 1] / (n as Float);
18        }
19        let func_int = cdf[n];
20        if func_int == 0.0 {
21            for i in 1..(n + 1) {
22                cdf[i] = (i as Float) / (n as Float);
23            }
24        } else {
25            for i in 1..(n + 1) {
26                cdf[i] /= func_int;
27            }
28        }
29        Distribution1D {
30            func,
31            cdf,
32            func_int,
33        }
34    }
35
36    pub fn count(&self) -> usize {
37        return self.func.len();
38    }
39
40    //value, pdf, offset
41    pub fn sample_continuous(&self, u: Float) -> (Float, Float, usize) {
42        let offset = find_interval(&self.cdf, &|vv, i| -> bool {
43            return vv[i] <= u;
44        });
45
46        let mut du = u - self.cdf[offset];
47        if (self.cdf[offset + 1] - self.cdf[offset]) > 0.0 {
48            du /= self.cdf[offset + 1] - self.cdf[offset];
49        }
50        // Compute PDF for sampled offset
51        let pdf = if self.func_int > 0.0 {
52            self.func[offset] / self.func_int
53        } else {
54            0.0
55        };
56        let r = ((offset as Float) + du) / (self.count() as Float);
57        return (r, pdf, offset);
58    }
59
60    //offset, pdf, remapped
61    pub fn sample_discrete(&self, u: Float) -> (usize, Float, Float) {
62        let offset = find_interval(&self.cdf, &|vv, i| -> bool {
63            return vv[i] <= u;
64        });
65        assert!(self.cdf[offset] <= u);
66        assert!(u <= self.cdf[offset + 1]);
67        let pdf = if self.func_int > 0.0 {
68            self.func[offset] / (self.func_int * self.count() as Float)
69        } else {
70            0.0
71        };
72        let remapped = (u - self.cdf[offset]) / (self.cdf[offset + 1] - self.cdf[offset]);
73        assert!(remapped >= 0.0 && remapped <= 1.0);
74        return (offset, pdf, remapped);
75    }
76
77    pub fn discrete_pdf(&self, index: usize) -> Float {
78        return self.func[index] / (self.func_int * self.count() as Float);
79    }
80}
81
82pub struct Distribution2D {
83    pub conditional_v: Vec<Box<Distribution1D>>,
84    pub marginal: Box<Distribution1D>,
85}
86
87impl Distribution2D {
88    pub fn new(data: &[Float], nu: usize, nv: usize) -> Self {
89        let mut conditional_v = Vec::with_capacity(nv);
90        for v in 0..nv {
91            let a = v * nu;
92            let b = a + nu;
93            let s = &data[a..b];
94            conditional_v.push(Box::new(Distribution1D::new(s)));
95        }
96        let mut marginal_func = Vec::with_capacity(nv);
97        for v in 0..nv {
98            marginal_func.push(conditional_v[v].func_int);
99        }
100        Distribution2D {
101            conditional_v,
102            marginal: Box::new(Distribution1D::new(&marginal_func)),
103        }
104    }
105
106    pub fn sample_continuous(&self, u: &Point2f) -> (Point2f, Float) {
107        let (d1, pdf1, v) = self.marginal.sample_continuous(u[1]);
108        let (d0, pdf0, _) = self.conditional_v[v].as_ref().sample_continuous(u[0]);
109        return (Point2f::new(d0, d1), pdf0 * pdf1);
110    }
111    pub fn pdf(&self, p: &Point2f) -> Float {
112        let ucount = self.conditional_v[0].count();
113        let vcount = self.marginal.count();
114        let iu = usize::clamp((p[0] * ucount as Float) as usize, 0, ucount - 1);
115        let iv = usize::clamp((p[1] * vcount as Float) as usize, 0, vcount - 1);
116        return self.conditional_v[iv].func[iu] / self.marginal.func_int;
117    }
118}