Support custom beam patterns in AnalyticBeam class

telegraphic commented 7 months ago

Extend AnalyticBeam class to support any function of alt, az and frequency

Description

The AnalyticBeam class has been extended to support beam calculations using custom functions.

To use, the type='func' argument is set, and the new optional argument func=my_custom_function is set:

def cos_squared(a, z, f):
    return np.cos(z)**2

beam = AnalyticBeam('func', func=cos_squared)

The function must have exactly three arguments, correspinding to az_array, za_array and freq_array, as used in the interp class method:

    def interp(self, az_array, za_array, freq_array, **kwargs):
        """
        Evaluate the primary beam at given sky coordinates and frequencies.

        (mocks :meth:`pyuvdata.UVBeam.interp`, but these are analytic, so no interpolation is done.)

        Parameters
        ----------
        az_array : array-like of float
            Azimuth values to evaluate at in radians. Should be a 1D array with the same
            length as `za_array`. The azimuth here has the :class:`pyuvdata.UVBeam` convention:
            North of East (East=0, North=pi/2)
        za_array : array-like of float
            Zenith angle values to evaluate at in radians. Should be a 1D array with the
            same length as `az_array`.
        freq_array : array-like of float
            Frequency values to evaluate at in Hz. Should be a 1D array.

Motivation and Context

The AnalyticBeam currently supports uniform, airy and Gaussian beam patterns. This PR allows any function of altitude, zenith angle, and frequency to be used, increasing flexibility.

My motivation was to add support for beams generated from simple combinations of Numpy ufuncs (e.g. cos, sin, exp) with corresponding analytical functions. However, as long as the function returns a value for any given alt/az/freq, it doesn't strictly need to be an analytical function.

I chose this approach over adding another simple analytic function, e.g. type='cos_squared', as it was far more flexible.

Types of changes

[ ] Bug fix (non-breaking change which fixes an issue)
[x] New feature (non-breaking change which adds functionality)
[ ] Breaking change (fix or feature that would cause existing functionality to change)
[ ] Reference simulation update or replacement
[ ] Documentation change (documentation changes only)
[ ] Version change
[ ] Build or continuous integration change

Checklist:

For all pull requests:

[x] I have read the contribution guide.
[x] My code follows the code style of this project.

New feature checklist:

[x] I have added or updated the docstrings associated with my feature using the numpy docstring format.
[ ] I have updated the documentation to highlight my new feature (if appropriate).
[x] I have added tests to cover my new feature.
[x] All new tests pass
[x] All existing tests pass
[x] updated the CHANGELOG

Things I haven't done:

[ ] Checked that I reproduce the reference simulations or if there are differences they are explained below (if appropriate). If there are changes that are correct, I will update the reference simulation files after this PR is merged. Not sure what this means.
[ ] checked (e.g., using the benchmarking tools) that this change does not significantly increase typical runtimes. If it does, I have included a justification in the comments on this PR. Have not done, but would be very surprised if it impacted runtimes
[ ] I have no intention of running linters or fixing any linting issues

Documentation change checklist:

[x] Any updated docstrings use the numpy docstring format.
[ ] If this is a significant change to the readme or other docs, I have checked that they are rendered properly on ReadTheDocs. (you may need help to get this branch to build on RTD, just ask!)

steven-murray commented 7 months ago

@telegraphic thank you for this. I think we need to think about this in relation to the work going on in https://github.com/RadioAstronomySoftwareGroup/pyuvdata/pull/1383, in which we are trying to generalize the UVBeam objects to have a more uniform interface between analaytic/simulated beams.

bhazelton commented 7 months ago

@telegraphic As Steven mentioned above, we are currently in the process of overhauling our approach to Analytic Beams, moving them into pyuvdata and trying to have a more flexible and useful API. We'd love your comments on that PR and would be happy to try to incorporate your ideas there as appropriate. The plan is to try to get them in a bit after our upcoming version 3.0 release of pyuvdata. We plan to make a branch of pyuvsim where we'll get the new analytic beams working to make sure everything works well before we merge it into pyuvdata (just as soon as we find the time...).

telegraphic commented 7 months ago

No worries, take your time, can reassess when beam approach has solidified.

bhazelton commented 2 months ago

@telegraphic I think we are getting close on the pyuvdata and pyuvsim PRs that will revamp the way the analytic beams are handled. The analytic beam structure was deliberately designed to make it easy for anyone to define their own beam objects by subclassing the AnalyticBeam abstract base class. We'd appreciate your thoughts on the PRs if you have a chance to take a look, both on the actual implementation and API and on the documentation.

The pyuvdata PR is here: https://github.com/RadioAstronomySoftwareGroup/pyuvdata/pull/1383 the pyuvsim PR is #487

RadioAstronomySoftwareGroup / pyuvsim