scirs2_vision/nerf/positional_encoding.rs
1//! Positional (Fourier) encoding for NeRF inputs.
2//!
3//! Implements the sinusoidal positional encoding from Mildenhall et al. 2020
4//! ("NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis").
5//!
6//! For a scalar `x` and `L` frequency levels, the encoding is:
7//!
8//! ```text
9//! γ(x) = [x, sin(2⁰ π x), cos(2⁰ π x), sin(2¹ π x), cos(2¹ π x), …,
10//! sin(2^{L-1} π x), cos(2^{L-1} π x)]
11//! ```
12//!
13//! giving `1 + 2L` output features per scalar.
14
15use std::f64::consts::PI;
16
17/// Encode a single scalar value with `n_freq` frequency bands.
18///
19/// Output length: `1 + 2 * n_freq`.
20///
21/// # Arguments
22///
23/// * `x` – scalar input value.
24/// * `n_freq` – number of frequency bands L.
25pub fn positional_encode(x: f64, n_freq: usize) -> Vec<f64> {
26 let out_len = 1 + 2 * n_freq;
27 let mut out = Vec::with_capacity(out_len);
28 out.push(x);
29 for k in 0..n_freq {
30 let freq = (1u64 << k) as f64 * PI; // 2^k · π
31 out.push((freq * x).sin());
32 out.push((freq * x).cos());
33 }
34 out
35}
36
37/// Encode a 3-D position vector with `n_freq` frequency bands per component.
38///
39/// Each of the three components is independently encoded; the results are
40/// concatenated:
41///
42/// ```text
43/// output = [γ(x), γ(y), γ(z)]
44/// ```
45///
46/// Output length: `3 * (1 + 2 * n_freq)`.
47///
48/// # Arguments
49///
50/// * `pos` – 3-D world-space position `[x, y, z]`.
51/// * `n_freq` – number of frequency bands L.
52pub fn encode_position(pos: &[f64; 3], n_freq: usize) -> Vec<f64> {
53 let component_len = 1 + 2 * n_freq;
54 let mut out = Vec::with_capacity(3 * component_len);
55 for &coord in pos.iter() {
56 out.extend_from_slice(&positional_encode(coord, n_freq));
57 }
58 out
59}
60
61/// Encode a unit viewing-direction vector with `n_freq` frequency bands per component.
62///
63/// Identical in structure to [`encode_position`] but intended for the lower-frequency
64/// direction branch of NeRF (typically `n_freq` = 4).
65///
66/// Output length: `3 * (1 + 2 * n_freq)`.
67///
68/// # Arguments
69///
70/// * `dir` – unit-length view direction `[dx, dy, dz]`.
71/// * `n_freq` – number of frequency bands L.
72pub fn encode_direction(dir: &[f64; 3], n_freq: usize) -> Vec<f64> {
73 encode_position(dir, n_freq)
74}
75
76/// Compute the expected output dimensionality for [`encode_position`] /
77/// [`encode_direction`] with a given number of frequency bands.
78#[inline]
79pub fn encoding_dim(n_freq: usize) -> usize {
80 3 * (1 + 2 * n_freq)
81}
82
83#[cfg(test)]
84mod tests {
85 use super::*;
86
87 #[test]
88 fn test_positional_encode_dim() {
89 for n_freq in [1_usize, 4, 10] {
90 let enc = positional_encode(0.5, n_freq);
91 assert_eq!(
92 enc.len(),
93 1 + 2 * n_freq,
94 "n_freq={n_freq}: expected dim {}, got {}",
95 1 + 2 * n_freq,
96 enc.len()
97 );
98 }
99 }
100
101 #[test]
102 fn test_encode_position_dim() {
103 for n_freq in [1_usize, 4, 10] {
104 let enc = encode_position(&[0.1, 0.2, 0.3], n_freq);
105 assert_eq!(enc.len(), 3 * (1 + 2 * n_freq));
106 }
107 }
108
109 #[test]
110 fn test_positional_encode_zero() {
111 // γ(0) = [0, sin(0), cos(0), sin(0), cos(0), …]
112 // = [0, 0, 1, 0, 1, …]
113 let enc = positional_encode(0.0, 3);
114 assert!((enc[0] - 0.0).abs() < 1e-12); // identity
115 for k in 0..3 {
116 let sin_idx = 1 + 2 * k;
117 let cos_idx = 2 + 2 * k;
118 assert!(
119 (enc[sin_idx] - 0.0).abs() < 1e-12,
120 "sin at k={k} should be 0, got {}",
121 enc[sin_idx]
122 );
123 assert!(
124 (enc[cos_idx] - 1.0).abs() < 1e-12,
125 "cos at k={k} should be 1, got {}",
126 enc[cos_idx]
127 );
128 }
129 }
130
131 #[test]
132 fn test_encoding_dim_helper() {
133 assert_eq!(encoding_dim(10), 3 * 21);
134 assert_eq!(encoding_dim(4), 3 * 9);
135 }
136}