Little improvements to angular neighbours

This commit is contained in:
rstiskalek 2023-10-20 18:36:18 +01:00
parent e90b8ea9be
commit 5090523a10

View file

@ -29,7 +29,8 @@ import numpy
from readfof import FoF_catalog from readfof import FoF_catalog
from sklearn.neighbors import NearestNeighbors from sklearn.neighbors import NearestNeighbors
from ..utils import (cartesian_to_radec, periodic_distance_two_points, from ..utils import (radec_to_cartesian, cartesian_to_radec,
periodic_distance_two_points,
real2redshift) real2redshift)
from .box_units import CSiBORGBox, QuijoteBox from .box_units import CSiBORGBox, QuijoteBox
from .paths import Paths from .paths import Paths
@ -159,7 +160,7 @@ class BaseCatalogue(ABC):
@property @property
def box(self): def box(self):
"""Box object.""" """Box object."""
pass return self._box
@box.setter @box.setter
def box(self, box): def box(self, box):
@ -279,7 +280,6 @@ class BaseCatalogue(ABC):
------- -------
:py:class:`sklearn.neighbors.NearestNeighbors` :py:class:`sklearn.neighbors.NearestNeighbors`
""" """
# TODO improve the caching
pos = self["lagpatch_pos"] if in_initial else self["cartesian_pos"] pos = self["lagpatch_pos"] if in_initial else self["cartesian_pos"]
L = self.box.boxsize L = self.box.boxsize
knn = NearestNeighbors( knn = NearestNeighbors(
@ -287,119 +287,96 @@ class BaseCatalogue(ABC):
knn.fit(pos) knn.fit(pos)
return knn return knn
# def nearest_neighbours(self, X, radius, in_initial, knearest=False, def nearest_neighbours(self, X, radius, in_initial, knearest=False):
# return_mass=False): r"""
# r""" Return nearest neighbours within `radius` of `X` from this catalogue.
# Return nearest neighbours within `radius` of `X` from this catalogue. Units of `X` are cMpc / h.
#
# Parameters
# ----------
# X : 2-dimensional array, shape `(n_queries, 3)`
# Query positions.
# radius : float or int
# Limiting distance or number of neighbours, depending on `knearest`.
# in_initial : bool
# Find nearest neighbours in the initial or final snapshot.
# knearest : bool, optional
# If True, `radius` is the number of neighbours to return.
# return_mass : bool, optional
# Return masses of the nearest neighbours.
#
# Returns
# -------
# dist : list of arrays
# Distances to the nearest neighbours for each query.
# indxs : list of arrays
# Indices of nearest neighbours for each query.
# mass (optional): list of arrays
# Masses of the nearest neighbours for each query.
# """
# if X.shape != (len(X), 3):
# raise ValueError("`X` must be of shape `(n_samples, 3)`.")
# if knearest and not isinstance(radius, int):
# raise ValueError("`radius` must be an integer if `knearest`.")
# # if return_mass and not mass_key:
# # raise ValueError("`mass_key` must be provided if `return_mass`.")
#
# knn = self.knn(in_initial, subtract_observer=False, periodic=True)
#
# if knearest:
# dist, indxs = knn.kneighbors(X, radius)
# else:
# dist, indxs = knn.radius_neighbors(X, radius, sort_results=True)
#
# if not return_mass:
# return dist, indxs
#
# mass = [self[self.mass_key][indx] for indx in indxs]
# return dist, indxs, mass
# def angular_neighbours(self, X, ang_radius, in_rsp, rad_tolerance=None): Parameters
# r""" ----------
# Find nearest neighbours within `ang_radius` of query points `X` in the X : 2-dimensional array, shape `(n_queries, 3)`
# final snaphot. Optionally applies radial distance tolerance, which is Query positions.
# expected to be in :math:`\mathrm{cMpc} / h`. radius : float or int
# Limiting distance or number of neighbours, depending on `knearest`.
# Parameters in_initial : bool
# ---------- Find nearest neighbours in the initial or final snapshot.
# X : 2-dimensional array of shape `(n_queries, 2)` or `(n_queries, 3)` knearest : bool, optional
# Query positions. Either RA/dec in degrees or dist/RA/dec with If True, `radius` is the number of neighbours to return.
# distance in :math:`\mathrm{cMpc} / h`. return_mass : bool, optional
# in_rsp : bool Return masses of the nearest neighbours.
# If True, use redshift space positions of haloes.
# ang_radius : float Returns
# Angular radius in degrees. -------
# rad_tolerance : float, optional dist : list of arrays
# Radial distance tolerance in :math:`\mathrm{cMpc} / h`. Distances to the nearest neighbours for each query.
# indxs : list of arrays
# Returns Indices of nearest neighbours for each query.
# ------- mass : list of arrays
# dist : array of 1-dimensional arrays of shape `(n_neighbours,)` Masses of the nearest neighbours for each query.
# Distance of each neighbour to the query point. """
# ind : array of 1-dimensional arrays of shape `(n_neighbours,)` if knearest and not isinstance(radius, int):
# Indices of each neighbour in this catalogue. raise ValueError("`radius` must be an integer if `knearest`.")
# """
# assert X.ndim == 2 knn = self.knn(in_initial)
# if knearest:
# # Get positions of haloes in this catalogue dist, indxs = knn.kneighbors(X, radius)
# if in_rsp: else:
# pos = self.redshift_space_position(cartesian=True, dist, indxs = knn.radius_neighbors(X, radius, sort_results=True)
# subtract_observer=True)
# else: return dist, indxs
# pos = self.position(in_initial=False, cartesian=True,
# subtract_observer=True) def angular_neighbours(self, X, in_rsp, angular_tolerance,
# radial_tolerance=None):
# # Convert halo positions to unit vectors. """
# raddist = numpy.linalg.norm(pos, axis=1) Find nearest angular neighbours of query points. Optionally applies
# pos /= raddist.reshape(-1, 1) radial distance tolerance. Units of `X` are cMpc / h and degrees.
#
# # Convert RA/dec query positions to unit vectors. If no radial Parameters
# # distance is provided artificially add it. ----------
# if X.shape[1] == 2: X : 2-dimensional array of shape `(n_queries, 3)`
# X = numpy.vstack([numpy.ones_like(X[:, 0]), X[:, 0], X[:, 1]]).T Query positions given as distance/RA/dec.
# radquery = None in_rsp : bool
# else: Whether to find neighbours in redshift space.
# radquery = X[:, 0] angular_tolerance : float
# X = radec_to_cartesian(X) Angular radius in degrees.
# radial_tolerance : float, optional
# # Find neighbours Radial tolerance.
# knn = NearestNeighbors(metric="cosine")
# knn.fit(pos) Returns
# metric_maxdist = 1 - numpy.cos(numpy.deg2rad(ang_radius)) -------
# dist, ind = knn.radius_neighbors(X, radius=metric_maxdist, dist : array of 1-dimensional arrays of shape `(n_neighbours,)`
# sort_results=True) Distance of each neighbour to the query point.
# ind : array of 1-dimensional arrays of shape `(n_neighbours,)`
# # Convert cosine difference to angular distance Indices of each neighbour in this catalogue.
# for i in range(X.shape[0]): """
# dist[i] = numpy.rad2deg(numpy.arccos(1 - dist[i])) if in_rsp:
# pos = self["cartesian_redshiftspace_pos"]
# # Apply radial tolerance else:
# if rad_tolerance and radquery: pos = self["cartesian_pos"]
# for i in range(X.shape[0]):
# mask = numpy.abs(raddist[ind[i]] - radquery[i]) < rad_tolerance radial_dist = numpy.linalg.norm(pos, axis=1)
# dist[i], ind[i] = dist[i][mask], ind[i][mask] pos = pos / radial_dist.reshape(-1, 1)
#
# return dist, ind knn = NearestNeighbors(metric="cosine")
knn.fit(pos)
metric_maxdist = 1 - numpy.cos(numpy.deg2rad(angular_tolerance))
dist, indxs = knn.radius_neighbors(radec_to_cartesian(X),
radius=metric_maxdist,
sort_results=True)
# Convert cosine difference to angular distance
for i in range(X.shape[0]):
dist[i] = numpy.rad2deg(numpy.arccos(1 - dist[i]))
if radial_tolerance is not None:
radial_query = numpy.linalg.norm(X, axis=1)
for i in range(X.shape[0]):
rad_sep = numpy.abs(radial_dist[indxs[i]] - radial_query[i])
mask = rad_sep < radial_tolerance
dist[i], indxs[i] = dist[i][mask], indxs[i][mask]
return dist, indxs
# def filter_data(self, data, bounds): # def filter_data(self, data, bounds):
# """ # """