diff --git a/src/pyuvdata/uvbeam/analytic_beam.py b/src/pyuvdata/uvbeam/analytic_beam.py index cc234183a1..b1f083b050 100644 --- a/src/pyuvdata/uvbeam/analytic_beam.py +++ b/src/pyuvdata/uvbeam/analytic_beam.py @@ -17,7 +17,7 @@ from ..docstrings import combine_docstrings from .uvbeam import UVBeam, _convert_feeds_to_pols -__all__ = ["AnalyticBeam", "GaussianBeam"] +__all__ = ["AnalyticBeam", "AiryBeam", "GaussianBeam", "ShortDipoleBeam", "UniformBeam"] class AnalyticBeam(ABC): @@ -184,8 +184,29 @@ def __eq__(self, other: Any, silent: bool = False) -> bool: print("Classes do not match") return False + def __ne__(self, other, *, check_extra=True, silent=True): + """ + Test if classes match and parameters are not equal. + + Parameters + ---------- + other : class + Other class instance to check + silent : bool + Option to turn off printing explanations of why equality fails. Useful to + prevent __ne__ from printing lots of messages. + + Returns + ------- + bool + True if the two instances are equivalent. + + """ + return not self.__eq__(other, silent=silent) + def _check_eval_inputs( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -211,6 +232,7 @@ def _get_empty_data_array( @abstractmethod def _efield_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -242,6 +264,7 @@ def _efield_eval( def efield_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -267,13 +290,18 @@ def efield_eval( (Naxes_vec, Nfeeds, freq_array.size, az_array.size) """ - self._check_eval_inputs(az_array, za_array, freq_array) + self._check_eval_inputs( + az_array=az_array, za_array=za_array, freq_array=freq_array + ) - return self._efield_eval(az_array, za_array, freq_array).astype(complex) + return self._efield_eval( + az_array=az_array, za_array=za_array, freq_array=freq_array + ).astype(complex) @abstractmethod def _power_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -305,6 +333,7 @@ def _power_eval( def power_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -332,7 +361,9 @@ def power_eval( float type. """ - self._check_eval_inputs(az_array, za_array, freq_array) + self._check_eval_inputs( + az_array=az_array, za_array=za_array, freq_array=freq_array + ) if self.Npols > self.Nfeeds: # cross pols are included @@ -340,7 +371,9 @@ def power_eval( else: expected_type = float - return self._power_eval(az_array, za_array, freq_array).astype(expected_type) + return self._power_eval( + az_array=az_array, za_array=za_array, freq_array=freq_array + ).astype(expected_type) @combine_docstrings(UVBeam.new) def to_uvbeam( @@ -495,8 +528,8 @@ class GaussianBeam(AnalyticBeam): 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_freq` must be set. Ignored if `sigma` is None. - reference_freq : float + `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 @@ -517,14 +550,17 @@ def __init__( sigma_type: Literal["efield", "power"] = "efield", diameter: float | None = None, spectral_index: float = 0.0, - reference_freq: float = None, + reference_frequency: float = None, feed_array: npt.NDArray[np.str] | None = None, include_cross_pols: bool = True, ): if (diameter is None and sigma is None) or ( diameter is not None and sigma is not None ): - raise ValueError("One of diameter or sigma must be set but not both.") + if diameter is None: + raise ValueError("Either diameter or sigma must be set.") + else: + raise ValueError("Only one of diameter or sigma can be set.") self.diameter = diameter @@ -536,18 +572,19 @@ def __init__( description_str = f", E-field sigma={self.sigma}" if spectral_index != 0.0: - if reference_freq is None: + if reference_frequency is None: raise ValueError( - "reference_freq must be set if `spectral_index` is not zero." + "reference_frequency must be set if `spectral_index` " + "is not zero." ) description_str += ( f", spectral index={spectral_index}, " - f"reference freq={reference_freq} Hz" + f"reference freq={reference_frequency} Hz" ) - if reference_freq is None: - reference_freq = 1.0 + if reference_frequency is None: + reference_frequency = 1.0 self.spectral_index = spectral_index - self.reference_freq = reference_freq + self.reference_frequency = reference_frequency else: description_str = f", equivalent diameter={self.diameter} m" @@ -575,12 +612,14 @@ def get_sigmas(self, freq_array: npt.NDArray[np.float]) -> npt.NDArray[np.float] sigmas = diameter_to_sigma(self.diameter, freq_array) elif self.sigma is not None: sigmas = ( - self.sigma * (freq_array / self.reference_freq) ** self.spectral_index + self.sigma + * (freq_array / self.reference_frequency) ** self.spectral_index ) return sigmas def _efield_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -604,6 +643,7 @@ def _efield_eval( def _power_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -659,6 +699,7 @@ def __init__( def _efield_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -687,6 +728,7 @@ def _efield_eval( def _power_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -748,6 +790,7 @@ def __init__( def _efield_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -769,6 +812,7 @@ def _efield_eval( def _power_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -821,6 +865,7 @@ def __init__( def _efield_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], @@ -838,6 +883,7 @@ def _efield_eval( def _power_eval( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float], diff --git a/src/pyuvdata/uvbeam/beam_interface.py b/src/pyuvdata/uvbeam/beam_interface.py index 786aec8758..d784309fa5 100644 --- a/src/pyuvdata/uvbeam/beam_interface.py +++ b/src/pyuvdata/uvbeam/beam_interface.py @@ -7,7 +7,7 @@ import copy import warnings -from typing import Literal +from typing import Any, Literal import numpy as np import numpy.typing as npt @@ -48,8 +48,11 @@ def __init__( beam_type: Literal["efield", "power"] | None = None, include_cross_pols: bool = True, ): - if not isinstance(beam, UVBeam) and not isinstance(beam, AnalyticBeam): - raise ValueError("beam must be a UVBeam or an AnalyticBeam instance.") + if not isinstance(beam, UVBeam) and not issubclass(type(beam), AnalyticBeam): + raise ValueError( + "beam must be a UVBeam or an AnalyticBeam instance, not a " + f"{type(beam)}." + ) self.beam = beam if isinstance(beam, UVBeam): self._isuvbeam = True @@ -73,15 +76,79 @@ def __init__( self._isuvbeam = False self.beam_type = beam_type + def __eq__(self, other: Any, silent: bool = False) -> bool: + """ + Test if classes match and parameters are equal. + + Parameters + ---------- + other : class + Other class instance to check + silent : bool + Option to turn off printing explanations of why equality fails. Useful to + prevent __ne__ from printing lots of messages. + + Returns + ------- + bool + True if the two instances are equivalent. + + """ + if isinstance(other, self.__class__): + # First check that the beam is the same + # If analytic, also check that the beam_type is the same + if self.beam.__ne__(other.beam, silent=silent): + if not silent: + print("Beams do not match. ") + return False + if not self._isuvbeam and self.beam_type != other.beam_type: + if not silent: + print( + "Beam types do not match. " + f"Left is {self.beam_type}," + f" right is {other.beam_type}." + ) + return False + else: + if not silent: + print("Classes do not match") + return False + + return True + + def __ne__(self, other, *, check_extra=True, silent=True): + """ + Test if classes match and parameters are not equal. + + Parameters + ---------- + other : class + Other class instance to check + silent : bool + Option to turn off printing explanations of why equality fails. Useful to + prevent __ne__ from printing lots of messages. + + Returns + ------- + bool + True if the two instances are equivalent. + + """ + return not self.__eq__(other, silent=silent) + def compute_response( self, + *, az_array: npt.NDArray[np.float], za_array: npt.NDArray[np.float], freq_array: npt.NDArray[np.float] | None, az_za_grid: bool = False, + interpolation_function=None, + freq_interp_kind=None, freq_interp_tol: float = 1.0, reuse_spline: bool = False, spline_opts: dict | None = None, + check_azza_domain: bool = True, ): """ Calculate beam responses, by interpolating UVBeams or evaluating AnalyticBeams. @@ -103,18 +170,32 @@ def compute_response( az_za_grid : bool Option to treat the `az_array` and `za_array` as the input vectors for points on a mesh grid. + interpolation_function : str, optional + Specify the interpolation function to use. Defaults to: "az_za_simple" for + objects with the "az_za" pixel_coordinate_system and "healpix_simple" for + objects with the "healpix" pixel_coordinate_system. Only applies if + beam is a UVBeam. + freq_interp_kind : str + Interpolation method to use frequency. See scipy.interpolate.interp1d + for details. Defaults to "cubic". freq_interp_tol : float Frequency distance tolerance [Hz] of nearest neighbors. If *all* elements in freq_array have nearest neighbor distances within the specified tolerance then return the beam at each nearest neighbor, - otherwise interpolate the beam. + otherwise interpolate the beam. Only applies if beam is a UVBeam. reuse_spline : bool - Save the interpolation functions for reuse. Only applies for - `az_za_simple` interpolation. + Save the interpolation functions for reuse. Only applies if beam is + a UVBeam and interpolation_function is "az_za_simple". spline_opts : dict Provide options to numpy.RectBivariateSpline. This includes spline - order parameters `kx` and `ky`, and smoothing parameter `s`. - Only applies for `az_za_simple` interpolation. + order parameters `kx` and `ky`, and smoothing parameter `s`. Only + applies if beam is a UVBeam and interpolation_function is "az_za_simple". + check_azza_domain : bool + Whether to check the domain of az/za to ensure that they are covered by the + intrinsic data array. Checking them can be quite computationally expensive. + Conversely, if the passed az/za are outside of the domain, they will be + silently extrapolated and the behavior is not well-defined. Only + applies if beam is a UVBeam and interpolation_function is "az_za_simple". Returns ------- @@ -133,9 +214,12 @@ def compute_response( za_array=za_array, az_za_grid=az_za_grid, freq_array=freq_array, + interpolation_function=interpolation_function, + freq_interp_kind=freq_interp_kind, freq_interp_tol=freq_interp_tol, reuse_spline=reuse_spline, spline_opts=spline_opts, + check_azza_domain=check_azza_domain, ) else: if not isinstance(freq_array, np.ndarray) or freq_array.ndim != 1: @@ -150,11 +234,11 @@ def compute_response( if self.beam_type == "efield": interp_data = self.beam.efield_eval( - az_array_use, za_array_use, freq_array + az_array=az_array_use, za_array=za_array_use, freq_array=freq_array ) else: interp_data = self.beam.power_eval( - az_array_use, za_array_use, freq_array + az_array=az_array_use, za_array=za_array_use, freq_array=freq_array ) return interp_data diff --git a/src/pyuvdata/uvbeam/uvbeam.py b/src/pyuvdata/uvbeam/uvbeam.py index 39d70bab60..c6988f6b24 100644 --- a/src/pyuvdata/uvbeam/uvbeam.py +++ b/src/pyuvdata/uvbeam/uvbeam.py @@ -1876,7 +1876,8 @@ def interp( Whether to check the domain of az/za to ensure that they are covered by the intrinsic data array. Checking them can be quite computationally expensive. Conversely, if the passed az/za are outside of the domain, they will be - silently extrapolated and the behavior is not well-defined. + silently extrapolated and the behavior is not well-defined. Only + applies for `az_za_simple` interpolation. Returns -------