Skip to main content

atfits_rs/
pixel.rs

1//! Pixel precision and section I/O on the primary image.
2use std::os::raw::{c_int, c_void};
3
4use fitsio::FitsFile;
5use fitsio::hdu::HduInfo;
6use fitsio::images::ImageType;
7use fitsio::sys::{self, LONGLONG};
8
9use crate::error::{AtfitsError, Result, ffcall};
10use crate::raw::read_flat;
11
12/// Precision a cube is processed in.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum PixelType {
15    F32,
16    F64,
17}
18
19impl PixelType {
20    /// `-64` maps to f64; `-32` and every integer BITPIX map to f32.
21    pub fn from_bitpix(bitpix: i64) -> Self {
22        if bitpix == -64 {
23            PixelType::F64
24        } else {
25            PixelType::F32
26        }
27    }
28}
29
30/// Element types (`f32`, `f64`) for reading and writing pixel ranges.
31///
32/// Ranges are flat pixel offsets `start..end` into the primary image. Integer
33/// images are scaled by BSCALE/BZERO; BLANK pixels read as NaN, and NaN is
34/// written as BLANK (an error if the header has no BLANK).
35pub trait CubeElem: Copy + Send + Sync + 'static {
36    /// BITPIX of a float image of this type.
37    const BITPIX: i64;
38    /// Read the whole primary image.
39    fn read_full(fptr: &mut FitsFile) -> Result<Vec<Self>>;
40    /// Read pixels `start..end`.
41    fn read_section(fptr: &mut FitsFile, start: usize, end: usize) -> Result<Vec<Self>>;
42    /// Read pixels `start..start + out.len()` into `out`.
43    fn read_section_into(fptr: &mut FitsFile, start: usize, out: &mut [Self]) -> Result<()> {
44        let data = Self::read_section(fptr, start, start + out.len())?;
45        out.copy_from_slice(&data);
46        Ok(())
47    }
48    /// Write `data` to pixels `start..end`; `data.len()` must equal `end - start`.
49    fn write_section(fptr: &mut FitsFile, start: usize, end: usize, data: &[Self]) -> Result<()>;
50}
51
52/// Read plane `index` of the primary image, where a plane is `plane_len` pixels.
53pub fn read_plane<T: CubeElem>(
54    fptr: &mut FitsFile,
55    plane_len: usize,
56    index: usize,
57) -> Result<Vec<T>> {
58    let (start, end) = plane_range(plane_len, index)?;
59    T::read_section(fptr, start, end)
60}
61
62/// Read plane `index` of the primary image into `out`, whose length is the
63/// plane size.
64pub fn read_plane_into<T: CubeElem>(
65    fptr: &mut FitsFile,
66    index: usize,
67    out: &mut [T],
68) -> Result<()> {
69    let (start, _) = plane_range(out.len(), index)?;
70    T::read_section_into(fptr, start, out)
71}
72
73/// Write `data` as plane `index` of the primary image.
74pub fn write_plane<T: CubeElem>(fptr: &mut FitsFile, index: usize, data: &[T]) -> Result<()> {
75    let (start, end) = plane_range(data.len(), index)?;
76    T::write_section(fptr, start, end, data)
77}
78
79fn plane_range(plane_len: usize, index: usize) -> Result<(usize, usize)> {
80    plane_len
81        .checked_mul(index)
82        .and_then(|start| Some((start, start.checked_add(plane_len)?)))
83        .ok_or_else(|| AtfitsError::Other(format!("plane {index} overflows usize")))
84}
85
86/// Pixel count of the primary image and whether it stores integers.
87fn primary_image(fptr: &mut FitsFile) -> Result<(usize, bool)> {
88    match fptr.primary_hdu()?.info {
89        HduInfo::ImageInfo { shape, image_type } => {
90            let npix = if shape.is_empty() {
91                0
92            } else {
93                shape.iter().product()
94            };
95            let integer = !matches!(image_type, ImageType::Float | ImageType::Double);
96            Ok((npix, integer))
97        }
98        _ => Err(AtfitsError::Other("primary HDU is not an image".into())),
99    }
100}
101
102fn check_range(start: usize, end: usize, npix: usize) -> Result<()> {
103    if start > end || end > npix {
104        return Err(range_error(start, end, Some(npix)));
105    }
106    Ok(())
107}
108
109fn range_error(start: usize, end: usize, npix: Option<usize>) -> AtfitsError {
110    match npix {
111        Some(n) => AtfitsError::Other(format!(
112            "pixel range {start}..{end} is outside the image ({n} pixels)"
113        )),
114        None => AtfitsError::Other(format!("pixel range {start}..{end} is reversed")),
115    }
116}
117
118macro_rules! impl_cube_elem {
119    ($t:ty, $bitpix:expr, $code:expr) => {
120        impl CubeElem for $t {
121            const BITPIX: i64 = $bitpix;
122
123            fn read_full(fptr: &mut FitsFile) -> Result<Vec<Self>> {
124                let (npix, _) = primary_image(fptr)?;
125                Self::read_section(fptr, 0, npix)
126            }
127
128            fn read_section(fptr: &mut FitsFile, start: usize, end: usize) -> Result<Vec<Self>> {
129                if start > end {
130                    return Err(range_error(start, end, None));
131                }
132                let mut out: Vec<$t> = vec![0.0; end - start];
133                Self::read_section_into(fptr, start, &mut out)?;
134                Ok(out)
135            }
136
137            fn read_section_into(
138                fptr: &mut FitsFile,
139                start: usize,
140                out: &mut [Self],
141            ) -> Result<()> {
142                let (npix, integer) = primary_image(fptr)?;
143                let end = start
144                    .checked_add(out.len())
145                    .ok_or_else(|| range_error(start, usize::MAX, Some(npix)))?;
146                check_range(start, end, npix)?;
147                // Float NaN reads as NaN anyway; a null value only costs time there.
148                read_flat(fptr, start, integer.then_some(<$t>::NAN), out)
149            }
150
151            fn write_section(
152                fptr: &mut FitsFile,
153                start: usize,
154                end: usize,
155                data: &[Self],
156            ) -> Result<()> {
157                let (npix, integer) = primary_image(fptr)?;
158                check_range(start, end, npix)?;
159                if data.len() != end - start {
160                    return Err(AtfitsError::Other(format!(
161                        "pixel range {start}..{end} needs {} values, got {}",
162                        end - start,
163                        data.len()
164                    )));
165                }
166                if !(integer && data.iter().any(|v| v.is_nan())) {
167                    // cfitsio only reads `data`; the C API lacks const.
168                    return ffcall(|s| unsafe {
169                        sys::ffppr(
170                            fptr.as_raw(),
171                            $code as c_int,
172                            start as LONGLONG + 1,
173                            data.len() as LONGLONG,
174                            data.as_ptr() as *mut c_void,
175                            s,
176                        )
177                    });
178                }
179                // cfitsio finds nulls by equality, which never holds for NaN.
180                let mut nulval = <$t>::NEG_INFINITY;
181                let mut buf: Vec<$t> = data
182                    .iter()
183                    .map(|&v| if v.is_nan() { nulval } else { v })
184                    .collect();
185                ffcall(|s| unsafe {
186                    sys::ffppn(
187                        fptr.as_raw(),
188                        $code as c_int,
189                        start as LONGLONG + 1,
190                        buf.len() as LONGLONG,
191                        buf.as_mut_ptr() as *mut c_void,
192                        &mut nulval as *mut $t as *mut c_void,
193                        s,
194                    )
195                })
196            }
197        }
198    };
199}
200impl_cube_elem!(f32, -32, sys::TFLOAT);
201impl_cube_elem!(f64, -64, sys::TDOUBLE);