Analytic Beam refactoring

[steven-murray]

I find the implementation of AnalyticBeam extremely restrictive and difficult to use for a number of reasons. I'll try to outline my ideas here, and open up discussion for ideas for improvements.

My frustrations lie along two major axes, both of which come from trying to use them in hera_sim, in the context of our implementations of PolyBeam and co.

1. It is not easily extendible. The AnalyticBeam class gets a keyword which tells it which kind of beam to evaluate (Gaussian or Airy etc.). New beams can inherit from AnalyticBeam, but it kinda breaks the way that it is meant to be used. The PolyBeam for instance can inherit the AnalyticBeam, but it is a new beam class entirely, and not specified via a new string option. There's no way for new code outside pyuvsim to "register" new analytic beams. This also means that for instance the simulation config code in pyuvsim is restricted to reading the actual AnalyticBeam class and the options it contains, rather than being able to read new derived beams as well. This is a problem for us in hera_sim.
2. The API for AnalyticBeam does not follow the UVBeam. Many of the methods/attributes in UVBeam are not available in AnalyticBeam. Some of these make sense since they're just not applicable for an analytic beam, but others would be helpful. For instance, the Nfeeds attribute makes sense in both contexts, but is not available for the AnalyticBeam. This means wrapper code that has to be able to deal with both has to make a lot of if/else branches to deal with each, which makes code brittle. As an example, the AnalyticBeam has no select() method.

There are a few ways to deal with this. I suggest something like the following API:

class BeamModel:
    def __init__(self, **beam_specific_kwargs):
        ... set all kwargs ...

    def __subclass_hook__(cls):
       <add subclass to global dict of beam models>

    def __call__(self, az, za,...):
        <compute the beam at az/za>

class AnalyticBeam:
    def __init__(self, beam_model: BeamModel, **other_generic_parameters):
        self.beam_model = beam_model
        ...

    def interp(self, ...):
        return self.beam_model(az, za)

    def other_methods_that_ensure_api_consistency(...):
        pass

class Gaussian(BeamModel):
    def __init__(self, width):
        self.width = width

    def __call__(self, az, za, ....):
       return np.exp(...)

In fact, with this kind of API it's VERY easy to add a YAML constructor for the BeamModel, so that you can specify a beam like this in YAML:

beam: !Gaussian
    width: 1.0

[piyanatk]

I do not think I understand how @steven-murray suggestion work completely, but two questions ...

Will this require defining a new yaml constructor for every new subclass of BeamModel?

And here a subclass of BeamModel does not have to necessary be an analytic representation of the beam as long as you can define the __call__ right? So how would you put the line between analytic vs non-analytic?

[steven-murray]

@piyanatk:

1. No, we can write the YAML constructor/representer as part of the base class (or as its own function).
2. Not sure if we really want to draw a sharp line between analytic and non-analytic. We want them to look the same!

[steven-murray]

Another useful thing to add to the AnalyticBeam here would be a to_uvbeam() method, to which you could pass an actual set of coordinates, and it would evaluate the beam there and create a new UVBeam object that knows nothing about the underlying analytic function, and has to do interpolation. This would be very useful for testing things.

[bhazelton]

We discussed this on a telecon on 8/15/22. Some thoughts:

• Inheritance feels a little wrong -- UVBeam needs to represent a bunch of arrays that are just not needed on an analytic beam. For those arrays, UVBeam really needs to inherit from UVBase, it's not clear if Analytic beam needs to inherit from UVBase at all.
• Look into composition. One idea would be to have UVBeam contain all the arrays and the code related to manipulating them (e.g. interpolation), then have another object that has methods that can either call that or the evaluation method for analytic beams. UVBeam may not need to change much at all for this.
• Analytic beams need to support polarized efield beams as well as power beams, so some representation of polarization will be needed both on the Analytic Beam object and probably on the higher level object.

[bhazelton]

This is being done in pyuvdata, see RadioAstronomySoftwareGroup/pyuvdata#1383