Merge pull request #26 from HERA-Team/fix_uvbeam_norm

Remove spurious beam normalisation logic in `uvbeam_to_lm`

.coveragerc

@@ -2,7 +2,7 @@
 [run]
 branch = True
 source = vis_cpu
-# omit = bad_file.py
+omit = */vis_gpu.py
 
 [paths]
 source =

CHANGELOG.rst

@@ -2,6 +2,25 @@
 Changelog
 =========
 
+Version 0.2.3
+=============
+
+Fixed
+-----
+
+- Fix issue with spurious beam normalization when a pixel beam
+  interpolation grid is generated from a UVBeam object
+- Fix bug where the imaginary part of complex pixel beams was
+  being discarded
+- Fix bug that was causing polarized calculations to fail with
+  ``simulate_vis``
+- CI paths fixed so coverage reports are linked properly
+
+Added
+-----
+
+- New units tests
+
 Version 0.2.2
 =============
 

ci/test-env.yml

@@ -13,6 +13,7 @@ dependencies:
   - matplotlib>=3.3.4,<4
   - ipython>=7.22,<8
   - h5py>=3.2,<4
+  - ffmpeg
   - pip:
     - pyuvsim[sim]>=1.2,<1.4
     - pyradiosky>=0.1.1,<0.3

codecov.yml

@@ -0,0 +1,2 @@
+fixes:
+  - "/home/runner/work/vis_cpu/::"

src/vis_cpu/conversions.py

@@ -266,6 +266,8 @@ def uvbeam_to_lm(uvbeam, freqs, n_pix_lm=63, polarized=False, **kwargs):
     drift-scan telescope). For a vector in East-North-Up (ENU) coordinates
     vec{p}, we therefore have ``l = vec{p}.hat{e}`` etc.
 
+    N.B. This function does not perform any beam normalization.
+
     Parameters
     ----------
     uvbeam : UVBeam object
@@ -301,19 +303,10 @@ def uvbeam_to_lm(uvbeam, freqs, n_pix_lm=63, polarized=False, **kwargs):
     else:
         bm = efield_beam[0, 0, 1, :, :]  # (phi, e) == 'xx' component
 
-    # Peak normalization and reshape output
+    # Reshape output
     if polarized:
         Naxes = bm.shape[0]  # polarization vector axes
         Nfeeds = bm.shape[1]  # polarized feeds
-
-        # Separately normalize each polarization channel
-        for i in range(Naxes):
-            for j in range(Nfeeds):
-                if np.max(bm[i, j]) > 0.0:
-                    bm /= np.max(bm[i, j])
         return bm.reshape((Naxes, Nfeeds, len(freqs), n_pix_lm, n_pix_lm))
     else:
-        # Normalize single polarization channel
-        if np.max(bm) > 0.0:
-            bm /= np.max(bm)
         return bm.reshape((len(freqs), n_pix_lm, n_pix_lm))

src/vis_cpu/plot.py

@@ -34,13 +34,12 @@ def _source_az_za_beam(
         Value of the beam (E-field, not power, unless the beam object contains
         only the power beam) for each source.
     """
-    # Equatorial to topocentric conversion at given LST
-    eq2tops = conversions.get_eq2tops(np.atleast_1d(lst), latitude=latitude)
-    eq2top = eq2tops[0]
+    # Get coordinate transforms as a function of LST
+    eq2top = conversions.eci_to_enu_matrix(lst, latitude)
 
-    # Get source az, za
+    # Get source az, za (note the azimuth convention used by UVBeam)
     tx, ty, tz = np.dot(eq2top, crd_eq)
-    az, za = conversions.lm_to_az_za(tx, ty)
+    az, za = conversions.enu_to_az_za(enu_e=tx, enu_n=ty, orientation="uvbeam")
 
     # Get beam values
     interp_beam = beam.interp(az, za, np.atleast_1d(ref_freq))[0]

src/vis_cpu/vis_cpu.py

@@ -163,7 +163,7 @@ def vis_cpu(
 
     if beam_list is None:
         bm_pix = bm_cube.shape[-1]
-        complex_bm_pix = np.iscomplex(bm_pix)
+        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) "
@@ -202,7 +202,7 @@ def vis_cpu(
     # Precompute splines using pixelized beams
     if beam_list is None:
         splines_re = construct_pixel_beam_spline(bm_cube.real)
-        if complex_bm_pix:
+        if complex_bm_cube:
             splines_im = construct_pixel_beam_spline(bm_cube.imag)
 
     # Loop over time samples
@@ -222,7 +222,7 @@ def vis_cpu(
                         # 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)
-                        if complex_bm_pix:
+                        if complex_bm_cube:
                             A_s[p1, p2, i] += 1.0j * splines_im[p1][p2][i](
                                 ty, tx, grid=False
                             )

src/vis_cpu/wrapper.py

@@ -87,8 +87,8 @@ def simulate_vis(
             naxes = beams[0].Naxes_vec
             nfeeds = beams[0].Nfeeds
         except AttributeError:
-            # If Naxes_vec and Nfeeds properties aren't set, assume no pol.
-            naxes = nfeeds = 1
+            # If Naxes_vec and Nfeeds properties aren't set, assume all pol.
+            naxes = nfeeds = 2
 
     # Antenna x,y,z positions
     antpos = np.array([ants[k] for k in ants.keys()])
@@ -122,18 +122,27 @@ def simulate_vis(
     for i in range(freqs.size):
 
         if pixel_beams:
-            vis[i] = vis_cpu(
+
+            # Get per-freq. pixel beam
+            bm = beam_cube[:, :, :, i, :, :] if polarized else beam_cube[:, i, :, :]
+
+            # Run vis_cpu
+            v = vis_cpu(
                 antpos,
                 freqs[i],
                 eq2tops,
                 crd_eq,
                 fluxes[:, i],
-                bm_cube=beam_cube[:, i, :, :],
+                bm_cube=bm,
                 precision=precision,
                 polarized=polarized,
             )
+            if polarized:
+                vis[:, :, i] = v  # v.shape: (nax, nfeed, ntimes, nant, nant)
+            else:
+                vis[i] = v  # v.shape: (ntimes, nant, nant)
         else:
-            vis[i] = vis_cpu(
+            v = vis_cpu(
                 antpos,
                 freqs[i],
                 eq2tops,
@@ -143,5 +152,9 @@ def simulate_vis(
                 precision=precision,
                 polarized=polarized,
             )
+            if polarized:
+                vis[:, :, i] = v  # v.shape: (nax, nfeed, ntimes, nant, nant)
+            else:
+                vis[i] = v  # v.shape: (ntimes, nant, nant)
 
     return vis

tests/test_compare_pyuvsim.py

@@ -50,6 +50,7 @@ def test_compare_pyuvsim():
         phase_type="drift",
         vis_units="Jy",
         complete=True,
+        write_files=False,
     )
     lsts = np.unique(uvdata.lst_array)
 

tests/test_plot.py

@@ -0,0 +1,62 @@
+"""Compare vis_cpu with pyuvsim visibilities."""
+import numpy as np
+from pyuvsim.analyticbeam import AnalyticBeam
+
+from vis_cpu import conversions, plot
+
+nsource = 10
+
+
+def test_source_az_za_beam():
+    """Test function that calculates the Az and ZA positions of sources."""
+    # Observation latitude and LST
+    hera_lat = -30.7215
+    lst = 0.78
+
+    # Add random sources
+    ra = np.random.uniform(low=0.0, high=360.0, size=nsource - 1)
+    dec = -30.72 + np.random.random(nsource - 1) * 10.0
+    ra = np.deg2rad(ra)
+    dec = np.deg2rad(dec)
+
+    # Point source coordinate transform, from equatorial to Cartesian
+    crd_eq = conversions.point_source_crd_eq(ra, dec)
+
+    # Beam model
+    beam = AnalyticBeam(type="gaussian", diameter=14.0)
+
+    # Calculate source locations and positions
+    az, za, beamval = plot._source_az_za_beam(
+        lst, crd_eq, beam, ref_freq=100.0e6, latitude=np.deg2rad(hera_lat)
+    )
+    assert np.all(np.isfinite(az))
+    assert np.all(np.isfinite(za))
+    # (Values of beamval should be NaN below the horizon)
+
+
+def test_animate_source_map():
+    """Test function that animates source positions vs LST."""
+    # Observation latitude and LSTs
+    hera_lat = -30.7215
+    lsts = np.linspace(0.0, 2.0 * np.pi, 5)
+
+    # Add random sources
+    ra = np.random.uniform(low=0.0, high=360.0, size=nsource - 1)
+    dec = -30.72 + np.random.random(nsource - 1) * 10.0
+    ra = np.deg2rad(ra)
+    dec = np.deg2rad(dec)
+
+    # Beam model
+    beam = AnalyticBeam(type="gaussian", diameter=14.0)
+
+    # Generate animation
+    anim = plot.animate_source_map(
+        ra,
+        dec,
+        lsts,
+        beam,
+        interval=200,
+        ref_freq=100.0e6,
+        latitude=np.deg2rad(hera_lat),
+    )
+    assert anim is not None

tests/test_vis_cpu.py

@@ -13,7 +13,7 @@
 NFREQ = 5
 NPTSRC = 20
 
-ants = {0: (0, 0, 0), 1: (1, 1, 0)}
+ants = {0: (0.0, 0.0, 0.0), 1: (20.0, 20.0, 0.0)}
 
 
 def test_vis_cpu():
@@ -102,6 +102,33 @@ def test_vis_cpu():
         )
     assert np.all(~np.isnan(_vis))  # check that there are no NaN values
 
+    # Check that errors are raised when beams are input incorrectly
+    with pytest.raises(RuntimeError):
+        vis_cpu(
+            antpos,
+            freq[0],
+            eq2tops,
+            crd_eq,
+            I_sky[:, i],
+            bm_cube=None,
+            beam_list=None,
+            precision=1,
+            polarized=False,
+        )
+
+    with pytest.raises(RuntimeError):
+        vis_cpu(
+            antpos,
+            freq[0],
+            eq2tops,
+            crd_eq,
+            I_sky[:, i],
+            bm_cube=beam_cube[:, i, :, :],
+            beam_list=[beam, beam],
+            precision=1,
+            polarized=False,
+        )
+
 
 def test_simulate_vis():
     """Test basic operation of simple wrapper around vis_cpu, `simulate_vis`."""
@@ -155,6 +182,38 @@ def test_simulate_vis():
     )
     assert np.all(~np.isnan(vis))  # check that there are no NaN values
 
+    # Run vis_cpu with UVBeam beams (polarized mode)
+    vis = simulate_vis(
+        ants,
+        I_sky,
+        ra,
+        dec,
+        freq,
+        lsts,
+        beams=[beam, beam],
+        pixel_beams=False,
+        polarized=True,
+        precision=1,
+        latitude=-30.7215 * np.pi / 180.0,
+    )
+    assert np.all(~np.isnan(vis))  # check that there are no NaN values
+
+    # Run vis_cpu with pixel beams (polarized mode)
+    vis = simulate_vis(
+        ants,
+        I_sky,
+        ra,
+        dec,
+        freq,
+        lsts,
+        beams=[beam, beam],
+        pixel_beams=True,
+        polarized=True,
+        precision=1,
+        latitude=-30.7215 * np.pi / 180.0,
+    )
+    assert np.all(~np.isnan(vis))  # check that there are no NaN values
+
 
 def test_construct_pixel_beam_spline():
     """Test pixel beam interpolation."""