astro py:docs

- Index
- Modules

### Navigation

- Astropy v5.2.1 »
- Module code »

# Source code for astropy.coordinates.matching

    # Licensed under a 3-clause BSD style license - see LICENSE.rst

    """
    This module contains functions for matching coordinate catalogs.
    """

    import numpy as np

    from astropy import units as u

    from . import Angle
    from .representation import UnitSphericalRepresentation
    from .sky_coordinate import SkyCoord

    __all__ = [
        "match_coordinates_3d",
        "match_coordinates_sky",
        "search_around_3d",
        "search_around_sky",
    ]

    def match_coordinates_3d(
        matchcoord, catalogcoord, nthneighbor=1, storekdtree="kdtree_3d"
    ):
        """
        Finds the nearest 3-dimensional matches of a coordinate or coordinates in
        a set of catalog coordinates.

        This finds the 3-dimensional closest neighbor, which is only different
        from the on-sky distance if ``distance`` is set in either ``matchcoord``
        or ``catalogcoord``.

        Parameters
        ----------
        matchcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
            The coordinate(s) to match to the catalog.
        catalogcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
            The base catalog in which to search for matches. Typically this will
            be a coordinate object that is an array (i.e.,
            ``catalogcoord.isscalar == False``)
        nthneighbor : int, optional
            Which closest neighbor to search for.  Typically ``1`` is desired here,
            as that is correct for matching one set of coordinates to another.
            The next likely use case is ``2``, for matching a coordinate catalog
            against *itself* (``1`` is inappropriate because each point will find
            itself as the closest match).
        storekdtree : bool or str, optional
            If a string, will store the KD-Tree used for the computation
            in the ``catalogcoord``, as in ``catalogcoord.cache`` with the
            provided name.  This dramatically speeds up subsequent calls with the
            same catalog. If False, the KD-Tree is discarded after use.

        Returns
        -------
        idx : int array
            Indices into ``catalogcoord`` to get the matched points for each
            ``matchcoord``. Shape matches ``matchcoord``.
        sep2d : `~astropy.coordinates.Angle`
            The on-sky separation between the closest match for each ``matchcoord``
            and the ``matchcoord``. Shape matches ``matchcoord``.
        dist3d : `~astropy.units.Quantity` ['length']
            The 3D distance between the closest match for each ``matchcoord`` and
            the ``matchcoord``. Shape matches ``matchcoord``.

        Notes
        -----
        This function requires `SciPy <https://www.scipy.org/>`_ to be installed
        or it will fail.
        """
        if catalogcoord.isscalar or len(catalogcoord) < 1:
            raise ValueError(
                "The catalog for coordinate matching cannot be a scalar or length-0."
            )

        kdt = _get_cartesian_kdtree(catalogcoord, storekdtree)

        # make sure coordinate systems match
        if isinstance(matchcoord, SkyCoord):
            matchcoord = matchcoord.transform_to(catalogcoord, merge_attributes=False)
        else:
            matchcoord = matchcoord.transform_to(catalogcoord)

        # make sure units match
        catunit = catalogcoord.cartesian.x.unit
        matchxyz = matchcoord.cartesian.xyz.to(catunit)

        matchflatxyz = matchxyz.reshape((3, np.prod(matchxyz.shape) // 3))
        # Querying NaN returns garbage
        if np.isnan(matchflatxyz.value).any():
            raise ValueError("Matching coordinates cannot contain NaN entries.")
        dist, idx = kdt.query(matchflatxyz.T, nthneighbor)

        if nthneighbor > 1:  # query gives 1D arrays if k=1, 2D arrays otherwise
            dist = dist[:, -1]
            idx = idx[:, -1]

        sep2d = catalogcoord[idx].separation(matchcoord)
        return (
            idx.reshape(matchxyz.shape[1:]),
            sep2d,
            dist.reshape(matchxyz.shape[1:]) * catunit,
        )

    def match_coordinates_sky(
        matchcoord, catalogcoord, nthneighbor=1, storekdtree="kdtree_sky"
    ):
        """
        Finds the nearest on-sky matches of a coordinate or coordinates in
        a set of catalog coordinates.

        This finds the on-sky closest neighbor, which is only different from the
        3-dimensional match if ``distance`` is set in either ``matchcoord``
        or ``catalogcoord``.

        Parameters
        ----------
        matchcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
            The coordinate(s) to match to the catalog.
        catalogcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
            The base catalog in which to search for matches. Typically this will
            be a coordinate object that is an array (i.e.,
            ``catalogcoord.isscalar == False``)
        nthneighbor : int, optional
            Which closest neighbor to search for.  Typically ``1`` is desired here,
            as that is correct for matching one set of coordinates to another.
            The next likely use case is ``2``, for matching a coordinate catalog
            against *itself* (``1`` is inappropriate because each point will find
            itself as the closest match).
        storekdtree : bool or str, optional
            If a string, will store the KD-Tree used for the computation
            in the ``catalogcoord`` in ``catalogcoord.cache`` with the
            provided name.  This dramatically speeds up subsequent calls with the
            same catalog. If False, the KD-Tree is discarded after use.

        Returns
        -------
        idx : int array
            Indices into ``catalogcoord`` to get the matched points for each
            ``matchcoord``. Shape matches ``matchcoord``.
        sep2d : `~astropy.coordinates.Angle`
            The on-sky separation between the closest match for each
            ``matchcoord`` and the ``matchcoord``. Shape matches ``matchcoord``.
        dist3d : `~astropy.units.Quantity` ['length']
            The 3D distance between the closest match for each ``matchcoord`` and
            the ``matchcoord``. Shape matches ``matchcoord``.  If either
            ``matchcoord`` or ``catalogcoord`` don't have a distance, this is the 3D
            distance on the unit sphere, rather than a true distance.

        Notes
        -----
        This function requires `SciPy <https://www.scipy.org/>`_ to be installed
        or it will fail.
        """
        if catalogcoord.isscalar or len(catalogcoord) < 1:
            raise ValueError(
                "The catalog for coordinate matching cannot be a scalar or length-0."
            )

        # send to catalog frame
        if isinstance(matchcoord, SkyCoord):
            newmatch = matchcoord.transform_to(catalogcoord, merge_attributes=False)
        else:
            newmatch = matchcoord.transform_to(catalogcoord)

        # strip out distance info
        match_urepr = newmatch.data.represent_as(UnitSphericalRepresentation)
        newmatch_u = newmatch.realize_frame(match_urepr)

        cat_urepr = catalogcoord.data.represent_as(UnitSphericalRepresentation)
        newcat_u = catalogcoord.realize_frame(cat_urepr)

        # Check for a stored KD-tree on the passed-in coordinate. Normally it will
        # have a distinct name from the "3D" one, so it's safe to use even though
        # it's based on UnitSphericalRepresentation.
        storekdtree = catalogcoord.cache.get(storekdtree, storekdtree)

        idx, sep2d, sep3d = match_coordinates_3d(
            newmatch_u, newcat_u, nthneighbor, storekdtree
        )
        # sep3d is *wrong* above, because the distance information was removed,
        # unless one of the catalogs doesn't have a real distance
        if not (
            isinstance(catalogcoord.data, UnitSphericalRepresentation)
            or isinstance(newmatch.data, UnitSphericalRepresentation)
        ):
            sep3d = catalogcoord[idx].separation_3d(newmatch)

        # update the kdtree on the actual passed-in coordinate
        if isinstance(storekdtree, str):
            catalogcoord.cache[storekdtree] = newcat_u.cache[storekdtree]
        elif storekdtree is True:
            # the old backwards-compatible name
            catalogcoord.cache["kdtree"] = newcat_u.cache["kdtree"]

        return idx, sep2d, sep3d

    def search_around_3d(coords1, coords2, distlimit, storekdtree="kdtree_3d"):
        """
        Searches for pairs of points that are at least as close as a specified
        distance in 3D space.

        This is intended for use on coordinate objects with arrays of coordinates,
        not scalars.  For scalar coordinates, it is better to use the
        ``separation_3d`` methods.

        Parameters
        ----------
        coords1 : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
            The first set of coordinates, which will be searched for matches from
            ``coords2`` within ``seplimit``.  Cannot be a scalar coordinate.
        coords2 : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
            The second set of coordinates, which will be searched for matches from
            ``coords1`` within ``seplimit``.  Cannot be a scalar coordinate.
        distlimit : `~astropy.units.Quantity` ['length']
            The physical radius to search within.
        storekdtree : bool or str, optional
            If a string, will store the KD-Tree used in the search with the name
            ``storekdtree`` in ``coords2.cache``. This speeds up subsequent calls
            to this function. If False, the KD-Trees are not saved.

        Returns
        -------
        idx1 : int array
            Indices into ``coords1`` that matches to the corresponding element of
            ``idx2``. Shape matches ``idx2``.
        idx2 : int array
            Indices into ``coords2`` that matches to the corresponding element of
            ``idx1``. Shape matches ``idx1``.
        sep2d : `~astropy.coordinates.Angle`
            The on-sky separation between the coordinates. Shape matches ``idx1``
            and ``idx2``.
        dist3d : `~astropy.units.Quantity` ['length']
            The 3D distance between the coordinates. Shape matches ``idx1`` and
            ``idx2``. The unit is that of ``coords1``.

        Notes
        -----
        This function requires `SciPy <https://www.scipy.org/>`_
        to be installed or it will fail.

        If you are using this function to search in a catalog for matches around
        specific points, the convention is for ``coords2`` to be the catalog, and
        ``coords1`` are the points to search around.  While these operations are
        mathematically the same if ``coords1`` and ``coords2`` are flipped, some of
        the optimizations may work better if this convention is obeyed.

        In the current implementation, the return values are always sorted in the
        same order as the ``coords1`` (so ``idx1`` is in ascending order).  This is
        considered an implementation detail, though, so it could change in a future
        release.
        """
        if not distlimit.isscalar:
            raise ValueError("distlimit must be a scalar in search_around_3d")

        if coords1.isscalar or coords2.isscalar:
            raise ValueError(
                "One of the inputs to search_around_3d is a scalar. search_around_3d is"
                " intended for use with array coordinates, not scalars.  Instead, use"
                " ``coord1.separation_3d(coord2) < distlimit`` to find the coordinates near"
                " a scalar coordinate."
            )

        if len(coords1) == 0 or len(coords2) == 0:
            # Empty array input: return empty match
            return (
                np.array([], dtype=int),
                np.array([], dtype=int),
                Angle([], u.deg),
                u.Quantity([], coords1.distance.unit),
            )

        kdt2 = _get_cartesian_kdtree(coords2, storekdtree)
        cunit = coords2.cartesian.x.unit

        # we convert coord1 to match coord2's frame.  We do it this way
        # so that if the conversion does happen, the KD tree of coord2 at least gets
        # saved. (by convention, coord2 is the "catalog" if that makes sense)
        coords1 = coords1.transform_to(coords2)

        kdt1 = _get_cartesian_kdtree(coords1, storekdtree, forceunit=cunit)

        # this is the *cartesian* 3D distance that corresponds to the given angle
        d = distlimit.to_value(cunit)

        idxs1 = []
        idxs2 = []
        for i, matches in enumerate(kdt1.query_ball_tree(kdt2, d)):
            for match in matches:
                idxs1.append(i)
                idxs2.append(match)
        idxs1 = np.array(idxs1, dtype=int)
        idxs2 = np.array(idxs2, dtype=int)

        if idxs1.size == 0:
            d2ds = Angle([], u.deg)
            d3ds = u.Quantity([], coords1.distance.unit)
        else:
            d2ds = coords1[idxs1].separation(coords2[idxs2])
            d3ds = coords1[idxs1].separation_3d(coords2[idxs2])

        return idxs1, idxs2, d2ds, d3ds

    def search_around_sky(coords1, coords2, seplimit, storekdtree="kdtree_sky"):
        """
        Searches for pairs of points that have an angular separation at least as
        close as a specified angle.

        This is intended for use on coordinate objects with arrays of coordinates,
        not scalars.  For scalar coordinates, it is better to use the ``separation``
        methods.

        Parameters
        ----------
        coords1 : coordinate-like
            The first set of coordinates, which will be searched for matches from
            ``coords2`` within ``seplimit``. Cannot be a scalar coordinate.
        coords2 : coordinate-like
            The second set of coordinates, which will be searched for matches from
            ``coords1`` within ``seplimit``. Cannot be a scalar coordinate.
        seplimit : `~astropy.units.Quantity` ['angle']
            The on-sky separation to search within.
        storekdtree : bool or str, optional
            If a string, will store the KD-Tree used in the search with the name
            ``storekdtree`` in ``coords2.cache``. This speeds up subsequent calls
            to this function. If False, the KD-Trees are not saved.

        Returns
        -------
        idx1 : int array
            Indices into ``coords1`` that matches to the corresponding element of
            ``idx2``. Shape matches ``idx2``.
        idx2 : int array
            Indices into ``coords2`` that matches to the corresponding element of
            ``idx1``. Shape matches ``idx1``.
        sep2d : `~astropy.coordinates.Angle`
            The on-sky separation between the coordinates. Shape matches ``idx1``
            and ``idx2``.
        dist3d : `~astropy.units.Quantity` ['length']
            The 3D distance between the coordinates. Shape matches ``idx1``
            and ``idx2``; the unit is that of ``coords1``.
            If either ``coords1`` or ``coords2`` don't have a distance,
            this is the 3D distance on the unit sphere, rather than a
            physical distance.

        Notes
        -----
        This function requires `SciPy <https://www.scipy.org/>`_
        to be installed or it will fail.

        In the current implementation, the return values are always sorted in the
        same order as the ``coords1`` (so ``idx1`` is in ascending order).  This is
        considered an implementation detail, though, so it could change in a future
        release.
        """
        if not seplimit.isscalar:
            raise ValueError("seplimit must be a scalar in search_around_sky")

        if coords1.isscalar or coords2.isscalar:
            raise ValueError(
                "One of the inputs to search_around_sky is a scalar. search_around_sky is"
                " intended for use with array coordinates, not scalars.  Instead, use"
                " ``coord1.separation(coord2) < seplimit`` to find the coordinates near a"
                " scalar coordinate."
            )

        if len(coords1) == 0 or len(coords2) == 0:
            # Empty array input: return empty match
            if coords2.distance.unit == u.dimensionless_unscaled:
                distunit = u.dimensionless_unscaled
            else:
                distunit = coords1.distance.unit
            return (
                np.array([], dtype=int),
                np.array([], dtype=int),
                Angle([], u.deg),
                u.Quantity([], distunit),
            )

        # we convert coord1 to match coord2's frame.  We do it this way
        # so that if the conversion does happen, the KD tree of coord2 at least gets
        # saved. (by convention, coord2 is the "catalog" if that makes sense)
        coords1 = coords1.transform_to(coords2)

        # strip out distance info
        urepr1 = coords1.data.represent_as(UnitSphericalRepresentation)
        ucoords1 = coords1.realize_frame(urepr1)

        kdt1 = _get_cartesian_kdtree(ucoords1, storekdtree)

        if storekdtree and coords2.cache.get(storekdtree):
            # just use the stored KD-Tree
            kdt2 = coords2.cache[storekdtree]
        else:
            # strip out distance info
            urepr2 = coords2.data.represent_as(UnitSphericalRepresentation)
            ucoords2 = coords2.realize_frame(urepr2)

            kdt2 = _get_cartesian_kdtree(ucoords2, storekdtree)
            if storekdtree:
                # save the KD-Tree in coords2, *not* ucoords2
                coords2.cache["kdtree" if storekdtree is True else storekdtree] = kdt2

        # this is the *cartesian* 3D distance that corresponds to the given angle
        r = (2 * np.sin(Angle(seplimit) / 2.0)).value

        idxs1 = []
        idxs2 = []
        for i, matches in enumerate(kdt1.query_ball_tree(kdt2, r)):
            for match in matches:
                idxs1.append(i)
                idxs2.append(match)
        idxs1 = np.array(idxs1, dtype=int)
        idxs2 = np.array(idxs2, dtype=int)

        if idxs1.size == 0:
            if coords2.distance.unit == u.dimensionless_unscaled:
                distunit = u.dimensionless_unscaled
            else:
                distunit = coords1.distance.unit
            d2ds = Angle([], u.deg)
            d3ds = u.Quantity([], distunit)
        else:
            d2ds = coords1[idxs1].separation(coords2[idxs2])
            try:
                d3ds = coords1[idxs1].separation_3d(coords2[idxs2])
            except ValueError:
                # they don't have distances, so we just fall back on the cartesian
                # distance, computed from d2ds
                d3ds = 2 * np.sin(d2ds / 2.0)

        return idxs1, idxs2, d2ds, d3ds

    def _get_cartesian_kdtree(coord, attrname_or_kdt="kdtree", forceunit=None):
        """
        This is a utility function to retrieve (and build/cache, if necessary)
        a 3D cartesian KD-Tree from various sorts of astropy coordinate objects.

        Parameters
        ----------
        coord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
            The coordinates to build the KD-Tree for.
        attrname_or_kdt : bool or str or KDTree
            If a string, will store the KD-Tree used for the computation in the
            ``coord``, in ``coord.cache`` with the provided name. If given as a
            KD-Tree, it will just be used directly.
        forceunit : unit or None
            If a unit, the cartesian coordinates will convert to that unit before
            being put in the KD-Tree.  If None, whatever unit it's already in
            will be used

        Returns
        -------
        kdt : `~scipy.spatial.cKDTree` or `~scipy.spatial.KDTree`
            The KD-Tree representing the 3D cartesian representation of the input
            coordinates.
        """
        from warnings import warn

        # without scipy this will immediately fail
        from scipy import spatial

        try:
            KDTree = spatial.cKDTree
        except Exception:
            warn(
                "C-based KD tree not found, falling back on (much slower) "
                "python implementation"
            )
            KDTree = spatial.KDTree

        if attrname_or_kdt is True:  # backwards compatibility for pre v0.4
            attrname_or_kdt = "kdtree"

        # figure out where any cached KDTree might be
        if isinstance(attrname_or_kdt, str):
            kdt = coord.cache.get(attrname_or_kdt, None)
            if kdt is not None and not isinstance(kdt, KDTree):
                raise TypeError(
                    f'The `attrname_or_kdt` "{attrname_or_kdt}" is not a scipy KD tree!'
                )
        elif isinstance(attrname_or_kdt, KDTree):
            kdt = attrname_or_kdt
            attrname_or_kdt = None
        elif not attrname_or_kdt:
            kdt = None
        else:
            raise TypeError(
                "Invalid `attrname_or_kdt` argument for KD-Tree:" + str(attrname_or_kdt)
            )

        if kdt is None:
            # need to build the cartesian KD-tree for the catalog
            if forceunit is None:
                cartxyz = coord.cartesian.xyz
            else:
                cartxyz = coord.cartesian.xyz.to(forceunit)
            flatxyz = cartxyz.reshape((3, np.prod(cartxyz.shape) // 3))
            # There should be no NaNs in the kdtree data.
            if np.isnan(flatxyz.value).any():
                raise ValueError("Catalog coordinates cannot contain NaN entries.")
            try:
                # Set compact_nodes=False, balanced_tree=False to use
                # "sliding midpoint" rule, which is much faster than standard for
                # many common use cases
                kdt = KDTree(flatxyz.value.T, compact_nodes=False, balanced_tree=False)
            except TypeError:
                # Python implementation does not take compact_nodes and balanced_tree
                # as arguments.  However, it uses sliding midpoint rule by default
                kdt = KDTree(flatxyz.value.T)

        if attrname_or_kdt:
            # cache the kdtree in `coord`
            coord.cache[attrname_or_kdt] = kdt

        return kdt