Skip to main content

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}