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

convolve_rs/
fits_io.rs
