Merge pull request #45 from HERA-Team/reuse-spline

perf: reuse splines on UVBeam objects

CHANGELOG.rst

@@ -2,10 +2,16 @@
 Changelog
 =========
 
-dev
-===
+Version 0.4.3
+=============
+
+Changed
+-------
+
+- Call ``UVBeam.interp`` with ``reuse_spline=True`` and ``check_azza_domain=False`` for
+  significantly faster performance when using ``beam_list``.
 
-Version 0.5.0
+Version 0.4.2
 =============
 
 Fixed

setup.cfg

@@ -32,7 +32,7 @@ install_requires =
     scipy
     astropy
     typing_extensions ; python_version<'3.8'
-    pyuvdata
+    pyuvdata>=2.2.8
 
 # The usage of test_requires is discouraged, see `Dependency Management` docs
 # tests_require = pytest; pytest-cov

src/vis_cpu/conversions.py

@@ -362,7 +362,7 @@ def uvbeam_to_lm(
         `polarized=True`.
     """
     # Define angle cosines
-    L = np.linspace(-1, 1, n_pix_lm)
+    L = bm_pix_to_lm(n_pix_lm)
     L, m = np.meshgrid(L, L)
     L = L.flatten()
     m = m.flatten()
@@ -392,3 +392,37 @@ def uvbeam_to_lm(
         return bm.reshape((Naxes, Nfeeds, len(freqs), n_pix_lm, n_pix_lm))
     else:
         return bm.reshape((len(freqs), n_pix_lm, n_pix_lm))
+
+
+def bm_pix_to_lm(bm_pix: int) -> np.ndarray:
+    """Return a sampling in projected coordinates.
+
+    This gives an array sampled from -1 to 1, with ``bm_pix`` points in it.
+    ``bm_pix`` must be an odd number of points since we ensure that zenith is
+    always sampled (i.e. l=0). We furthermore always put two points very close to
+    zenith (at 1e-18) to ensure that we can "interpolate away" from zenith properly.
+
+    Since we interpolate in real/imag, the phase flip at zenith will always mean that
+    there is a point with magnitude of zero between zenith and one of the points on the
+    circle closest to zenith. By making this circle very small (1e-18) we reduce the
+    chance that this artifact has any effect.
+    """
+    if not bm_pix % 2:
+        raise ValueError("Must use odd bm_pix so that you sample zenith!")
+
+    if bm_pix < 5:
+        raise ValueError("can't use bm_pix less than 5!")
+
+    lm = np.linspace(-1, 1, bm_pix - 2)
+
+    # The smallest value of lm should always be zero, however numerically it doesn't
+    # always quite work out, so we force it to be exactly zero.
+    lm[np.argmin(np.abs(lm))] = 0.0
+    lm = np.insert(
+        lm,
+        bm_pix // 2 - 1,
+        -1e-18,
+    )
+    lm = np.insert(lm, bm_pix // 2 + 1, 1e-18)
+
+    return lm

src/vis_cpu/vis_cpu.py

@@ -47,8 +47,7 @@ def construct_pixel_beam_spline(bm_cube):
     nax, nfeed, nbeam, bm_pix, _ = bm_cube.shape
 
     # x and y coordinates of beam
-    bm_pix_x = np.linspace(-1, 1, bm_pix)
-    bm_pix_y = np.linspace(-1, 1, bm_pix)
+    lm = conversions.bm_pix_to_lm(bm_pix)
 
     # Construct splines for each polarization (pol. vector axis + feed) and
     # antenna. The `splines` list has shape (Naxes, Nfeeds, Nants).
@@ -61,9 +60,7 @@ def construct_pixel_beam_spline(bm_cube):
             # 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
-                )
+                spl = RectBivariateSpline(lm, lm, bm_cube[p1, p2, i], kx=1, ky=1)
                 spl_feeds.append(spl)
             spl_axes.append(spl_feeds)
         splines.append(spl_axes)
@@ -117,7 +114,8 @@ def vis_cpu(
         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.
+        provided also. Note that the projected coordinates corresponding to the bm_cube
+        MUST be equivalent to those returned by :func:`~conversions.bm_pix_to_lm`.
     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
@@ -282,9 +280,15 @@ def vis_cpu(
 
             # 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(nbeam):
-                interp_beam = beam_list[i].interp(
-                    az_array=az, za_array=za, freq_array=np.atleast_1d(freq)
+            for i, bm in enumerate(beam_list):
+                kw = (
+                    {"reuse_spline": True, "check_azza_domain": False}
+                    if isinstance(bm, UVBeam)
+                    else {}
+                )
+
+                interp_beam = bm.interp(
+                    az_array=az, za_array=za, freq_array=np.atleast_1d(freq), **kw
                 )[0]
 
                 if polarized:
@@ -324,7 +328,4 @@ def vis_cpu(
             )
 
     # Return visibilities with or without multiple polarization channels
-    if polarized:
-        return vis
-    else:
-        return vis[0, 0]
+    return vis if polarized else vis[0, 0]

tests/test_beams.py

@@ -185,7 +185,7 @@ def beam_list_pol() -> List[EllipticalBeam]:
 def beam_cube(beam_list_unpol, freq) -> np.ndarray:
     """Construct pixel beam from analytic beam."""
     beam_pix = conversions.uvbeam_to_lm(
-        beam_list_unpol[0], freq, n_pix_lm=1000, polarized=False
+        beam_list_unpol[0], freq, n_pix_lm=1001, polarized=False
     )
     return np.array([beam_pix, beam_pix])
 
@@ -279,14 +279,14 @@ def test_beam_interpolation_pol():
 
     # Construct pixel beam from analytic beam
     beam_pix = conversions.uvbeam_to_lm(
-        beam_analytic, freq, n_pix_lm=20, polarized=True
+        beam_analytic, freq, n_pix_lm=21, polarized=True
     )
     assert np.all(~np.isnan(beam_pix))
     assert np.all(~np.isinf(beam_pix))
 
     # Check that unpolarized beam pixelization has 2 fewer dimensions
     beam_pix_unpol = conversions.uvbeam_to_lm(
-        beam_analytic, freq, n_pix_lm=200, polarized=False
+        beam_analytic, freq, n_pix_lm=201, polarized=False
     )
     assert len(beam_pix.shape) == len(beam_pix_unpol.shape) + 2
 
@@ -379,7 +379,7 @@ def test_prepare_beam_unpol_uvbeam_pol_no_exist():
 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
+        beam_list_unpol[0], freq, n_pix_lm=1001, polarized=False
     )
 
     antpos = np.array([[0, 0, 0], [1, 1, 0], [-1, 1, 0]])
@@ -448,7 +448,7 @@ def test_unique_beam_passed(beam_list_unpol, freq, sky_flux, crd_eq, eq2tops):
 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
+        beam_list_unpol[0], freq, n_pix_lm=1001, polarized=False
     )
 
     antpos = np.array([[0, 0, 0], [1, 1, 0], [-1, 1, 0]])

tests/test_conversions.py

@@ -182,3 +182,21 @@ def test_equatorial_to_eci_coords():
         _ra, _dec = conversions.equatorial_to_eci_coords(
             ra, dec, obstime, (-30.7, 21.4, 1073.0), unit="rad", frame="icrs"
         )
+
+
+def test_bm_pix():
+    """Test the bm_pix_to_lm function."""
+    with pytest.raises(ValueError, match="can't use bm_pix less than 5!"):
+        conversions.bm_pix_to_lm(3)
+
+    with pytest.raises(
+        ValueError, match="Must use odd bm_pix so that you sample zenith!"
+    ):
+        conversions.bm_pix_to_lm(10)
+
+    for n in [101, 201, 153, 17]:
+        lm = conversions.bm_pix_to_lm(n)
+        assert 0.0 in lm.tolist()
+        assert np.all(np.diff(lm)) > 0.0
+        assert 1e-18 in lm.tolist()
+        assert -1e-18 in lm.tolist()

tests/test_vis_cpu.py

@@ -180,7 +180,7 @@ def test_construct_pixel_beam_spline():
     beams = [beam, beam, beam]  # 3 ants
 
     beam_pix = [
-        conversions.uvbeam_to_lm(bm, freqs, n_pix_lm=20, polarized=False)
+        conversions.uvbeam_to_lm(bm, freqs, n_pix_lm=21, polarized=False)
         for bm in beams
     ]
     beam_cube = np.array(beam_pix)