From 2a5d6a03b61e7ec21865ad496036c067091b0566 Mon Sep 17 00:00:00 2001
From: rstiskalek <richard.stiskalek@protonmail.com>
Date: Fri, 25 Nov 2022 16:45:27 +0000
Subject: [PATCH] add distance cut

---
 csiborgtools/read/make_cat.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/csiborgtools/read/make_cat.py b/csiborgtools/read/make_cat.py
index 3a9d269..90d56a3 100644
--- a/csiborgtools/read/make_cat.py
+++ b/csiborgtools/read/make_cat.py
@@ -34,9 +34,11 @@ class HaloCatalogue:
     ----------
     paths : py:class:`csiborgtools.read.CSiBORGPaths`
         CSiBORG paths-handling object with set `n_sim` and `n_snap`.
-    minimum_m500 : float, optional
+    min_m500 : float, optional
         The minimum :math:`M_{rm 500c} / M_\odot` mass. By default no
         threshold.
+    max_dist : float, optional
+        The maximum comoving distance of a halo. By default no upper limit.
     """
     _box = None
     _paths = None
@@ -44,11 +46,12 @@ class HaloCatalogue:
     _knn = None
     _positions = None
 
-    def __init__(self, paths, minimum_m500=None):
+    def __init__(self, paths, min_m500=None, max_dist=None):
         self._box = BoxUnits(paths)
-        minimum_m500 = 0 if minimum_m500 is None else minimum_m500
+        min_m500 = 0 if min_m500 is None else min_m500
+        max_dist = numpy.infty if max_dist is None else max_dist
         self._paths = paths
-        self._set_data(minimum_m500)
+        self._set_data(min_m500, max_dist)
         # Initialise the KNN
         knn = NearestNeighbors()
         knn.fit(self.positions)
@@ -122,7 +125,7 @@ class HaloCatalogue:
         """
         return self.paths.n_sim
 
-    def _set_data(self, minimum_m500):
+    def _set_data(self, min_m500, max_dist):
         """
         Loads the data, merges with mmain, does various coordinate transforms.
         """
@@ -157,12 +160,15 @@ class HaloCatalogue:
         data = self.box.convert_from_boxunits(data, convert_cols)
 
         # Cut on mass. Note that this is in Msun
-        data = data[data["m500"] > minimum_m500]
+        data = data[data["m500"] > min_m500]
 
         # Now calculate spherical coordinates
         d, ra, dec = cartesian_to_radec(data)
         data = add_columns(data, [d, ra, dec], ["dist", "ra", "dec"])
 
+        # Cut on separation
+        data = data[data["dist"] < max_dist]
+
         # Pre-allocate the positions array
         self._positions = numpy.vstack(
             [data["peak_{}".format(p)] for p in ("x", "y", "z")]).T
@@ -265,9 +271,11 @@ class CombinedHaloCatalogue:
     paths : py:class`csiborgtools.read.CSiBORGPaths`
         CSiBORG paths-handling object. Doest not have to have set set `n_sim`
         and `n_snap`.
-    minimum_m500 : float, optional
+    min_m500 : float, optional
         The minimum :math:`M_{rm 500c} / M_\odot` mass. By default no
         threshold.
+    max_dist : float, optional
+        The maximum comoving distance of a halo. By default no upper limit.
     verbose : bool, optional
         Verbosity flag for reading the catalogues.
     """
@@ -275,7 +283,7 @@ class CombinedHaloCatalogue:
     _n_snaps = None
     _cats = None
 
-    def __init__(self, paths, minimum_m500=None, verbose=True):
+    def __init__(self, paths, min_m500=None, max_dist=None, verbose=True):
         # Read simulations and their maximum snapshots
         # NOTE remove this later and take all cats
         self._n_sims = paths.ic_ids[:10]
@@ -286,7 +294,7 @@ class CombinedHaloCatalogue:
         for i in trange(self.N) if verbose else range(self.N):
             paths = deepcopy(paths)
             paths.set_info(self.n_sims[i], self.n_snaps[i])
-            cats[i] = HaloCatalogue(paths, minimum_m500)
+            cats[i] = HaloCatalogue(paths, min_m500, max_dist)
         self._cats = cats
 
     @property