more tutorial!

docs/analytic_beam_tutorial.rst

@@ -16,8 +16,8 @@ To evaluate an analytic beam at one or more frequencies and in in one or more
 directions, use either the :meth:`pyuvdata.analytic_beam.AnalyticBeam.efield_eval`
 or :meth:`pyuvdata.analytic_beam.AnalyticBeam.power_eval` methods as appropriate.
 
-a) Evaluating an Airy Beam power response
-*****************************************
+Evaluating an Airy Beam power response
+**************************************
 
 .. code-block:: python
 
@@ -78,8 +78,8 @@ a) Evaluating an Airy Beam power response
   :width: 600
 
 
-a) Evaluating a Short Dipole Beam E-Field response
-**************************************************
+Evaluating a Short Dipole Beam E-Field response
+***********************************************
 
 .. code-block:: python
 
@@ -240,12 +240,16 @@ actually calculate the response of the new beam:
     ``Npols`` is equal to either ``Nfeeds`` squared if ``include_cross_pols``
     is True (the default) or ``Nfeeds`` if ``include_cross_pols`` is False.
 
+Note that if you need to do some manipulation or validation of the fields after
+they are specified by the user, you can use the dataclass's ``__post_init__``
+method to do that, being sure to call the super class ``__post_init__`` as well.
+The gaussian beam example below shows how this can be done.
 
 Below we provide some examples of beams defined in pyuvdata to make this more
 concrete.
 
-a) Defining an Airy Beam
-************************
+Defining a simple unpolarized beam
+**********************************
 
 Airy beams are unpolarized but frequency dependent and require one parameter,
 the dish diameter in meters.
@@ -367,8 +371,8 @@ the dish diameter in meters.
 
           return data_array
 
-a) Defining a Short Dipole Beam
-*******************************
+Defining a simple polarized beam
+********************************
 
 Short (Hertzian) dipole beams are polarized but frequency independent and do not
 require any parameters.
@@ -475,3 +479,248 @@ require any parameters.
               data_array[0, 3] = data_array[0, 2]
 
           return data_array
+
+
+Defining a beam with post init validation
+*****************************************
+
+The gaussian beam defined in pyuvdata has several optional configurations that
+require some validation, which we do using the dataclass ``__post_init__`` method.
+Note that we call the ``super().__post_init__`` within that method to ensure that
+all the normal AnalyticBeam setup has been done.
+
+.. code-block:: python
+  :linenos:
+
+  import dataclasses
+  from dataclasses import InitVar, dataclass, field
+  from typing import Literal
+
+  import numpy as np
+  import numpy.typing as npt
+  from astropy.constants import c as speed_of_light
+  from pyuvdata.analytic_beam import AnalyticBeam
+
+  def diameter_to_sigma(diameter: float, freq_array: npt.NDArray[float]) -> float:
+      """
+      Find the sigma that gives a beam width similar to an Airy disk.
+
+      Find the stddev of a gaussian with fwhm equal to that of
+      an Airy disk's main lobe for a given diameter.
+
+      Parameters
+      ----------
+      diameter : float
+          Antenna diameter in meters
+      freq_array : array of float
+          Frequencies in Hz
+
+      Returns
+      -------
+      sigma : float
+          The standard deviation in zenith angle radians for a Gaussian beam
+          with FWHM equal to that of an Airy disk's main lobe for an aperture
+          with the given diameter.
+
+      """
+      wavelengths = speed_of_light.to("m/s").value / freq_array
+
+      scalar = 2.2150894  # Found by fitting a Gaussian to an Airy disk function
+
+      sigma = np.arcsin(scalar * wavelengths / (np.pi * diameter)) * 2 / 2.355
+
+      return sigma
+
+
+  @dataclass(kw_only=True)
+  class GaussianBeam(AnalyticBeam):
+      """
+      A circular, zenith pointed Gaussian beam.
+
+      Requires either a dish diameter in meters or a standard deviation sigma in
+      radians. Gaussian beams specified by a diameter will have their width
+      matched to an Airy beam at each simulated frequency, so are inherently
+      chromatic. For Gaussian beams specified with sigma, the sigma_type defines
+      whether the width specified by sigma specifies the width of the E-Field beam
+      (default) or power beam in zenith angle. If only sigma is specified, the
+      beam is achromatic, optionally both the spectral_index and reference_frequency
+      parameters can be set to generate a chromatic beam with standard deviation
+      defined by a power law:
+
+      stddev(f) = sigma * (f/ref_freq)**(spectral_index)
+
+      The unpolarized nature leads to some results that may be
+      surprising to radio astronomers: if two feeds are specified they will have
+      identical responses and the cross power beam between the two feeds will be
+      identical to the power beam for a single feed.
+
+      Attributes
+      ----------
+      sigma : float
+          Standard deviation in radians for the gaussian beam. Only one of sigma
+          and diameter should be set.
+      sigma_type : str
+          Either "efield" or "power" to indicate whether the sigma specifies the size of
+          the efield or power beam. Ignored if `sigma` is None.
+      diameter : float
+          Dish diameter in meters to use to define the size of the gaussian beam, by
+          matching the FWHM of the gaussian to the FWHM of an Airy disk. This will result
+          in a frequency dependent beam.  Only one of sigma and diameter should be set.
+      spectral_index : float
+          Option to scale the gaussian beam width as a power law with frequency. If set
+          to anything other than zero, the beam will be frequency dependent and the
+          `reference_frequency` must be set. Ignored if `sigma` is None.
+      reference_frequency : float
+          The reference frequency for the beam width power law, required if `sigma` is not
+          None and `spectral_index` is not zero. Ignored if `sigma` is None.
+      feed_array : np.ndarray of str
+          Feeds to define this beam for, e.g. x & y or n & e (for "north" and "east")
+          or r & l.
+      x_orientation : str
+          Physical orientation of the feed for the x feed. Not meaningful for
+          GaussianBeams, which are unpolarized.
+
+      Parameters
+      ----------
+      sigma : float
+          Standard deviation in radians for the gaussian beam. Only one of sigma
+          and diameter should be set.
+      sigma_type : str
+          Either "efield" or "power" to indicate whether the sigma specifies the size of
+          the efield or power beam. Ignored if `sigma` is None.
+      diameter : float
+          Dish diameter in meters to use to define the size of the gaussian beam, by
+          matching the FWHM of the gaussian to the FWHM of an Airy disk. This will result
+          in a frequency dependent beam.  Only one of sigma and diameter should be set.
+      spectral_index : float
+          Option to scale the gaussian beam width as a power law with frequency. If set
+          to anything other than zero, the beam will be frequency dependent and the
+          `reference_frequency` must be set. Ignored if `sigma` is None.
+      reference_frequency : float
+          The reference frequency for the beam width power law, required if `sigma` is not
+          None and `spectral_index` is not zero. Ignored if `sigma` is None.
+      feed_array : np.ndarray of str
+          Feeds to define this beam for, e.g. n & e or x & y or r & l.
+      x_orientation : str
+          Physical orientation of the feed for the x feed. Not meaningful for
+          GaussianBeams, which are unpolarized.
+      include_cross_pols : bool
+          Option to include the cross polarized beams (e.g. xy and yx or en and ne) for
+          the power beam.
+
+      """
+
+      sigma: float | None = None
+      sigma_type: Literal["efield", "power"] = "efield"
+      diameter: float | None = None
+      spectral_index: float = 0.0
+      reference_frequency: float = None
+
+      feed_array: npt.NDArray[str] | None = field(default=None, repr=False, compare=False)
+      x_orientation: Literal["east", "north"] = field(
+          default="east", repr=False, compare=False
+      )
+
+      include_cross_pols: InitVar[bool] = True
+
+      basis_vector_type = "az_za"
+
+      def __post_init__(self, include_cross_pols):
+          """
+          Post-initialization validation and conversions.
+
+          Parameters
+          ----------
+          include_cross_pols : bool
+              Option to include the cross polarized beams (e.g. xy and yx or en and ne)
+              for the power beam.
+
+          """
+          if (self.diameter is None and self.sigma is None) or (
+              self.diameter is not None and self.sigma is not None
+          ):
+              if self.diameter is None:
+                  raise ValueError("Either diameter or sigma must be set.")
+              else:
+                  raise ValueError("Only one of diameter or sigma can be set.")
+
+          if self.sigma is not None:
+              if self.sigma_type != "efield":
+                  self.sigma = np.sqrt(2) * self.sigma
+
+              if self.spectral_index != 0.0 and self.reference_frequency is None:
+                  raise ValueError(
+                      "reference_frequency must be set if `spectral_index` is not zero."
+                  )
+              if self.reference_frequency is None:
+                  self.reference_frequency = 1.0
+
+          super().__post_init__(include_cross_pols=include_cross_pols)
+
+      def get_sigmas(self, freq_array: npt.NDArray[float]) -> npt.NDArray[float]:
+          """
+          Get the sigmas for the gaussian beam using the diameter (if defined).
+
+          Parameters
+          ----------
+          freq_array : array of floats
+              Frequency values to get the sigmas for in Hertz.
+
+          Returns
+          -------
+          sigmas : array_like of float
+              Beam sigma values as a function of frequency. Size will match the
+              freq_array size.
+
+          """
+          if self.diameter is not None:
+              sigmas = diameter_to_sigma(self.diameter, freq_array)
+          elif self.sigma is not None:
+              sigmas = (
+                  self.sigma
+                  * (freq_array / self.reference_frequency) ** self.spectral_index
+              )
+          return sigmas
+
+      def _efield_eval(
+          self,
+          *,
+          az_array: npt.NDArray[float],
+          za_array: npt.NDArray[float],
+          freq_array: npt.NDArray[float],
+      ) -> npt.NDArray[float]:
+          """Evaluate the efield at the given coordinates."""
+          sigmas = self.get_sigmas(freq_array)
+
+          values = np.exp(
+              -(za_array[np.newaxis, ...] ** 2) / (2 * sigmas[:, np.newaxis] ** 2)
+          )
+          data_array = self._get_empty_data_array(az_array.size, freq_array.size)
+          for fn in np.arange(self.Nfeeds):
+              data_array[0, fn, :, :] = values / np.sqrt(2.0)
+              data_array[1, fn, :, :] = values / np.sqrt(2.0)
+
+          return data_array
+
+      def _power_eval(
+          self,
+          *,
+          az_array: npt.NDArray[float],
+          za_array: npt.NDArray[float],
+          freq_array: npt.NDArray[float],
+      ) -> npt.NDArray[float]:
+          """Evaluate the power at the given coordinates."""
+          sigmas = self.get_sigmas(freq_array)
+
+          values = np.exp(
+              -(za_array[np.newaxis, ...] ** 2) / (sigmas[:, np.newaxis] ** 2)
+          )
+          data_array = self._get_empty_data_array(
+              az_array.size, freq_array.size, beam_type="power"
+          )
+          for fn in np.arange(self.Npols):
+              # For power beams the first axis is shallow because we don't have to worry
+              # about polarization.
+              data_array[0, fn, :, :] = values
+
+          return data_array