Periodic neighbours (#84)

* Edit the HMF plot

* Add periodic dist 2 points

* Add boxsize to RVSSphere

* Add periodic distance

* Adding periodic distance

* Add imports

* Change arguments

* Update bounds

* Lower min number of particles

* Change kwargs

* Add paths overlap quijote

* Add some comments

csiborgtools/__init__.py

@@ -13,7 +13,9 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 from csiborgtools import clustering, field, match, read  # noqa
-from .utils import (center_of_mass, delta2ncells, periodic_distance, number_counts)  # noqa
+
+from .utils import (center_of_mass, delta2ncells, number_counts,
+                    periodic_distance, periodic_distance_two_points)  # noqa
 
 # Arguments to csiborgtools.read.Paths.
 paths_glamdring = {"srcdir": "/mnt/extraspace/hdesmond/",

csiborgtools/clustering/utils.py

@@ -51,16 +51,20 @@ class BaseRVS(ABC):
 class RVSinsphere(BaseRVS):
     """
     Generator of uniform RVS in a sphere of radius `R` in Cartesian
-    coordinates centered at the origin.
+    coordinates centered at the centre of the box.
 
     Parameters
     ----------
     R : float
         Radius of the sphere.
+    boxsize : float
+        Box size
     """
-    def __init__(self, R):
+    def __init__(self, R, boxsize):
         assert R > 0, "Radius must be positive."
+        assert boxsize > 0, "Box size must be positive."
         self.R = R
+        self.boxsize = boxsize
         BaseRVS.__init__(self)
 
     def __call__(self, nsamples, random_state=42, dtype=numpy.float32):
@@ -73,7 +77,7 @@ class RVSinsphere(BaseRVS):
         x = r * numpy.sin(theta) * numpy.cos(phi)
         y = r * numpy.sin(theta) * numpy.sin(phi)
         z = r * numpy.cos(theta)
-        return numpy.vstack([x, y, z]).T
+        return numpy.vstack([x, y, z]).T + self.boxsize / 2
 
 
 class RVSinbox(BaseRVS):

csiborgtools/match/match.py

@@ -239,7 +239,8 @@ class RealisationsMatcher(BaseMatcher):
         if verbose:
             print(f"{datetime.now()}: querying the KNN.", flush=True)
         match_indxs = radius_neighbours(
-            catx.knn(in_initial=True), cat0.position(in_initial=True),
+            catx.knn(in_initial=True, subtract_observer=False, periodic=True),
+            cat0.position(in_initial=True),
             radiusX=cat0["lagpatch_size"], radiusKNN=catx["lagpatch_size"],
             nmult=self.nmult, enforce_int32=True, verbose=verbose)
 
@@ -515,7 +516,7 @@ class ParticleOverlap(BaseMatcher):
             Minimun and maximum cell numbers along each dimension of `halo1`.
             Optional.
         mins2, maxs2 : 1-dimensional arrays of shape `(3,)`
-            Minimun and maximum cell numbers along each dimension of `halo2`.
+            Minimum and maximum cell numbers along each dimension of `halo2`.
             Optional.
         smooth_kwargs : kwargs, optional
             Kwargs to be passed to :py:func:`scipy.ndimage.gaussian_filter`.
@@ -1014,7 +1015,8 @@ def radius_neighbours(knn, X, radiusX, radiusKNN, nmult=1.0,
 def find_neighbour(nsim0, cats):
     """
     Find the nearest neighbour of halos from a reference catalogue indexed
-    `nsim0` in the remaining simulations.
+    `nsim0` in the remaining simulations. Note that this must be the same
+    simulation suite.
 
     Parameters
     ----------
@@ -1030,8 +1032,10 @@ def find_neighbour(nsim0, cats):
     cross_hindxs : 2-dimensional array of shape `(nhalos, len(cats) - 1)`
         Halo indices of the nearest neighbour.
     """
+    assert all(isinstance(cat, type(cats[nsim0])) for cat in cats.values())
+
     cat0 = cats[nsim0]
-    X = cat0.position(in_initial=False, subtract_observer=True)
+    X = cat0.position(in_initial=False)
 
     nhalos = X.shape[0]
     num_cats = len(cats) - 1

csiborgtools/read/halo_cat.py

@@ -33,6 +33,7 @@ from .paths import Paths
 from .readsim import CSiBORGReader
 from .utils import (add_columns, cartesian_to_radec, cols_to_structured,
                     flip_cols, radec_to_cartesian, real2redshift)
+from ..utils import periodic_distance_two_points
 
 
 class BaseCatalogue(ABC):
@@ -73,6 +74,17 @@ class BaseCatalogue(ABC):
         """
         pass
 
+    @abstractproperty
+    def simname(self):
+        """
+        Simulation name.
+
+        Returns
+        -------
+        simname : str
+        """
+        pass
+
     @property
     def paths(self):
         """
@@ -327,7 +339,7 @@ class BaseCatalogue(ABC):
         return numpy.vstack([self["L{}".format(p)] for p in ("x", "y", "z")]).T
 
     @lru_cache(maxsize=2)
-    def knn(self, in_initial):
+    def knn(self, in_initial, subtract_observer, periodic):
         r"""
         kNN object for catalogue objects with caching. Positions are centered
         on the observer.
@@ -336,14 +348,32 @@ class BaseCatalogue(ABC):
         ----------
         in_initial : bool
             Whether to define the kNN on the initial or final snapshot.
+        subtract_observer : bool
+            Whether to subtract the observer's location from the positions.
+        periodic : bool
+            Whether to use periodic boundary conditions.
 
         Returns
         -------
         knn : :py:class:`sklearn.neighbors.NearestNeighbors`
             kNN object fitted with object positions.
         """
-        pos = self.position(in_initial=in_initial)
-        return NearestNeighbors().fit(pos)
+        if subtract_observer and periodic:
+            raise ValueError("Subtracting observer is not supported for "
+                             "periodic boundary conditions.")
+
+        pos = self.position(in_initial=in_initial,
+                            subtract_observer=subtract_observer)
+
+        if periodic:
+            L = self.box.boxsize
+            knn = NearestNeighbors(
+                metric=lambda a, b: periodic_distance_two_points(a, b, L))
+        else:
+            knn = NearestNeighbors()
+
+        knn.fit(pos)
+        return knn
 
     def nearest_neighbours(self, X, radius, in_initial, knearest=False,
                            return_mass=False, mass_key=None):
@@ -382,7 +412,8 @@ class BaseCatalogue(ABC):
         if return_mass and not mass_key:
             raise ValueError("`mass_key` must be provided if `return_mass`.")
 
-        knn = self.knn(in_initial)
+        knn = self.knn(in_initial, subtract_observer=False, periodic=True)
+
         if knearest:
             dist, indxs = knn.kneighbors(X, radius)
         else:
@@ -564,6 +595,10 @@ class CSiBORGHaloCatalogue(BaseCatalogue):
         """
         return CSiBORGBox(self.nsnap, self.nsim, self.paths)
 
+    @property
+    def simname(self):
+        return "csiborg"
+
 
 ###############################################################################
 #                         Quijote halo catalogue                              #
@@ -661,6 +696,10 @@ class QuijoteHaloCatalogue(BaseCatalogue):
         assert nsnap in [0, 1, 2, 3, 4]
         self._nsnap = nsnap
 
+    @property
+    def simname(self):
+        return "quijote"
+
     @property
     def redshift(self):
         """

csiborgtools/read/overlap_summary.py

@@ -47,6 +47,10 @@ class PairOverlap:
     _data = None
 
     def __init__(self, cat0, catx, paths, maxdist=None):
+        if cat0.simname != catx.simname:
+            raise ValueError("The two catalogues must be from the same "
+                             "simulation.")
+
         self._cat0 = cat0
         self._catx = catx
         self.load(cat0, catx, paths, maxdist)
@@ -77,8 +81,8 @@ class PairOverlap:
 
         # We first load in the output files. We need to find the right
         # combination of the reference and cross simulation.
-        fname = paths.overlap(nsim0, nsimx, smoothed=False)
-        fname_inv = paths.overlap(nsimx, nsim0, smoothed=False)
+        fname = paths.overlap(cat0.simname, nsim0, nsimx, smoothed=False)
+        fname_inv = paths.overlap(cat0.simname, nsimx, nsim0, smoothed=False)
         if isfile(fname):
             data_ngp = numpy.load(fname, allow_pickle=True)
             to_invert = False
@@ -89,7 +93,8 @@ class PairOverlap:
         else:
             raise FileNotFoundError(f"No file found for {nsim0} and {nsimx}.")
 
-        fname_smooth = paths.overlap(cat0.nsim, catx.nsim, smoothed=True)
+        fname_smooth = paths.overlap(cat0.simname, cat0.nsim, catx.nsim,
+                                     smoothed=True)
         data_smooth = numpy.load(fname_smooth, allow_pickle=True)
 
         # Create mapping from halo indices to array positions in the catalogue.
@@ -764,13 +769,15 @@ class NPairsOverlap:
 ###############################################################################
 
 
-def get_cross_sims(nsim0, paths, smoothed):
+def get_cross_sims(simname, nsim0, paths, smoothed):
     """
     Get the list of cross simulations for a given reference simulation for
     which the overlap has been calculated.
 
     Parameters
     ----------
+    simname : str
+        Simulation name.
     nsim0 : int
         Reference simulation number.
     paths : :py:class:`csiborgtools.paths.Paths`
@@ -782,8 +789,8 @@ def get_cross_sims(nsim0, paths, smoothed):
     for nsimx in paths.get_ics("csiborg"):
         if nsimx == nsim0:
             continue
-        f1 = paths.overlap(nsim0, nsimx, smoothed)
-        f2 = paths.overlap(nsimx, nsim0, smoothed)
+        f1 = paths.overlap(simname, nsim0, nsimx, smoothed)
+        f2 = paths.overlap(simname, nsimx, nsim0, smoothed)
         if isfile(f1) or isfile(f2):
             nsimxs.append(nsimx)
     return nsimxs

csiborgtools/read/paths.py

@@ -462,12 +462,14 @@ class Paths:
         fname = f"out_{str(nsim).zfill(5)}_{str(nsnap).zfill(5)}.npy"
         return join(fdir, fname)
 
-    def overlap(self, nsim0, nsimx, smoothed):
+    def overlap(self, simname, nsim0, nsimx, smoothed):
         """
         Path to the overlap files between two CSiBORG simulations.
 
         Parameters
         ----------
+        simname : str
+            Simulation name. Must be one of `csiborg` or `quijote`.
         nsim0 : int
             IC realisation index of the first simulation.
         nsimx : int
@@ -479,7 +481,12 @@ class Paths:
         -------
         path : str
         """
-        fdir = join(self.postdir, "overlap")
+        if simname == "csiborg":
+            fdir = join(self.postdir, "overlap")
+        elif simname == "quijote":
+            fdir = join(self.quijote_dir, "overlap")
+        else:
+            raise ValueError(f"Unknown simulation name `{simname}`.")
 
         try_create_directory(fdir)
 

csiborgtools/utils.py

@@ -92,6 +92,39 @@ def periodic_distance(points, reference, boxsize):
     return dist
 
 
+@jit(nopython=True, fastmath=True, boundscheck=False)
+def periodic_distance_two_points(p1, p2, boxsize):
+    """
+    Compute the 3D distance between two points using periodic boundary
+    conditions. This is an optimized JIT implementation.
+
+    Parameters
+    ----------
+    p1 : 1-dimensional array of shape `(3, )`
+        First point.
+    p2 : 1-dimensional array of shape `(3, )`
+        Second point.
+    boxsize : float
+        Box size.
+
+    Returns
+    -------
+    dist : 1-dimensional array of shape `(n_points, )`
+    """
+    half_box = boxsize / 2
+
+    dist = 0
+    for i in range(3):
+        dist_1d = abs(p1[i] - p2[i])
+
+        if dist_1d > (half_box):
+            dist_1d = boxsize - dist_1d
+
+        dist += dist_1d**2
+
+    return dist**0.5
+
+
 @jit(nopython=True, fastmath=True, boundscheck=False)
 def delta2ncells(delta):
     """