Better support for FITS TOAST

[pkgw]

Now that the engine can render TOAST FITS, Toasty should get some better support for creating FITS TOAST data sets.

As of #79 and #80 we can create FITS TOAST pyramids from HEALPix files, which is a good start. But we should also plug into the general FITS tiling framework that's been built. Ideally we should be able to automagically tile into TOAST when needed instead of using hipsgen.

The base need here is a "sampler" for sampling into TOAST given a generic data+WCS image. I have first-draft code to do that, pasted below. One thing that hasn't been worked out super well is that most of the TOAST sampling code assumes that there will never be any masked pixels, but that's not something that we can/should assume for a WCS input. Unfortunately the sampling framework doesn't have a convenient way for plugging into the existing image.py tools for managing image formats and how to mask.

On top of that, I think we need some code to do "filtered" TOAST sampling with an input ImageCollection. This would basically be doing the same thing as the hacky ChunkedPlateCareeSampler code:

toasty/toasty/samplers.py

Line 434 in 5054088

class ChunkedPlateCarreeSampler(object):

See also:

toasty/toasty/tests/test_toast.py

Line 232 in 5054088

def test_earth_plate_carree_jpeg2000_chunked_planetary(self):

I wrote this to deal with a special JPEG2000 input data set. The code is basically recreating much of the infrastructure associated with the ImageCollection framework, and it would probably be best to centralize it. In particular, the latlon_tile_filter implementation was quite tricky to get right and the best approach is almost definitely to reuse it for FITS WCS tiling as well. The JPEG2000 stuff — and in fact the plate carrée samplers in general — could be handled using ImageCollection and WCS ... I've just never bothered to go to the effort of expressing the plate carrée coordinate systems in a WCS data structure.

As with the multi-WCS FITS sampler, you run into issues regarding what to do when multiple inputs overlap each other; and as with the existing implementation, I think our best bet is to just have various inputs overwrite each other and worry about the details later.

CC @imbasimba — hopefully we can talk about this more on Monday.

---

Here's the draft WCS sampler function for toasty/samplers.py:

def wcs_sampler(data, wcs):
    """Create a sampler for data with a generic WCS specification.

    Parameters
    ----------
    data : array
        The image data
    wcs : :class:`astropy.wcs.WCS`
        A WCS data structure describing the data.

    Returns
    -------
    A function that samples the data; the call signature is ``vec2pix(lon, lat)
    -> data``, where the inputs and output are 2D arrays and *lon* and *lat* are
    in radians.

    """
    from astropy.coordinates import SkyCoord
    from astropy import units as u

    def vec2pix(lon, lat):
        c = SkyCoord(
            lon * u.rad,
            lat * u.rad,
            frame='icrs'
        )
        idx = wcs.world_to_array_index(c)

        # Flag out-of-bounds pixels. Currently (April 2022) world_to_array_index
        # returns a tuple of lists, so we need to array-ify.

        idx = tuple(np.asarray(i) for i in idx)
        bad = (idx[0] < 0) | (idx[0] >= data.shape[0])

        for i, i_idx in enumerate(idx[1:]):
            bad |= (i_idx < 0)
            bad |= (i_idx >= data.shape[i+1]) # note that i is off by 1 here

        for i_idx in idx:
            i_idx[bad] = 0

        samp = data[idx]
        samp[bad] = np.nan
        return samp

    return vec2pix