convolve_rs/
cube_io.rs

1//! FITS spectral cube reading and writing with per-channel beams.
2//!
3//! Cubes are 3D (RA, Dec, FREQ) or 4D with one Stokes plane, in either axis
4//! order. Per-channel beams come from, in order of preference:
5//!   1. a CASA BEAMS binary table (`CASAMBM = T`),
6//!   2. a beamlog text file `{dir}/beamlog.{stem}.txt`,
7//!   3. BMAJ/BMIN/BPA in the primary header, used for every channel.
8use std::path::{Path, PathBuf};
9
10use atfits_rs::{key_bool, pixel_scale_deg, update_key_f64, update_key_logical};
11use fitsio::{
12    FitsFile,
13    tables::{ColumnDataType, ColumnDescription},
14};
15use ndarray::Array2;
16use thiserror::Error;
17
18use crate::beam::{Beam, BeamError};
19use crate::convolve_uv::FftFloat;
20use crate::fits_io::{FitsError, create_output, read_header_beam, read_unit};
21use crate::smooth::BrightnessUnit;
22
23pub use atfits_rs::PixelType;
24
25/// Pixel types a plane is read, convolved and written in: `f32` or `f64`.
26pub trait CubeElem: atfits_rs::CubeElem + FftFloat {}
27
28impl<T: atfits_rs::CubeElem + FftFloat> CubeElem for T {}
29
30#[derive(Debug, Error)]
31pub enum CubeError {
32    #[error("FITS I/O error: {0}")]
33    Fits(#[from] fitsio::errors::Error),
34    #[error(transparent)]
35    Header(#[from] FitsError),
36    #[error("I/O error: {0}")]
37    Io(#[from] std::io::Error),
38    #[error("shape error: {0}")]
39    Shape(#[from] ndarray::ShapeError),
40    #[error("invalid beam: {0}")]
41    Beam(#[from] BeamError),
42    #[error("unsupported NAXIS={0} (expected 3 or 4)")]
43    UnsupportedNaxis(i64),
44    #[error("missing header keyword: {0}")]
45    MissingKeyword(String),
46    #[error("channel count mismatch in beam table: expected {expected}, got {got}")]
47    BeamCountMismatch { expected: usize, got: usize },
48    #[error("beamlog parse error at line {line}: {msg}")]
49    BeamlogParse { line: usize, msg: String },
50    #[error("no per-channel beam source found (no CASAMBM, no beamlog, no header beam)")]
51    NoBeams,
52}
53
54impl From<atfits_rs::AtfitsError> for CubeError {
55    fn from(e: atfits_rs::AtfitsError) -> Self {
56        CubeError::Header(e.into())
57    }
58}
59
60/// Metadata of a FITS spectral cube.
61#[derive(Debug)]
62pub struct CubeMeta {
63    pub path: PathBuf,
64    /// NAXIS1 (RA).
65    pub nx: usize,
66    /// NAXIS2 (Dec).
67    pub ny: usize,
68    pub nfreq: usize,
69    pub nstokes: usize,
70    /// Pixel size along x in degrees.
71    pub dx_deg: f64,
72    /// Pixel size along y in degrees.
73    pub dy_deg: f64,
74    /// CRPIX of the spectral axis, truncated to a whole pixel as racs_tools
75    /// does. Picks the channel whose beam goes in the output header; may lie
76    /// outside `1..=nfreq`.
77    pub crpix_freq: f64,
78    /// Per-channel beams; `None` for a masked channel.
79    pub beams: Vec<Option<Beam>>,
80    pub is_4d: bool,
81    pub unit: BrightnessUnit,
82    pub dtype: PixelType,
83}
84
85impl CubeMeta {
86    /// Flat pixel range `[start, end)` of channel `chan` (first Stokes plane).
87    /// With one Stokes plane both axis orders put channel `c` at `c·ny·nx`.
88    pub fn channel_range(&self, chan: usize) -> (usize, usize) {
89        let plane = self.ny * self.nx;
90        (chan * plane, (chan + 1) * plane)
91    }
92}
93
94/// The beamlog next to a FITS file: `{dir}/beamlog.{stem}.txt`.
95pub fn beamlog_path(fits: &Path) -> PathBuf {
96    let dir = fits.parent().unwrap_or(Path::new("."));
97    let stem = fits.file_stem().unwrap_or_default();
98    dir.join(format!("beamlog.{}.txt", stem.to_string_lossy()))
99}
100
101/// Read the shape, pixel scale and per-channel beams of a cube.
102pub fn read_cube_meta(path: &Path) -> Result<CubeMeta, CubeError> {
103    let mut fptr = FitsFile::open(path)?;
104    let hdu = fptr.primary_hdu()?;
105
106    let naxis: i64 = hdu.read_key(&mut fptr, "NAXIS")?;
107    if naxis != 3 && naxis != 4 {
108        return Err(CubeError::UnsupportedNaxis(naxis));
109    }
110    let mut dims = Vec::with_capacity(naxis as usize);
111    for i in 1..=naxis {
112        dims.push(hdu.read_key::<i64>(&mut fptr, &format!("NAXIS{i}"))? as usize);
113    }
114
115    // CASA writes RA, DEC, STOKES, FREQ; ASKAP writes RA, DEC, FREQ, STOKES.
116    let ctype3 = hdu
117        .read_key::<String>(&mut fptr, "CTYPE3")
118        .map(|s| s.trim().to_ascii_uppercase())
119        .unwrap_or_default();
120    let (nstokes, nfreq, freq_axis) = match dims[..] {
121        [_, _, n3, n4] if ctype3 == "STOKES" => (n3, n4, 4),
122        [_, _, n3, n4] => (n4, n3, 3),
123        _ => (1, dims[2], 3),
124    };
125
126    let (dx_deg, dy_deg) = pixel_scale_deg(&mut fptr)?;
127    let crpix_freq = atfits_rs::key_f64(&mut fptr, &format!("CRPIX{freq_axis}"))?
128        .unwrap_or(1.0)
129        .floor();
130    let bitpix: i64 = hdu.read_key(&mut fptr, "BITPIX")?;
131    let unit = read_unit(&mut fptr, &hdu, path);
132
133    // CASA writes CASAMBM as a logical, but some tools write a string.
134    let casambm = key_bool(&mut fptr, "CASAMBM")?.unwrap_or(false);
135
136    let beamlog = beamlog_path(path);
137    let beams = if casambm {
138        read_casambm_beams(&mut fptr, nfreq, path)?
139    } else if beamlog.exists() {
140        let parsed = read_beamlog(&beamlog)?;
141        if parsed.len() != nfreq {
142            return Err(CubeError::BeamCountMismatch {
143                expected: nfreq,
144                got: parsed.len(),
145            });
146        }
147        parsed.into_iter().map(Some).collect()
148    } else {
149        let beam = read_header_beam(&mut fptr)?.ok_or(CubeError::NoBeams)?;
150        vec![Some(beam); nfreq]
151    };
152
153    Ok(CubeMeta {
154        path: path.to_path_buf(),
155        nx: dims[0],
156        ny: dims[1],
157        nfreq,
158        nstokes,
159        dx_deg,
160        dy_deg,
161        crpix_freq,
162        beams,
163        is_4d: naxis == 4,
164        unit,
165        dtype: PixelType::from_bitpix(bitpix),
166    })
167}
168
169/// Per-channel beams from the CASA BEAMS table (BMAJ, BMIN in arcsec; BPA in
170/// degrees).
171fn read_casambm_beams(
172    fptr: &mut FitsFile,
173    nfreq: usize,
174    path: &Path,
175) -> Result<Vec<Option<Beam>>, CubeError> {
176    let hdu = fptr
177        .hdu("BEAMS")
178        .map_err(|_| CubeError::MissingKeyword("BEAMS extension".into()))?;
179    let bmaj: Vec<f32> = hdu.read_col(fptr, "BMAJ")?;
180    let bmin: Vec<f32> = hdu.read_col(fptr, "BMIN")?;
181    let bpa: Vec<f32> = hdu.read_col(fptr, "BPA")?;
182    if bmaj.len() != nfreq || bmin.len() != nfreq || bpa.len() != nfreq {
183        return Err(CubeError::BeamCountMismatch {
184            expected: nfreq,
185            got: bmaj.len(),
186        });
187    }
188    Ok((0..nfreq)
189        .map(|c| {
190            beam_from_row(
191                f64::from(bmaj[c]),
192                f64::from(bmin[c]),
193                f64::from(bpa[c]),
194                c,
195                path,
196            )
197        })
198        .collect())
199}
200
201/// Beam from one row of a beam table (arcsec, arcsec, deg).
202///
203/// `None` for a masked channel: a non-finite value, or an axis no larger than
204/// the `f32::MIN_POSITIVE` written for masked channels. Swapped axes are
205/// fixed with a warning.
206fn beam_from_row(maj: f64, min: f64, pa: f64, chan: usize, path: &Path) -> Option<Beam> {
207    let tiny = f64::from(f32::MIN_POSITIVE);
208    if !(maj.is_finite() && min.is_finite() && pa.is_finite()) || maj <= tiny || min <= tiny {
209        return None;
210    }
211    if min > maj {
212        tracing::warn!(
213            "{}: channel {chan} has BMIN > BMAJ ({min}\" > {maj}\"); swapping the axes",
214            path.display()
215        );
216        let pa = if pa > 0.0 { pa - 90.0 } else { pa + 90.0 };
217        return Beam::from_arcsec(min, maj, pa).ok();
218    }
219    Beam::from_arcsec(maj, min, pa).ok()
220}
221
222/// An open input cube for reading channels. cfitsio handles are not `Send`,
223/// so each thread opens its own.
224pub struct CubeReader {
225    fptr: FitsFile,
226}
227
228impl CubeReader {
229    pub fn open(path: &Path) -> Result<Self, CubeError> {
230        Ok(Self {
231            fptr: FitsFile::open(path)?,
232        })
233    }
234
235    /// Channel `chan` (first Stokes plane) as an `ny × nx` array.
236    pub fn read_channel<T: CubeElem>(
237        &mut self,
238        chan: usize,
239        meta: &CubeMeta,
240    ) -> Result<Array2<T>, CubeError> {
241        let (start, end) = meta.channel_range(chan);
242        let data = T::read_section(&mut self.fptr, start, end)?;
243        Ok(Array2::from_shape_vec((meta.ny, meta.nx), data)?)
244    }
245}
246
247/// Mode for choosing the common beam of a cube.
248#[derive(Debug, Clone, Copy, PartialEq, Eq)]
249pub enum CubeMode {
250    /// A common beam per channel, written to a BEAMS table.
251    Natural,
252    /// One common beam for every channel, written to the primary header.
253    Total,
254}
255
256/// An output cube held open while channels are written, so cfitsio writes the
257/// data unit once, on [`CubeWriter::finish`]. Use from a single thread.
258pub struct CubeWriter {
259    fptr: FitsFile,
260    /// Natural mode's BEAMS table. Creating it makes cfitsio flush the primary
261    /// data unit, so it is written after the channels.
262    pending_beams: Option<Vec<Option<Beam>>>,
263}
264
265impl CubeWriter {
266    /// Create `output_path` with the header of `input_path` and the beam
267    /// keywords for `target_beams`.
268    pub fn create(
269        input_path: &Path,
270        output_path: &Path,
271        target_beams: &[Option<Beam>],
272        mode: CubeMode,
273        meta: &CubeMeta,
274    ) -> Result<Self, CubeError> {
275        let mut fptr = create_output(input_path, output_path, None)?;
276        write_primary_beam_keys(
277            &mut fptr,
278            ref_beam_for(target_beams, meta),
279            mode == CubeMode::Natural,
280        )?;
281        Ok(Self {
282            fptr,
283            pending_beams: (mode == CubeMode::Natural).then(|| target_beams.to_vec()),
284        })
285    }
286
287    pub fn write_channel<T: CubeElem>(
288        &mut self,
289        chan: usize,
290        data: &Array2<T>,
291        meta: &CubeMeta,
292    ) -> Result<(), CubeError> {
293        let (start, end) = meta.channel_range(chan);
294        let flat = data.as_standard_layout();
295        let slice = flat
296            .as_slice()
297            .ok_or_else(|| std::io::Error::other("plane is not contiguous"))?;
298        T::write_section(&mut self.fptr, start, end, slice)?;
299        Ok(())
300    }
301
302    /// Append the BEAMS table (natural mode) and close the file. Channels never
303    /// written read as zero.
304    pub fn finish(mut self) -> Result<(), CubeError> {
305        if let Some(beams) = self.pending_beams.take() {
306            write_beams_table(&mut self.fptr, &beams)?;
307        }
308        Ok(())
309    }
310}
311
312/// The beam for the primary header: the target at the spectral CRPIX (clamped
313/// to the channel range), else the first valid one, else zero.
314fn ref_beam_for(target_beams: &[Option<Beam>], meta: &CubeMeta) -> Beam {
315    let last = meta.nfreq.saturating_sub(1) as f64;
316    let ref_idx = (meta.crpix_freq - 1.0).clamp(0.0, last) as usize;
317    target_beams[ref_idx]
318        .or_else(|| target_beams.iter().find_map(|b| *b))
319        .unwrap_or(Beam::zero())
320}
321
322/// Set BMAJ/BMIN/BPA and CASAMBM in the primary header, replacing the copied
323/// cards. CASAMBM must be a logical: casacore reads it with `asBool`.
324fn write_primary_beam_keys(
325    fptr: &mut FitsFile,
326    ref_beam: Beam,
327    natural: bool,
328) -> Result<(), CubeError> {
329    fptr.primary_hdu()?;
330    update_key_f64(fptr, "BMAJ", ref_beam.major_deg())?;
331    update_key_f64(fptr, "BMIN", ref_beam.minor_deg())?;
332    update_key_f64(fptr, "BPA", ref_beam.pa_deg())?;
333    update_key_logical(fptr, "CASAMBM", natural)?;
334    Ok(())
335}
336
337/// Append a CASA BEAMS table. Masked channels get `f32::MIN_POSITIVE`, which
338/// [`read_cube_meta`] reads back as masked.
339fn write_beams_table(fptr: &mut FitsFile, target_beams: &[Option<Beam>]) -> Result<(), CubeError> {
340    let nfreq = target_beams.len();
341    let tiny = f32::MIN_POSITIVE;
342    let column = |f: fn(&Beam) -> f64| -> Vec<f32> {
343        target_beams
344            .iter()
345            .map(|b| b.map_or(tiny, |b| f(&b) as f32))
346            .collect()
347    };
348    let bmaj = column(Beam::major_arcsec);
349    let bmin = column(Beam::minor_arcsec);
350    let bpa = column(Beam::pa_deg);
351    let chan: Vec<i32> = (0..nfreq as i32).collect();
352    let pol = vec![0i32; nfreq];
353
354    let col = |name: &str, ty: ColumnDataType| ColumnDescription::new(name).with_type(ty).create();
355    let cols = [
356        col("BMAJ", ColumnDataType::Float)?,
357        col("BMIN", ColumnDataType::Float)?,
358        col("BPA", ColumnDataType::Float)?,
359        col("CHAN", ColumnDataType::Int)?,
360        col("POL", ColumnDataType::Int)?,
361    ];
362    let table = fptr.create_table("BEAMS", &cols)?;
363    table.write_col(fptr, "BMAJ", &bmaj)?;
364    table.write_col(fptr, "BMIN", &bmin)?;
365    table.write_col(fptr, "BPA", &bpa)?;
366    table.write_col(fptr, "CHAN", &chan)?;
367    table.write_col(fptr, "POL", &pol)?;
368
369    // casacore needs the column units.
370    let table = fptr.hdu("BEAMS")?;
371    table.write_key(fptr, "TUNIT1", "arcsec")?;
372    table.write_key(fptr, "TUNIT2", "arcsec")?;
373    table.write_key(fptr, "TUNIT3", "deg")?;
374    table.write_key(fptr, "NCHAN", nfreq as i64)?;
375    table.write_key(fptr, "NPOL", 1i64)?;
376    Ok(())
377}
378
379/// Read a beamlog as written by RACS-tools or [`write_beamlog`]:
380///
381/// ```text
382/// # Channel BMAJ[arcsec] BMIN[arcsec] BPA[deg]
383/// 0  20.0  10.0  10.0
384/// ```
385///
386/// Masked rows come back as [`Beam::zero`].
387pub fn read_beamlog(path: &Path) -> Result<Vec<Beam>, CubeError> {
388    let content = std::fs::read_to_string(path)?;
389    let mut beams = Vec::new();
390    for (i, line) in content.lines().enumerate() {
391        let trimmed = line.trim();
392        if trimmed.is_empty() || trimmed.starts_with('#') {
393            continue;
394        }
395        let fields: Vec<&str> = trimmed.split_whitespace().collect();
396        if fields.len() < 4 {
397            return Err(CubeError::BeamlogParse {
398                line: i + 1,
399                msg: format!("expected 4 fields, got {}", fields.len()),
400            });
401        }
402        let parse = |s: &str, n: &str| -> Result<f64, CubeError> {
403            s.parse::<f64>().map_err(|_| CubeError::BeamlogParse {
404                line: i + 1,
405                msg: format!("cannot parse {n}={s:?} as float"),
406            })
407        };
408        let (bmaj, bmin, bpa) = (
409            parse(fields[1], "BMAJ")?,
410            parse(fields[2], "BMIN")?,
411            parse(fields[3], "BPA")?,
412        );
413        let chan = beams.len();
414        beams.push(beam_from_row(bmaj, bmin, bpa, chan, path).unwrap_or(Beam::zero()));
415    }
416    Ok(beams)
417}
418
419/// Write per-channel beams as a beamlog; masked channels are written as NaN.
420pub fn write_beamlog(path: &Path, beams: &[Option<Beam>]) -> Result<(), CubeError> {
421    let mut out = String::from("# Channel BMAJ[arcsec] BMIN[arcsec] BPA[deg]\n");
422    for (i, b) in beams.iter().enumerate() {
423        out += &match b {
424            Some(b) => format!(
425                "{i} {} {} {}\n",
426                b.major_arcsec(),
427                b.minor_arcsec(),
428                b.pa_deg()
429            ),
430            None => format!("{i} nan nan nan\n"),
431        };
432    }
433    std::fs::write(path, out)?;
434    Ok(())
435}
convolve_rs/cube_io.rs

convolve_rs/
cube_io.rs
