Add AnalyticBeam objects and a unified interface for UVBeams and AnalyticBeams

[bhazelton]

Description

This is a refactor of some code that has existed in pyuvsim for some time with some new functionality to make it easier to work with analytic beams.

Major changes:

• Add a new AnalyticBeam base class and several specific analytic beam objects (Airy, Gaussian, Short dipole, Uniform) that support evaluating either the E-field or power response in any direction and frequency.
• Add a BeamInterface object which can have either a UVBeam or an Analytic Beam attached to it and which can be used to get beam responses in any direction and frequency via a unified method that either calls UVBeam.interp or the appropriate AnalyticBeam eval method.

There are also a couple of minor changes:

• Fix a bug in the the new beam initializer method where it makes data_arrays with the wrong type when cross pols are included.
• Small refactoring in UVBeam

Motivation and Context

This will substantially clean up code in pyuvsim, make the analytic beams more accessible for other users and make it easier for simulators by providing a unified interface for UVBeam and Analytic Beam objects.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation change (documentation changes only)
• Version change
• Build or continuous integration change

Checklist:

• I have read the contribution guide.
• My code follows the code style of this project.

New feature checklist:

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

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.93%. Comparing base (72dd3bc) to head (7ed37e9).
Report is 167 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff            @@
##             main    #1383    +/-   ##
========================================
  Coverage   99.92%   99.93%            
========================================
  Files          61       62     +1     
  Lines       21337    21743   +406     
========================================
+ Hits        21322    21728   +406     
  Misses         15       15            

Flag	Coverage Δ
	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[steven-murray]

Thanks @bhazelton, this is looking really nice. I've made a few comments on things I think could be a little tighter.

Some extra stuff:

1. I think the analytic beams could easily support YAML interfaces. There are two ways to do this. The first would be to use the __init_subclass__ method on the base class, and fill up a private global dictionary with any subclasses, i.e.:

def __init_subclass__(cls, is_abstract=False, **kwargs):
        """Provide plugin capablity and do some verification of a defined plugin."""
        super().__init_subclass__(**kwargs)

        # Plugin framework
        if not is_abstract:
            cls._plugins[cls.__name__] = cls

This is nice because now we globally know the possible list of AnalyticBeams. However, to be added to the list, the class has to be imported, which can't easily be ensured from a YAML file (unless you have some syntax where you have a top-level list like:

plugins:
   - custom.module
   - other_package.submodule

And your YAML reader always makes sure to import those modules before doing any other loading. If you do all this, you can simply write a YAML constructor on the base class whose name is the subclass. So you could then write YAML as eg.:

beam: !GaussianBeam
   diameter: 4.0

On yaml.load(fl) you'd get an instantiated GaussianBeam, and any custom-written classes would also be fine, AS LONG AS THEY ARE IMPORTED FIRST.

If you don't want to rely on things being imported first, then you could instead allow behaviour like this:

beam: !AnalyticBeam
   class: module.CustomBeam
   diameter: 5.0

To do this, you'd simply use yaml.add_constructor("!AnalyticBeam", func_that_turns_data_into_beam, Loader=yaml.FullLoader). In that func, you'd do an import of the class and use it to construct the object. You could also support both ways if you wanted.

2. We should add some docs on how to use the new AnalyticBeam's and the beam interface.

[bhazelton]

Ugg, h5py 3.12 was just released to pypi (it's not on conda-forge yet). It's breaking the Windows build wheels CI.

There's now an issue reporting the problem: h5py/h5py#2505

Edit: They've released a fix, re-running now.

[telegraphic]

Hi @bhazelton, finally got a chance to look at this. It's great! A few comments/questions:

• I think the tutorials should mention the BeamInterface (I started writing a comment about UVBeam and AnalyticBeam both being 'beamy', so unifying them would be nice -- then I noticed beam_interface.py)
• Could (should?) analytic_beam.py live in the uvbeam directory?
• I know feed rotation is already on your radars (Generalize x_orientation concept #854 and Add support for feed position angles #1470), but noting x_orientation is still limited to either E or N. Don't let this be a hold-up though.

Building on #1470, for the lazy among us, would it be possible to include a UserFunctionBeam(AnalyticBeam), e.g.

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

def sin_cubed(a, z, f):
    return np.sin(z)**3

beam = UserFunctionBeam(efield_func=cos_squared, power_func=sin_cubed)

I'm not 100% sure how to do this with dataclass shenanigans, can we overwrite the _efield_eval and _power_eval functions via the __post_init__()? E.g.

import types

@dataclass(kw_only=True)
class UserFunctionBeam(AnalyticBeam):
    efield_func: types.FunctionType
    power_func: types.FunctionType

    def __post_init__(self):
        self._efield_eval = efield_func
        self._power_eval = power_func

This might overlap with @steven-murray's YAML plugin idea above though?

[steven-murray]

Thanks for having a look @telegraphic!

As to your suggestion of a UserFunctionBeam, I think the simplest way to do this would simply be to subclass AnalyticBeam yourself -- that is what is intended. So:

from pyuvdata.analytic_beam import AnalyticBeam

class MyBeam(AnalyticBeam):
    def _efield_eval(self, *, az_array, za_array, freq_array):
        return np.cos(za_array)**2

    def _power_eval(self, *, az_array, za_array, freq_array):
        return np.sin(za_array)**2

IMHO this is not much more writing than you need to create a new "UserFunctionBeam" as you have suggested. We could, if you really think it useful, provide a factory function, e.g. make_custom_beam which would construct a new class and return it, given two functions, so e.g.:

def make_custom_beam(name: str, efield: types.FunctionType, power types.FunctionType):
    @dataclass
    class CustomBeam(AnalyticBeam):
        __name__ = name
        def _efield_eval(self, *, az, za, freq):
            return efield(az, za, freq)
        def _power_eval(self, *az, za, freq):
            return power(az, za, freq)

    return CustomBeam

which could be called by the user like so:

my_new_beam_class = make_custom_beam(name='cosbeam', efield=lambda az, za, freq: np.cos(za)**2, power=lambda az,za,freq: np.sin(za)**2)

[bhazelton]

@telegraphic Thanks for your comments! As @steven-murray outlined above, the idea was to make it easy to make new AnalyticBeams by subclassing that object, and I tried to write the tutorial to outline how to do that. I'm willing to entertain Steven's idea of a factory function, but I'd need to explore it more to understand it. Note that we did implement the plugin model Steven suggested, so if the analytic beam class is imported it can be instantiated from a yaml.

I will add some tutorial stuff on BeamInterface to better demonstrate how to use it.

I initially had both analytic_beam.py and beam_interface.py in the UVBeam folder, but I decided to move them out because I felt like that was confusing since they aren't based on UVBeam.

[bhazelton]

@telegraphic I added a BeamInterface tutorial, comments welcome!

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite	Current: 8f3d7d9	Previous: 72dd3bc	Ratio
tests/utils/test_bls.py::test_bls_to_ant[min=4259840-len=1]	22411.74124582377 iter/sec (stddev: 0.00002546482354602611)	45253.21348892038 iter/sec (stddev: 0.000006309466103868728)	2.02

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite	Current: 8f3d7d9	Previous: 72dd3bc	Ratio
tests/utils/test_bls.py::test_ants_to_bls[min=65536-len=1]	7802.570390038221 iter/sec (stddev: 0.00030422167319555456)	19018.809758528256 iter/sec (stddev: 0.000020634598156091437)	2.44

This comment was automatically generated by workflow using github-action-benchmark.

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite	Current: 9455ec5	Previous: 72dd3bc	Ratio
tests/utils/test_bls.py::test_bls_to_ant[min=0-len=10000]	8113.563431727591 iter/sec (stddev: 0.0002847113722937477)	17153.133943111716 iter/sec (stddev: 0.000014328094102494816)	2.11

This comment was automatically generated by workflow using github-action-benchmark.

[bhazelton]

@telegraphic We did some work to try to make defining new beams easier and I've re-written the docs. Can you please take a look and see how you feel about it?

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite	Current: 0d0e092	Previous: 72dd3bc	Ratio
tests/utils/test_bls.py::test_bls_to_ant[min=65536-len=10]	15878.88875763105 iter/sec (stddev: 0.0005017992393593569)	39574.76477815693 iter/sec (stddev: 0.000014323943066984319)	2.49
tests/utils/test_bls.py::test_bls_to_ant[min=65536-len=100000]	1205.2620433393151 iter/sec (stddev: 0.00014395155612953415)	2697.029604108284 iter/sec (stddev: 0.00006945868233318298)	2.24
tests/utils/test_bls.py::test_bls_to_ant[min=4259840-len=100]	20168.228751046987 iter/sec (stddev: 0.00005135060970793413)	42670.934804243894 iter/sec (stddev: 0.000009504441735601591)	2.12
tests/utils/test_bls.py::test_bls_to_ant[min=4259840-len=1000]	17056.375855360315 iter/sec (stddev: 0.000031597962152010776)	35820.826656128345 iter/sec (stddev: 0.000008388776850375593)	2.10
tests/utils/test_bls.py::test_bls_to_ant[min=4259840-len=10000]	7210.711976161784 iter/sec (stddev: 0.00016902262084636491)	17546.28040765295 iter/sec (stddev: 0.000009184979574146035)	2.43
tests/utils/test_bls.py::test_ants_to_bls[min=65536-len=1]	9188.911561814726 iter/sec (stddev: 0.00008982676050676747)	19018.809758528256 iter/sec (stddev: 0.000020634598156091437)	2.07
tests/utils/test_bls.py::test_ants_to_bls[min=65536-len=100]	7783.648994006051 iter/sec (stddev: 0.00029543439140781476)	16502.48961933711 iter/sec (stddev: 0.000017937603983490586)	2.12
tests/utils/test_bls.py::test_ants_to_bls[min=65536-len=1000]	4592.552795705242 iter/sec (stddev: 0.0005951607196643278)	10081.853191575563 iter/sec (stddev: 0.00010024060490679263)	2.20
tests/utils/test_bls.py::test_ants_to_bls[min=65536-len=10000]	1040.8907637089126 iter/sec (stddev: 0.0004308662912820643)	2626.1443811390286 iter/sec (stddev: 0.00004557742609505649)	2.52
tests/utils/test_bls.py::test_ants_to_bls[min=4259840-len=100]	6804.738969574004 iter/sec (stddev: 0.0002964374977119043)	15497.255840555703 iter/sec (stddev: 0.00038930621398815333)	2.28
tests/utils/test_bls.py::test_ants_to_bls[min=4259840-len=10000]	1382.0636218817967 iter/sec (stddev: 0.00019983296867285546)	3000.14029555561 iter/sec (stddev: 0.00018125242381257953)	2.17
tests/uvdata/test_mwa_corr_fits.py::test_van_vleck[cheby=False]	0.06298907376585466 iter/sec (stddev: 1.5100134105019933)	0.13610131863522423 iter/sec (stddev: 0.4555432816545101)	2.16

This comment was automatically generated by workflow using github-action-benchmark.

[steven-murray]

Two small comments, but otherwise approved.

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite	Current: 8b6ee44	Previous: 72dd3bc	Ratio
tests/uvdata/test_mwa_corr_fits.py::test_read_mwax	3.2821206155192035 iter/sec (stddev: 0.14039880006939515)	7.004348042731254 iter/sec (stddev: 0.013840146951652157)	2.13

This comment was automatically generated by workflow using github-action-benchmark.

[steven-murray]

Wonderful, thanks @bhazelton --- this was a lot of work and will be so useful!