isolatedStarAssociation.py

Go to the documentation of this file.

# This file is part of pipe_tasks.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (https://www.lsst.org).
# See the COPYRIGHT file at the top-level directory of this distribution
# for details of code ownership.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
__all__ = ['IsolatedStarAssociationConnections',
           'IsolatedStarAssociationConfig',
           'IsolatedStarAssociationTask']
 
import esutil
import hpgeom as hpg
import numpy as np
import pandas as pd
from smatch.matcher import Matcher
 
import lsst.pex.config as pexConfig
import lsst.pipe.base as pipeBase
from lsst.skymap import BaseSkyMap
from lsst.meas.algorithms.sourceSelector import sourceSelectorRegistry
 
from .healSparseMapping import _is_power_of_two
 
 
class IsolatedStarAssociationConnections(pipeBase.PipelineTaskConnections,
                                         dimensions=('instrument', 'tract', 'skymap',),
                                         defaultTemplates={}):
    source_table_visit = pipeBase.connectionTypes.Input(
        doc='Source table in parquet format, per visit',
        name='preSourceTable_visit',
        storageClass='DataFrame',
        dimensions=('instrument', 'visit'),
        deferLoad=True,
        multiple=True,
    )
    skymap = pipeBase.connectionTypes.Input(
        doc="Input definition of geometry/bbox and projection/wcs for warped exposures",
        name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
        storageClass='SkyMap',
        dimensions=('skymap',),
    )
    isolated_star_sources = pipeBase.connectionTypes.Output(
        doc='Catalog of individual sources for the isolated stars',
        name='isolated_star_presources',
        storageClass='DataFrame',
        dimensions=('instrument', 'tract', 'skymap'),
    )
    isolated_star_cat = pipeBase.connectionTypes.Output(
        doc='Catalog of isolated star positions',
        name='isolated_star_presource_associations',
        storageClass='DataFrame',
        dimensions=('instrument', 'tract', 'skymap'),
    )
 
 
class IsolatedStarAssociationConfig(pipeBase.PipelineTaskConfig,
                                    pipelineConnections=IsolatedStarAssociationConnections):
    """Configuration for IsolatedStarAssociationTask."""
 
    inst_flux_field = pexConfig.Field(
        doc=('Full name of instFlux field to use for s/n selection and persistence. '
             'The associated flag will be implicity included in bad_flags. '
             'Note that this is expected to end in ``instFlux``.'),
        dtype=str,
        default='normCompTophatFlux_instFlux',
    )
    match_radius = pexConfig.Field(
        doc='Match radius (arcseconds)',
        dtype=float,
        default=1.0,
    )
    nside_split = pexConfig.Field(
        doc="HealPix nside to use for splitting tract for reduced memory. Must be power of 2.",
        dtype=int,
        default=32,
        check=_is_power_of_two,
    )
    isolation_radius = pexConfig.Field(
        doc=('Isolation radius (arcseconds).  Any stars with average centroids '
             'within this radius of another star will be rejected from the final '
             'catalog.  This radius should be at least 2x match_radius.'),
        dtype=float,
        default=2.0,
    )
    band_order = pexConfig.ListField(
        doc=(('Ordered list of bands to use for matching/storage. '
              'Any bands not listed will not be matched.')),
        dtype=str,
        default=['i', 'z', 'r', 'g', 'y', 'u'],
    )
    id_column = pexConfig.Field(
        doc='Name of column with source id.',
        dtype=str,
        default='sourceId',
    )
    ra_column = pexConfig.Field(
        doc='Name of column with right ascension.',
        dtype=str,
        default='ra',
    )
    dec_column = pexConfig.Field(
        doc='Name of column with declination.',
        dtype=str,
        default='dec',
    )
    physical_filter_column = pexConfig.Field(
        doc='Name of column with physical filter name',
        dtype=str,
        default='physical_filter',
    )
    band_column = pexConfig.Field(
        doc='Name of column with band name',
        dtype=str,
        default='band',
    )
    extra_columns = pexConfig.ListField(
        doc='Extra names of columns to read and persist (beyond instFlux and error).',
        dtype=str,
        default=['x',
                 'y',
                 'xErr',
                 'yErr',
                 'apFlux_12_0_instFlux',
                 'apFlux_12_0_instFluxErr',
                 'apFlux_12_0_flag',
                 'apFlux_17_0_instFlux',
                 'apFlux_17_0_instFluxErr',
                 'apFlux_17_0_flag',
                 'localBackground_instFlux',
                 'localBackground_flag',
                 'ixx',
                 'iyy',
                 'ixy',]
    )
    source_selector = sourceSelectorRegistry.makeField(
        doc='How to select sources.  Under normal usage this should not be changed.',
        default='science'
    )
 
    def setDefaults(self):
        super().setDefaults()
 
        source_selector = self.source_selector['science']
        source_selector.setDefaults()
 
        source_selector.doFlags = True
        source_selector.doUnresolved = True
        source_selector.doSignalToNoise = True
        source_selector.doIsolated = True
        source_selector.doRequireFiniteRaDec = True
        source_selector.doRequirePrimary = True
 
        source_selector.signalToNoise.minimum = 8.0
        source_selector.signalToNoise.maximum = 1000.0
 
        flux_flag_name = self.inst_flux_field.replace("instFlux", "flag")
 
        source_selector.flags.bad = ['pixelFlags_edge',
                                     'pixelFlags_interpolatedCenter',
                                     'pixelFlags_saturatedCenter',
                                     'pixelFlags_crCenter',
                                     'pixelFlags_bad',
                                     'pixelFlags_interpolated',
                                     'pixelFlags_saturated',
                                     'centroid_flag',
                                     flux_flag_name]
 
        source_selector.signalToNoise.fluxField = self.inst_flux_field
        source_selector.signalToNoise.errField = self.inst_flux_field + 'Err'
 
        source_selector.isolated.parentName = 'parentSourceId'
        source_selector.isolated.nChildName = 'deblend_nChild'
 
        source_selector.unresolved.name = 'sizeExtendedness'
 
        source_selector.requireFiniteRaDec.raColName = self.ra_column
        source_selector.requireFiniteRaDec.decColName = self.dec_column
 
 
class IsolatedStarAssociationTask(pipeBase.PipelineTask):
    """Associate sources into isolated star catalogs.
    """
    ConfigClass = IsolatedStarAssociationConfig
    _DefaultName = 'isolatedStarAssociation'
 
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
 
        self.makeSubtask('source_selector')
        # Only log warning and fatal errors from the source_selector
        self.source_selector.log.setLevel(self.source_selector.log.WARN)
 
    def runQuantum(self, butlerQC, inputRefs, outputRefs):
        input_ref_dict = butlerQC.get(inputRefs)
 
        tract = butlerQC.quantum.dataId['tract']
 
        source_table_refs = input_ref_dict['source_table_visit']
 
        self.log.info('Running with %d source_table_visit dataRefs',
                      len(source_table_refs))
 
        source_table_ref_dict_temp = {source_table_ref.dataId['visit']: source_table_ref for
                                      source_table_ref in source_table_refs}
 
        bands = {source_table_ref.dataId['band'] for source_table_ref in source_table_refs}
        for band in bands:
            if band not in self.config.band_order:
                self.log.warning('Input data has data from band %s but that band is not '
                                 'configured for matching', band)
 
        # TODO: Sort by visit until DM-31701 is done and we have deterministic
        # dataset ordering.
        source_table_ref_dict = {visit: source_table_ref_dict_temp[visit] for
                                 visit in sorted(source_table_ref_dict_temp.keys())}
 
        struct = self.run(input_ref_dict['skymap'], tract, source_table_ref_dict)
 
        butlerQC.put(pd.DataFrame(struct.star_source_cat),
                     outputRefs.isolated_star_sources)
        butlerQC.put(pd.DataFrame(struct.star_cat),
                     outputRefs.isolated_star_cat)
 
    def run(self, skymap, tract, source_table_ref_dict):
        """Run the isolated star association task.
 
        Parameters
        ----------
        skymap : `lsst.skymap.SkyMap`
            Skymap object.
        tract : `int`
            Tract number.
        source_table_ref_dict : `dict`
            Dictionary of source_table refs.  Key is visit, value is dataref.
 
        Returns
        -------
        struct : `lsst.pipe.base.struct`
            Struct with outputs for persistence.
        """
        star_source_cat = self._make_all_star_sources(skymap[tract], source_table_ref_dict)
 
        primary_bands = self.config.band_order
 
        # Do primary matching
        primary_star_cat = self._match_primary_stars(primary_bands, star_source_cat)
 
        if len(primary_star_cat) == 0:
            return pipeBase.Struct(star_source_cat=np.zeros(0, star_source_cat.dtype),
                                   star_cat=np.zeros(0, primary_star_cat.dtype))
 
        # Remove neighbors
        primary_star_cat = self._remove_neighbors(primary_star_cat)
 
        if len(primary_star_cat) == 0:
            return pipeBase.Struct(star_source_cat=np.zeros(0, star_source_cat.dtype),
                                   star_cat=np.zeros(0, primary_star_cat.dtype))
 
        # Crop to inner tract region
        inner_tract_ids = skymap.findTractIdArray(primary_star_cat[self.config.ra_column],
                                                  primary_star_cat[self.config.dec_column],
                                                  degrees=True)
        use = (inner_tract_ids == tract)
        self.log.info('Total of %d isolated stars in inner tract.', use.sum())
 
        primary_star_cat = primary_star_cat[use]
 
        if len(primary_star_cat) == 0:
            return pipeBase.Struct(star_source_cat=np.zeros(0, star_source_cat.dtype),
                                   star_cat=np.zeros(0, primary_star_cat.dtype))
 
        # Set the unique ids.
        primary_star_cat['isolated_star_id'] = self._compute_unique_ids(skymap,
                                                                        tract,
                                                                        len(primary_star_cat))
 
        # Match to sources.
        star_source_cat, primary_star_cat = self._match_sources(primary_bands,
                                                                star_source_cat,
                                                                primary_star_cat)
 
        return pipeBase.Struct(star_source_cat=star_source_cat,
                               star_cat=primary_star_cat)
 
    def _make_all_star_sources(self, tract_info, source_table_ref_dict):
        """Make a catalog of all the star sources.
 
        Parameters
        ----------
        tract_info : `lsst.skymap.TractInfo`
            Information about the tract.
        source_table_ref_dict : `dict`
            Dictionary of source_table refs.  Key is visit, value is dataref.
 
        Returns
        -------
        star_source_cat : `np.ndarray`
            Catalog of star sources.
        """
        # Internally, we use a numpy recarray, they are by far the fastest
        # option in testing for relatively narrow tables.
        # (have not tested wide tables)
        all_columns, persist_columns = self._get_source_table_visit_column_names()
        poly = tract_info.outer_sky_polygon
 
        tables = []
        for visit in source_table_ref_dict:
            source_table_ref = source_table_ref_dict[visit]
            df = source_table_ref.get(parameters={'columns': all_columns})
            df.reset_index(inplace=True)
 
            goodSrc = self.source_selector.selectSources(df)
 
            table = df[persist_columns][goodSrc.selected].to_records()
 
            # Append columns that include the row in the source table
            # and the matched object index (to be filled later).
            table = np.lib.recfunctions.append_fields(table,
                                                      ['source_row',
                                                       'obj_index'],
                                                      [np.where(goodSrc.selected)[0],
                                                       np.zeros(goodSrc.selected.sum(), dtype=np.int32)],
                                                      dtypes=['i4', 'i4'],
                                                      usemask=False)
 
            # We cut to the outer tract polygon to ensure consistent matching
            # from tract to tract.
            tract_use = poly.contains(np.deg2rad(table[self.config.ra_column]),
                                      np.deg2rad(table[self.config.dec_column]))
 
            tables.append(table[tract_use])
 
        # Combine tables
        star_source_cat = np.concatenate(tables)
 
        return star_source_cat
 
    def _get_source_table_visit_column_names(self):
        """Get the list of sourceTable_visit columns from the config.
 
        Returns
        -------
        all_columns : `list` [`str`]
            All columns to read
        persist_columns : `list` [`str`]
            Columns to persist (excluding selection columns)
        """
        columns = [self.config.id_column,
                   'visit', 'detector',
                   self.config.ra_column, self.config.dec_column,
                   self.config.physical_filter_column, self.config.band_column,
                   self.config.inst_flux_field, self.config.inst_flux_field + 'Err']
        columns.extend(self.config.extra_columns)
 
        all_columns = columns.copy()
        if self.source_selector.config.doFlags:
            all_columns.extend(self.source_selector.config.flags.bad)
        if self.source_selector.config.doUnresolved:
            all_columns.append(self.source_selector.config.unresolved.name)
        if self.source_selector.config.doIsolated:
            all_columns.append(self.source_selector.config.isolated.parentName)
            all_columns.append(self.source_selector.config.isolated.nChildName)
        if self.source_selector.config.doRequirePrimary:
            all_columns.append(self.source_selector.config.requirePrimary.primaryColName)
 
        return all_columns, columns
 
    def _match_primary_stars(self, primary_bands, star_source_cat):
        """Match primary stars.
 
        Parameters
        ----------
        primary_bands : `list` [`str`]
            Ordered list of primary bands.
        star_source_cat : `np.ndarray`
            Catalog of star sources.
 
        Returns
        -------
        primary_star_cat : `np.ndarray`
            Catalog of primary star positions
        """
        ra_col = self.config.ra_column
        dec_col = self.config.dec_column
 
        dtype = self._get_primary_dtype(primary_bands)
 
        # Figure out nside for computing boundary pixels.
        # We want to be 10x the match radius to be conservative.
        nside_split = self.config.nside_split
        nside_boundary = nside_split
        res = hpg.nside_to_resolution(nside_boundary)*3600.
        while res > 10.*self.config.match_radius:
            nside_boundary *= 2
            res = hpg.nside_to_resolution(nside_boundary)*3600.
 
        # Cancel out last jump.
        nside_boundary //= 2
 
        # Compute healpix nest bit shift between resolutions.
        bit_shift = 2*int(np.round(np.log2(nside_boundary / nside_split)))
 
        primary_star_cat = None
        for primary_band in primary_bands:
            use = (star_source_cat['band'] == primary_band)
 
            ra = star_source_cat[ra_col][use]
            dec = star_source_cat[dec_col][use]
 
            hpix_split = np.unique(hpg.angle_to_pixel(nside_split, ra, dec))
            hpix_highres = hpg.angle_to_pixel(nside_boundary, ra, dec)
 
            band_cats = []
            for pix in hpix_split:
                self.log.info("Matching primary stars in %s band in pixel %d", primary_band, pix)
                # We take all the high resolution pixels in this pixel, and
                # add in all the boundary pixels to ensure objects do not get split.
                pixels_highres = np.left_shift(pix, bit_shift) + np.arange(2**bit_shift)
                pixels_highres = np.unique(hpg.neighbors(nside_boundary, pixels_highres).ravel())
                a, b = esutil.numpy_util.match(pixels_highres, hpix_highres)
 
                _ra = ra[b]
                _dec = dec[b]
 
                with Matcher(_ra, _dec) as matcher:
                    idx = matcher.query_groups(self.config.match_radius/3600., min_match=1)
 
                count = len(idx)
 
                if count == 0:
                    continue
 
                band_cat = np.zeros(count, dtype=dtype)
                band_cat['primary_band'] = primary_band
 
                # If the tract cross ra=0 (that is, it has both low ra and high ra)
                # then we need to remap all ra values from [0, 360) to [-180, 180)
                # before doing any position averaging.
                remapped = False
                if _ra.min() < 60.0 and _ra.max() > 300.0:
                    _ra_temp = (_ra + 180.0) % 360. - 180.
                    remapped = True
                else:
                    _ra_temp = _ra
 
                # Compute mean position for each primary star
                for i, row in enumerate(idx):
                    row = np.array(row)
                    band_cat[ra_col][i] = np.mean(_ra_temp[row])
                    band_cat[dec_col][i] = np.mean(_dec[row])
 
                if remapped:
                    # Remap ra back to [0, 360)
                    band_cat[ra_col] %= 360.0
 
                # And cut down to objects that are not in the boundary region
                # after association.
                inpix = (hpg.angle_to_pixel(nside_split, band_cat[ra_col], band_cat[dec_col]) == pix)
                if inpix.sum() > 0:
                    band_cats.append(band_cat[inpix])
 
            if len(band_cats) == 0:
                self.log.info('Found 0 primary stars in %s band.', primary_band)
                continue
 
            band_cat = np.concatenate(band_cats)
 
            # Match to previous band catalog(s), and remove duplicates.
            if primary_star_cat is None or len(primary_star_cat) == 0:
                primary_star_cat = band_cat
            else:
                with Matcher(band_cat[ra_col], band_cat[dec_col]) as matcher:
                    idx = matcher.query_radius(primary_star_cat[ra_col],
                                               primary_star_cat[dec_col],
                                               self.config.match_radius/3600.)
                # Any object with a match should be removed.
                match_indices = np.array([i for i in range(len(idx)) if len(idx[i]) > 0])
                if len(match_indices) > 0:
                    band_cat = np.delete(band_cat, match_indices)
 
                primary_star_cat = np.append(primary_star_cat, band_cat)
            self.log.info('Found %d primary stars in %s band.', len(band_cat), primary_band)
 
        # If everything was cut, we still want the correct datatype.
        if primary_star_cat is None:
            primary_star_cat = np.zeros(0, dtype=dtype)
 
        return primary_star_cat
 
    def _remove_neighbors(self, primary_star_cat):
        """Remove neighbors from the primary star catalog.
 
        Parameters
        ----------
        primary_star_cat : `np.ndarray`
            Primary star catalog.
 
        Returns
        -------
        primary_star_cat_cut : `np.ndarray`
            Primary star cat with neighbors removed.
        """
        ra_col = self.config.ra_column
        dec_col = self.config.dec_column
 
        with Matcher(primary_star_cat[ra_col], primary_star_cat[dec_col]) as matcher:
            # By setting min_match=2 objects that only match to themselves
            # will not be recorded.
            try:
                # New smatch API
                idx = matcher.query_groups(self.config.isolation_radius/3600., min_match=2)
            except AttributeError:
                # Old smatch API
                idx = matcher.query_self(self.config.isolation_radius/3600., min_match=2)
 
        try:
            neighbor_indices = np.concatenate(idx)
        except ValueError:
            neighbor_indices = np.zeros(0, dtype=int)
 
        if len(neighbor_indices) > 0:
            neighbored = np.unique(neighbor_indices)
            self.log.info('Cutting %d objects with close neighbors.', len(neighbored))
            primary_star_cat = np.delete(primary_star_cat, neighbored)
 
        return primary_star_cat
 
    def _match_sources(self, bands, star_source_cat, primary_star_cat):
        """Match individual sources to primary stars.
 
        Parameters
        ----------
        bands : `list` [`str`]
            List of bands.
        star_source_cat : `np.ndarray`
            Array of star sources.
        primary_star_cat : `np.ndarray`
            Array of primary stars.
 
        Returns
        -------
        star_source_cat_sorted : `np.ndarray`
            Sorted and cropped array of star sources.
        primary_star_cat : `np.ndarray`
            Catalog of isolated stars, with indexes to star_source_cat_cut.
        """
        ra_col = self.config.ra_column
        dec_col = self.config.dec_column
 
        # We match sources per-band because it allows us to have sorted
        # sources for easy retrieval of per-band matches.
        n_source_per_band_per_obj = np.zeros((len(bands),
                                              len(primary_star_cat)),
                                             dtype=np.int32)
        band_uses = []
        idxs = []
        with Matcher(primary_star_cat[ra_col], primary_star_cat[dec_col]) as matcher:
            for b, band in enumerate(bands):
                band_use, = np.where(star_source_cat['band'] == band)
 
                idx = matcher.query_radius(star_source_cat[ra_col][band_use],
                                           star_source_cat[dec_col][band_use],
                                           self.config.match_radius/3600.)
                n_source_per_band_per_obj[b, :] = np.array([len(row) for row in idx])
                idxs.append(idx)
                band_uses.append(band_use)
 
        n_source_per_obj = np.sum(n_source_per_band_per_obj, axis=0)
 
        primary_star_cat['nsource'] = n_source_per_obj
        primary_star_cat['source_cat_index'][1:] = np.cumsum(n_source_per_obj)[:-1]
 
        n_tot_source = primary_star_cat['source_cat_index'][-1] + primary_star_cat['nsource'][-1]
 
        # Temporary arrays until we crop/sort the source catalog
        source_index = np.zeros(n_tot_source, dtype=np.int32)
        obj_index = np.zeros(n_tot_source, dtype=np.int32)
 
        ctr = 0
        for i in range(len(primary_star_cat)):
            obj_index[ctr: ctr + n_source_per_obj[i]] = i
            for b in range(len(bands)):
                source_index[ctr: ctr + n_source_per_band_per_obj[b, i]] = band_uses[b][idxs[b][i]]
                ctr += n_source_per_band_per_obj[b, i]
 
        source_cat_index_band_offset = np.cumsum(n_source_per_band_per_obj, axis=0)
 
        for b, band in enumerate(bands):
            primary_star_cat[f'nsource_{band}'] = n_source_per_band_per_obj[b, :]
            if b == 0:
                # The first band listed is the same as the overall star
                primary_star_cat[f'source_cat_index_{band}'] = primary_star_cat['source_cat_index']
            else:
                # Other band indices are offset from the previous band
                primary_star_cat[f'source_cat_index_{band}'] = (primary_star_cat['source_cat_index']
                                                                + source_cat_index_band_offset[b - 1, :])
 
        star_source_cat = star_source_cat[source_index]
        star_source_cat['obj_index'] = obj_index
 
        return star_source_cat, primary_star_cat
 
    def _compute_unique_ids(self, skymap, tract, nstar):
        """Compute unique star ids.
 
        This is a simple hash of the tract and star to provide an
        id that is unique for a given processing.
 
        Parameters
        ----------
        skymap : `lsst.skymap.Skymap`
            Skymap object.
        tract : `int`
            Tract id number.
        nstar : `int`
            Number of stars.
 
        Returns
        -------
        ids : `np.ndarray`
            Array of unique star ids.
        """
        # The end of the id will be big enough to hold the tract number
        mult = 10**(int(np.log10(len(skymap))) + 1)
 
        return (np.arange(nstar) + 1)*mult + tract
 
    def _get_primary_dtype(self, primary_bands):
        """Get the numpy datatype for the primary star catalog.
 
        Parameters
        ----------
        primary_bands : `list` [`str`]
            List of primary bands.
 
        Returns
        -------
        dtype : `numpy.dtype`
            Datatype of the primary catalog.
        """
        max_len = max([len(primary_band) for primary_band in primary_bands])
 
        dtype = [('isolated_star_id', 'i8'),
                 (self.config.ra_column, 'f8'),
                 (self.config.dec_column, 'f8'),
                 ('primary_band', f'U{max_len}'),
                 ('source_cat_index', 'i4'),
                 ('nsource', 'i4')]
 
        for band in primary_bands:
            dtype.append((f'source_cat_index_{band}', 'i4'))
            dtype.append((f'nsource_{band}', 'i4'))
 
        return dtype