convolve_rs/
fits_io.rs

1//! FITS image reading and writing.
2//!
3//! Handles both 2D (NAXIS=2, shape `[ny, nx]`) and 4D (NAXIS=4, shape
4//! `[1,1,ny,nx]`) images as produced by ASKAP/CASA.
5use std::path::{Path, PathBuf};
6
7use atfits_rs::copy_header_only;
8use fitsio::FitsFile;
9use ndarray::Array2;
10use thiserror::Error;
11
12use crate::beam::Beam;
13use crate::smooth::BrightnessUnit;
14
15#[derive(Debug, Error)]
16pub enum FitsError {
17    #[error("FITS I/O error: {0}")]
18    Fitsio(#[from] fitsio::errors::Error),
19    #[error("missing keyword: {0}")]
20    MissingKeyword(String),
21    #[error("unsupported NAXIS={0} (expected 2 or 4)")]
22    UnsupportedNaxis(i64),
23    #[error("I/O error: {0}")]
24    Io(#[from] std::io::Error),
25}
26
27/// Map the shared [`atfits_rs::AtfitsError`] onto this module's error type.
28impl From<atfits_rs::AtfitsError> for FitsError {
29    fn from(e: atfits_rs::AtfitsError) -> Self {
30        use atfits_rs::AtfitsError as A;
31        match e {
32            A::Fits(e) => FitsError::Fitsio(e),
33            A::Io(e) => FitsError::Io(e),
34            A::MissingKeyword(s) | A::TargetAxisMissing(s) => FitsError::MissingKeyword(s),
35            A::UnsupportedNaxis(n) => FitsError::UnsupportedNaxis(n),
36            A::Other(s) => FitsError::Io(std::io::Error::other(s)),
37        }
38    }
39}
40
41pub struct FitsImageData {
42    pub path: PathBuf,
43    pub image: Array2<f32>,
44    pub is_4d: bool,
45    /// |CDELT1| in degrees (x / RA pixel size)
46    pub dx_deg: f64,
47    /// |CDELT2| in degrees (y / Dec pixel size)
48    pub dy_deg: f64,
49    pub beam: Beam,
50    /// Brightness unit from BUNIT (defaults to Jy/beam if absent).
51    pub unit: BrightnessUnit,
52    /// Full header keyword list (key, value strings) for re-writing.
53    pub header_cards: Vec<(String, String)>,
54}
55
56/// Read a 2D or 4D FITS image.
57pub fn read_fits(path: &Path) -> Result<FitsImageData, FitsError> {
58    let path_str = path.to_string_lossy().into_owned();
59    let mut fptr = FitsFile::open(&path_str)?;
60    let hdu = fptr.primary_hdu()?;
61
62    let naxis: i64 = hdu.read_key(&mut fptr, "NAXIS")?;
63    let naxis1: i64 = hdu.read_key(&mut fptr, "NAXIS1")?; // x / RA (cols)
64    let naxis2: i64 = hdu.read_key(&mut fptr, "NAXIS2")?; // y / Dec (rows)
65
66    if naxis != 2 && naxis != 4 {
67        return Err(FitsError::UnsupportedNaxis(naxis));
68    }
69
70    let nx = naxis1 as usize;
71    let ny = naxis2 as usize;
72
73    // Read flat image data (works for both 2D and 4D when extra axes are size 1).
74    let data: Vec<f32> = hdu.read_image(&mut fptr)?;
75    let image = Array2::from_shape_vec((ny, nx), data)
76        .map_err(|e| FitsError::Io(std::io::Error::other(e.to_string())))?;
77
78    // Pixel scales — use absolute values since CDELT1 is negative for RA.
79    let cdelt1: f64 = hdu.read_key(&mut fptr, "CDELT1")?;
80    let cdelt2: f64 = hdu.read_key(&mut fptr, "CDELT2")?;
81    let dx_deg = cdelt1.abs();
82    let dy_deg = cdelt2.abs();
83
84    // Beam.
85    let bmaj: f64 = hdu
86        .read_key(&mut fptr, "BMAJ")
87        .map_err(|_| FitsError::MissingKeyword("BMAJ".into()))?;
88    let bmin: f64 = hdu.read_key(&mut fptr, "BMIN").unwrap_or(bmaj);
89    let bpa: f64 = hdu.read_key(&mut fptr, "BPA").unwrap_or(0.0);
90
91    let beam = Beam::new(bmaj, bmin, bpa)
92        .map_err(|e| FitsError::Io(std::io::Error::other(e.to_string())))?;
93
94    // Brightness unit (BUNIT); warn and default to Jy/beam when absent.
95    let unit = match hdu.read_key::<String>(&mut fptr, "BUNIT") {
96        Ok(s) => BrightnessUnit::from_bunit(&s),
97        Err(_) => {
98            tracing::warn!(
99                "No BUNIT keyword in {}; assuming Jy/beam (flux scaling applied).",
100                path.display()
101            );
102            BrightnessUnit::default()
103        }
104    };
105
106    Ok(FitsImageData {
107        path: path.to_path_buf(),
108        image,
109        is_4d: naxis == 4,
110        dx_deg,
111        dy_deg,
112        beam,
113        unit,
114        header_cards: vec![],
115    })
116}
117
118/// Write a smoothed image to `out_path`.
119///
120/// Copies `template_path` to `out_path`, then overwrites the pixel data and
121/// updates BMAJ/BMIN/BPA in the header.  This preserves all other keywords
122/// (WCS, HISTORY, etc.) from the original file.
123pub fn write_fits(
124    image: &Array2<f32>,
125    out_path: &Path,
126    template_path: &Path,
127    new_beam: &Beam,
128    _was_4d: bool,
129) -> Result<(), FitsError> {
130    // Initialise the destination cheaply: copy only the template's primary-HDU
131    // header (preserving WCS/HISTORY/etc.), not its pixel data.  The data unit is
132    // defined by the copied NAXIS keywords; we overwrite every pixel below, so
133    // reading the template's data via `std::fs::copy` would be wasted I/O.
134    copy_header_only(template_path, out_path)?;
135
136    let out_str = out_path.to_string_lossy().into_owned();
137    let mut fptr = FitsFile::edit(&out_str)?;
138    let hdu = fptr.primary_hdu()?;
139
140    // Flatten to Vec<f32> in row-major order (C order = FITS row-major).
141    let flat: Vec<f32> = image.iter().cloned().collect();
142    hdu.write_image(&mut fptr, &flat)?;
143
144    // Update beam keywords.
145    hdu.write_key(&mut fptr, "BMAJ", new_beam.major_deg)?;
146    hdu.write_key(&mut fptr, "BMIN", new_beam.minor_deg)?;
147    hdu.write_key(&mut fptr, "BPA", new_beam.pa_deg)?;
148
149    Ok(())
150}
151
152/// Build the output path from the input path with an optional suffix/prefix/outdir
153/// (re-exported from [`atfits_rs`]).
154pub use atfits_rs::output_path;
convolve_rs/fits_io.rs

convolve_rs/
fits_io.rs
