convolve-rs 1.2.1

import astropy.units as u
import numpy as np
import pytest
import radio_beam as rb
from astropy.io import fits

from convolve_rs import Beam, common_beam, smooth

ARCSEC = 1.0 / 3600.0  # degrees


def _rb(major_as, minor_as, pa_deg):
    return rb.Beam(major_as * u.arcsec, minor_as * u.arcsec, pa_deg * u.deg)


# ── Beam construction ─────────────────────────────────────────────────────────


class TestBeamConstructor:
    def test_basic(self):
        b = Beam(10 * ARCSEC, 8 * ARCSEC, 30.0)
        assert b.major_deg == pytest.approx(10 * ARCSEC)
        assert b.minor_deg == pytest.approx(8 * ARCSEC)
        assert b.pa_deg == pytest.approx(30.0)

    def test_arcsec_properties(self):
        b = Beam.from_arcsec(10.0, 8.0, 30.0)
        assert b.major_arcsec == pytest.approx(10.0)
        assert b.minor_arcsec == pytest.approx(8.0)

    def test_invalid_axes_raises(self):
        with pytest.raises(ValueError):
            Beam(5 * ARCSEC, 10 * ARCSEC, 0.0)  # minor > major

    def test_nonfinite_raises(self):
        with pytest.raises(ValueError):
            Beam(float("nan"), 8 * ARCSEC, 0.0)


class TestBeamClassmethods:
    def test_from_fits_header(self):
        hdr = fits.Header()
        hdr["BMAJ"] = 10 * ARCSEC
        hdr["BMIN"] = 8 * ARCSEC
        hdr["BPA"] = 30.0
        b = Beam.from_fits_header(hdr)
        assert b.major_deg == pytest.approx(10 * ARCSEC)
        assert b.minor_deg == pytest.approx(8 * ARCSEC)
        assert b.pa_deg == pytest.approx(30.0)

    def test_from_fits_header_matches_radio_beam(self):
        hdr = fits.Header()
        hdr["BMAJ"] = 10 * ARCSEC
        hdr["BMIN"] = 8 * ARCSEC
        hdr["BPA"] = 30.0
        ours = Beam.from_fits_header(hdr)
        theirs = rb.Beam.from_fits_header(hdr)
        assert ours.major_arcsec == pytest.approx(theirs.major.to(u.arcsec).value)
        assert ours.minor_arcsec == pytest.approx(theirs.minor.to(u.arcsec).value)
        assert ours.pa_deg == pytest.approx(theirs.pa.to(u.deg).value)

    def test_from_fits_header_missing_bpa_defaults_zero(self):
        hdr = fits.Header()
        hdr["BMAJ"] = 10 * ARCSEC
        hdr["BMIN"] = 8 * ARCSEC
        assert Beam.from_fits_header(hdr).pa_deg == 0.0

    def test_from_radio_beam(self):
        rb_beam = _rb(10.0, 8.0, 30.0)
        b = Beam.from_radio_beam(rb_beam)
        assert b.major_arcsec == pytest.approx(rb_beam.major.to(u.arcsec).value)
        assert b.minor_arcsec == pytest.approx(rb_beam.minor.to(u.arcsec).value)
        assert b.pa_deg == pytest.approx(rb_beam.pa.to(u.deg).value)

    def test_from_arcsec_matches_from_radio_beam(self):
        ours = Beam.from_arcsec(10.0, 8.0, 30.0)
        rb_beam = _rb(10.0, 8.0, 30.0)
        via_rb = Beam.from_radio_beam(rb_beam)
        assert ours.major_deg == pytest.approx(via_rb.major_deg)
        assert ours.minor_deg == pytest.approx(via_rb.minor_deg)
        assert ours.pa_deg == pytest.approx(via_rb.pa_deg)


# ── Beam math vs radio_beam ───────────────────────────────────────────────────


class TestBeamMath:
    def test_convolve_matches_radio_beam(self):
        b1 = Beam.from_arcsec(10.0, 8.0, 30.0)
        b2 = Beam.from_arcsec(12.0, 6.0, 60.0)
        ours = b1.convolve(b2)
        theirs = _rb(10.0, 8.0, 30.0).convolve(_rb(12.0, 6.0, 60.0))
        assert ours.major_arcsec == pytest.approx(
            theirs.major.to(u.arcsec).value, rel=1e-6
        )
        assert ours.minor_arcsec == pytest.approx(
            theirs.minor.to(u.arcsec).value, rel=1e-6
        )
        assert ours.pa_deg == pytest.approx(theirs.pa.to(u.deg).value, abs=1e-6)

    def test_deconvolve_matches_radio_beam(self):
        b1 = Beam.from_arcsec(15.0, 10.0, 45.0)
        b2 = Beam.from_arcsec(10.0, 8.0, 30.0)
        ours = b1.deconvolve(b2)
        theirs = _rb(15.0, 10.0, 45.0).deconvolve(_rb(10.0, 8.0, 30.0))
        assert ours.major_arcsec == pytest.approx(
            theirs.major.to(u.arcsec).value, rel=1e-5
        )
        assert ours.minor_arcsec == pytest.approx(
            theirs.minor.to(u.arcsec).value, rel=1e-5
        )

    def test_deconvolve_fails_when_psf_larger(self):
        small = Beam.from_arcsec(5.0, 5.0, 0.0)
        large = Beam.from_arcsec(10.0, 10.0, 0.0)
        with pytest.raises(ValueError):
            small.deconvolve(large)

    def test_convolve_deconvolve_roundtrip(self):
        a = Beam.from_arcsec(10.0, 8.0, 30.0)
        b = Beam.from_arcsec(6.0, 5.0, 15.0)
        recovered = a.convolve(b).deconvolve(a)
        assert recovered.major_arcsec == pytest.approx(b.major_arcsec, rel=1e-6)
        assert recovered.minor_arcsec == pytest.approx(b.minor_arcsec, rel=1e-6)

    def test_area_sr(self):
        b = Beam.from_arcsec(10.0, 10.0, 0.0)
        fwhm_rad = (10.0 * ARCSEC) * np.pi / 180.0
        expected = np.pi / (4 * np.log(2)) * fwhm_rad**2
        assert b.area_sr() == pytest.approx(expected, rel=1e-10)


# ── common_beam vs radio_beam ─────────────────────────────────────────────────


class TestCommonBeam:
    def test_two_beams_matches_radio_beam(self):
        params = [(10.0, 8.0, 30.0), (12.0, 6.0, 60.0)]
        beams = [Beam.from_arcsec(*p) for p in params]
        theirs = _rb(*params[0]).commonbeam_with(_rb(*params[1]))
        ours = common_beam(beams)
        assert ours.major_arcsec == pytest.approx(
            theirs.major.to(u.arcsec).value, rel=1e-4
        )
        assert ours.minor_arcsec == pytest.approx(
            theirs.minor.to(u.arcsec).value, rel=1e-4
        )

    def test_many_beams_matches_radio_beam(self):
        params = [
            (10.0, 8.0, 30.0),
            (12.0, 6.0, 60.0),
            (11.0, 9.0, 45.0),
            (9.0, 7.0, 15.0),
        ]
        beams = [Beam.from_arcsec(*p) for p in params]
        rb_beams = rb.Beams(
            major=[p[0] for p in params] * u.arcsec,
            minor=[p[1] for p in params] * u.arcsec,
            pa=[p[2] for p in params] * u.deg,
        )
        ours = common_beam(beams)
        theirs = rb_beams.common_beam()
        assert ours.major_arcsec == pytest.approx(
            theirs.major.to(u.arcsec).value, rel=1e-3
        )
        assert ours.minor_arcsec == pytest.approx(
            theirs.minor.to(u.arcsec).value, rel=1e-3
        )

    def test_identical_beams(self):
        b = Beam.from_arcsec(10.0, 8.0, 30.0)
        result = common_beam([b, b, b])
        assert result.major_arcsec == pytest.approx(b.major_arcsec, rel=1e-6)
        assert result.minor_arcsec == pytest.approx(b.minor_arcsec, rel=1e-6)

    def test_single_beam_returns_itself(self):
        b = Beam.from_arcsec(10.0, 8.0, 30.0)
        result = common_beam([b])
        assert result.major_arcsec == pytest.approx(b.major_arcsec)
        assert result.minor_arcsec == pytest.approx(b.minor_arcsec)

    def test_common_beam_contains_all_inputs(self):
        params = [(10.0, 8.0, 30.0), (12.0, 6.0, 60.0), (11.0, 9.0, 45.0)]
        beams = [Beam.from_arcsec(*p) for p in params]
        result = common_beam(beams)
        for b in beams:
            result.deconvolve(b)  # must not raise

    def test_circular_plus_symmetric_elongated(self):
        # Circular beam + two elongated beams at equal and opposite PAs.
        # Both implementations find a valid common beam but the iterative MVE
        # algorithm can converge to slightly different solutions, so we check
        # correctness properties rather than exact numerical agreement.
        beams = [Beam(10, 10, 0), Beam(20, 10, 30), Beam(20, 10, -30)]
        result = common_beam(beams)
        # Must be larger than the largest input beam
        assert result.major_deg >= max(b.major_deg for b in beams)
        # Must contain every input beam
        for b in beams:
            result.deconvolve(b)  # must not raise

    def test_empty_raises(self):
        with pytest.raises(ValueError):
            common_beam([])


# ── smooth ────────────────────────────────────────────────────────────────────


class TestSmooth:
    """Flux-scaling behaviour of smooth() for different brightness units."""

    OLD = Beam(10 * ARCSEC, 10 * ARCSEC, 0.0)
    NEW = Beam(20 * ARCSEC, 20 * ARCSEC, 0.0)
    PIX = 2.5 * ARCSEC

    def _flat(self):
        return np.ones((32, 32), dtype=np.float32)

    def test_default_is_jy_per_beam(self):
        # Convolving a flat Jy/beam image scales it by Ω_new/Ω_old = 4.
        out = smooth(self._flat(), self.OLD, self.NEW, self.PIX, self.PIX)
        assert out[16, 16] == pytest.approx(4.0, abs=1e-3)

    def test_explicit_jy_per_beam(self):
        out = smooth(
            self._flat(), self.OLD, self.NEW, self.PIX, self.PIX, bunit="Jy/beam"
        )
        assert out[16, 16] == pytest.approx(4.0, abs=1e-3)

    def test_kelvin_conserves_surface_brightness(self):
        # Brightness temperature: a flat image stays flat (no flux scaling).
        out = smooth(self._flat(), self.OLD, self.NEW, self.PIX, self.PIX, bunit="K")
        assert out[16, 16] == pytest.approx(1.0, abs=1e-3)

    def test_unrecognised_bunit_warns_and_assumes_jy_per_beam(self):
        with pytest.warns(UserWarning, match="Could not determine brightness unit"):
            out = smooth(
                self._flat(), self.OLD, self.NEW, self.PIX, self.PIX, bunit="furlongs"
            )
        assert out[16, 16] == pytest.approx(4.0, abs=1e-3)

    def test_recognised_bunit_does_not_warn(self):
        import warnings  # noqa: PLC0415

        with warnings.catch_warnings():
            warnings.simplefilter("error")
            smooth(self._flat(), self.OLD, self.NEW, self.PIX, self.PIX, bunit="K")
            smooth(
                self._flat(), self.OLD, self.NEW, self.PIX, self.PIX, bunit="Jy/beam"
            )


class TestSmoothDtype:
    """smooth() runs at the input precision and returns the matching dtype."""

    OLD = Beam(10 * ARCSEC, 10 * ARCSEC, 0.0)
    NEW = Beam(20 * ARCSEC, 20 * ARCSEC, 0.0)
    PIX = 2.5 * ARCSEC

    def test_float32_returns_float32(self):
        img = np.ones((32, 32), dtype=np.float32)
        out = smooth(img, self.OLD, self.NEW, self.PIX, self.PIX)
        assert out.dtype == np.float32
        assert out[16, 16] == pytest.approx(4.0, abs=1e-3)

    def test_float64_returns_float64(self):
        img = np.ones((32, 32), dtype=np.float64)
        out = smooth(img, self.OLD, self.NEW, self.PIX, self.PIX)
        assert out.dtype == np.float64
        assert out[16, 16] == pytest.approx(4.0, abs=1e-3)

    def test_float32_and_float64_agree(self):
        rng = np.random.default_rng(0)
        img64 = rng.normal(size=(48, 48))
        img32 = img64.astype(np.float32)
        out64 = smooth(img64, self.OLD, self.NEW, self.PIX, self.PIX)
        out32 = smooth(img32, self.OLD, self.NEW, self.PIX, self.PIX)
        np.testing.assert_allclose(out32, out64, atol=1e-3)

    def test_non_float_array_raises(self):
        img = np.ones((32, 32), dtype=np.int32)
        with pytest.raises(ValueError):
            smooth(img, self.OLD, self.NEW, self.PIX, self.PIX)


class TestSmoothInputs:
    """smooth() takes arrays as they come from FITS and from other threads."""

    OLD = Beam.from_arcsec(12.0, 10.0, 20.0)
    NEW = Beam.from_arcsec(18.0, 15.0, 35.0)
    PIX = 2.5 * ARCSEC

    def _image(self, shape=(65, 48), dtype=np.float32, nans=True):
        rng = np.random.default_rng(1234)
        img = rng.normal(size=shape).astype(dtype)
        img[shape[0] // 2 :, : shape[1] // 3] += 5.0
        if nans:
            img[:20, :25] = np.nan
        return img

    @pytest.mark.parametrize("dtype", [np.float32, np.float64])
    def test_big_endian_fits_data(self, tmp_path, dtype):
        # astropy hands back FITS data big-endian (">f4"/">f8"); it must be
        # accepted and give the same answer as the native-order array.
        path = tmp_path / "plane.fits"
        fits.PrimaryHDU(data=self._image(dtype=dtype)).writeto(path)
        data = fits.getdata(path)
        assert not data.dtype.isnative

        out = smooth(data, self.OLD, self.NEW, self.PIX, self.PIX)
        native = smooth(data.astype(dtype), self.OLD, self.NEW, self.PIX, self.PIX)
        assert out.dtype == np.dtype(dtype)
        np.testing.assert_array_equal(out, native)

    def test_non_contiguous_input(self):
        base = self._image(shape=(96, 130))
        for view in (base.T, base[::2, 1::3], np.asfortranarray(base)):
            out = smooth(view, self.OLD, self.NEW, self.PIX, self.PIX)
            expected = smooth(
                np.ascontiguousarray(view), self.OLD, self.NEW, self.PIX, self.PIX
            )
            np.testing.assert_array_equal(out, expected)

    def test_input_is_left_alone(self):
        for nans in (False, True):
            img = self._image(nans=nans)
            before = img.copy()
            smooth(img, self.OLD, self.NEW, self.PIX, self.PIX)
            np.testing.assert_array_equal(img, before)

    def test_concurrent_threads_match_sequential(self):
        # The GIL is released during the convolution, so a thread-pool executor
        # (racs_tools' default) runs planes in parallel; results must not mix.
        from concurrent.futures import ThreadPoolExecutor  # noqa: PLC0415

        images = [self._image(shape=(200, 180 + k)) * (k + 1) for k in range(6)]
        expected = [smooth(im, self.OLD, self.NEW, self.PIX, self.PIX) for im in images]
        with ThreadPoolExecutor(max_workers=4) as pool:
            got = list(
                pool.map(
                    lambda im: smooth(im, self.OLD, self.NEW, self.PIX, self.PIX),
                    images,
                )
            )
        for g, e in zip(got, expected, strict=True):
            np.testing.assert_array_equal(g, e)

    @pytest.mark.parametrize("shape", [(128, 96), (129, 97)])
    def test_matches_racs_tools_robust(self, shape):
        # Parity with the Python implementation this ports, when it is installed.
        convolve_uv = pytest.importorskip("racs_tools.convolve_uv")
        img = self._image(shape=shape)
        old = rb.Beam(12 * u.arcsec, 10 * u.arcsec, 20 * u.deg)
        new = rb.Beam(18 * u.arcsec, 15 * u.arcsec, 35 * u.deg)
        theirs = convolve_uv.smooth(
            img, old, new, self.PIX * u.deg, self.PIX * u.deg, conv_mode="robust"
        )
        ours = smooth(
            img,
            Beam.from_radio_beam(old),
            Beam.from_radio_beam(new),
            self.PIX,
            self.PIX,
        )
        # convolve-rs blanks where the convolved NaN mask reaches 0.99 of the
        # filter gain rather than 1, so f32 round-off cannot un-blank a masked
        # region; that adds a thin rim of blanked pixels but never removes one.
        assert np.isnan(ours)[np.isnan(theirs)].all()
        assert (
            np.isnan(ours).sum() - np.isnan(theirs).sum()
            < 0.02 * np.isnan(theirs).sum()
        )
        finite = ~np.isnan(ours)
        np.testing.assert_allclose(
            ours[finite], theirs[finite], rtol=0, atol=1e-4 * theirs[finite].std()
        )