Skip to main content

leodos_protocols/application/compression/
spectral.rs

1//! CCSDS 122.1-B-1 Spectral Preprocessing Transform.
2//!
3//! Provides spectral decorrelation for multispectral and
4//! hyperspectral images before 2D compression (CCSDS 122.0).
5//!
6//! The pipeline is:
7//!   Input → Upshift → Spectral Transform → Downshift → 2D Encoder
8//!
9//! Three spectral transforms are defined:
10//! - Identity (passthrough)
11//! - IWT (Integer Wavelet Transform, CDF 5/3, 5 levels)
12//! - POT (Pairwise Orthogonal Transform, approximates KLT)
13//!
14//! The IWT is applied independently along the spectral axis (z)
15//! for each spatial pixel (x, y).
16
17/// Upshift: multiply each sample by 2^u (Section 3.2.2, Eq. 11).
18pub fn upshift(image: &mut [i32], u: u32) {
19    for sample in image.iter_mut() {
20        *sample <<= u;
21    }
22}
23
24/// Downshift: multiply each sample by 2^-d, rounding to nearest
25/// integer (Section 3.2.3, Eq. 13).
26pub fn downshift(image: &mut [i32], d: u32) {
27    if d == 0 {
28        return;
29    }
30    let half = 1i32 << (d - 1);
31    for sample in image.iter_mut() {
32        *sample = (*sample + half) >> d;
33    }
34}
35
36/// Number of IWT decomposition levels (Section 4.4.4).
37pub const IWT_LEVELS: usize = 5;
38
39/// Single-level forward IWT decomposition (CDF 5/3 lifting).
40///
41/// Splits input `x` of length `n` into low-frequency `l` and
42/// high-frequency `h` coefficients (Section 4.4.3.1, Eq. 22-23).
43pub fn iwt_forward_single(x: &[i32], l: &mut [i32], h: &mut [i32]) {
44    let n = x.len();
45    if n <= 1 {
46        if n == 1 {
47            l[0] = x[0];
48        }
49        return;
50    }
51
52    let p = (n + 1) / 2;
53    let q = n / 2;
54
55    for k in 0..q {
56        let is_last = k == q - 1 && n % 2 == 0;
57        h[k] = if is_last {
58            x[2 * k + 1] - x[2 * k]
59        } else {
60            x[2 * k + 1] - ((x[2 * k] + x[2 * k + 2]) / 2)
61        };
62    }
63
64    for k in 0..p {
65        l[k] = if k == 0 {
66            x[0] + ((h[0] + 1) / 2)
67        } else if k == p - 1 && n % 2 != 0 {
68            x[2 * k] + ((h[k - 1] + 1) / 2)
69        } else {
70            x[2 * k] + ((h[k - 1] + h[k] + 2) / 4)
71        };
72    }
73}
74
75/// Single-level inverse IWT decomposition (Section 4.4.3.2, Eq. 25-26).
76pub fn iwt_inverse_single(l: &[i32], h: &[i32], x: &mut [i32]) {
77    let p = l.len();
78    let q = h.len();
79    let n = p + q;
80
81    if n <= 1 {
82        if n == 1 {
83            x[0] = l[0];
84        }
85        return;
86    }
87
88    for k in 0..p {
89        x[2 * k] = if k == 0 {
90            l[0] - ((h[0] + 1) / 2)
91        } else if k == p - 1 && n % 2 != 0 {
92            l[k] - ((h[k - 1] + 1) / 2)
93        } else {
94            l[k] - ((h[k - 1] + h[k] + 2) / 4)
95        };
96    }
97
98    for k in 0..q {
99        let is_last = k == q - 1 && n % 2 == 0;
100        x[2 * k + 1] = if is_last {
101            h[k] + x[2 * k]
102        } else {
103            h[k] + ((x[2 * k] + x[2 * k + 2]) / 2)
104        };
105    }
106}
107
108/// Five-level forward IWT decomposition (Section 4.4.4).
109///
110/// Output order: L5, H5, H4, H3, H2, H1 (Eq. 27).
111/// `buf` must have length >= `n`. `scratch` must have length >= `n`.
112pub fn iwt_forward_5(input: &[i32], output: &mut [i32], scratch: &mut [i32]) {
113    let n = input.len();
114    output[..n].copy_from_slice(input);
115
116    let mut current_len = n;
117    for _ in 0..IWT_LEVELS {
118        if current_len <= 1 {
119            break;
120        }
121        let p = (current_len + 1) / 2;
122        let q = current_len / 2;
123
124        let (src, _) = output.split_at(current_len);
125        let src_copy: heapless::Vec<i32, 4096> = src.iter().copied().collect();
126
127        let (scratch_l, scratch_h) = scratch.split_at_mut(p);
128        iwt_forward_single(&src_copy[..current_len], scratch_l, &mut scratch_h[..q]);
129
130        output[..p].copy_from_slice(&scratch[..p]);
131        output[p..p + q].copy_from_slice(&scratch[p..p + q]);
132
133        current_len = p;
134    }
135}
136
137/// Five-level inverse IWT decomposition (Section 4.4.4.2).
138pub fn iwt_inverse_5(input: &[i32], output: &mut [i32], scratch: &mut [i32]) {
139    let n = input.len();
140    output[..n].copy_from_slice(input);
141
142    let mut sizes = [0usize; IWT_LEVELS + 1];
143    sizes[0] = n;
144    for i in 0..IWT_LEVELS {
145        sizes[i + 1] = (sizes[i] + 1) / 2;
146        if sizes[i] <= 1 {
147            break;
148        }
149    }
150
151    let levels_used = sizes.iter().take_while(|&&s| s > 1).count().min(IWT_LEVELS);
152
153    for level in (0..levels_used).rev() {
154        let current_len = sizes[level];
155        let p = (current_len + 1) / 2;
156        let q = current_len / 2;
157
158        scratch[..p].copy_from_slice(&output[..p]);
159        scratch[p..p + q].copy_from_slice(&output[p..p + q]);
160
161        iwt_inverse_single(&scratch[..p], &scratch[p..p + q], &mut output[..current_len]);
162    }
163}
164
165/// Spectral transform selection.
166#[derive(Debug, Copy, Clone, Eq, PartialEq)]
167pub enum SpectralTransform {
168    /// No spectral decorrelation (Section 4.3).
169    Identity,
170    /// Integer Wavelet Transform, 5 levels (Section 4.4).
171    Iwt,
172}
173
174/// Apply the spectral transform along the z-axis for a 3D image.
175///
176/// Image layout: `image[z * nx * ny + y * nx + x]` (band-interleaved by pixel).
177/// `nx` × `ny` spatial, `nz` spectral bands.
178pub fn transform_spectral(
179    image: &mut [i32],
180    nx: usize,
181    ny: usize,
182    nz: usize,
183    transform: SpectralTransform,
184) {
185    match transform {
186        SpectralTransform::Identity => {}
187        SpectralTransform::Iwt => {
188            let mut col = [0i32; 4096];
189            let mut out = [0i32; 4096];
190            let mut scratch = [0i32; 4096];
191            assert!(nz <= 4096);
192
193            for y in 0..ny {
194                for x in 0..nx {
195                    for z in 0..nz {
196                        col[z] = image[z * nx * ny + y * nx + x];
197                    }
198                    iwt_forward_5(&col[..nz], &mut out[..nz], &mut scratch[..nz]);
199                    for z in 0..nz {
200                        image[z * nx * ny + y * nx + x] = out[z];
201                    }
202                }
203            }
204        }
205    }
206}
207
208/// Apply the inverse spectral transform along the z-axis.
209pub fn inverse_transform_spectral(
210    image: &mut [i32],
211    nx: usize,
212    ny: usize,
213    nz: usize,
214    transform: SpectralTransform,
215) {
216    match transform {
217        SpectralTransform::Identity => {}
218        SpectralTransform::Iwt => {
219            let mut col = [0i32; 4096];
220            let mut out = [0i32; 4096];
221            let mut scratch = [0i32; 4096];
222            assert!(nz <= 4096);
223
224            for y in 0..ny {
225                for x in 0..nx {
226                    for z in 0..nz {
227                        col[z] = image[z * nx * ny + y * nx + x];
228                    }
229                    iwt_inverse_5(&col[..nz], &mut out[..nz], &mut scratch[..nz]);
230                    for z in 0..nz {
231                        image[z * nx * ny + y * nx + x] = out[z];
232                    }
233                }
234            }
235        }
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242
243    #[test]
244    fn upshift_downshift_roundtrip() {
245        let original = [10, -5, 127, 0, -128];
246        let mut data: [i32; 5] = original;
247        upshift(&mut data, 3);
248        for (o, d) in original.iter().zip(data.iter()) {
249            assert_eq!(*d, *o << 3);
250        }
251        downshift(&mut data, 3);
252        assert_eq!(data, original);
253    }
254
255    #[test]
256    fn downshift_rounds() {
257        let mut data = [3i32, 5, 7, -3];
258        downshift(&mut data, 1);
259        assert_eq!(data, [2, 3, 4, -1]);
260    }
261
262    #[test]
263    fn iwt_single_level_roundtrip() {
264        let input = [10, 20, 30, 40, 50, 60, 70, 80];
265        let n = input.len();
266        let p = (n + 1) / 2;
267        let q = n / 2;
268        let mut l = [0i32; 4];
269        let mut h = [0i32; 4];
270        iwt_forward_single(&input, &mut l[..p], &mut h[..q]);
271
272        let mut recovered = [0i32; 8];
273        iwt_inverse_single(&l[..p], &h[..q], &mut recovered[..n]);
274        assert_eq!(recovered, input);
275    }
276
277    #[test]
278    fn iwt_single_level_odd_length() {
279        let input = [10, 20, 30, 40, 50];
280        let n = input.len();
281        let p = (n + 1) / 2;
282        let q = n / 2;
283        let mut l = [0i32; 3];
284        let mut h = [0i32; 2];
285        iwt_forward_single(&input, &mut l[..p], &mut h[..q]);
286
287        let mut recovered = [0i32; 5];
288        iwt_inverse_single(&l[..p], &h[..q], &mut recovered[..n]);
289        assert_eq!(recovered, input);
290    }
291
292    #[test]
293    fn iwt_five_level_roundtrip() {
294        let input: [i32; 32] = core::array::from_fn(|i| (i as i32) * 7 - 100);
295        let mut output = [0i32; 32];
296        let mut scratch = [0i32; 32];
297        iwt_forward_5(&input, &mut output, &mut scratch);
298
299        let mut recovered = [0i32; 32];
300        iwt_inverse_5(&output, &mut recovered, &mut scratch);
301        assert_eq!(recovered, input);
302    }
303
304    #[test]
305    fn iwt_five_level_short_sequence() {
306        let input = [42i32, -7, 100];
307        let mut output = [0i32; 3];
308        let mut scratch = [0i32; 3];
309        iwt_forward_5(&input, &mut output, &mut scratch);
310
311        let mut recovered = [0i32; 3];
312        iwt_inverse_5(&output, &mut recovered, &mut scratch);
313        assert_eq!(recovered, input);
314    }
315
316    #[test]
317    fn spectral_transform_identity() {
318        let mut image = [1, 2, 3, 4, 5, 6, 7, 8];
319        let original = image;
320        transform_spectral(&mut image, 2, 2, 2, SpectralTransform::Identity);
321        assert_eq!(image, original);
322    }
323
324    #[test]
325    fn spectral_transform_iwt_roundtrip() {
326        let nx = 2;
327        let ny = 2;
328        let nz = 8;
329        let mut image: [i32; 32] = core::array::from_fn(|i| (i as i32) * 3 + 1);
330        let original = image;
331
332        transform_spectral(&mut image, nx, ny, nz, SpectralTransform::Iwt);
333        assert_ne!(image, original);
334
335        inverse_transform_spectral(&mut image, nx, ny, nz, SpectralTransform::Iwt);
336        assert_eq!(image, original);
337    }
338
339    #[test]
340    fn iwt_decorrelates_constant_bands() {
341        let input = [42i32; 16];
342        let mut output = [0i32; 16];
343        let mut scratch = [0i32; 16];
344        iwt_forward_5(&input, &mut output, &mut scratch);
345        assert_eq!(output[0], 42);
346        for &h in &output[1..] {
347            assert_eq!(h, 0);
348        }
349    }
350}