Add nearest neighbour search
- Dominant language
- Python
- Stars
- 5.3k
- Forks
- 2.2k
- Avg merge
- 1d 18h
- Merged PRs (30d)
- 74
Description
### What is the problem this feature will solve?
Hi all,
I was recently in need of finding all matches between a set of catalogues (even a self-match) up to a Nth nearest neighbour.
`astropy.coordinates.match_coordinates_sky` provides some of this functionality, but only for a single value of `N`. I've put together a little function that mirrors the behaviour of `astropy.coordinates.search_around_sky` but uses an Nth neighbour limit, rather than a sky separation limit.
If we think this would be broadly useful, I'm happy to open a PR
### Describe the desired outcome
Here is my current draft implmentation. I've tried to mirror most of the API and checks in `search_around_sky`. I'll note that the KDTree caching is currently absent, but I think that could be easily placed in with a call to `_get_cartesian_kdtree`.
```python
def nearest_neighbour_sky(
coords1: SkyCoord,
coords2: SkyCoord,
nthneighbor_max: int,
) -> CoordinateSearchResult:
"""
Searches for pairs of points that are at most ``nthneighbor_max+1`` nearest neighbours.
For finding _just_ the nearest neighbour see ``match_coordinates_sky``.
``nthneighbor_max+1`` is needed to find ``nthneighbor_max`` neighbours excluding a self-match.
Parameters
----------
coords1 : SkyCoord
The first set of coordinates, which will be searched for matches from
``coords2`` within ``nthneighbor_max+1`` neighbours.
Must be a one-dimensional coordinate array.
coords2 : SkyCoord
The first set of coordinates, which will be searched for matches from
``coords1`` within ``nthneighbor_max+1`` neighbours.
Must be a one-dimensional coordinate array.
nthneighbor_max : int
The maximum nearest neighbour to be searched.
Returns
-------
CoordinateSearchResult
A `~typing.NamedTuple` with attributes representing the indices
of the elements of found pairs in both source sets and angular
and physical separations of the pairs. If either set of sources
lack distances, the physical separation is the 3D distance on
the unit sphere, rather than a true distance.
"""
if coords1.ndim != 1 or coords2.ndim != 1:
msg = "nearest_neighbour_sky only supports 1-dimensional coordinate arrays."
if coords1.isscalar or coords2.isscalar:
msg += " The catalog for coordinate matching cannot be a scalar or length-0"
raise ValueError(msg)
xyz1 = coords1.cartesian.xyz.value.T
xyz2 = coords2.cartesian.xyz.value.T
xyz_unit = coords2.cartesian.xyz.unit
tree = KDTree(xyz2, compact_nodes=False, balanced_tree=False)
# Increase nthneighbor_max by one since the nearest neighbour may be itself
dist, idx2 = tree.query(xyz1, k=nthneighbor_max + 1)
# ensure 2D arrays if k=1
if nthneighbor_max+1 == 1:
dist = dist[:, None]
idx2 = idx2[:, None]
n_points, n_neighbours = idx2.shape
# `idx1` is just the simple index into coords1
# `idx2` is the index from the tree
idx1 = np.repeat(np.arange(n_points), repeats=n_neighbours)
idx2 = idx2.ravel()
# Convert cartesian distance to angular separation
sep = Angle(2 * np.arcsin(dist.ravel() / 2) * u.rad)
return CoordinateSearchResult(
idx1,
idx2,
sep,
dist * xyz_unit
)
```
### Additional context
_No response_
Contributor guide
Research direction
Start by comparing the proposed nearest_neighbour_sky API with astropy.coordinates.search_around_sky and match_coordinates_sky, then inspect _get_cartesian_kdtree for the mentioned caching path. The issue names no file or test; done means the coordinate matching behavior, validation, neighbour indexing, self-match handling, and returned separations are covered consistently with the existing APIs.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100