Skip to main content

ad_core/
ndarray.rs

1use crate::attributes::NDAttributeList;
2use crate::codec::Codec;
3use crate::error::{ADError, ADResult};
4use crate::timestamp::EpicsTimestamp;
5
6/// NDArray data types matching areaDetector NDDataType_t.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8#[repr(u8)]
9pub enum NDDataType {
10    Int8 = 0,
11    UInt8 = 1,
12    Int16 = 2,
13    UInt16 = 3,
14    Int32 = 4,
15    UInt32 = 5,
16    Int64 = 6,
17    UInt64 = 7,
18    Float32 = 8,
19    Float64 = 9,
20}
21
22impl NDDataType {
23    pub fn element_size(&self) -> usize {
24        match self {
25            Self::Int8 | Self::UInt8 => 1,
26            Self::Int16 | Self::UInt16 => 2,
27            Self::Int32 | Self::UInt32 | Self::Float32 => 4,
28            Self::Int64 | Self::UInt64 | Self::Float64 => 8,
29        }
30    }
31
32    pub fn from_ordinal(v: u8) -> Option<Self> {
33        match v {
34            0 => Some(Self::Int8),
35            1 => Some(Self::UInt8),
36            2 => Some(Self::Int16),
37            3 => Some(Self::UInt16),
38            4 => Some(Self::Int32),
39            5 => Some(Self::UInt32),
40            6 => Some(Self::Int64),
41            7 => Some(Self::UInt64),
42            8 => Some(Self::Float32),
43            9 => Some(Self::Float64),
44            _ => None,
45        }
46    }
47}
48
49/// Typed buffer for NDArray data.
50#[derive(Debug, Clone)]
51pub enum NDDataBuffer {
52    I8(Vec<i8>),
53    U8(Vec<u8>),
54    I16(Vec<i16>),
55    U16(Vec<u16>),
56    I32(Vec<i32>),
57    U32(Vec<u32>),
58    I64(Vec<i64>),
59    U64(Vec<u64>),
60    F32(Vec<f32>),
61    F64(Vec<f64>),
62}
63
64impl NDDataBuffer {
65    pub fn zeros(data_type: NDDataType, count: usize) -> Self {
66        match data_type {
67            NDDataType::Int8 => Self::I8(vec![0; count]),
68            NDDataType::UInt8 => Self::U8(vec![0; count]),
69            NDDataType::Int16 => Self::I16(vec![0; count]),
70            NDDataType::UInt16 => Self::U16(vec![0; count]),
71            NDDataType::Int32 => Self::I32(vec![0; count]),
72            NDDataType::UInt32 => Self::U32(vec![0; count]),
73            NDDataType::Int64 => Self::I64(vec![0; count]),
74            NDDataType::UInt64 => Self::U64(vec![0; count]),
75            NDDataType::Float32 => Self::F32(vec![0.0; count]),
76            NDDataType::Float64 => Self::F64(vec![0.0; count]),
77        }
78    }
79
80    pub fn data_type(&self) -> NDDataType {
81        match self {
82            Self::I8(_) => NDDataType::Int8,
83            Self::U8(_) => NDDataType::UInt8,
84            Self::I16(_) => NDDataType::Int16,
85            Self::U16(_) => NDDataType::UInt16,
86            Self::I32(_) => NDDataType::Int32,
87            Self::U32(_) => NDDataType::UInt32,
88            Self::I64(_) => NDDataType::Int64,
89            Self::U64(_) => NDDataType::UInt64,
90            Self::F32(_) => NDDataType::Float32,
91            Self::F64(_) => NDDataType::Float64,
92        }
93    }
94
95    pub fn len(&self) -> usize {
96        match self {
97            Self::I8(v) => v.len(),
98            Self::U8(v) => v.len(),
99            Self::I16(v) => v.len(),
100            Self::U16(v) => v.len(),
101            Self::I32(v) => v.len(),
102            Self::U32(v) => v.len(),
103            Self::I64(v) => v.len(),
104            Self::U64(v) => v.len(),
105            Self::F32(v) => v.len(),
106            Self::F64(v) => v.len(),
107        }
108    }
109
110    pub fn is_empty(&self) -> bool {
111        self.len() == 0
112    }
113
114    pub fn total_bytes(&self) -> usize {
115        self.len() * self.data_type().element_size()
116    }
117
118    /// Capacity of the underlying Vec in bytes.
119    pub fn capacity_bytes(&self) -> usize {
120        let cap = match self {
121            Self::I8(v) => v.capacity(),
122            Self::U8(v) => v.capacity(),
123            Self::I16(v) => v.capacity(),
124            Self::U16(v) => v.capacity(),
125            Self::I32(v) => v.capacity(),
126            Self::U32(v) => v.capacity(),
127            Self::I64(v) => v.capacity(),
128            Self::U64(v) => v.capacity(),
129            Self::F32(v) => v.capacity(),
130            Self::F64(v) => v.capacity(),
131        };
132        cap * self.data_type().element_size()
133    }
134
135    /// Resize the buffer, zeroing new elements if growing.
136    pub fn resize(&mut self, new_len: usize) {
137        match self {
138            Self::I8(v) => v.resize(new_len, 0),
139            Self::U8(v) => v.resize(new_len, 0),
140            Self::I16(v) => v.resize(new_len, 0),
141            Self::U16(v) => v.resize(new_len, 0),
142            Self::I32(v) => v.resize(new_len, 0),
143            Self::U32(v) => v.resize(new_len, 0),
144            Self::I64(v) => v.resize(new_len, 0),
145            Self::U64(v) => v.resize(new_len, 0),
146            Self::F32(v) => v.resize(new_len, 0.0),
147            Self::F64(v) => v.resize(new_len, 0.0),
148        }
149    }
150
151    /// View the underlying data as a byte slice.
152    pub fn as_u8_slice(&self) -> &[u8] {
153        match self {
154            Self::I8(v) => unsafe {
155                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len())
156            },
157            Self::U8(v) => v.as_slice(),
158            Self::I16(v) => unsafe {
159                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2)
160            },
161            Self::U16(v) => unsafe {
162                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2)
163            },
164            Self::I32(v) => unsafe {
165                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
166            },
167            Self::U32(v) => unsafe {
168                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
169            },
170            Self::I64(v) => unsafe {
171                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
172            },
173            Self::U64(v) => unsafe {
174                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
175            },
176            Self::F32(v) => unsafe {
177                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4)
178            },
179            Self::F64(v) => unsafe {
180                std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8)
181            },
182        }
183    }
184
185    /// Get element at index as f64.
186    pub fn get_as_f64(&self, index: usize) -> Option<f64> {
187        match self {
188            Self::I8(v) => v.get(index).map(|&x| x as f64),
189            Self::U8(v) => v.get(index).map(|&x| x as f64),
190            Self::I16(v) => v.get(index).map(|&x| x as f64),
191            Self::U16(v) => v.get(index).map(|&x| x as f64),
192            Self::I32(v) => v.get(index).map(|&x| x as f64),
193            Self::U32(v) => v.get(index).map(|&x| x as f64),
194            Self::I64(v) => v.get(index).map(|&x| x as f64),
195            Self::U64(v) => v.get(index).map(|&x| x as f64),
196            Self::F32(v) => v.get(index).map(|&x| x as f64),
197            Self::F64(v) => v.get(index).copied(),
198        }
199    }
200
201    /// Set element at index from f64 value.
202    pub fn set_from_f64(&mut self, index: usize, value: f64) {
203        match self {
204            Self::I8(v) => { if let Some(e) = v.get_mut(index) { *e = value as i8; } }
205            Self::U8(v) => { if let Some(e) = v.get_mut(index) { *e = value as u8; } }
206            Self::I16(v) => { if let Some(e) = v.get_mut(index) { *e = value as i16; } }
207            Self::U16(v) => { if let Some(e) = v.get_mut(index) { *e = value as u16; } }
208            Self::I32(v) => { if let Some(e) = v.get_mut(index) { *e = value as i32; } }
209            Self::U32(v) => { if let Some(e) = v.get_mut(index) { *e = value as u32; } }
210            Self::I64(v) => { if let Some(e) = v.get_mut(index) { *e = value as i64; } }
211            Self::U64(v) => { if let Some(e) = v.get_mut(index) { *e = value as u64; } }
212            Self::F32(v) => { if let Some(e) = v.get_mut(index) { *e = value as f32; } }
213            Self::F64(v) => { if let Some(e) = v.get_mut(index) { *e = value; } }
214        }
215    }
216}
217
218/// A single dimension of an NDArray.
219#[derive(Debug, Clone)]
220pub struct NDDimension {
221    pub size: usize,
222    pub offset: usize,
223    pub binning: usize,
224    pub reverse: bool,
225}
226
227impl NDDimension {
228    pub fn new(size: usize) -> Self {
229        Self {
230            size,
231            offset: 0,
232            binning: 1,
233            reverse: false,
234        }
235    }
236}
237
238/// Computed info about an NDArray's layout.
239#[derive(Debug, Clone)]
240pub struct NDArrayInfo {
241    pub total_bytes: usize,
242    pub bytes_per_element: usize,
243    pub num_elements: usize,
244    pub x_size: usize,
245    pub y_size: usize,
246    pub color_size: usize,
247}
248
249/// N-dimensional array with typed data buffer.
250#[derive(Debug, Clone)]
251pub struct NDArray {
252    pub unique_id: i32,
253    pub timestamp: EpicsTimestamp,
254    pub dims: Vec<NDDimension>,
255    pub data: NDDataBuffer,
256    pub attributes: NDAttributeList,
257    pub codec: Option<Codec>,
258}
259
260impl NDArray {
261    /// Create a new NDArray with zeroed buffer matching dimensions.
262    pub fn new(dims: Vec<NDDimension>, data_type: NDDataType) -> Self {
263        let num_elements: usize = dims.iter().map(|d| d.size).product();
264        Self {
265            unique_id: 0,
266            timestamp: EpicsTimestamp::default(),
267            dims,
268            data: NDDataBuffer::zeros(data_type, num_elements),
269            attributes: NDAttributeList::new(),
270            codec: None,
271        }
272    }
273
274    /// Compute layout info for this array.
275    pub fn info(&self) -> NDArrayInfo {
276        let bytes_per_element = self.data.data_type().element_size();
277        let num_elements = self.data.len();
278        let total_bytes = num_elements * bytes_per_element;
279
280        let ndims = self.dims.len();
281        let (x_size, y_size, color_size) = match ndims {
282            0 => (0, 0, 0),
283            1 => (self.dims[0].size, 1, 1),
284            2 => (self.dims[0].size, self.dims[1].size, 1),
285            _ => {
286                // 3D: interpret as color image
287                // NDColorMode convention: dim[0]=color for RGB1
288                (self.dims[1].size, self.dims[2].size, self.dims[0].size)
289            }
290        };
291
292        NDArrayInfo {
293            total_bytes,
294            bytes_per_element,
295            num_elements,
296            x_size,
297            y_size,
298            color_size,
299        }
300    }
301
302    /// Validate that buffer length matches dimension product.
303    pub fn validate(&self) -> ADResult<()> {
304        let expected: usize = if self.dims.is_empty() {
305            0
306        } else {
307            self.dims.iter().map(|d| d.size).product()
308        };
309        if self.data.len() != expected {
310            return Err(ADError::BufferSizeMismatch {
311                expected,
312                actual: self.data.len(),
313            });
314        }
315        Ok(())
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn test_element_size_all_types() {
325        assert_eq!(NDDataType::Int8.element_size(), 1);
326        assert_eq!(NDDataType::UInt8.element_size(), 1);
327        assert_eq!(NDDataType::Int16.element_size(), 2);
328        assert_eq!(NDDataType::UInt16.element_size(), 2);
329        assert_eq!(NDDataType::Int32.element_size(), 4);
330        assert_eq!(NDDataType::UInt32.element_size(), 4);
331        assert_eq!(NDDataType::Int64.element_size(), 8);
332        assert_eq!(NDDataType::UInt64.element_size(), 8);
333        assert_eq!(NDDataType::Float32.element_size(), 4);
334        assert_eq!(NDDataType::Float64.element_size(), 8);
335    }
336
337    #[test]
338    fn test_from_ordinal_roundtrip() {
339        for i in 0..10u8 {
340            let dt = NDDataType::from_ordinal(i).unwrap();
341            assert_eq!(dt as u8, i);
342        }
343        assert!(NDDataType::from_ordinal(10).is_none());
344    }
345
346    #[test]
347    fn test_buffer_zeros_type_and_len() {
348        let buf = NDDataBuffer::zeros(NDDataType::UInt16, 100);
349        assert_eq!(buf.data_type(), NDDataType::UInt16);
350        assert_eq!(buf.len(), 100);
351        assert_eq!(buf.total_bytes(), 200);
352    }
353
354    #[test]
355    fn test_buffer_zeros_all_types() {
356        for i in 0..10u8 {
357            let dt = NDDataType::from_ordinal(i).unwrap();
358            let buf = NDDataBuffer::zeros(dt, 10);
359            assert_eq!(buf.data_type(), dt);
360            assert_eq!(buf.len(), 10);
361            assert_eq!(buf.total_bytes(), 10 * dt.element_size());
362        }
363    }
364
365    #[test]
366    fn test_buffer_as_u8_slice() {
367        let buf = NDDataBuffer::U8(vec![1, 2, 3]);
368        assert_eq!(buf.as_u8_slice(), &[1, 2, 3]);
369    }
370
371    #[test]
372    fn test_ndarray_new_allocates() {
373        let dims = vec![NDDimension::new(256), NDDimension::new(256)];
374        let arr = NDArray::new(dims, NDDataType::UInt8);
375        assert_eq!(arr.data.len(), 256 * 256);
376        assert_eq!(arr.data.data_type(), NDDataType::UInt8);
377    }
378
379    #[test]
380    fn test_ndarray_validate_ok() {
381        let dims = vec![NDDimension::new(10), NDDimension::new(20)];
382        let arr = NDArray::new(dims, NDDataType::Float64);
383        arr.validate().unwrap();
384    }
385
386    #[test]
387    fn test_ndarray_validate_mismatch() {
388        let mut arr = NDArray::new(
389            vec![NDDimension::new(10), NDDimension::new(20)],
390            NDDataType::UInt8,
391        );
392        arr.data = NDDataBuffer::U8(vec![0; 100]);
393        assert!(arr.validate().is_err());
394    }
395
396    #[test]
397    fn test_ndarray_info_2d_mono() {
398        let dims = vec![NDDimension::new(640), NDDimension::new(480)];
399        let arr = NDArray::new(dims, NDDataType::UInt16);
400        let info = arr.info();
401        assert_eq!(info.x_size, 640);
402        assert_eq!(info.y_size, 480);
403        assert_eq!(info.color_size, 1);
404        assert_eq!(info.num_elements, 640 * 480);
405        assert_eq!(info.bytes_per_element, 2);
406        assert_eq!(info.total_bytes, 640 * 480 * 2);
407    }
408
409    #[test]
410    fn test_ndarray_info_3d_rgb() {
411        let dims = vec![
412            NDDimension::new(3),
413            NDDimension::new(640),
414            NDDimension::new(480),
415        ];
416        let arr = NDArray::new(dims, NDDataType::UInt8);
417        let info = arr.info();
418        assert_eq!(info.color_size, 3);
419        assert_eq!(info.x_size, 640);
420        assert_eq!(info.y_size, 480);
421        assert_eq!(info.num_elements, 3 * 640 * 480);
422    }
423
424    #[test]
425    fn test_ndarray_info_1d() {
426        let dims = vec![NDDimension::new(1024)];
427        let arr = NDArray::new(dims, NDDataType::Float64);
428        let info = arr.info();
429        assert_eq!(info.x_size, 1024);
430        assert_eq!(info.y_size, 1);
431        assert_eq!(info.color_size, 1);
432    }
433
434    #[test]
435    fn test_buffer_is_empty() {
436        let buf = NDDataBuffer::zeros(NDDataType::UInt8, 0);
437        assert!(buf.is_empty());
438        let buf2 = NDDataBuffer::zeros(NDDataType::UInt8, 1);
439        assert!(!buf2.is_empty());
440    }
441
442    #[test]
443    fn test_codec_field_preserved() {
444        let mut arr = NDArray::new(vec![NDDimension::new(10)], NDDataType::UInt8);
445        arr.codec = Some(Codec { name: crate::codec::CodecName::JPEG, compressed_size: 42 });
446        let cloned = arr.clone();
447        assert_eq!(cloned.codec.as_ref().unwrap().compressed_size, 42);
448    }
449}