Merge pull request #99 from HERA-Team/fix-analytic-beam

maint: update API to conform to new pyuvdata conventions.

Author: steven-murray

docs/tutorials/matvis_tutorial.ipynb

@@ -48,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -59,7 +59,7 @@
     "from astropy.time import Time\n",
     "\n",
     "from matvis import simulate_vis\n",
-    "from pyuvsim.analyticbeam import AnalyticBeam"
+    "from pyuvdata.analytic_beam import GaussianBeam, UniformBeam"
    ]
   },
   {
@@ -107,17 +107,18 @@
    "source": [
     "Now, let's define the beams for each of these antennas. `matvis` allows us to specify\n",
     "different beams for each antenna, but we need only specify *unique* beams. Each beam\n",
-    "must be either a `UVBeam` or `AnalyticBeam`. Here, for simplicity we just use the \n",
-    "`AnalyticBeam` class. We do this as follows:"
+    "must be either a `UVBeam` or `AnalyticBeam`. Here, for simplicity we just use some\n",
+    "subclasses of the `AnalyticBeam` class. You can create your own custom `AnalyticBeam`\n",
+    "classes to use as well (see the [pyuvdata tutorial](https://pyuvdata.readthedocs.io/en/latest/analytic_beam_tutorial.html#defining-new-analytic-beams))."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "beams = [AnalyticBeam(\"gaussian\", sigma=0.5), AnalyticBeam(\"uniform\")]\n",
+    "beams = [GaussianBeam(sigma=0.5), UniformBeam()]\n",
     "beam_idx = [0, 0, 1]"
    ]
   },
@@ -135,9 +136,7 @@
    "metadata": {},
    "source": [
     "We also need to tell `matvis` the observational configuration, such as the frequency\n",
-    "channels and times to use. In this example, we don't worry too much about the exact date\n",
-    "of observation, but rather the LST. In general, the exact date of observation does make\n",
-    "a little difference."
+    "channels and times to use:"
    ]
   },
   {
@@ -422,7 +421,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -436,12 +435,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
-  },
-  "vscode": {
-   "interpreter": {
-    "hash": "ad4f6c198aebbf84691339a9a95c2c9f143d88558d502a312c8b0f70f8363bfd"
-   }
+   "version": "3.12.6"
   }
  },
  "nbformat": 4,

setup.cfg

@@ -25,10 +25,10 @@ install_requires =
     line-profiler
     numpy>=2.0
     psutil
-    pyuvdata>=3.1.0
+    pyuvdata>=3.1.2
     rich
     scipy
-python_requires = >=3.9
+python_requires = >=3.10
 include_package_data = True
 package_dir =
     =src

tests/__init__.py src/matvis/_test_utils.pytests/__init__.py renamed to src/matvis/_test_utils.py

@@ -5,14 +5,17 @@
 from astropy.coordinates import EarthLocation, Latitude, Longitude
 from astropy.time import Time
 from astropy.units import Quantity
+from dataclasses import replace
 from pathlib import Path
 from pyradiosky import SkyModel
 from pyuvdata import UVBeam
+from pyuvdata.analytic_beam import GaussianBeam
+from pyuvdata.beam_interface import BeamInterface
 from pyuvdata.telescopes import get_telescope
-from pyuvsim import AnalyticBeam, simsetup
+from pyuvsim import simsetup
 from pyuvsim.telescope import BeamList
 
-from matvis import DATA_PATH, coordinates
+from matvis import DATA_PATH
 
 nfreq = 1
 ntime = 5  # 20
@@ -39,34 +42,36 @@ def get_standard_sim_params(
     # Beam model
     if use_analytic_beam:
         n_freq = nfreq
-        beam = AnalyticBeam("gaussian", diameter=14.0)
+        beam = BeamInterface(
+            GaussianBeam(diameter=14.0), beam_type="efield" if polarized else "power"
+        )
     else:
         n_freq = min(nfreq, 2)
         # This is a peak-normalized e-field beam file at 100 and 101 MHz,
         # downsampled to roughly 4 square-degree resolution.
         beam = UVBeam()
         beam.read_beamfits(beam_file)
-        if not polarized:
-            uvsim_beam = beam.copy()
-            beam.efield_to_power(calc_cross_pols=False, inplace=True)
-            beam.select(polarizations=["xx"], inplace=True)
+        # Even though we sometimes want a power beam for matvis, we always need
+        # an efield beam for pyuvsim, so we create an efield beam here, and let matvis
+        # take care of conversion later.
+        beam = BeamInterface(beam, beam_type="efield")
 
         # Now, the beam we have on file doesn't actually properly wrap around in azimuth.
         # This is fine -- the uvbeam.interp() handles the interpolation well. However, when
         # comparing to the GPU interpolation, which first has to interpolate to a regular
         # grid that ends right at 2pi, it's better to compare like for like, so we
         # interpolate to such a grid here.
-        beam = beam.interp(
-            az_array=np.linspace(0, 2 * np.pi, 181),
-            za_array=np.linspace(0, np.pi / 2, 46),
-            az_za_grid=True,
-            new_object=True,
+        beam = beam.clone(
+            beam=beam.beam.interp(
+                az_array=np.linspace(0, 2 * np.pi, 181),
+                za_array=np.linspace(0, np.pi / 2, 46),
+                az_za_grid=True,
+                new_object=True,
+            )
         )
 
-    if polarized or use_analytic_beam:
-        uvsim_beams = BeamList([beam])
-    else:
-        uvsim_beams = BeamList([uvsim_beam])
+    # The UVSim beams must be efield all the time.
+    uvsim_beams = BeamList([beam])
 
     beam_dict = {str(i): 0 for i in range(nants)}
 

src/matvis/core/beams.py

@@ -3,43 +3,41 @@
 import numpy as np
 from abc import ABC, abstractmethod
 from copy import deepcopy
+from dataclasses import replace
 from pyuvdata import UVBeam
+from pyuvdata.analytic_beam import AnalyticBeam
+from pyuvdata.beam_interface import BeamInterface
 from pyuvdata.utils.pol import polstr2num
 from typing import Any, Literal
 
 
-def prepare_beam_unpolarized(uvbeam, use_pol: Literal["xy", "xx", "yx", "yy"] = "xx"):
-    """Given a UVBeam or AnalyticBeam, prepare it for an un-polarized simulation."""
-    if uvbeam.beam_type == "power" and getattr(uvbeam, "Npols", 1) == 1:
-        return uvbeam
+def prepare_beam_unpolarized(
+    beam: BeamInterface,
+    use_feed: Literal["x", "y"] = "x",
+    allow_beam_mutation: bool = False,
+) -> BeamInterface:
+    """Given a BeamInterface, prepare it for an un-polarized simulation."""
+    if beam.beam_type == "power" and beam.Npols == 1:
+        return beam
 
-    uvbeam_ = uvbeam.copy() if isinstance(uvbeam, UVBeam) else deepcopy(uvbeam)
-
-    if uvbeam_.beam_type == "efield":
-        if isinstance(uvbeam, UVBeam):
-            uvbeam_.efield_to_power(calc_cross_pols=False)
-        else:
-            uvbeam_.efield_to_power()
-
-    if getattr(uvbeam_, "Npols", 1) > 1:
-        pol = polstr2num(use_pol)
+    if beam.beam_type == "efield":
+        beam = beam.as_power_beam(
+            include_cross_pols=False, allow_beam_mutation=allow_beam_mutation
+        )
 
-        if pol not in uvbeam_.polarization_array:
-            raise ValueError(
-                f"You want to use {use_pol} pol, but it does not exist in the UVBeam"
-            )
-        uvbeam_.select(polarizations=[pol])
+    if beam.Npols > 1:
+        beam = beam.with_feeds([use_feed])
 
-    return uvbeam_
+    return beam
 
 
 def _wrangle_beams(
     beam_idx: np.ndarray | None,
-    beam_list: list[UVBeam],
+    beam_list: list[BeamInterface | UVBeam | AnalyticBeam],
     polarized: bool,
     nant: int,
     freq: float,
-) -> tuple[list[UVBeam], int, np.ndarray]:
+) -> tuple[list[BeamInterface], int, np.ndarray]:
     """Perform all the operations and checks on the input beams.
 
     Checks that the beam indices match the number of antennas, pre-interpolates to the
@@ -61,6 +59,10 @@ def _wrangle_beams(
     """
     # Get the number of unique beams
     nbeam = len(beam_list)
+    beam_list = [BeamInterface(beam) for beam in beam_list]
+
+    if not polarized:
+        beam_list = [prepare_beam_unpolarized(beam) for beam in beam_list]
 
     # Check the beam indices
     if beam_idx is None and nbeam not in (1, nant):
@@ -78,8 +80,12 @@ def _wrangle_beams(
     # make sure we interpolate to the right frequency first.
     beam_list = [
         (
-            bm.interp(freq_array=np.array([freq]), new_object=True, run_check=False)
-            if isinstance(bm, UVBeam)
+            bm.clone(
+                beam=bm.beam.interp(
+                    freq_array=np.array([freq]), new_object=True, run_check=False
+                )
+            )
+            if bm._isuvbeam
             else bm
         )
         for bm in beam_list
@@ -89,19 +95,6 @@ def _wrangle_beams(
         if any(b.beam_type != "efield" for b in beam_list):
             raise ValueError("beam type must be efield if using polarized=True")
 
-    # The following applies if we're not polarized
-    elif 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"
-        )
-
     return beam_list, nbeam, beam_idx
 
 
@@ -132,7 +125,7 @@ class BeamInterpolator(ABC):
 
     def __init__(
         self,
-        beam_list: list[UVBeam],
+        beam_list: list[BeamInterface],
         beam_idx: np.ndarray,
         polarized: bool,
         nant: int,
@@ -142,6 +135,7 @@ def __init__(
         precision: int = 1,
     ):
         self.polarized = polarized
+
         self.beam_list, self.nbeam, self.beam_idx = _wrangle_beams(
             beam_idx=beam_idx,
             beam_list=beam_list,

src/matvis/cpu/beams.py

@@ -24,24 +24,19 @@ def interp(self, tx: np.ndarray, ty: np.ndarray, out: np.ndarray) -> np.ndarray:
         # Primary beam pattern using direct interpolation of UVBeam object
         az, za = enu_to_az_za(enu_e=tx, enu_n=ty, orientation="uvbeam")
 
+        kw = {
+            "reuse_spline": True,
+            "check_azza_domain": False,
+            "interpolation_function": "az_za_map_coordinates",
+            "spline_opts": self.spline_opts,
+        }
         for i, bm in enumerate(self.beam_list):
-            kw = (
-                {
-                    "reuse_spline": True,
-                    "check_azza_domain": False,
-                    "interpolation_function": "az_za_map_coordinates",
-                    "spline_opts": self.spline_opts,
-                }
-                if isinstance(bm, UVBeam)
-                else {}
-            )
-
-            interp_beam = bm.interp(
+            interp_beam = bm.compute_response(
                 az_array=az,
                 za_array=za,
                 freq_array=np.array([self.freq]),
                 **kw,
-            )[0]
+            )
 
             if self.polarized:
                 interp_beam = interp_beam[:, :, 0, :].transpose((1, 0, 2))

src/matvis/cpu/cpu.py

@@ -11,6 +11,8 @@
 from astropy.time import Time
 from collections.abc import Sequence
 from pyuvdata import UVBeam
+from pyuvdata.analytic_beam import AnalyticBeam
+from pyuvdata.beam_interface import BeamInterface
 from typing import Callable, Literal
 
 from .._utils import get_desired_chunks, get_dtypes, log_progress, logdebug, memtrace
@@ -32,7 +34,7 @@ def simulate(
     skycoords: SkyCoord,
     telescope_loc: EarthLocation,
     I_sky: np.ndarray,
-    beam_list: Sequence[UVBeam | Callable] | None,
+    beam_list: Sequence[UVBeam | AnalyticBeam | BeamInterface] | None,
     antpairs: np.ndarray | list[tuple[int, int]] | None = None,
     precision: int = 1,
     polarized: bool = False,

src/matvis/gpu/beams.py

@@ -11,7 +11,7 @@
 from ..cpu.beams import UVBeamInterpolator
 
 
-def prepare_for_map_coords(uvbeam):
+def prepare_for_map_coords(uvbeam: UVBeam):
     """Obtain coordinates for doing map_coordinates interpolation from a UVBeam."""
     d0, az, za = uvbeam._prepare_coordinate_data(uvbeam.data_array)
     d0 = d0[:, :, 0]  # only one frequency
@@ -30,8 +30,8 @@ def setup(self):
         Decides if the beam_list is a list of UVBeam objects or AnalyticBeam objects,
         and dispatches accordingly.
         """
-        self.use_interp = isinstance(self.beam_list[0], UVBeam)
-        if self.use_interp and not all(isinstance(b, UVBeam) for b in self.beam_list):
+        self.use_interp = self.beam_list[0]._isuvbeam
+        if self.use_interp and not all(b._isuvbeam for b in self.beam_list):
             raise ValueError(
                 "GPUBeamInterpolator only supports beam_lists with either all UVBeam or all AnalyticBeam objects."
             )
@@ -48,15 +48,15 @@ def setup(self):
             # we might want to do cubic interpolation with pyuvbeam onto a much higher-res
             # grid, then use linear interpolation on the GPU with that high-res grid.
             # We can explore this later...
-            if any(bm.pixel_coordinate_system != "az_za" for bm in self.beam_list):
+            if any(bm.beam.pixel_coordinate_system != "az_za" for bm in self.beam_list):
                 raise ValueError('pixel coordinate system must be "az_za"')
 
             self.daz = np.zeros(len(self.beam_list))
             self.dza = np.zeros(len(self.beam_list))
             self.azmin = np.zeros(len(self.beam_list))
 
             d0, self.daz[0], self.dza[0], self.azmin[0] = prepare_for_map_coords(
-                self.beam_list[0]
+                self.beam_list[0].beam
             )
 
             self.beam_data = cp.zeros(
@@ -68,7 +68,7 @@ def setup(self):
             if len(self.beam_list) > 1:
                 for i, b in enumerate(self.beam_list[1:]):
                     d, self.daz[i + 1], self.dza[i + 1], self.azmin[i + 1] = (
-                        prepare_for_map_coords(b)
+                        prepare_for_map_coords(b.beam)
                     )
                     self.beam_data[i + 1].set(d)
         else: