HERA-Team / hera_pspec

HERA power spectrum estimation code and data formats
http://hera-pspec.readthedocs.io/en/latest/
BSD 3-Clause "New" or "Revised" License
5 stars 3 forks source link

Implement psuedostokes that uses pol_convention appropriately #403

Closed jsdillon closed 1 month ago

jsdillon commented 1 month ago

When https://github.com/RadioAstronomySoftwareGroup/pyuvdata/pull/1455 is merged in, UVData and UVCal objects will know about the pol_convention which is either 'sum' or 'avg'. hera_pspec currently assumes the average convention, but we're trying to move to sum as a collaboration.

https://github.com/HERA-Team/hera_pspec/blob/main/hera_pspec/pstokes.py needs to be modified to be responsive to the pol convention. I think it's fine to keep using the average convention when no pol_convention is provided, but there should definitely be a warning. If uvd1 and uvd2 disagree on the pol_convention, that should definitely be an error.

Also, when we do this, I'd like there to be a lower-level function (or set of functions) that is called by _combine_pol() and that operates on numpy arrays, rather than UVData objects. Right now, in https://github.com/HERA-Team/hera_notebook_templates/blob/master/notebooks/single_baseline_postprocessing_and_pspec.ipynb this is done with the following function:

def form_pseudostokes(data=None, flags=None, nsamples=None):
    '''This function creates pseudo-Stokes I and Q versions corresponding to the ee and nn entries
    in the input data containers, and stores them in the same data containers.'''
    # data are averaged (or diffed in the case of pQ)
    if data is not None:
        for antpair in data.antpairs():
            data[antpair + ('pI',)] = (data[antpair + ('ee',)] + data[antpair + ('nn',)]) / 2
            data[antpair + ('pQ',)] = (data[antpair + ('ee',)] - data[antpair + ('nn',)]) / 2

    # flags are ORed
    if flags is not None:
        for antpair in flags.antpairs():            
            flags[antpair + ('pI',)] = flags[antpair + ('ee',)] | flags[antpair + ('nn',)]
            flags[antpair + ('pQ',)] = flags[antpair + ('ee',)] | flags[antpair + ('nn',)]

    # nsamples are combined to produce the correct variance for pI and pQ. 
    # In the limit where the two pols have the same nsamples, it's just the sum.
    if nsamples is not None:
        for antpair in nsamples.antpairs():
            nsamples[antpair + ('pI',)] = 4 * (nsamples[antpair + ('ee',)]**-1 + nsamples[antpair + ('nn',)]**-1)**-1
            nsamples[antpair + ('pQ',)] = 4 * (nsamples[antpair + ('ee',)]**-1 + nsamples[antpair + ('nn',)]**-1)**-1

This needs to be updated to be pol_covention-aware too. It also needs to be able to do psuedo U and pseudo V as soon as cross-pols are available. But ideally the lower-level functions in hera_pspec would be callable here so that we can avoid code duplication.