1use 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#[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
28pub 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
41pub fn parse_beams(file_list: &[PathBuf]) -> Result<Vec<Beam>> {
43 file_list.iter().map(|p| read_beam(p)).collect()
44}
45
46fn 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
55pub 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
71pub 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
80pub 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}