conspire/geometry/grid/
mod.rs

1#[cfg(test)]
2mod test;
3
4mod defeature;
5mod from;
6mod read;
7mod write;
8
9pub use self::{read::Input, write::Output};
10
11use std::{
12    array,
13    borrow::Cow,
14    iter,
15    ops::{Index, IndexMut, Range},
16};
17
18pub type Pixels<T> = Grid<2, T>;
19pub type Voxels<T> = Grid<3, T>;
20
21pub struct Grid<const D: usize, T> {
22    data: Vec<T>,
23    nel: [usize; D],
24    strides: [usize; D],
25}
26
27fn col_major_strides<const D: usize>(nel: [usize; D]) -> [usize; D] {
28    let mut strides = [1usize; D];
29    for axis in 1..D {
30        strides[axis] = strides[axis - 1] * nel[axis - 1];
31    }
32    strides
33}
34
35fn row_major_strides<const D: usize>(nel: [usize; D]) -> [usize; D] {
36    let mut strides = [1usize; D];
37    for axis in (0..D.saturating_sub(1)).rev() {
38        strides[axis] = strides[axis + 1] * nel[axis + 1];
39    }
40    strides
41}
42
43impl<const D: usize, T> Grid<D, T> {
44    pub fn new(data: Vec<T>, nel: [usize; D]) -> Self {
45        assert_eq!(
46            data.len(),
47            nel.iter().product::<usize>(),
48            "voxel data length must equal the product of nel"
49        );
50        Self {
51            data,
52            nel,
53            strides: col_major_strides(nel),
54        }
55    }
56    pub fn new_row_major(data: Vec<T>, nel: [usize; D]) -> Self {
57        assert_eq!(
58            data.len(),
59            nel.iter().product::<usize>(),
60            "voxel data length must equal the product of nel"
61        );
62        Self {
63            data,
64            nel,
65            strides: row_major_strides(nel),
66        }
67    }
68    pub fn data(&self) -> &[T] {
69        &self.data
70    }
71    pub fn flat(&self, index: [usize; D]) -> usize {
72        index
73            .iter()
74            .zip(&self.strides)
75            .map(|(&i, &stride)| i * stride)
76            .sum()
77    }
78    fn axes_by_stride(&self) -> [usize; D] {
79        let mut order: [usize; D] = array::from_fn(|axis| axis);
80        order.sort_by_key(|&axis| self.strides[axis]);
81        order
82    }
83    pub fn logical_iter<'a>(&'a self) -> impl Iterator<Item = ([usize; D], &'a T)> + 'a {
84        let nel = self.nel;
85        let order = self.axes_by_stride();
86        let mut index = [0usize; D];
87        let mut offset = 0usize;
88        iter::from_fn(move || {
89            if offset == self.data.len() {
90                return None;
91            }
92            let item = (index, &self.data[offset]);
93            offset += 1;
94            for &axis in &order {
95                index[axis] += 1;
96                if index[axis] < nel[axis] {
97                    break;
98                }
99                index[axis] = 0;
100            }
101            Some(item)
102        })
103    }
104    pub fn nel(&self) -> &[usize; D] {
105        &self.nel
106    }
107    pub fn is_col_major(&self) -> bool {
108        self.strides == col_major_strides(self.nel)
109    }
110    pub fn len(&self) -> usize {
111        self.data.len()
112    }
113    pub fn is_empty(&self) -> bool {
114        self.data.is_empty()
115    }
116}
117
118impl<const D: usize, T: Copy> Grid<D, T> {
119    pub fn data_col_major(&self) -> Cow<'_, [T]> {
120        if self.strides == col_major_strides(self.nel) {
121            Cow::Borrowed(&self.data)
122        } else {
123            Cow::Owned(read::transpose(self.data.clone(), self.nel))
124        }
125    }
126}
127
128impl<const D: usize, T: Copy> Grid<D, T> {
129    pub fn extract(&self, ranges: [Range<usize>; D]) -> Self {
130        let nel: [usize; D] = array::from_fn(|axis| ranges[axis].len());
131        let total: usize = nel.iter().product();
132        let mut data = Vec::with_capacity(total);
133        let mut index = [0usize; D];
134        for offset in 0..total {
135            let mut remainder = offset;
136            for axis in 0..D {
137                index[axis] = ranges[axis].start + remainder % nel[axis];
138                remainder /= nel[axis];
139            }
140            data.push(self[index]);
141        }
142        Self::new(data, nel)
143    }
144}
145
146impl<const D: usize, T: PartialEq> Grid<D, T> {
147    pub fn diff(&self, other: &Self) -> Grid<D, u8> {
148        assert_eq!(
149            self.nel(),
150            other.nel(),
151            "grids do not have the same dimensions"
152        );
153        if self.strides == other.strides {
154            let data = self
155                .data
156                .iter()
157                .zip(&other.data)
158                .map(|(a, b)| (a != b) as u8)
159                .collect();
160            Grid {
161                data,
162                nel: self.nel,
163                strides: self.strides,
164            }
165        } else {
166            let nel = self.nel;
167            let data = (0..self.data.len())
168                .map(|offset| {
169                    let mut index = [0usize; D];
170                    let mut remainder = offset;
171                    for axis in 0..D {
172                        index[axis] = remainder % nel[axis];
173                        remainder /= nel[axis];
174                    }
175                    (self[index] != other[index]) as u8
176                })
177                .collect();
178            Grid::new(data, nel)
179        }
180    }
181}
182
183impl<const D: usize, T> Index<[usize; D]> for Grid<D, T> {
184    type Output = T;
185    fn index(&self, index: [usize; D]) -> &T {
186        &self.data[self.flat(index)]
187    }
188}
189
190impl<const D: usize, T> IndexMut<[usize; D]> for Grid<D, T> {
191    fn index_mut(&mut self, index: [usize; D]) -> &mut T {
192        let offset = self.flat(index);
193        &mut self.data[offset]
194    }
195}
conspire/geometry/grid/mod.rs

conspire/geometry/grid/
mod.rs
