Add AnalyticBeam objects and a unified interface for UVBeams and AnalyticBeams

.coveragerc

@@ -5,6 +5,7 @@ parallel = true
 omit =
   */tests/*
   */docs/*
+  */data/*
   setup.py
 source =
   pyuvdata
@@ -13,6 +14,7 @@ source =
 omit =
   */tests/*
   */docs/*
+  */data/*
   setup.py
 
 show_missing = true

CHANGELOG.md

@@ -4,8 +4,13 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
-- New optional spatial interpolation method, ``interpolation_function="az_za_map_coordinates"`` that improves the linear
-interpolation speed for data in ``az_za`` coordinates.
+- New analytic beam classes: AiryBeam, GaussianBeam, ShortDipoleBeam, UniformBeam
+that support evaluating either the efield or power beam in any direction and frequency.
+- A new BeamInterface class that provides a unified interface to UVBeam and analytic
+beam objects to get beam responses in any direction and frequency (via
+interpolation or evaluation as appropriate).
+- New optional spatial interpolation method, ``interpolation_function="az_za_map_coordinates"``
+that improves the linear interpolation speed for data in ``az_za`` coordinates.
 - New UVParameter `pol_convention` on `UVData` and `UVCal`. This specifies the convention
 assumed for converting linear to stokes polarizations -- either "sum" or "avg". Also
 added to `uvcalibrate` to apply from the `UVCal` to the `UVData`.

docs/analytic_beams.rst

@@ -0,0 +1,78 @@
+Analytic Beams
+==============
+
+pyuvdata defines several analytic primary beams for radio telescopes. While these
+are not realistic models for true antennas (like those represented in
+:class:`pyuvdata.UVBeam`), they can be useful in simulation because they are
+lightweight and fast to evaluate (as opposed to having to interpolate).
+
+The analytic beams defined in pyuvdata are based on a base class,
+:class:`pyuvdata.analytic_beam.AnalyticBeam`, which ensures a standard interface
+and can be used to define other analytic beams in a consistent way (see the
+:ref:`analytic beam tutorial <analytic_beam_tutorial>`). To evaluate analytic
+beams in particular directions at particular frequencies, use the
+:meth:`pyuvdata.analytic_beam.AnalyticBeam.efield_eval`
+or :meth:`pyuvdata.analytic_beam.AnalyticBeam.power_eval` methods as appropriate.
+
+The ``AnalyticBeam`` base class also provides a yaml constructor that can enable
+analytic beams to be instantiated directly from yaml files (see
+:ref:`yaml_constructors`, similar constructors are also available for UVBeam
+objects) and a plugin infrastructure that can automatically include any imported
+subclass even if they are defined in other packages. This infrastructure, along
+with the :class:`pyuvdata.BeamInterface` class, can simplify the setup for
+simulations.
+
+.. autoclass:: pyuvdata.analytic_beam.AnalyticBeam
+   :members:
+
+
+.. autoclass:: pyuvdata.AiryBeam
+   :members:
+
+.. autoclass:: pyuvdata.GaussianBeam
+   :members:
+
+.. autoclass:: pyuvdata.ShortDipoleBeam
+   :members:
+
+.. autoclass:: pyuvdata.UniformBeam
+   :members:
+
+
+.. _yaml_constructors:
+
+yaml constructors
+-----------------
+
+Analytic beams can be instantiated directly from yaml files using the
+``!AnalyticBeam`` tag. The ``class`` parameter must be specified and it can be
+set to one of the pyuvdata provided analytic beams or to any AnalyticBeam
+subclass. If the subclass is not defined in pyuvdata, either the subclass must
+already be imported or it must be specified with the dot-pathed modules included
+(i.e. ``my_module.MyAnalyticBeam``). Some analytic beams have required parameters
+(e.g. ``diameter`` for AiryBeams), which must also be provided, see the object
+definitions for details.
+
+Some examples are provided below, note that the node key can be anything, it
+does not need to be ``beam``:
+
+A 16 meter diameter Airy beam::
+
+    beam: !AnalyticBeam
+        class: AiryBeam
+        diameter: 16
+
+A classical short dipole beam (the dot-pathed module notation is not required
+for pyvudata beams but is shown here as an example)::
+
+    beam: !AnalyticBeam
+        class: pyuvdata.ShortDipoleBeam
+
+A gaussian shaped beam with an E-Field beam sigma of 0.26 radians that has
+width that scales as a power law with frequency::
+
+    beam: !AnalyticBeam
+        class: GaussianBeam
+        reference_frequency: 120000000.
+        spectral_index: -1.5
+        sigma: 0.26

docs/beam_interface.rst

@@ -0,0 +1,13 @@
+Beam Interface
+==============
+
+The BeamInterface object is designed to provide a unified interface for UVBeam
+and AnalyticBeam objects to compute beam response values. It can be constructed
+with either a :class:`pyuvdata.UVBeam` or :class:`AnalyticBeam` and it provides
+the :meth:`pyuvdata.BeamInterface.compute_response` method, which can be used to
+either evaluate the response of an AnalyticBeam or interpolate a UVBeam to get
+the beam response. This interface provides a simplified way for simulators to
+get beam responses from either UVBeams or analytic beams with the same code.
+
+.. autoclass:: pyuvdata.BeamInterface
+   :members:

docs/beam_interface_tutorial.rst

@@ -0,0 +1,103 @@
+--------------
+Beam Interface
+--------------
+
+The BeamInterface object is designed to provide a unified interface for UVBeam
+and AnalyticBeam objects to compute beam response values. It can be constructed
+with either a :class:`pyuvdata.UVBeam` or :class:`AnalyticBeam` and the beam
+response can be calculated using the :meth:`pyuvdata.BeamInterface.compute_response`
+method.
+
+Using BeamInterface
+-------------------
+
+The following code shows how to set up two BeamInterface objects, one with an
+analytic beam and one with a UVBeam. Then each is evalated at the same frequency
+and directions using the same call to :meth:`pyuvdata.BeamInterface.compute_response`
+and the results are plotted. The value of the BeamInterface object is that it
+unifies the interface so the code calling it does not need to know if the beam
+that is attached to it is an analytic beam or a UVBeam.
+
+.. code-block:: python
+
+  >>> import os
+  >>> import matplotlib.pyplot as plt
+  >>> import numpy as np
+  >>> from matplotlib.colors import LogNorm
+
+  >>> from pyuvdata import ShortDipoleBeam, BeamInterface, UVBeam
+  >>> from pyuvdata.data import DATA_PATH
+
+  >>> filename = os.path.join(DATA_PATH, "mwa_full_EE_test.h5")
+
+  >>> dipole_beam = BeamInterface(ShortDipoleBeam(), beam_type="power")
+  >>> mwa_beam = BeamInterface(UVBeam.from_file(filename, pixels_per_deg=1), beam_type="power")
+
+  >>> # set up zenith angle, azimuth and frequency arrays to evaluate with
+  >>> # make a regular grid in direction cosines for nice plots
+  >>> n_vals = 100
+  >>> zmax = np.radians(90)  # Degrees
+  >>> axis_arr = np.arange(-n_vals/2., n_vals/2.) / float(n_vals/2.)
+  >>> l_arr, m_arr = np.meshgrid(axis_arr, axis_arr)
+  >>> radius = np.sqrt(l_arr**2 + m_arr**2)
+  >>> za_array = radius * zmax
+  >>> az_array = np.arctan2(m_arr, l_arr)
+
+  >>> # Wrap the azimuth array to [0, 2pi] to match the extent of the UVBeam azimuth
+  >>> where_neg_az = np.nonzero(az_array < 0)
+  >>> az_array[where_neg_az] = az_array[where_neg_az] + np.pi * 2.
+  >>> az_array = az_array.flatten()
+  >>> za_array = za_array.flatten()
+
+  >>> # find the values above the horizon so we don't try to interpolate the MWA beam
+  >>> # beyond the horizon
+  >>> above_hor = np.nonzero(za_array <= np.pi / 2.)[0]
+
+  >>> # set up output arrays that matches the expected shape, except that they
+  >>> # include the points beyond the horizon, and fill them with infinity.
+  >>> # Then we will set the points above the horizon to the computed responses.
+  >>> dipole_beam_vals = np.full((1, 4, 1, n_vals * n_vals), np.inf, dtype=complex)
+  >>> mwa_beam_vals = np.full((1, 4, 1, n_vals * n_vals), np.inf, dtype=complex)
+
+  >>> # The MWA beam we have in our test data is small, it only has 3 frequencies,
+  >>> # so we will just get the value at one of those frequencies rather than
+  >>> # trying to interpolate to a new frequency.
+  >>> freqs = np.array([mwa_beam.beam.freq_array[-1]])
+
+  >>> dipole_beam_vals[:, :, :, above_hor] = dipole_beam.compute_response(
+  ...     az_array=az_array[above_hor], za_array=za_array[above_hor], freq_array=freqs
+  ... )
+  >>> dipole_beam_vals = dipole_beam_vals.reshape(4, n_vals, n_vals)
+
+  >>> mwa_beam_vals[:, :, :, above_hor] = mwa_beam.compute_response(
+  ...     az_array=az_array[above_hor], za_array=za_array[above_hor], freq_array=freqs
+  ... )
+  >>> mwa_beam_vals = mwa_beam_vals.reshape(4, n_vals, n_vals)
+
+  >>> fig, ax = plt.subplots(1, 2)
+  >>> bp_dip = ax[0].imshow(
+  ...   dipole_beam_vals[0].real,
+  ...   norm=LogNorm(vmin = 1e-4, vmax =1),
+  ...   extent=[np.min(l_arr), np.max(l_arr), np.min(m_arr), np.max(m_arr)],
+  ... )
+  >>> _ = ax[0].set_title(f"E/W Dipole power beam")
+  >>> _ = ax[0].set_xlabel("direction cosine l")
+  >>> _ = ax[0].set_ylabel("direction cosine m")
+  >>> _ = fig.colorbar(bp_dip, ax=ax[0], fraction=0.046, pad=0.04)
+
+  >>> bp_mwa = ax[1].imshow(
+  ...   mwa_beam_vals[0].real,
+  ...   norm=LogNorm(vmin = 1e-4, vmax =1),
+  ...   extent=[np.min(l_arr), np.max(l_arr), np.min(m_arr), np.max(m_arr)],
+  ... )
+  >>> _ = ax[1].set_title(f"MWA E/W power beam")
+  >>> _ = ax[1].set_xlabel("direction cosine l")
+  >>> _ = ax[1].set_ylabel("direction cosine m")
+  >>> _ = fig.colorbar(bp_mwa, ax=ax[1], fraction=0.046, pad=0.04)
+  >>> fig.tight_layout()
+  >>> plt.show()  # doctest: +SKIP
+  >>> plt.savefig("Images/dipole_mwa_power.png", bbox_inches='tight')
+  >>> plt.clf()
+
+.. image:: Images/dipole_mwa_power.png
+  :width: 600

docs/cst_settings_yaml.rst

@@ -7,8 +7,7 @@ via keywords when the files are read in, but it is better for the metadata to be
 specified once and carried with the data files. To that end, we developed a yaml
 settings file specification to carry all the metadata. This format is very human
 readable and writeable and we encourage using such a file as the best way to
-ensure the metadata is preserved. Note that reading a yaml settings file into
-UVBeam requires that pyyaml is installed.
+ensure the metadata is preserved.
 
 Required Fields
 ***************

docs/make_index.py

@@ -49,6 +49,8 @@ def write_index_rst(readme_file=None, write_file=None):
         "   uvbeam\n"
         "   uvflag\n"
         "   telescope\n"
+        "   analytic_beams\n"
+        "   beam_interface\n"
         "   fast_uvh5_meta\n"
         "   fast_calh5_meta\n"
         "   utility_functions\n"

docs/make_uvbeam.py

@@ -21,7 +21,9 @@ def write_uvbeam_rst(write_file=None):
         "transforming the data (interpolating/regridding, selecting, converting\n"
         "types) and can be interacted with directly.\n\n"
         "Note that there are some tricks that can help with reading in CST beam\n"
-        "simulation files in `CST Settings Files`_.\n\n"
+        "simulation files in `CST Settings Files`_.\n"
+        "UVBeam also provides a yaml constructor that can enable beams to be\n"
+        "instantiated directly from yaml files (see: `UVbeam yaml constructor`_)\n\n"
         "Attributes\n----------\n"
         "The attributes on UVBeam hold all of the metadata and data required to\n"
         "describe primary beam models. Under the hood, the attributes are implemented\n"
@@ -67,6 +69,33 @@ def write_uvbeam_rst(write_file=None):
 
     out += "Methods\n-------\n.. autoclass:: pyuvdata.UVBeam\n  :members:\n\n"
 
+    out += """
+UVBeam yaml constructor
+-----------------------
+
+UVbeams can be instantiated directly from yaml files using the
+``!UVBeam`` tag. The ``filename`` parameter must be specified and
+any other parameter that can be passed to the :meth:`pyuvdata.UVBeam.read`
+method can also be specified.
+
+Some examples are provided below, note that the node key can be anything,
+it does not need to be ``beam``:
+
+A simple UVBeam specification::
+
+    beam: !UVBeam
+        filename: hera.beamfits
+
+An UVbeam specification with some keywords to pass to ``UVBeam.read``::
+
+    beam: !UVBeam
+        filename: mwa_full_EE_test.h5
+        pixels_per_deg: 1
+        freq_range: [100.e+6, 200.e+6]
+
+\n\n
+"""
+
     with open("cst_settings_yaml.rst", encoding="utf-8") as cst_settings_file:
         cst_setting_text = cst_settings_file.read()
 

docs/tutorial.rst

@@ -16,3 +16,5 @@ Tutorials are available for each major user class:
     uvdata_tutorial.rst
     uvcal_tutorial.rst
     uvbeam_tutorial.rst
+    analytic_beam_tutorial.rst
+    beam_interface_tutorial.rst

docs/uvcal_tutorial.rst

@@ -324,6 +324,7 @@ that should be noted:
   on the object itself).
 * Regardless of the value of ``undo``, the convention that is inferred for the
   calibration solutions is determined as follows:
+
   * If neither ``uvc_pol_convention`` nor ``uvcal.pol_convention`` are specified, a
     a warning is raised (since the resulting calibrated data is not well-determined),
     and it is *assumed* that the solutions have the same convention as the ``UVData``
@@ -332,7 +333,9 @@ that should be noted:
     corrections are applied, and the result is ambiguous.
   * If both ``uvc_pol_convention`` and ``uvcal.pol_convention`` are specified and are
     different, an error is raised.
+
 * When **calibrating** in :func:`pyuvdata.utils.uvcalibrate` (i.e. ``undo=False``):
+
   * If ``uvdata.pol_convention`` is specified, an error is raised, because you are
     trying to calibrate already-calibrated data.
   * The convention applied to the resulting ``UVData`` object is inferred in the
@@ -341,14 +344,17 @@ that should be noted:
     or ``uvcal.pol_convention``, see above), (iii) if still unspecified, no convention
     will be used and a warning will be raised. This was always the behaviour in earlier
     versions of ``pyuvdata`` (pre-v3).
+
 * When **un-calibrating** with :func:`pyuvdata.utils.uvcalibrate` (i.e. ``undo=True``):
+
   * If both ``uvd_pol_convention`` and ``uvdata.pol_convention`` are defined and
     are different, an error is raised.
   * If neither are set, a warning is raised, since the resulting un-calibrated values
     may not be the same as the original values before calibration (since a different
     convention could have been used to calibrate originally than is being used to
     de-calibrate). However, calibration will continue, assuming that the ``UVData``
     object has the same convention as the ``UVCal`` object used to de-calibrate.
+
 * It is not supported to have ``pol_convention`` set on ``UVCal``, but *not*
   ``gain_scale``. A ``pol_convention`` only makes sense in the context of having a
   scale for the gains.

src/pyuvdata/__init__.py

@@ -45,6 +45,8 @@ def branch_scheme(version):  # pragma: nocover
 warnings.filterwarnings("ignore", message="numpy.dtype size changed")
 warnings.filterwarnings("ignore", message="numpy.ufunc size changed")
 
+from .analytic_beam import AiryBeam, GaussianBeam, ShortDipoleBeam, UniformBeam  # noqa
+from .beam_interface import BeamInterface  # noqa
 from .telescopes import Telescope  # noqa
 from .telescopes import get_telescope  # noqa  # NB: get_telescopes is deprecated
 from .uvbeam import UVBeam  # noqa
@@ -59,6 +61,11 @@ def branch_scheme(version):  # pragma: nocover
     "UVCal",
     "UVFlag",
     "UVBeam",
+    "BeamInterface",
+    "AiryBeam",
+    "GaussianBeam",
+    "ShortDipoleBeam",
+    "UniformBeam",
     "Telescope",
     "get_telescope",
 ]