1use rand::Rng;
2use std::f32;
3
4fn get_coefficient_count(order: i32) -> usize {
5 ((order + 1) * (order + 1)) as usize
6}
7
8fn get_idx(l: i32, m: i32) -> usize {
9 (l * (l + 1) + m) as usize
10}
11
12fn img_y_to_theta(y: u32, h: u32) -> f32 {
13 f32::consts::PI * (y as f32 + 0.5) / h as f32
14}
15
16fn img_x_to_phi(x: u32, w: u32) -> f32 {
17 2.0 * f32::consts::PI * (x as f32 + 0.5) / w as f32
18}
19
20fn factorial(x: i32) -> f32 {
21 const CACHE: [f32; 16] = [
22 1.0,
23 1.0,
24 2.0,
25 6.0,
26 24.0,
27 120.0,
28 720.0,
29 5040.0,
30 40320.0,
31 362880.0,
32 3628800.0,
33 39916800.0,
34 479001600.0,
35 6227020800.0,
36 87178291200.0,
37 1307674368000.0,
38 ];
39
40 let idx = x as usize;
41
42 if idx < CACHE.len() {
43 CACHE[idx]
44 } else {
45 let mut s = 1.0;
46 for n in 2..=x {
47 s *= n as f32;
48 }
49 s
50 }
51}
52
53fn to_spherical(dir: &(f32, f32, f32)) -> (f32, f32) {
54 (dir.1.atan2(dir.0), dir.2.max(0.0).min(1.0).acos())
55}
56
57fn eval_legendre_polynomial(l: i32, m: i32, x: f32) -> f32 {
58 let mut pmm = 1.0;
59 if m > 0 {
60 let sign = if m % 2 == 0 { 1.0 } else { -1.0 };
61 pmm = sign * factorial(2 * m - 1) * (1.0 - x * x).powf((m / 2) as f32);
62 }
63
64 if l as i32 == m {
65 return pmm;
66 }
67
68 let mut pmm1 = x * ((2 * m + 1) as f32) * pmm;
69 if l as i32 == m + 1 {
70 return pmm1;
71 }
72
73 for n in m + 2..=l as i32 {
74 let pmn = (x * ((2 * n - 1) as f32) * pmm1 - ((n + m - 1) as f32) * pmm) / ((n - m) as f32);
75 pmm = pmm1;
76 pmm1 = pmn;
77 }
78
79 return pmm1;
80}
81
82fn eval_sh_slow(l: i32, m: i32, phi: f32, theta: f32) -> f32 {
83 let klm = (((2.0 * l as f32 + 1.0) * factorial(l - m.abs()))
84 / (4.0 * f32::consts::PI * factorial(l + m.abs())))
85 .sqrt();
86
87 if m > 0 {
88 2.0f32.sqrt() * klm * (m as f32 * phi).cos() * eval_legendre_polynomial(l, m, theta.cos())
89 } else if m < 0 {
90 2.0f32.sqrt() * klm * (-m as f32 * phi).sin() * eval_legendre_polynomial(l, -m, theta.cos())
91 } else {
92 klm * eval_legendre_polynomial(l, 0, theta.cos())
93 }
94}
95
96fn eval_sh(l: i32, m: i32, phi: f32, theta: f32) -> f32 {
97 eval_sh_slow(l, m, phi, theta)
98}
99
100pub fn project_env(order: i32, img: &Image) -> Vec<(f32, f32, f32)> {
101 let mut coeffs = vec![(0.0f32, 0.0f32, 0.0f32); get_coefficient_count(order)];
102
103 let pixel_area =
104 (2.0 * f32::consts::PI / img.width() as f32) * (f32::consts::PI / img.height() as f32);
105
106 for t in 0..img.height() {
107 let theta = img_y_to_theta(t, img.height());
108 let weight = pixel_area * theta.sin();
109
110 for p in 0..img.width() {
111 let phi = img_x_to_phi(p, img.width());
112 let c = img.get_pixel(p, t);
113
114 for l in 0..=order {
115 for m in -(l as i32)..=l as i32 {
116 let idx = get_idx(l, m);
117 let sh = eval_sh(l, m, phi, theta);
118 let r = (c.0 * sh * weight, c.1 * sh * weight, c.2 * sh * weight);
119
120 let ref mut c = &mut coeffs[idx];
121 c.0 += r.0;
122 c.1 += r.1;
123 c.2 += r.2;
124 }
125 }
126 }
127 }
128
129 coeffs
130}
131
132pub fn project_fn<R: Rng, F: Fn(f32, f32) -> f32>(
133 order: i32,
134 sample_count: u32,
135 rng: &mut R,
136 f: F,
137) -> Vec<f32> {
138 let mut coeffs = vec![0.0f32; get_coefficient_count(order)];
139
140 let sample_side = (sample_count as f32).sqrt().floor() as i32;
141
142 for t in 0..sample_side {
143 for p in 0..sample_side {
144 let alpha = (t as f32 + rng.gen::<f32>()) / sample_side as f32;
145 let beta = (p as f32 + rng.gen::<f32>()) / sample_side as f32;
146
147 let phi = 2.0 * f32::consts::PI * beta;
148 let theta = (2.0 * alpha - 1.0).acos();
149
150 let val = f(phi, theta);
151
152 for l in 0..=order {
153 for m in -(l as i32)..=l as i32 {
154 let idx = get_idx(l, m);
155 let sh = eval_sh(l, m, phi, theta);
156 coeffs[idx] += val * sh;
157 }
158 }
159 }
160 }
161
162 let w = 4.0 * f32::consts::PI / (sample_side * sample_side) as f32;
163 for c in &mut coeffs {
164 *c *= w;
165 }
166
167 coeffs
168}
169
170pub fn project_sparse_samples(
171 order: i32,
172 dirs: &Vec<(f32, f32, f32)>,
173 values: &Vec<f32>,
174) -> Vec<f32> {
175 use nalgebra::*;
176
177 let mut basis_values =
178 DMatrix::<f32>::from_element(dirs.len(), get_coefficient_count(order), 0.0);
179 let func_values = DVector::<f32>::from_iterator(values.len(), values.iter().cloned());
180
181 for (idx, d) in dirs.iter().enumerate() {
182 let (phi, theta) = to_spherical(d);
183 for l in 0..=order {
184 for m in -(l as i32)..=l as i32 {
185 basis_values[(idx, get_idx(l, m))] = eval_sh(l, m, phi, theta);
186 }
187 }
188 }
189
190 let coeffs = basis_values
191 .svd(true, true)
192 .solve(&func_values, f32::EPSILON)
193 .unwrap();
194
195 coeffs.iter().map(|x| *x).collect::<Vec<_>>()
196}
197
198pub struct Image {
199 w: u32,
200 h: u32,
201 data: Vec<(f32, f32, f32)>,
202}
203
204impl Image {
205 fn width(&self) -> u32 {
206 self.w
207 }
208
209 fn height(&self) -> u32 {
210 self.h
211 }
212
213 fn get_pixel(&self, x: u32, y: u32) -> (f32, f32, f32) {
214 self.data[(x + y * self.w) as usize]
215 }
216}
217
218