Merge branch 'main' into pre-commit-ci-update-config

CHANGELOG.rst

@@ -2,6 +2,15 @@
 Changelog
 =========
 
+dev
+===
+
+Changed
+-------
+
+- Enhanced performance by allowing unique beams only to be passed (no breaking API
+  change).
+
 Version 0.2.3
 =============
 

src/vis_cpu/__init__.py

@@ -6,7 +6,7 @@
     # Change here if project is renamed and does not equal the package name
     dist_name = __name__
     __version__ = get_distribution(dist_name).version
-except DistributionNotFound:
+except DistributionNotFound:  # pragma: no cover
     __version__ = "unknown"
 finally:
     del get_distribution, DistributionNotFound

src/vis_cpu/vis_cpu.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from astropy.constants import c
+from pyuvdata import UVBeam
 from scipy.interpolate import RectBivariateSpline
 from typing import Optional, Sequence
 
@@ -32,12 +33,8 @@ def construct_pixel_beam_spline(bm_cube):
             "imaginary components separately."
         )
 
-    if len(bm_cube.shape) == 5:
-        # Polarized beam
-        nax, nfeed, nant, bm_pix, _ = bm_cube.shape
-    else:
-        nax = nfeed = 1
-        nant, bm_pix, _ = bm_cube.shape
+    # Polarized beam
+    nax, nfeed, nbeam, bm_pix, _ = bm_cube.shape
 
     # x and y coordinates of beam
     bm_pix_x = np.linspace(-1, 1, bm_pix)
@@ -51,8 +48,8 @@ def construct_pixel_beam_spline(bm_cube):
         for p2 in range(nfeed):
             spl_feeds = []
 
-            # Loop over antennas
-            for i in range(nant):
+            # Loop over beams/antennas
+            for i in range(nbeam):
                 # Linear interpolation of primary beam pattern.
                 spl = RectBivariateSpline(
                     bm_pix_y, bm_pix_x, bm_cube[p1, p2, i], kx=1, ky=1
@@ -70,9 +67,10 @@ def vis_cpu(
     crd_eq: np.ndarray,
     I_sky: np.ndarray,
     bm_cube: Optional[np.ndarray] = None,
-    beam_list: Optional[Sequence[np.ndarray]] = None,
+    beam_list: Optional[Sequence[UVBeam]] = None,
     precision: int = 1,
     polarized: bool = False,
+    beam_idx: Optional[np.ndarray] = None,
 ):
     """
     Calculate visibility from an input intensity map and beam model.
@@ -105,9 +103,10 @@ def vis_cpu(
         Shape=(NSRCS,).
     bm_cube : array_like, optional
         Pixelized beam maps for each antenna. If ``polarized=False``,
-        shape=``(NANT, BM_PIX, BM_PIX)``, otherwise
-        shape=``(NAX, NFEED, NANT, BM_PIX, BM_PIX)``. Only one of ``bm_cube`` and
-        ``beam_list`` should be provided.
+        shape=``(NBEAMS, BM_PIX, BM_PIX)``, otherwise
+        shape=``(NAX, NFEED, NBEAMS, BM_PIX, BM_PIX)``. Only one of ``bm_cube`` and
+        ``beam_list`` should be provided. If NBEAMS != NANT, then `beam_idx` must be
+        provided also.
     beam_list : list of UVBeam, optional
         If specified, evaluate primary beam values directly using UVBeam
         objects instead of using pixelized beam maps. Only one of ``bm_cube`` and
@@ -123,6 +122,10 @@ def vis_cpu(
         Whether to simulate a full polarized response in terms of nn, ne, en,
         ee visibilities. See Eq. 6 of Kohn+ (arXiv:1802.04151) for notation.
         Default: False.
+    beam_idx
+        Optional length-NANT array specifying a beam index for each antenna.
+        By default, either a single beam is assumed to apply to all antennas or
+        each antenna gets its own beam.
 
     Returns
     -------
@@ -160,39 +163,58 @@ def vis_cpu(
         I_sky.ndim == 1 and I_sky.shape[0] == nsrcs
     ), "I_sky must have shape (NSRCS,)."
 
+    # Get the number of unique beams
+    if bm_cube is not None:
+        nbeam = bm_cube.shape[2 if polarized else 0]
+    else:
+        nbeam = len(beam_list)
+
+    # Check the beam indices
+    if beam_idx is None:
+        if nbeam == 1:
+            beam_idx = np.zeros(nant, dtype=int)
+        elif nbeam == nant:
+            beam_idx = np.arange(nant, dtype=int)
+        else:
+            raise ValueError(
+                "If number of beams provided is not 1 or nant, beam_idx must be provided."
+            )
+    else:
+        assert beam_idx.shape == (nant,), "beam_idx must be length nant"
+        assert all(
+            0 <= i < nbeam for i in beam_idx
+        ), "beam_idx contains indices greater than the number of beams"
+
     if beam_list is None:
         bm_pix = bm_cube.shape[-1]
         complex_bm_cube = np.any(np.iscomplex(bm_cube))
         if polarized:
-            assert bm_cube.shape == (nax, nfeed, nant, bm_pix, bm_pix), (
-                "bm_cube must have shape (NAXES, NFEEDS, NANTS, BM_PIX, BM_PIX) if "
+            assert bm_cube.shape == (nax, nfeed, nbeam, bm_pix, bm_pix), (
+                "bm_cube must have shape (NAXES, NFEEDS, NBEAMS, BM_PIX, BM_PIX) if "
                 f"polarized=True. Shape wanted: {(nax, nfeed, nant, bm_pix, bm_pix)}; "
                 f"shape given: {bm_cube.shape}"
             )
-        elif bm_cube.shape != (1, 1, nant, bm_pix, bm_pix):
-            assert bm_cube.shape == (nant, bm_pix, bm_pix), (
-                "bm_cube must have shape (NANTS, BM_PIX, BM_PIX) "
-                "or (1, 1, nant, bm_pix, bm_pix) if polarized=False. "
-                f"Shape wanted: {(nant, bm_pix, bm_pix)}; "
+        elif bm_cube.shape != (1, 1, nbeam, bm_pix, bm_pix):
+            assert bm_cube.shape == (nbeam, bm_pix, bm_pix), (
+                "bm_cube must have shape (NBEAMS, BM_PIX, BM_PIX) "
+                "or (1, 1, nbeam, bm_pix, bm_pix) if polarized=False. "
+                f"Shape wanted: {(nbeam, bm_pix, bm_pix)}; "
                 f"shape given: {bm_cube.shape}"
             )
             bm_cube = bm_cube[np.newaxis, np.newaxis]
-    else:
-        if polarized and any(b.beam_type != "efield" for b in beam_list):
-            raise ValueError("beam type must be efield if using polarized=True")
-        elif not polarized and any(
-            (
-                b.beam_type != "power"
-                or getattr(b, "Npols", 1) > 1
-                or b.polarization_array[0] not in [-5, -6]
-            )
-            for b in beam_list
-        ):
-            raise ValueError(
-                "beam type must be power and have only one pol (either xx or yy) if polarized=False"
-            )
-
-        assert len(beam_list) == nant, "beam_list must have length nant"
+    elif polarized and any(b.beam_type != "efield" for b in beam_list):
+        raise ValueError("beam type must be efield if using polarized=True")
+    elif not polarized and any(
+        (
+            b.beam_type != "power"
+            or getattr(b, "Npols", 1) > 1
+            or b.polarization_array[0] not in [-5, -6]
+        )
+        for b in beam_list
+    ):
+        raise ValueError(
+            "beam type must be power and have only one pol (either xx or yy) if polarized=False"
+        )
 
     # Intensity distribution (sqrt) and antenna positions. Does not support
     # negative sky. Factor of 0.5 accounts for splitting Stokes I between
@@ -215,6 +237,8 @@ def vis_cpu(
         if complex_bm_cube:
             splines_im = construct_pixel_beam_spline(bm_cube.imag)
 
+    im = 0
+
     # Loop over time samples
     for t, eq2top in enumerate(eq2tops.astype(real_dtype)):
         # Dot product converts ECI cosines (i.e. from RA and Dec) into ENU
@@ -225,32 +249,36 @@ def vis_cpu(
         # Primary beam response
         if beam_list is None:
             # Primary beam pattern using pixelized primary beam
-            for i in range(nant):
+            for i in range(nbeam):
                 # Extract requested polarizations
                 for p1 in range(nax):
                     for p2 in range(nfeed):
                         # The beam pixel grid has been reshaped in the order
                         # ty,tx, which implies m,l order
-                        A_s[p1, p2, i] = splines_re[p1][p2][i](ty, tx, grid=False)
+                        re = splines_re[p1][p2][i](ty, tx, grid=False)
+
                         if complex_bm_cube:
-                            A_s[p1, p2, i] += 1.0j * splines_im[p1][p2][i](
-                                ty, tx, grid=False
-                            )
+                            im = 1.0j * splines_im[p1][p2][i](ty, tx, grid=False)
+
+                        A_s[p1, p2, beam_idx == i] = re + im
         else:
 
             # Primary beam pattern using direct interpolation of UVBeam object
             az, za = conversions.enu_to_az_za(enu_e=tx, enu_n=ty, orientation="uvbeam")
-            for i in range(nant):
+            for i in range(nbeam):
                 interp_beam = beam_list[i].interp(
                     az_array=az, za_array=za, freq_array=np.atleast_1d(freq)
                 )[0]
 
                 if polarized:
-                    A_s[:, :, i] = interp_beam[:, 0, :, 0, :]  # spw=0 and freq=0
+                    for ant in range(nant):
+                        A_s[:, :, beam_idx[ant]] = interp_beam[
+                            :, 0, :, 0, :
+                        ]  # spw=0 and freq=0
                 else:
                     # Here we have already asserted that the beam is a power beam and
                     # has only one polarization, so we just evaluate that one.
-                    A_s[:, :, i] = np.sqrt(interp_beam[0, 0, 0, :, :])
+                    A_s[:, :, beam_idx == i] = np.sqrt(interp_beam[0, 0, 0, 0, :])
 
         # Horizon cut
         A_s = np.where(tz > 0, A_s, 0)

tests/test_beams.py

@@ -20,7 +20,7 @@
 cst_file = Path(DATA_PATH) / "NicCSTbeams" / "HERA_NicCST_150MHz.txt"
 
 
-class EllipticalBeam(object):
+class EllipticalBeam:
     """Add ellipticity/shearing to an existing UVBeam/AnalyticBeam object."""
 
     def __init__(self, base_beam, xstretch=1.0, ystretch=1.0, rotation=0.0):
@@ -372,3 +372,96 @@ def test_prepare_beam_unpol_uvbeam_pol_no_exist():
         ValueError, match="You want to use x feed, but it does not exist in the UVBeam"
     ):
         conversions.prepare_beam(beam, polarized=False, use_feed="x")
+
+
+def test_unique_beam_passed(beam_list_unpol, freq, sky_flux, crd_eq, eq2tops):
+    """Test passing different numbers of beams than nant."""
+    beam_pix = conversions.uvbeam_to_lm(
+        beam_list_unpol[0], freq, n_pix_lm=1000, polarized=False
+    )
+
+    antpos = np.array([[0, 0, 0], [1, 1, 0], [-1, 1, 0]])
+
+    bm_cube = np.array([beam_pix, beam_pix, beam_pix])
+
+    for i in range(freq.size):
+        # Pixel beams
+        vis_pix = vis_cpu(
+            antpos,
+            freq[i],
+            eq2tops,
+            crd_eq,
+            sky_flux[:, i],
+            bm_cube=bm_cube[:1, i, :, :],
+            precision=2,
+            polarized=False,
+        )
+
+        vis_pix2 = vis_cpu(
+            antpos,
+            freq[i],
+            eq2tops,
+            crd_eq,
+            sky_flux[:, i],
+            bm_cube=bm_cube[:2, i, :, :],
+            precision=2,
+            polarized=False,
+            beam_idx=np.array([0, 1, 0]),
+        )
+
+        vis_pix3 = vis_cpu(
+            antpos,
+            freq[i],
+            eq2tops,
+            crd_eq,
+            sky_flux[:, i],
+            bm_cube=bm_cube[:, i, :, :],
+            precision=2,
+            polarized=False,
+        )
+
+        # Analytic beams
+        vis_analytic = vis_cpu(
+            antpos,
+            freq[i],
+            eq2tops,
+            crd_eq,
+            sky_flux[:, i],
+            beam_list=beam_list_unpol[:1],
+            precision=2,
+            polarized=False,
+        )
+
+        assert np.all(~np.isnan(vis_pix))  # check that there are no NaN values
+        assert np.all(~np.isnan(vis_pix2))  # check that there are no NaN values
+        assert np.all(~np.isnan(vis_analytic))
+
+        # Check that results are close (they should be for 1000^2 pixel-beams
+        # if the elliptical beams are both oriented the same way)
+        np.testing.assert_allclose(vis_pix, vis_analytic, rtol=1e-5, atol=1e-5)
+        np.testing.assert_allclose(vis_pix, vis_pix2, rtol=1e-5, atol=1e-5)
+        np.testing.assert_allclose(vis_pix3, vis_pix2, rtol=1e-5, atol=1e-5)
+
+
+def test_wrong_numbeams_passed(beam_list_unpol, freq, sky_flux, crd_eq, eq2tops):
+    """Test passing different numbers of beams than nant."""
+    beam_pix = conversions.uvbeam_to_lm(
+        beam_list_unpol[0], freq, n_pix_lm=1000, polarized=False
+    )
+
+    antpos = np.array([[0, 0, 0], [1, 1, 0], [-1, 1, 0]])
+
+    bm_cube = np.array([beam_pix, beam_pix, beam_pix])
+
+    # Pixel beams
+    with pytest.raises(ValueError, match="beam_idx must be provided"):
+        vis_cpu(
+            antpos,
+            freq[0],
+            eq2tops,
+            crd_eq,
+            sky_flux[:, 0],
+            bm_cube=bm_cube[:2, 0, :, :],
+            precision=2,
+            polarized=False,
+        )