From d2c49c78a0f79f247fc6784fdd002892c8a8aaee Mon Sep 17 00:00:00 2001 From: Bryna Hazelton Date: Tue, 20 Aug 2024 12:47:05 -0700 Subject: [PATCH] updates for pyuvsim --- src/pyuvdata/uvbeam/analytic_beam.py | 82 +++++++++++++++----- src/pyuvdata/uvbeam/beam_interface.py | 104 +++++++++++++++++++++++--- src/pyuvdata/uvbeam/uvbeam.py | 3 +- tests/uvbeam/test_analytic_beam.py | 20 +++-- tests/uvbeam/test_beam_interface.py | 54 ++++++++++++- 5 files changed, 220 insertions(+), 43 deletions(-) diff --git a/src/pyuvdata/uvbeam/analytic_beam.py b/src/pyuvdata/uvbeam/analytic_beam.py index cc234183a1..ef23c18b09 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( @@ -434,7 +467,9 @@ def to_uvbeam( else: eval_function = "power_eval" - data_array = getattr(self, eval_function)(az_array, za_array, freq_array) + data_array = getattr(self, eval_function)( + az_array=az_array, za_array=za_array, freq_array=freq_array + ) if uvb.pixel_coordinate_system == "az_za": data_array = data_array.reshape(uvb.data_array.shape) @@ -495,8 +530,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 +552,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 +574,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 +614,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 +645,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 +701,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 +730,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 +792,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 +814,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 +867,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 +885,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 1bf7bc2060..518f94c64f 100644 --- a/src/pyuvdata/uvbeam/uvbeam.py +++ b/src/pyuvdata/uvbeam/uvbeam.py @@ -1877,7 +1877,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 ------- diff --git a/tests/uvbeam/test_analytic_beam.py b/tests/uvbeam/test_analytic_beam.py index 8e29fedf4f..d38ed817fa 100644 --- a/tests/uvbeam/test_analytic_beam.py +++ b/tests/uvbeam/test_analytic_beam.py @@ -115,15 +115,16 @@ def test_chromatic_gaussian(): # Error if trying to define chromatic beam without a reference frequency with pytest.raises( - ValueError, match="reference_freq must be set if `spectral_index` is not zero." + ValueError, + match="reference_frequency must be set if `spectral_index` is not zero.", ): GaussianBeam(sigma=sigma, spectral_index=alpha) - beam = GaussianBeam(sigma=sigma, reference_freq=freqs[0], spectral_index=alpha) + beam = GaussianBeam(sigma=sigma, reference_frequency=freqs[0], spectral_index=alpha) # Get the widths at each frequency. - vals = beam.efield_eval(az, za, freqs) + vals = beam.efield_eval(az_array=az, za_array=za, freq_array=freqs) # pick out a single polarization direction and feed vals = vals[0, 0] @@ -241,8 +242,8 @@ def test_power_analytic_beam(beam_obj, kwargs, xy_grid): beam = beam_obj(**kwargs) - efield_vals = beam.efield_eval(az_vals, za_vals, freqs) - power_vals = beam.power_eval(az_vals, za_vals, freqs) + efield_vals = beam.efield_eval(az_array=az_vals, za_array=za_vals, freq_array=freqs) + power_vals = beam.power_eval(az_array=az_vals, za_array=za_vals, freq_array=freqs) # check power beams are peak normalized assert np.max(power_vals) == 1.0 @@ -297,7 +298,7 @@ def test_eval_errors(az_za_deg_grid): [UniformBeam(), False, None], [GaussianBeam(sigma=0.05), False, None], [ - GaussianBeam(sigma=0.02, reference_freq=100e8, spectral_index=-1.5), + GaussianBeam(sigma=0.02, reference_frequency=100e8, spectral_index=-1.5), False, None, ], @@ -339,11 +340,8 @@ def test_comparison(compare_beam, equality, operation): {"feed_array": "w", "diameter": 5}, re.escape("Feeds must be one of: ['n', 'e', 'x', 'y', 'r', 'l']"), ], - [{}, "One of diameter or sigma must be set but not both."], - [ - {"diameter": 5, "sigma": 0.2}, - "One of diameter or sigma must be set but not both.", - ], + [{}, "Either diameter or sigma must be set."], + [{"diameter": 5, "sigma": 0.2}, "Only one of diameter or sigma can be set."], ], ) def test_beamerrs(beam_kwargs, err_msg): diff --git a/tests/uvbeam/test_beam_interface.py b/tests/uvbeam/test_beam_interface.py index f67adbbc58..533055511b 100644 --- a/tests/uvbeam/test_beam_interface.py +++ b/tests/uvbeam/test_beam_interface.py @@ -137,11 +137,17 @@ def test_beam_interface( az_za_grid = False analytic_data = bi_analytic.compute_response( - az_array, za_array, freq_array, az_za_grid=az_za_grid + az_array=az_array, + za_array=za_array, + freq_array=freq_array, + az_za_grid=az_za_grid, ) uvb_data = bi_uvbeam.compute_response( - az_array, za_array, freq_array, az_za_grid=az_za_grid + az_array=az_array, + za_array=za_array, + freq_array=freq_array, + az_za_grid=az_za_grid, ) np.testing.assert_allclose(analytic_data, uvb_data, rtol=0, atol=1e-14) @@ -150,9 +156,11 @@ def test_beam_interface( # larger differences of course az_vals, za_vals, freqs = xy_grid_coarse analytic_data = bi_analytic.compute_response( - az_vals, za_vals, freqs, az_za_grid=True + az_array=az_vals, za_array=za_vals, freq_array=freqs, az_za_grid=True + ) + uvb_data = bi_uvbeam.compute_response( + az_array=az_vals, za_array=za_vals, freq_array=freqs, az_za_grid=True ) - uvb_data = bi_uvbeam.compute_response(az_vals, za_vals, freqs, az_za_grid=True) if not (coord_sys == "healpix" and "dipole" in beam_obj.name.lower()): np.testing.assert_allclose(analytic_data, uvb_data, rtol=0, atol=1e-1) @@ -171,6 +179,44 @@ def test_beam_interface( ) +@pytest.mark.parametrize( + ["bi1", "bi2", "equality"], + [ + [ + BeamInterface(ShortDipoleBeam(), beam_type="efield"), + BeamInterface(ShortDipoleBeam(), beam_type="efield"), + True, + ], + [ + BeamInterface(ShortDipoleBeam(), beam_type="efield"), + BeamInterface(ShortDipoleBeam(), beam_type="power"), + False, + ], + [ + BeamInterface(ShortDipoleBeam(), beam_type="efield"), + BeamInterface(AiryBeam(diameter=14.0), beam_type="efield"), + False, + ], + [ + BeamInterface(AiryBeam(diameter=12.0), beam_type="efield"), + BeamInterface(AiryBeam(diameter=14.0), beam_type="efield"), + False, + ], + [ + BeamInterface(ShortDipoleBeam(), beam_type="efield"), + ShortDipoleBeam(), + False, + ], + ], +) +def test_beam_interface_equality(bi1, bi2, equality): + if equality: + assert bi1 == bi2 + else: + assert bi1 != bi2 + assert not bi1 == bi2 # noqa SIM201 + + def test_beam_interface_errors(): with pytest.raises( ValueError, match="beam must be a UVBeam or an AnalyticBeam instance."