ndwt 0.1.1

High-performance discrete and lifting wavelet transforms for 1-D and N-D signals, with SIMD acceleration, adjoint operations, and 8 boundary conditions.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

use aligned_vec::{AVec, avec};
use itertools::Itertools;
use ndwt::iter::LanesIterator;
use ndwt::utils::deinterleave_nd;

fn ref_deinterleave_nd<T: Clone + num_traits::Zero>(x: &[T], shape: &[usize]) -> Vec<T> {
    let mut out = x.to_vec();

    for ax in 0..shape.len() {
        let n = shape[ax];
        let ne = (n + 1) / 2;
        let no = n / 2;

        let mut work_e = vec![T::zero(); ne];
        let mut work_o = vec![T::zero(); no];

        for mut lane in out.iter_lanes_mut(shape, ax) {
            lane.deinterleave(&mut work_e, &mut work_o);
            lane.stack(&work_e, &work_o);
        }
    }
    out
}

fn ref_deinterleave_chunk_nd<T: Clone + num_traits::Zero>(x: &[T], shape: &[usize]) -> Vec<T> {
    let mut out = x.to_vec();

    const N: usize = 4;

    for ax in 0..shape.len() {
        let n = shape[ax];
        let ne = (n + 1) / 2;
        let no = n / 2;

        let (chunks, rem) = out.iter_lane_chunks_mut::<4>(shape, ax);

        let mut work_e: [AVec<T>; N] = core::array::from_fn(|_| avec![T::zero(); ne]);
        let mut work_o: [AVec<T>; N] = core::array::from_fn(|_| avec![T::zero(); no]);

        for mut chunk in chunks {
            chunk.deinterleave(&mut work_e, &mut work_o);
            chunk.stack(&work_e, &work_o);
        }

        let mut work_e = vec![T::zero(); ne];
        let mut work_o = vec![T::zero(); no];

        for mut lane in rem {
            lane.deinterleave(&mut work_e, &mut work_o);
            lane.stack(&work_e, &work_o);
        }
    }
    out
}

#[test]
fn test_deinterleave_2d() {
    let ns = [10, 11];

    for dim in [1, 2, 3, 4, 5] {
        for n in ns {
            let shape = vec![n; dim];
            let n_total = shape.iter().product();
            let x = (0..n_total).collect_vec();

            let mut out = vec![0; n_total];
            deinterleave_nd(&x, &mut out, &shape);

            let desired = ref_deinterleave_nd(&x, &shape);

            assert_eq!(out, desired);
        }
    }
}

#[test]
fn test_deinterleave_chunked_2d() {
    let ns = [10, 11];

    for dim in [1, 2, 3, 4, 5] {
        for n in ns {
            let shape = vec![n; dim];
            let n_total = shape.iter().product();
            let x = (0..n_total).collect_vec();

            let mut out = vec![0; n_total];
            deinterleave_nd(&x, &mut out, &shape);

            let desired = ref_deinterleave_chunk_nd(&x, &shape);

            assert_eq!(out, desired);
        }
    }
}