Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools_public.git synced 2025-05-20 17:41:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix angular

Browse source
This commit is contained in:
rstiskalek 2023-10-23 15:59:42 +01:00
parent 1b498c939c
commit 45f6584493
1 changed files with 22 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

48
csiborgtools/read/halo_cat.py
View file

@ -25,7 +25,8 @@ from h5py import File
from sklearn.neighbors import NearestNeighbors from sklearn.neighbors import NearestNeighbors
from ..utils import (cartesian_to_radec, fprint, periodic_distance_two_points, from ..utils import (cartesian_to_radec, fprint, periodic_distance_two_points,
radec_to_cartesian, real2redshift, number_counts) radec_to_cartesian, real2redshift, number_counts,
great_circle_distance)
from .box_units import CSiBORGBox, QuijoteBox from .box_units import CSiBORGBox, QuijoteBox
from .paths import Paths from .paths import Paths
from .readsim import load_halo_particles, make_halomap_dict from .readsim import load_halo_particles, make_halomap_dict
@ -299,8 +300,8 @@ class BaseCatalogue(ABC):
key = f"snapshot_{'initial' if in_initial else 'final'}/{kind}" key = f"snapshot_{'initial' if in_initial else 'final'}/{kind}"
return load_halo_particles(hid, self[key], self["particle_offset"]) return load_halo_particles(hid, self[key], self["particle_offset"])
@lru_cache(maxsize=2) @lru_cache(maxsize=4)
def knn(self, in_initial): def knn(self, in_initial, angular=False):
r""" r"""
Periodic kNN object in real space in Cartesian coordinates trained on Periodic kNN object in real space in Cartesian coordinates trained on
`self["cartesian_pos"]`. `self["cartesian_pos"]`.
@ -309,15 +310,23 @@ class BaseCatalogue(ABC):
---------- ----------
in_initial : bool in_initial : bool
Whether to define the kNN on the initial or final snapshot. Whether to define the kNN on the initial or final snapshot.
angular : bool, optional
Whether to define the kNN in RA/dec.
Returns Returns
------- -------
:py:class:`sklearn.neighbors.NearestNeighbors` :py:class:`sklearn.neighbors.NearestNeighbors`
""" """
pos = self["lagpatch_pos"] if in_initial else self["cartesian_pos"] if angular:
L = self.box.boxsize assert not in_initial, "Angular kNN not available for initial."
knn = NearestNeighbors( pos = self["spherical_pos"][:, 1:]
metric=lambda a, b: periodic_distance_two_points(a, b, L)) knn = NearestNeighbors(metric=great_circle_distance)
else:
pos = self["lagpatch_pos"] if in_initial else self["cartesian_pos"]
L = self.box.boxsize
knn = NearestNeighbors(
metric=lambda a, b: periodic_distance_two_points(a, b, L))
knn.fit(pos) knn.fit(pos)
return knn return knn
@ -383,28 +392,15 @@ class BaseCatalogue(ABC):
ind : array of 1-dimensional arrays of shape `(n_neighbours,)` ind : array of 1-dimensional arrays of shape `(n_neighbours,)`
Indices of each neighbour in this catalogue. Indices of each neighbour in this catalogue.
""" """
if in_rsp: knn = self.knn(in_initial=False, angular=True)
pos = self["cartesian_redshiftspace_pos"] dist, indxs = knn.radius_neighbors(X[:, 1:],
else: radius=angular_tolerance,
pos = self["cartesian_pos"]
radial_dist = numpy.linalg.norm(pos, axis=1)
pos = pos / radial_dist.reshape(-1, 1)
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) 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: if radial_tolerance is not None:
radial_query = numpy.linalg.norm(X, axis=1) radial_dist = self["redshift_dist"] if in_rsp else self["dist"]
radial_query = X[:, 0]
for i in range(X.shape[0]): for i in range(X.shape[0]):
rad_sep = numpy.abs(radial_dist[indxs[i]] - radial_query[i]) rad_sep = numpy.abs(radial_dist[indxs[i]] - radial_query[i])
mask = rad_sep < radial_tolerance mask = rad_sep < radial_tolerance