Skip to main content

fitscube_rs/
beams.rs

1//! Per-channel restoring-beam handling.
2//!
3//! Port of the beam logic in `fitscube.combine_fits`: read BMAJ/BMIN/BPA (all in
4//! degrees) from each input header, decide whether they are all the same, and —
5//! when they vary — emit a CASA-style `BEAMS` binary-table extension.
6use std::path::{Path, PathBuf};
7
8use fitsio::FitsFile;
9use fitsio::tables::{ColumnDataType, ColumnDescription};
10
11use crate::error::Result;
12use crate::fits_io::{find_target_axis, read_key_f64};
13
14/// A restoring beam in degrees. Any field may be NaN (no beam in header).
15#[derive(Debug, Clone, Copy, PartialEq)]
16pub struct Beam {
17    pub major_deg: f64,
18    pub minor_deg: f64,
19    pub pa_deg: f64,
20}
21
22impl Beam {
23    pub fn is_nan(&self) -> bool {
24        self.major_deg.is_nan() || self.minor_deg.is_nan() || self.pa_deg.is_nan()
25    }
26}
27
28/// Read BMAJ/BMIN/BPA (degrees) from a header. Missing → a NaN beam, matching
29/// the `NoBeamException` fallback in the original.
30pub fn read_beam(path: &Path) -> Result<Beam> {
31    let major_deg = read_key_f64(path, "BMAJ")?.unwrap_or(f64::NAN);
32    let minor_deg = read_key_f64(path, "BMIN")?.unwrap_or(f64::NAN);
33    let pa_deg = read_key_f64(path, "BPA")?.unwrap_or(f64::NAN);
34    Ok(Beam {
35        major_deg,
36        minor_deg,
37        pa_deg,
38    })
39}
40
41/// Read one beam per file.
42pub fn parse_beams(file_list: &[PathBuf]) -> Result<Vec<Beam>> {
43    file_list.iter().map(|p| read_beam(p)).collect()
44}
45
46/// `numpy.isclose` with the default tolerances (`rtol = 1e-5`, `atol = 1e-8`).
47/// NaNs are never close (matching numpy with `equal_nan=False`).
48fn isclose(a: f64, b: f64) -> bool {
49    if a.is_nan() || b.is_nan() {
50        return false;
51    }
52    (a - b).abs() <= 1e-8 + 1e-5 * b.abs()
53}
54
55/// Whether every beam matches the first on major, minor and PA.
56///
57/// Mirrors the `single_beam` test in the original (compares shape, not area).
58/// If beams are all-NaN this returns `false`, since NaN is never close to NaN.
59pub fn is_single_beam(beams: &[Beam]) -> bool {
60    if beams.is_empty() {
61        return false;
62    }
63    let first = beams[0];
64    beams.iter().all(|b| {
65        isclose(first.major_deg, b.major_deg)
66            && isclose(first.minor_deg, b.minor_deg)
67            && isclose(first.pa_deg, b.pa_deg)
68    })
69}
70
71/// FITS polarisation index: `CRPIX − 1` of the STOKES axis, or 0 if absent.
72/// Mirrors `get_polarisation`.
73pub fn get_polarisation(path: &Path) -> Result<i32> {
74    match find_target_axis(path, "STOKES") {
75        Ok(axis) => Ok((axis.crpix - 1.0) as i32),
76        Err(_) => Ok(0),
77    }
78}
79
80/// Append a CASA `BEAMS` binary-table extension to an open output cube.
81///
82/// Columns: BMAJ/BMIN [arcsec], BPA [deg], CHAN, POL. NaN beam values are
83/// replaced with `f32::MIN_POSITIVE` (numpy `finfo(float32).tiny`) to keep CASA
84/// happy, exactly as `make_beam_table` does.
85pub fn write_beam_table(fptr: &mut FitsFile, beams: &[Beam], pol: i32) -> Result<()> {
86    let tiny = f32::MIN_POSITIVE;
87    let nchan = beams.len();
88
89    let nan_to_tiny = |v: f32| if v.is_nan() { tiny } else { v };
90
91    let bmaj: Vec<f32> = beams
92        .iter()
93        .map(|b| nan_to_tiny((b.major_deg * 3600.0) as f32))
94        .collect();
95    let bmin: Vec<f32> = beams
96        .iter()
97        .map(|b| nan_to_tiny((b.minor_deg * 3600.0) as f32))
98        .collect();
99    let bpa: Vec<f32> = beams.iter().map(|b| nan_to_tiny(b.pa_deg as f32)).collect();
100    let chan: Vec<i32> = (0..nchan as i32).collect();
101    let pols: Vec<i32> = vec![pol; nchan];
102
103    let col_bmaj = ColumnDescription::new("BMAJ")
104        .with_type(ColumnDataType::Float)
105        .create()?;
106    let col_bmin = ColumnDescription::new("BMIN")
107        .with_type(ColumnDataType::Float)
108        .create()?;
109    let col_bpa = ColumnDescription::new("BPA")
110        .with_type(ColumnDataType::Float)
111        .create()?;
112    let col_chan = ColumnDescription::new("CHAN")
113        .with_type(ColumnDataType::Int)
114        .create()?;
115    let col_pol = ColumnDescription::new("POL")
116        .with_type(ColumnDataType::Int)
117        .create()?;
118
119    let table_hdu =
120        fptr.create_table("BEAMS", &[col_bmaj, col_bmin, col_bpa, col_chan, col_pol])?;
121    table_hdu.write_col(fptr, "BMAJ", &bmaj)?;
122    table_hdu.write_col(fptr, "BMIN", &bmin)?;
123    table_hdu.write_col(fptr, "BPA", &bpa)?;
124    table_hdu.write_col(fptr, "CHAN", &chan)?;
125    table_hdu.write_col(fptr, "POL", &pols)?;
126
127    let beam_hdu = fptr.hdu("BEAMS")?;
128    beam_hdu.write_key(fptr, "NCHAN", nchan as i64)?;
129    beam_hdu.write_key(fptr, "NPOL", 1i64)?;
130
131    Ok(())
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137
138    fn b(maj: f64, min: f64, pa: f64) -> Beam {
139        Beam {
140            major_deg: maj,
141            minor_deg: min,
142            pa_deg: pa,
143        }
144    }
145
146    #[test]
147    fn identical_beams_are_single() {
148        let beams = vec![b(1.0, 0.5, 10.0); 4];
149        assert!(is_single_beam(&beams));
150    }
151
152    #[test]
153    fn varying_beams_are_not_single() {
154        let beams = vec![b(1.0, 0.5, 10.0), b(1.1, 0.5, 10.0)];
155        assert!(!is_single_beam(&beams));
156    }
157
158    #[test]
159    fn all_nan_is_not_single() {
160        let beams = vec![b(f64::NAN, f64::NAN, f64::NAN); 3];
161        assert!(!is_single_beam(&beams));
162    }
163}