convolve_rs/
smooth.rs

1//! High-level smoothing: convolve + apply Jy/beam flux scaling.
2use ndarray::Array2;
3#[cfg(feature = "python")]
4use ndarray::ArrayView2;
5use thiserror::Error;
6
7use crate::beam::Beam;
8use crate::convolve_uv::{
9    ConvolveError, FftFloat, FftPlans, OutputGain, PlaneInput, convolve_plane,
10};
11
12#[derive(Debug, Error)]
13pub enum SmoothError {
14    #[error("convolution failed: {0}")]
15    Convolve(#[from] ConvolveError),
16}
17
18/// Brightness unit of an image, determining whether flux scaling applies after
19/// convolution to a larger beam.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
21pub enum BrightnessUnit {
22    /// Jy/beam (or any per-beam flux density) — apply the Gaussian flux-scaling
23    /// factor so the output stays in the same units.
24    #[default]
25    JyPerBeam,
26    /// Kelvin (brightness temperature) — surface brightness is conserved under
27    /// convolution, so the scaling factor is 1.
28    Kelvin,
29}
30
31impl BrightnessUnit {
32    /// Classify a FITS `BUNIT` string, returning `None` if the unit is not
33    /// recognised (neither a Kelvin nor a Jy/beam form).
34    ///
35    /// # Examples
36    ///
37    /// ```
38    /// use convolve_rs::BrightnessUnit;
39    ///
40    /// assert_eq!(BrightnessUnit::parse("Jy/beam"), Some(BrightnessUnit::JyPerBeam));
41    /// assert_eq!(BrightnessUnit::parse(" K "), Some(BrightnessUnit::Kelvin));
42    /// assert_eq!(BrightnessUnit::parse("Jy/pixel"), None);
43    /// ```
44    pub fn parse(bunit: &str) -> Option<Self> {
45        let u = bunit.trim().trim_matches('\'').trim().to_ascii_uppercase();
46        match u.as_str() {
47            "K" | "KELVIN" => Some(BrightnessUnit::Kelvin),
48            "JY/BEAM" | "JY BEAM-1" | "JY/BM" | "JYBEAM" => Some(BrightnessUnit::JyPerBeam),
49            _ => None,
50        }
51    }
52
53    /// Classify a FITS `BUNIT` string.  Anything recognised as a brightness
54    /// temperature (Kelvin) skips flux scaling; recognised Jy/beam forms get the
55    /// Gaussian factor.  Unrecognised units cannot be determined automatically:
56    /// a warning is emitted and Jy/beam is assumed.
57    pub fn from_bunit(bunit: &str) -> Self {
58        match Self::parse(bunit) {
59            Some(unit) => unit,
60            None => {
61                tracing::warn!(
62                    "Could not determine brightness unit from BUNIT={bunit:?}; \
63                     assuming Jy/beam (flux scaling applied). Pass a recognised \
64                     unit (e.g. 'Jy/beam' or 'K') to silence this warning."
65                );
66                BrightnessUnit::JyPerBeam
67            }
68        }
69    }
70}
71
72/// Smooth `image` from `old_beam` to `new_beam`.
73///
74/// `dx_deg` / `dy_deg` are pixel sizes in degrees.  `unit` selects the flux
75/// scaling: [`BrightnessUnit::JyPerBeam`] applies the Gaussian factor,
76/// [`BrightnessUnit::Kelvin`] leaves the data unscaled (factor 1).  The
77/// convolution runs in the image's element type `T` (`f32` or `f64`); the
78/// result has the same dtype and pixel shape, ready to write back to FITS.
79///
80/// # Examples
81///
82/// Smoothing a flat image from a 10″ to a 20″ circular beam:
83///
84/// ```
85/// use convolve_rs::{Beam, BrightnessUnit, smooth};
86/// use ndarray::Array2;
87///
88/// let old = Beam::from_arcsec(10.0, 10.0, 0.0)?;
89/// let new = Beam::from_arcsec(20.0, 20.0, 0.0)?;
90/// let image = Array2::<f32>::from_elem((64, 64), 1.0);
91/// let dx = 2.5 / 3600.0;
92///
93/// // Jy/beam: pixel values scale by the beam-area ratio Ω_new/Ω_old = 4.
94/// let jy = smooth(&image, &old, &new, dx, dx, None, BrightnessUnit::JyPerBeam)?;
95/// assert!((jy[(32, 32)] - 4.0).abs() < 1e-3);
96///
97/// // Kelvin: surface brightness is conserved, so a flat image stays at 1.
98/// let k = smooth(&image, &old, &new, dx, dx, None, BrightnessUnit::Kelvin)?;
99/// assert!((k[(32, 32)] - 1.0).abs() < 1e-3);
100/// # Ok::<(), Box<dyn std::error::Error>>(())
101/// ```
102pub fn smooth<T: FftFloat>(
103    image: &Array2<T>,
104    old_beam: &Beam,
105    new_beam: &Beam,
106    dx_deg: f64,
107    dy_deg: f64,
108    cutoff_arcsec: Option<f64>,
109    unit: BrightnessUnit,
110) -> Result<Array2<T>, SmoothError> {
111    let (nrows, ncols) = image.dim();
112    let plans = FftPlans::<T>::new(nrows, ncols);
113    smooth_with_plans(
114        image,
115        old_beam,
116        new_beam,
117        dx_deg,
118        dy_deg,
119        cutoff_arcsec,
120        unit,
121        &plans,
122    )
123}
124
125/// Like [`smooth`], but reuses pre-built [`FftPlans`] across calls (e.g. every
126/// channel of a cube), avoiding per-call FFT planning. `plans` must match the
127/// image dimensions.
128#[allow(clippy::too_many_arguments)]
129pub fn smooth_with_plans<T: FftFloat>(
130    image: &Array2<T>,
131    old_beam: &Beam,
132    new_beam: &Beam,
133    dx_deg: f64,
134    dy_deg: f64,
135    cutoff_arcsec: Option<f64>,
136    unit: BrightnessUnit,
137    plans: &FftPlans<T>,
138) -> Result<Array2<T>, SmoothError> {
139    smooth_plane(
140        PlaneInput::Borrowed(image.view()),
141        old_beam,
142        new_beam,
143        dx_deg,
144        dy_deg,
145        cutoff_arcsec,
146        unit,
147        plans,
148    )
149}
150
151/// Like [`smooth_with_plans`], but takes ownership of `image` and smooths it in
152/// its own buffer, so a NaN-free plane needs no image-sized allocation at all
153/// (one with NaNs needs one, for the mask). Prefer this when the input is not
154/// needed afterwards, such as a cube channel just read from disk.
155#[allow(clippy::too_many_arguments)]
156pub fn smooth_owned_with_plans<T: FftFloat>(
157    image: Array2<T>,
158    old_beam: &Beam,
159    new_beam: &Beam,
160    dx_deg: f64,
161    dy_deg: f64,
162    cutoff_arcsec: Option<f64>,
163    unit: BrightnessUnit,
164    plans: &FftPlans<T>,
165) -> Result<Array2<T>, SmoothError> {
166    smooth_plane(
167        PlaneInput::Owned(image),
168        old_beam,
169        new_beam,
170        dx_deg,
171        dy_deg,
172        cutoff_arcsec,
173        unit,
174        plans,
175    )
176}
177
178/// [`smooth`] on a borrowed view (any memory layout), planning for its size.
179#[cfg(feature = "python")]
180#[allow(clippy::too_many_arguments)]
181pub(crate) fn smooth_view<T: FftFloat>(
182    image: ArrayView2<'_, T>,
183    old_beam: &Beam,
184    new_beam: &Beam,
185    dx_deg: f64,
186    dy_deg: f64,
187    cutoff_arcsec: Option<f64>,
188    unit: BrightnessUnit,
189) -> Result<Array2<T>, SmoothError> {
190    let (nrows, ncols) = image.dim();
191    let plans = FftPlans::<T>::new(nrows, ncols);
192    smooth_plane(
193        PlaneInput::Borrowed(image),
194        old_beam,
195        new_beam,
196        dx_deg,
197        dy_deg,
198        cutoff_arcsec,
199        unit,
200        &plans,
201    )
202}
203
204#[allow(clippy::too_many_arguments)]
205fn smooth_plane<T: FftFloat>(
206    image: PlaneInput<'_, T>,
207    old_beam: &Beam,
208    new_beam: &Beam,
209    dx_deg: f64,
210    dy_deg: f64,
211    cutoff_arcsec: Option<f64>,
212    unit: BrightnessUnit,
213    plans: &FftPlans<T>,
214) -> Result<Array2<T>, SmoothError> {
215    // The convolution already bakes one g_ratio (= √(Ω_new/Ω_old)) into the image.
216    // Jy/beam needs the full beam-area ratio Ω_new/Ω_old = g_ratio², so multiply
217    // by g_ratio once more. Kelvin conserves surface brightness, so the image
218    // must be flux-normalised — divide the baked-in g_ratio back out. Either is
219    // folded into the filter, so it costs no extra pass over the image.
220    let gain = match unit {
221        BrightnessUnit::JyPerBeam => OutputGain::Ratio,
222        BrightnessUnit::Kelvin => OutputGain::InverseRatio,
223    };
224    let result = convolve_plane(
225        image,
226        old_beam,
227        new_beam,
228        dx_deg,
229        dy_deg,
230        cutoff_arcsec,
231        gain,
232        plans,
233    )?;
234    Ok(result.image)
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use ndarray::Array2;
241
242    #[test]
243    fn test_parse_recognised() {
244        assert_eq!(BrightnessUnit::parse("K"), Some(BrightnessUnit::Kelvin));
245        assert_eq!(BrightnessUnit::parse(" k "), Some(BrightnessUnit::Kelvin));
246        assert_eq!(BrightnessUnit::parse("'K'"), Some(BrightnessUnit::Kelvin));
247        assert_eq!(
248            BrightnessUnit::parse("Kelvin"),
249            Some(BrightnessUnit::Kelvin)
250        );
251        assert_eq!(
252            BrightnessUnit::parse("Jy/beam"),
253            Some(BrightnessUnit::JyPerBeam)
254        );
255        assert_eq!(
256            BrightnessUnit::parse("JY BEAM-1"),
257            Some(BrightnessUnit::JyPerBeam)
258        );
259    }
260
261    #[test]
262    fn test_parse_unrecognised_is_none() {
263        // Unknown / ambiguous units cannot be determined automatically.
264        assert_eq!(BrightnessUnit::parse(""), None);
265        assert_eq!(BrightnessUnit::parse("Jy/pixel"), None);
266        assert_eq!(BrightnessUnit::parse("mJy"), None);
267    }
268
269    #[test]
270    fn test_from_bunit_falls_back_to_jy_per_beam() {
271        // Recognised forms classify directly; unrecognised forms warn and
272        // assume Jy/beam.
273        assert_eq!(BrightnessUnit::from_bunit("K"), BrightnessUnit::Kelvin);
274        assert_eq!(
275            BrightnessUnit::from_bunit("Jy/beam"),
276            BrightnessUnit::JyPerBeam
277        );
278        assert_eq!(BrightnessUnit::from_bunit("wat"), BrightnessUnit::JyPerBeam);
279    }
280
281    #[test]
282    fn test_kelvin_skips_flux_scaling() {
283        let old = Beam::new(10.0 / 3600.0, 10.0 / 3600.0, 0.0).unwrap();
284        let new = Beam::new(20.0 / 3600.0, 20.0 / 3600.0, 0.0).unwrap();
285        let img = Array2::from_elem((32, 32), 1.0_f32);
286        let dx = 2.5 / 3600.0;
287
288        let jy = smooth(&img, &old, &new, dx, dx, None, BrightnessUnit::JyPerBeam).unwrap();
289        let k = smooth(&img, &old, &new, dx, dx, None, BrightnessUnit::Kelvin).unwrap();
290
291        // Jy/beam scales flux up (Ω_new/Ω_old = 4); Kelvin leaves it unscaled.
292        let center = (16, 16);
293        assert!(jy[center] > k[center] * 1.5, "Jy/beam should be scaled up");
294        // Kelvin output is the pure convolution: a flat image stays ~1.
295        assert!(
296            (k[center] - 1.0).abs() < 1e-3,
297            "Kelvin center = {}",
298            k[center]
299        );
300    }
301}
convolve_rs/smooth.rs

convolve_rs/
smooth.rs
