convolve_rs/

common_beam.rs

//! The smallest beam a set of beams can be convolved to, as radio_beam finds
//! it: the CASA formula for two beams, otherwise a minimum-volume ellipse
//! around points sampled on every beam (Khachiyan's algorithm).
use crate::beam::{Beam, BeamError};
use thiserror::Error;

const DEG2RAD: f64 = std::f64::consts::PI / 180.0;

#[derive(Debug, Error)]
pub enum CommonBeamError {
    #[error("no beams provided")]
    NoBeams,
    #[error("all beams are invalid or flagged")]
    AllFlagged,
    #[error("Khachiyan algorithm did not converge after {0} iterations")]
    NoConvergence(usize),
    #[error("no common beam found that contains every input (epsilon reached {0:.2e})")]
    NotFound(f64),
    #[error("singular matrix in the minimum-volume ellipse fit")]
    Singular,
    #[error("nsamps must be at least 1")]
    NoSamples,
    #[error("beam error: {0}")]
    Beam(#[from] BeamError),
}

/// The smallest beam every beam in `beams` can be convolved to.
///
/// `tolerance`, `nsamps` and `epsilon` are as in
/// `radio_beam.Beams.common_beam(method='pts')`.
///
/// ```
/// use convolve_rs::{Beam, common_beam};
///
/// let b1 = Beam::from_arcsec(10.0, 8.0, 30.0)?;
/// let b2 = Beam::from_arcsec(12.0, 6.0, 60.0)?;
/// let common = common_beam(&[b1, b2], 1e-4, 200, 5e-4)?;
///
/// assert!(common.area_sr() >= b1.area_sr());
/// assert!(common.area_sr() >= b2.area_sr());
/// assert!(common.major_arcsec() >= 12.0);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub fn common_beam(
    beams: &[Beam],
    tolerance: f64,
    nsamps: usize,
    epsilon: f64,
) -> Result<Beam, CommonBeamError> {
    if beams.is_empty() {
        return Err(CommonBeamError::NoBeams);
    }
    if nsamps == 0 {
        return Err(CommonBeamError::NoSamples);
    }
    if beams.len() == 1 {
        return Ok(beams[0]);
    }

    let largest = largest_beam(beams);
    if fits_in_beam(beams, &largest) {
        return Ok(largest);
    }

    if beams.len() == 2
        && let Ok(b) = find_commonbeam_between(&beams[0], &beams[1])
        && fits_in_beam(beams, &b)
    {
        return Ok(b);
    }

    common_manybeams_mve(beams, tolerance, nsamps, epsilon)
}

/// The common beam of two beams by the CASA `ia.commonbeam` formula. The
/// result is not checked; [`common_beam`] checks it.
pub fn find_commonbeam_between(beam1: &Beam, beam2: &Beam) -> Result<Beam, CommonBeamError> {
    if beam1.approx_eq(beam2) {
        return Ok(*beam1);
    }

    let (large_beam, small_beam) = if beam1.area_sr() >= beam2.area_sr() {
        (beam1, beam2)
    } else {
        (beam2, beam1)
    };

    if large_beam.contains(small_beam) {
        return Ok(*large_beam);
    }

    let large_major = large_beam.major_arcsec();
    let large_minor = large_beam.minor_arcsec();
    let small_major = small_beam.major_arcsec();

    // If the small beam is circular the minor axis is the circle radius.
    if small_beam.is_circular(1e-6) {
        let beam = Beam::from_arcsec(large_major, small_major, large_beam.pa_deg())?;
        return Ok(beam);
    }

    let pa_diff_rad = ((small_beam.pa_deg() - large_beam.pa_deg()) * DEG2RAD
        + std::f64::consts::FRAC_PI_2
        + std::f64::consts::PI)
        .rem_euclid(std::f64::consts::PI)
        - std::f64::consts::FRAC_PI_2;

    // Perpendicular beams: the axes are the two major axes.
    if (pa_diff_rad.abs() - std::f64::consts::FRAC_PI_2).abs() < 1e-9 {
        let (major, minor) = if large_major >= small_major {
            (large_major, small_major)
        } else {
            (small_major, large_major)
        };
        let pa = if large_major >= small_major {
            large_beam.pa_deg()
        } else {
            small_beam.pa_deg()
        };
        return Beam::from_arcsec(major, minor, pa).map_err(Into::into);
    }

    // Work in the frame where the large beam is circular.
    let major_comb = (large_major * small_major).sqrt();
    let p = major_comb / large_major;
    let q = major_comb / large_minor;

    let (trans_maj_sc, _trans_min_sc, trans_pa_sc) =
        transform_ellipse_arcsec(small_major, small_beam.minor_arcsec(), pa_diff_rad, p, q);

    // CASA sets the minor axis to the geometric mean of the majors.
    let trans_min_sc = major_comb;

    let (trans_maj_unsc, trans_min_unsc, trans_pa_unsc) =
        transform_ellipse_arcsec(trans_maj_sc, trans_min_sc, trans_pa_sc, 1.0 / p, 1.0 / q);

    let final_pa_deg = trans_pa_unsc.to_degrees() + large_beam.pa_deg();

    let eps = 100.0 * f64::EPSILON;
    let beam = Beam::from_arcsec(trans_maj_unsc + eps, trans_min_unsc + eps, final_pa_deg)?;

    Ok(beam)
}

/// Scale an ellipse (arcsec axes, PA in radians) by `(x_scale, y_scale)`, as
/// radio_beam's `transform_ellipse`. Returns `(major, minor, pa_rad)`.
fn transform_ellipse_arcsec(
    major: f64,
    minor: f64,
    pa: f64,
    x_scale: f64,
    y_scale: f64,
) -> (f64, f64, f64) {
    let cospa = pa.cos();
    let sinpa = pa.sin();
    let cos2pa = cospa * cospa;
    let sin2pa = sinpa * sinpa;
    let major2 = major * major;
    let minor2 = minor * minor;

    let a = cos2pa / major2 + sin2pa / minor2;
    let b = -2.0 * cospa * sinpa * (1.0 / major2 - 1.0 / minor2);
    let c = sin2pa / major2 + cos2pa / minor2;

    let x2 = x_scale * x_scale;
    let y2 = y_scale * y_scale;

    let r = a / x2;
    let s = b * b / (4.0 * x2 * y2);
    let t = c / y2;

    let udiff = r - t;
    let f1 = udiff * udiff + 4.0 * s;
    let f2 = f1.sqrt() * udiff.abs();

    let j1 = (f2 + f1) / (2.0 * f1);
    let j2 = (f1 - f2) / (2.0 * f1);

    let k1 = (j1 * (r + t) - t) / (2.0 * j1 - 1.0);
    let k2 = (j2 * (r + t) - t) / (2.0 * j2 - 1.0);

    let c1 = 1.0 / k1.sqrt();
    let c2 = 1.0 / k2.sqrt();

    let pa_sign = if pa >= 0.0 { 1.0 } else { -1.0 };

    if (c1 - c2).abs() < f64::EPSILON {
        (1.0 / c1, 1.0 / c1, 0.0)
    } else if c1 > c2 {
        (c1, c2, pa_sign * j1.sqrt().acos())
    } else {
        (c2, c1, pa_sign * j2.sqrt().acos())
    }
}

/// The minimum-volume ellipse around `nsamps` points on each beam. If it does
/// not contain every beam, the points are pushed out by a growing `epsilon`
/// and the fit repeated, as radio_beam does.
pub fn common_manybeams_mve(
    beams: &[Beam],
    tolerance: f64,
    nsamps: usize,
    epsilon: f64,
) -> Result<Beam, CommonBeamError> {
    let max_iter = 10;
    let max_epsilon = 1e-3_f64;
    let mut eps = epsilon;

    for step in 0..=max_iter {
        let all_pts = collect_ellipse_points(beams, nsamps, eps);
        let hull_pts = convex_hull_2d(&all_pts);

        let (radii, rotation) = min_vol_ellipse(&hull_pts, tolerance)?;

        // radio_beam's rotation matrix convention: ((sin θ, cos θ), (cos θ, -sin θ)).
        let pa = (-rotation[0][0]).atan2(rotation[1][0]);
        let pa = if pa.abs() == std::f64::consts::PI {
            0.0
        } else {
            pa
        };

        let (major_deg, minor_deg) = if radii[0] >= radii[1] {
            (radii[0], radii[1])
        } else {
            (radii[1], radii[0])
        };
        let com_beam = Beam::new(major_deg, minor_deg, pa.to_degrees())?;
        if fits_in_beam(beams, &com_beam) {
            return Ok(com_beam);
        }

        if step < max_iter {
            eps += (step as f64 + 1.0) * (max_epsilon - eps) / max_iter as f64;
        }
    }
    Err(CommonBeamError::NotFound(eps))
}

/// Sample `nsamps` points on the edge of each beam's ellipse (scaled by `1 + epsilon`).
fn collect_ellipse_points(beams: &[Beam], nsamps: usize, epsilon: f64) -> Vec<[f64; 2]> {
    let mut pts = Vec::with_capacity(beams.len() * nsamps);
    for beam in beams {
        let bpa = beam.pa_deg() * DEG2RAD;
        let major = beam.major_deg() * (1.0 + epsilon);
        let minor = beam.minor_deg() * (1.0 + epsilon);
        for k in 0..nsamps {
            let phi = 2.0 * std::f64::consts::PI * k as f64 / nsamps as f64;
            let x = major * phi.cos();
            let y = minor * phi.sin();
            let xr = x * bpa.cos() - y * bpa.sin();
            let yr = x * bpa.sin() + y * bpa.cos();
            pts.push([xr, yr]);
        }
    }
    pts
}

/// 2D convex hull via Graham scan. Returns hull vertices.
fn convex_hull_2d(pts: &[[f64; 2]]) -> Vec<[f64; 2]> {
    if pts.len() <= 3 {
        return pts.to_vec();
    }

    let pivot = pts
        .iter()
        .enumerate()
        .min_by(|(_, a), (_, b)| a[1].total_cmp(&b[1]).then(a[0].total_cmp(&b[0])))
        .map(|(i, _)| i)
        .unwrap();

    let pivot_pt = pts[pivot];

    let mut sorted: Vec<[f64; 2]> = pts.iter().filter(|&&p| p != pivot_pt).cloned().collect();

    sorted.sort_by(|a, b| {
        let angle_a = (a[1] - pivot_pt[1]).atan2(a[0] - pivot_pt[0]);
        let angle_b = (b[1] - pivot_pt[1]).atan2(b[0] - pivot_pt[0]);
        angle_a.total_cmp(&angle_b)
    });

    let mut hull: Vec<[f64; 2]> = vec![pivot_pt];
    for &p in &sorted {
        while hull.len() > 1 {
            let n = hull.len();
            let cross = cross2d(hull[n - 2], hull[n - 1], p);
            if cross <= 0.0 {
                hull.pop();
            } else {
                break;
            }
        }
        hull.push(p);
    }
    hull
}

fn cross2d(o: [f64; 2], a: [f64; 2], b: [f64; 2]) -> f64 {
    (a[0] - o[0]) * (b[1] - o[1]) - (a[1] - o[1]) * (b[0] - o[0])
}

/// Khachiyan's minimum-volume enclosing ellipse of 2D points, following
/// radio_beam's `getMinVolEllipse`. Returns the semi-axes and the rotation
/// matrix whose columns are the axes.
fn min_vol_ellipse(
    pts: &[[f64; 2]],
    tolerance: f64,
) -> Result<([f64; 2], [[f64; 2]; 2]), CommonBeamError> {
    let n = pts.len();
    let d = 2_usize;

    let q: Vec<[f64; 3]> = pts.iter().map(|p| [p[0], p[1], 1.0]).collect(); // (N, 3)

    let mut u = vec![1.0 / n as f64; n];

    let max_iter = 100_000;
    let mut err = 1.0_f64;
    let mut iter = 0;

    while err > tolerance {
        // V = Q.T * diag(u) * Q  (3x3 matrix)
        let v = matmul_qt_diag_q(&q, &u); // 3x3
        let v_inv = mat3_inv(v)?;

        // M[i] = q[i].T * V_inv * q[i]
        let m: Vec<f64> = q.iter().map(|qi| quadratic_form_3(qi, &v_inv)).collect();

        let j = m
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.total_cmp(b))
            .map(|(i, _)| i)
            .unwrap();
        let maximum = m[j];

        let step = (maximum - d as f64 - 1.0) / ((d as f64 + 1.0) * (maximum - 1.0));
        let new_u: Vec<f64> = u.iter().map(|&ui| (1.0 - step) * ui).collect();
        err = new_u
            .iter()
            .zip(u.iter())
            .map(|(a, b)| (a - b).powi(2))
            .sum::<f64>()
            .sqrt();
        u = new_u;
        u[j] += step;

        iter += 1;
        if iter >= max_iter {
            return Err(CommonBeamError::NoConvergence(max_iter));
        }
    }

    let center = [
        pts.iter()
            .zip(u.iter())
            .map(|(p, &ui)| p[0] * ui)
            .sum::<f64>(),
        pts.iter()
            .zip(u.iter())
            .map(|(p, &ui)| p[1] * ui)
            .sum::<f64>(),
    ];

    // A = inv(P.T * diag(u) * P - center*center.T) / d  (2x2)
    let ptdp = matmul_pt_diag_p(pts, &u); // 2x2
    let cc = [
        [center[0] * center[0], center[0] * center[1]],
        [center[1] * center[0], center[1] * center[1]],
    ];
    let inner = [
        [ptdp[0][0] - cc[0][0], ptdp[0][1] - cc[0][1]],
        [ptdp[1][0] - cc[1][0], ptdp[1][1] - cc[1][1]],
    ];
    let a = mat2_scale(mat2_inv(inner)?, 1.0 / d as f64);

    // A is symmetric positive definite, so its SVD is its eigendecomposition.
    let (eigenvalues, eigenvectors) = symmetric_2x2_eig(a);

    let radii = [
        (1.0 / eigenvalues[0].abs().sqrt()) * (1.0 + tolerance),
        (1.0 / eigenvalues[1].abs().sqrt()) * (1.0 + tolerance),
    ];

    Ok((radii, eigenvectors))
}

/// Q.T * diag(u) * Q where Q is (N, 3). Returns 3x3.
fn matmul_qt_diag_q(q: &[[f64; 3]], u: &[f64]) -> [[f64; 3]; 3] {
    let mut v = [[0.0_f64; 3]; 3];
    for (qi, &ui) in q.iter().zip(u.iter()) {
        for r in 0..3 {
            for c in 0..3 {
                v[r][c] += qi[r] * ui * qi[c];
            }
        }
    }
    v
}

/// P.T * diag(u) * P where P is (N, 2). Returns 2x2.
fn matmul_pt_diag_p(p: &[[f64; 2]], u: &[f64]) -> [[f64; 2]; 2] {
    let mut m = [[0.0_f64; 2]; 2];
    for (pi, &ui) in p.iter().zip(u.iter()) {
        for r in 0..2 {
            for c in 0..2 {
                m[r][c] += pi[r] * ui * pi[c];
            }
        }
    }
    m
}

/// Quadratic form x.T * M * x for 3-vector x and 3x3 matrix M.
fn quadratic_form_3(x: &[f64; 3], m: &[[f64; 3]; 3]) -> f64 {
    let mut acc = 0.0_f64;
    for r in 0..3 {
        for c in 0..3 {
            acc += x[r] * m[r][c] * x[c];
        }
    }
    acc
}

/// Invert a 3x3 matrix (Cramer's rule).
fn mat3_inv(m: [[f64; 3]; 3]) -> Result<[[f64; 3]; 3], CommonBeamError> {
    let det = m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
        - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
        + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]);

    if det.abs() < f64::EPSILON {
        return Err(CommonBeamError::Singular);
    }
    let inv_det = 1.0 / det;

    Ok([
        [
            inv_det * (m[1][1] * m[2][2] - m[1][2] * m[2][1]),
            inv_det * (m[0][2] * m[2][1] - m[0][1] * m[2][2]),
            inv_det * (m[0][1] * m[1][2] - m[0][2] * m[1][1]),
        ],
        [
            inv_det * (m[1][2] * m[2][0] - m[1][0] * m[2][2]),
            inv_det * (m[0][0] * m[2][2] - m[0][2] * m[2][0]),
            inv_det * (m[0][2] * m[1][0] - m[0][0] * m[1][2]),
        ],
        [
            inv_det * (m[1][0] * m[2][1] - m[1][1] * m[2][0]),
            inv_det * (m[0][1] * m[2][0] - m[0][0] * m[2][1]),
            inv_det * (m[0][0] * m[1][1] - m[0][1] * m[1][0]),
        ],
    ])
}

/// Invert a symmetric 2x2 matrix.
fn mat2_inv(m: [[f64; 2]; 2]) -> Result<[[f64; 2]; 2], CommonBeamError> {
    let det = m[0][0] * m[1][1] - m[0][1] * m[1][0];
    if det.abs() < f64::EPSILON {
        return Err(CommonBeamError::Singular);
    }
    let inv_det = 1.0 / det;
    Ok([
        [inv_det * m[1][1], -inv_det * m[0][1]],
        [-inv_det * m[1][0], inv_det * m[0][0]],
    ])
}

fn mat2_scale(m: [[f64; 2]; 2], s: f64) -> [[f64; 2]; 2] {
    [[m[0][0] * s, m[0][1] * s], [m[1][0] * s, m[1][1] * s]]
}

/// Eigenvalues and eigenvectors (as columns) of a symmetric 2x2 matrix.
fn symmetric_2x2_eig(m: [[f64; 2]; 2]) -> ([f64; 2], [[f64; 2]; 2]) {
    let a = m[0][0];
    let b = m[0][1]; // = m[1][0] since symmetric
    let c = m[1][1];

    let trace = a + c;
    let det = a * c - b * b;
    let disc = ((trace / 2.0).powi(2) - det).max(0.0).sqrt();

    let lam1 = trace / 2.0 + disc;
    let lam2 = trace / 2.0 - disc;

    let (v1, v2) = if b.abs() > f64::EPSILON {
        let v1 = normalise([lam1 - c, b]);
        let v2 = normalise([lam2 - c, b]);
        (v1, v2)
    } else if a >= c {
        ([1.0, 0.0], [0.0, 1.0])
    } else {
        ([0.0, 1.0], [1.0, 0.0])
    };

    // Columns are eigenvectors, as from radio_beam's SVD.
    let rotation = [[v1[0], v2[0]], [v1[1], v2[1]]];

    ([lam1, lam2], rotation)
}

fn normalise(v: [f64; 2]) -> [f64; 2] {
    let len = (v[0] * v[0] + v[1] * v[1]).sqrt();
    if len < f64::EPSILON {
        return v;
    }
    [v[0] / len, v[1] / len]
}

pub fn largest_beam(beams: &[Beam]) -> Beam {
    beams
        .iter()
        .max_by(|a, b| a.area_sr().total_cmp(&b.area_sr()))
        .copied()
        .unwrap()
}

/// True if every beam in `beams` can be convolved up to `large_beam`.
pub fn fits_in_beam(beams: &[Beam], large_beam: &Beam) -> bool {
    beams.iter().all(|b| large_beam.contains(b))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_common_beam_identical() {
        let b = Beam::new(10.0 / 3600.0, 8.0 / 3600.0, 30.0).unwrap();
        let result = common_beam(&[b, b], 1e-4, 200, 5e-4).unwrap();
        assert!(result.approx_eq(&b) || result.major_deg() >= b.major_deg());
    }

    #[test]
    fn test_common_beam_two_different() {
        let b1 = Beam::new(10.0 / 3600.0, 8.0 / 3600.0, 30.0).unwrap();
        let b2 = Beam::new(12.0 / 3600.0, 6.0 / 3600.0, 60.0).unwrap();
        let result = common_beam(&[b1, b2], 1e-4, 200, 5e-4).unwrap();
        assert!(
            result.major_deg() >= b1.major_deg().max(b2.major_deg())
                || fits_in_beam(&[b1, b2], &result)
        );
    }

    /// A target whose major axis equals an input's major axis still contains
    /// it; the kernel is one-dimensional.
    #[test]
    fn test_equal_major_axis_fits() {
        let target = Beam::from_arcsec(20.0, 20.0, 0.0).unwrap();
        for pa in [0.0, 37.0, 90.0, -45.0] {
            let b = Beam::from_arcsec(20.0, 10.0, pa).unwrap();
            assert!(fits_in_beam(&[b], &target), "PA {pa}");
        }
        let common = common_beam(
            &[target, Beam::from_arcsec(20.0, 10.0, 0.0).unwrap()],
            1e-4,
            200,
            5e-4,
        )
        .unwrap();
        assert!(common.approx_eq(&target), "got {common}");
    }

    #[test]
    fn test_zero_samples_is_an_error() {
        let b = Beam::from_arcsec(10.0, 8.0, 0.0).unwrap();
        let c = Beam::from_arcsec(9.0, 9.0, 0.0).unwrap();
        assert!(common_beam(&[b, c, b], 1e-4, 0, 5e-4).is_err());
    }

    /// The CASA two-beam solution is only used if it contains both inputs.
    #[test]
    fn test_two_beam_result_contains_inputs() {
        for (a, b) in [
            ((10.0, 8.0, 30.0), (12.0, 6.0, 60.0)),
            // The CASA formula gives 19.6255"x8.0238" here, which does
            // not contain the first beam.
            (
                (17.134150314045915, 7.278818100904679, -4.1357805889849),
                (17.230792491993967, 4.380913928098483, 14.758744881833891),
            ),
        ] {
            let b1 = Beam::from_arcsec(a.0, a.1, a.2).unwrap();
            let b2 = Beam::from_arcsec(b.0, b.1, b.2).unwrap();
            let common = common_beam(&[b1, b2], 1e-4, 200, 5e-4).unwrap();
            assert!(fits_in_beam(&[b1, b2], &common), "{common}");
        }
    }
}