convolve_rs/
fits_io.rs

1//! FITS image reading and writing.
2//!
3//! Handles 2D images (NAXIS = 2) and 4D images (NAXIS = 4). Only the first
4//! plane of a 4D image is used, as in racs_tools.
5use std::path::{Path, PathBuf};
6
7use atfits_rs::{HeaderGeom, PixelType, create_cube_open, update_key_f64};
8use fitsio::FitsFile;
9use fitsio::hdu::{FitsHdu, HduInfo};
10use fitsio::images::ImageType;
11use ndarray::Array2;
12use thiserror::Error;
13
14use crate::beam::{Beam, BeamError};
15use crate::cube_io::CubeElem;
16use crate::smooth::BrightnessUnit;
17
18pub use atfits_rs::output_path;
19
20#[derive(Debug, Error)]
21pub enum FitsError {
22    #[error("FITS I/O error: {0}")]
23    Fitsio(#[from] fitsio::errors::Error),
24    #[error("missing keyword: {0}")]
25    MissingKeyword(String),
26    #[error("unsupported NAXIS={0} (expected 2 or 4)")]
27    UnsupportedNaxis(i64),
28    #[error("invalid header beam: {0}")]
29    Beam(#[from] BeamError),
30    #[error("shape error: {0}")]
31    Shape(#[from] ndarray::ShapeError),
32    #[error("I/O error: {0}")]
33    Io(#[from] std::io::Error),
34}
35
36impl From<atfits_rs::AtfitsError> for FitsError {
37    fn from(e: atfits_rs::AtfitsError) -> Self {
38        use atfits_rs::AtfitsError as A;
39        match e {
40            A::Fits(e) => FitsError::Fitsio(e),
41            A::Io(e) => FitsError::Io(e),
42            A::MissingKeyword(s) | A::TargetAxisMissing(s) => FitsError::MissingKeyword(s),
43            A::UnsupportedNaxis(n) => FitsError::UnsupportedNaxis(n),
44            A::Other(s) => FitsError::Io(std::io::Error::other(s)),
45        }
46    }
47}
48
49/// Header metadata of a 2D image, or of the first plane of a 4D one.
50#[derive(Debug, Clone)]
51pub struct FitsImageHeader {
52    pub path: PathBuf,
53    pub nx: usize,
54    pub ny: usize,
55    /// Planes beyond the first (1 for a plain image).
56    pub nplanes: usize,
57    /// |CDELT1| in degrees.
58    pub dx_deg: f64,
59    /// |CDELT2| in degrees.
60    pub dy_deg: f64,
61    pub beam: Beam,
62    pub unit: BrightnessUnit,
63    pub dtype: PixelType,
64}
65
66pub struct FitsImageData<T = f32> {
67    pub header: FitsImageHeader,
68    pub image: Array2<T>,
69}
70
71/// Read the shape, pixel scale, beam and unit of a 2D or 4D image.
72pub fn read_fits_header(path: &Path) -> Result<FitsImageHeader, FitsError> {
73    let mut fptr = FitsFile::open(path)?;
74    let geom = HeaderGeom::from_file(&mut fptr)?;
75    if geom.naxis != 2 && geom.naxis != 4 {
76        return Err(FitsError::UnsupportedNaxis(geom.naxis as i64));
77    }
78    let nplanes = geom.dims[2..].iter().product();
79    if nplanes > 1 {
80        tracing::warn!(
81            "{}: {nplanes} planes beyond the image axes; only the first is smoothed",
82            path.display()
83        );
84    }
85    let hdu = fptr.primary_hdu()?;
86    let (dx_deg, dy_deg) = pixel_scales(&mut fptr, &hdu)?;
87    let beam = read_header_beam(&mut fptr)?.ok_or(FitsError::MissingKeyword("BMAJ".into()))?;
88    Ok(FitsImageHeader {
89        path: path.to_path_buf(),
90        nx: geom.dims[0],
91        ny: geom.dims[1],
92        nplanes,
93        dx_deg,
94        dy_deg,
95        beam,
96        unit: read_unit(&mut fptr, &hdu, path),
97        dtype: PixelType::from_bitpix(geom.bitpix),
98    })
99}
100
101/// Read a 2D image, or the first plane of a 4D one, in precision `T`.
102pub fn read_fits<T: CubeElem>(path: &Path) -> Result<FitsImageData<T>, FitsError> {
103    let header = read_fits_header(path)?;
104    let image = read_fits_plane(&header)?;
105    Ok(FitsImageData { header, image })
106}
107
108/// Read the first plane of the image `header` describes.
109pub fn read_fits_plane<T: CubeElem>(header: &FitsImageHeader) -> Result<Array2<T>, FitsError> {
110    let mut fptr = FitsFile::open(&header.path)?;
111    let data = read_pixels(&mut fptr, 0, header.nx * header.ny)?;
112    Ok(Array2::from_shape_vec((header.ny, header.nx), data)?)
113}
114
115/// Write `image` to `out_path` with the header of `template_path`, the beam
116/// set to `new_beam`, and any planes beyond the first dropped.
117pub fn write_fits<T: CubeElem>(
118    image: &Array2<T>,
119    out_path: &Path,
120    template_path: &Path,
121    new_beam: &Beam,
122) -> Result<(), FitsError> {
123    let (ny, nx) = image.dim();
124    let mut dims = vec![1; HeaderGeom::read(template_path)?.naxis];
125    dims[0] = nx;
126    dims[1] = ny;
127    let mut fptr = create_output(template_path, out_path, Some(&dims))?;
128    update_key_f64(&mut fptr, "BMAJ", new_beam.major_deg())?;
129    update_key_f64(&mut fptr, "BMIN", new_beam.minor_deg())?;
130    update_key_f64(&mut fptr, "BPA", new_beam.pa_deg())?;
131
132    let flat = image.as_standard_layout();
133    let slice = flat
134        .as_slice()
135        .ok_or_else(|| FitsError::Io(std::io::Error::other("image is not contiguous")))?;
136    T::write_section(&mut fptr, 0, nx * ny, slice)?;
137    Ok(())
138}
139
140/// Create `output` with `template`'s header and shape `dims` (default: the
141/// template's), open for writing. Integer templates give a float image, since
142/// smoothed values are fractional and blanked pixels are NaN.
143pub(crate) fn create_output(
144    template: &Path,
145    output: &Path,
146    dims: Option<&[usize]>,
147) -> Result<FitsFile, atfits_rs::AtfitsError> {
148    let geom = HeaderGeom::read(template)?;
149    let bitpix = if geom.bitpix > 0 { -32 } else { geom.bitpix };
150    create_cube_open(template, output, bitpix, dims.unwrap_or(&geom.dims))
151}
152
153/// Pixels `start..end` of the primary image.
154///
155/// Float images are read through fitsio: atfits 1.1 has cfitsio null-check
156/// every float pixel, which doubles the read time (1.0 s against 0.46 s for a
157/// 1 GB cube). Integer images need atfits' BLANK handling.
158pub(crate) fn read_pixels<T: CubeElem>(
159    fptr: &mut FitsFile,
160    start: usize,
161    end: usize,
162) -> Result<Vec<T>, FitsError> {
163    let hdu = fptr.primary_hdu()?;
164    if let HduInfo::ImageInfo {
165        image_type: ImageType::Float | ImageType::Double,
166        ..
167    } = hdu.info
168    {
169        Ok(T::read_float(fptr, &hdu, start, end)?)
170    } else {
171        Ok(T::read_section(fptr, start, end)?)
172    }
173}
174
175/// `(|CDELT1|, |CDELT2|)` in degrees.
176pub(crate) fn pixel_scales(fptr: &mut FitsFile, hdu: &FitsHdu) -> Result<(f64, f64), FitsError> {
177    let mut read = |key: &str| {
178        hdu.read_key::<f64>(fptr, key).map(f64::abs).map_err(|_| {
179            FitsError::MissingKeyword(if hdu.read_key::<f64>(fptr, "CD1_1").is_ok() {
180                format!("{key} (CD-matrix headers are not supported; add CDELT1/CDELT2)")
181            } else {
182                key.to_string()
183            })
184        })
185    };
186    Ok((read("CDELT1")?, read("CDELT2")?))
187}
188
189/// BMAJ/BMIN/BPA from the primary header, or `None` without BMAJ. BMIN
190/// defaults to BMAJ and BPA to 0.
191pub(crate) fn read_header_beam(fptr: &mut FitsFile) -> Result<Option<Beam>, FitsError> {
192    let Some(bmaj) = atfits_rs::key_f64(fptr, "BMAJ")? else {
193        return Ok(None);
194    };
195    let bmin = atfits_rs::key_f64(fptr, "BMIN")?.unwrap_or(bmaj);
196    let bpa = atfits_rs::key_f64(fptr, "BPA")?.unwrap_or(0.0);
197    Ok(Some(Beam::new(bmaj, bmin, bpa)?))
198}
199
200/// The brightness unit from BUNIT, warning and assuming Jy/beam if it is
201/// missing or not recognised.
202pub(crate) fn read_unit(fptr: &mut FitsFile, hdu: &FitsHdu, path: &Path) -> BrightnessUnit {
203    match hdu.read_key::<String>(fptr, "BUNIT") {
204        Ok(s) => BrightnessUnit::from_bunit(&s),
205        Err(_) => {
206            tracing::warn!(
207                "No BUNIT keyword in {}; assuming Jy/beam (flux scaling applied).",
208                path.display()
209            );
210            BrightnessUnit::default()
211        }
212    }
213}
convolve_rs/fits_io.rs

convolve_rs/
fits_io.rs
