From c7b600d0ad8aa0cae00d8360defc5e62714b2190 Mon Sep 17 00:00:00 2001
From: Richard Stiskalek <richard.stiskalek@protonmail.com>
Date: Tue, 8 Aug 2023 12:19:40 +0200
Subject: [PATCH] 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             |  4 ++-
 csiborgtools/clustering/utils.py     | 10 ++++--
 csiborgtools/match/match.py          | 12 ++++---
 csiborgtools/read/halo_cat.py        | 47 +++++++++++++++++++++++++---
 csiborgtools/read/overlap_summary.py | 19 +++++++----
 csiborgtools/read/paths.py           | 11 +++++--
 csiborgtools/utils.py                | 33 +++++++++++++++++++
 scripts/cluster_knn_auto.py          | 12 +++----
 scripts/cluster_tpcf_auto.py         |  6 ++--
 scripts/fit_init.py                  |  9 +++---
 scripts/match_singlematch.py         | 20 +++++++-----
 scripts_plots/plot_data.py           | 15 +++++----
 scripts_plots/plot_match.py          | 40 ++++++++++++++---------
 setup.py                             | 19 +++++++++++
 14 files changed, 196 insertions(+), 61 deletions(-)

diff --git a/csiborgtools/__init__.py b/csiborgtools/__init__.py
index a019822..b21249a 100644
--- a/csiborgtools/__init__.py
+++ b/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/",
diff --git a/csiborgtools/clustering/utils.py b/csiborgtools/clustering/utils.py
index 7f6d89e..db44b5e 100644
--- a/csiborgtools/clustering/utils.py
+++ b/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):
diff --git a/csiborgtools/match/match.py b/csiborgtools/match/match.py
index b9c6fad..aaac705 100644
--- a/csiborgtools/match/match.py
+++ b/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
diff --git a/csiborgtools/read/halo_cat.py b/csiborgtools/read/halo_cat.py
index d90e43e..241be99 100644
--- a/csiborgtools/read/halo_cat.py
+++ b/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):
         """
diff --git a/csiborgtools/read/overlap_summary.py b/csiborgtools/read/overlap_summary.py
index 422ce13..fac2bef 100644
--- a/csiborgtools/read/overlap_summary.py
+++ b/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
diff --git a/csiborgtools/read/paths.py b/csiborgtools/read/paths.py
index 03716c5..4ae2ffa 100644
--- a/csiborgtools/read/paths.py
+++ b/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)
 
diff --git a/csiborgtools/utils.py b/csiborgtools/utils.py
index 2ea50f9..bf1be68 100644
--- a/csiborgtools/utils.py
+++ b/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):
     """
diff --git a/scripts/cluster_knn_auto.py b/scripts/cluster_knn_auto.py
index ea02ed5..04b8b50 100644
--- a/scripts/cluster_knn_auto.py
+++ b/scripts/cluster_knn_auto.py
@@ -59,10 +59,10 @@ def do_auto(args, config, cats, nsim, paths):
     -------
     None
     """
-    rvs_gen = csiborgtools.clustering.RVSinsphere(args.Rmax)
-    knncdf = csiborgtools.clustering.kNN_1DCDF()
     cat = cats[nsim]
-    knn = cat.knn(in_initial=False)
+    rvs_gen = csiborgtools.clustering.RVSinsphere(args.Rmax, cat.boxsize)
+    knncdf = csiborgtools.clustering.kNN_1DCDF()
+    knn = cat.knn(in_initial=False, subtract_observer=False, periodic=True)
     rs, cdf = knncdf(
         knn, rvs_gen=rvs_gen, nneighbours=config["nneighbours"],
         rmin=config["rmin"], rmax=config["rmax"],
@@ -97,9 +97,9 @@ def do_cross_rand(args, config, cats, nsim, paths):
     -------
     None
     """
-    rvs_gen = csiborgtools.clustering.RVSinsphere(args.Rmax)
     cat = cats[nsim]
-    knn1 = cat.knn(in_initial=False)
+    rvs_gen = csiborgtools.clustering.RVSinsphere(args.Rmax, cat.boxsize)
+    knn1 = cat.knn(in_initial=False, subtract_observer=False, periodic=True)
 
     knn2 = NearestNeighbors()
     pos2 = rvs_gen(len(cat).shape[0])
@@ -126,7 +126,7 @@ if __name__ == "__main__":
                         help="Simulation name")
     parser.add_argument("--nsims", type=int, nargs="+", default=None,
                         help="Indices of simulations to cross. If `-1` processes all simulations.")  # noqa
-    parser.add_argument("--Rmax", type=float, default=155/0.705,
+    parser.add_argument("--Rmax", type=float, default=155,
                         help="High-resolution region radius")  # noqa
     parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
                         default=False)
diff --git a/scripts/cluster_tpcf_auto.py b/scripts/cluster_tpcf_auto.py
index 546f80e..d9263df 100644
--- a/scripts/cluster_tpcf_auto.py
+++ b/scripts/cluster_tpcf_auto.py
@@ -37,12 +37,12 @@ except ModuleNotFoundError:
 
 
 def do_auto(args, config, cats, nsim, paths):
+    cat = cats[nsim]
     tpcf = csiborgtools.clustering.Mock2PCF()
-    rvs_gen = csiborgtools.clustering.RVSinsphere(args.Rmax)
+    rvs_gen = csiborgtools.clustering.RVSinsphere(args.Rmax, cat.boxsize)
     bins = numpy.logspace(
         numpy.log10(config["rpmin"]), numpy.log10(config["rpmax"]),
         config["nrpbins"] + 1,)
-    cat = cats[nsim]
 
     pos = cat.position(in_initial=False, cartesian=True)
     nrandom = int(config["randmult"] * pos.shape[0])
@@ -59,7 +59,7 @@ if __name__ == "__main__":
                         help="Simulation name")
     parser.add_argument("--nsims", type=int, nargs="+", default=None,
                         help="Indices of simulations to cross. If `-1` processes all simulations.")  # noqa
-    parser.add_argument("--Rmax", type=float, default=155/0.705,
+    parser.add_argument("--Rmax", type=float, default=155,
                         help="High-resolution region radius.")
     parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
                         default=False, help="Verbosity flag.")
diff --git a/scripts/fit_init.py b/scripts/fit_init.py
index e816e84..8f69258 100644
--- a/scripts/fit_init.py
+++ b/scripts/fit_init.py
@@ -14,8 +14,9 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 """
 Script to calculate the particle centre of mass, Lagrangian patch size in the
-initial snapshot. The initial snapshot particles are read from the sorted
-files.
+initial snapshot.
+
+The initial snapshot particles are read from the sorted files.
 """
 from argparse import ArgumentParser
 from datetime import datetime
@@ -74,7 +75,7 @@ def _main(nsim, simname, verbose):
 
     # Initialise the overlapper.
     if simname == "csiborg":
-        kwargs = {"box_size": 2048, "bckg_halfsize": 475}
+        kwargs = {"box_size": 2048, "bckg_halfsize": 512}
     else:
         kwargs = {"box_size": 512, "bckg_halfsize": 256}
     overlapper = csiborgtools.match.ParticleOverlap(**kwargs)
@@ -86,7 +87,7 @@ def _main(nsim, simname, verbose):
                                                      hid2map)
 
         # Skip if the halo has no particles or is too small.
-        if part is None or part.size < 100:
+        if part is None or part.size < 40:
             continue
 
         pos, mass = part[:, :3], part[:, 3]
diff --git a/scripts/match_singlematch.py b/scripts/match_singlematch.py
index b1b911b..47a462d 100644
--- a/scripts/match_singlematch.py
+++ b/scripts/match_singlematch.py
@@ -57,12 +57,13 @@ def pair_match(nsim0, nsimx, simname, sigma, verbose):
     None
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    smooth_kwargs = {"sigma": sigma, "mode": "wrap"}
+    smooth_kwargs = {"sigma": sigma, "mode": "constant", "cval": 0}
 
     if simname == "csiborg":
-        overlapper_kwargs = {"box_size": 2048, "bckg_halfsize": 475}
+        overlapper_kwargs = {"box_size": 2048, "bckg_halfsize": 512}
         mass_kind = "fof_totpartmass"
-        bounds = {mass_kind: (1e13, None)}
+        bounds = {"dist": (0, 155), mass_kind: (10**13.25, None)}
+
         cat0 = csiborgtools.read.CSiBORGHaloCatalogue(
             nsim0, paths, bounds=bounds, load_fitted=False,
             with_lagpatch=True)
@@ -72,11 +73,14 @@ def pair_match(nsim0, nsimx, simname, sigma, verbose):
     elif simname == "quijote":
         overlapper_kwargs = {"box_size": 512, "bckg_halfsize": 256}
         mass_kind = "group_mass"
-        bounds = {mass_kind: (1e14, None)}
+        bounds = {mass_kind: (10**13.25, None)}
+
         cat0 = csiborgtools.read.QuijoteHaloCatalogue(
-            nsim0, paths, 4, load_fitted=False, with_lagpatch=True)
+            nsim0, paths, 4, bounds=bounds, load_fitted=False,
+            with_lagpatch=True)
         catx = csiborgtools.read.QuijoteHaloCatalogue(
-            nsimx, paths, 4, load_fitted=False, with_lagpatch=True)
+            nsimx, paths, 4, bounds=bounds, load_fitted=False,
+            with_lagpatch=True)
     else:
         raise ValueError(f"Unknown simulation name: `{simname}`.")
 
@@ -116,7 +120,7 @@ def pair_match(nsim0, nsimx, simname, sigma, verbose):
         for j, match in enumerate(matches):
             match_hids[i][j] = catx["index"][match]
 
-    fout = paths.overlap(nsim0, nsimx, smoothed=False)
+    fout = paths.overlap(simname, nsim0, nsimx, smoothed=False)
     if verbose:
         print(f"{datetime.now()}: saving to ... `{fout}`.", flush=True)
     numpy.savez(fout, ref_hids=cat0["index"], match_hids=match_hids,
@@ -135,7 +139,7 @@ def pair_match(nsim0, nsimx, simname, sigma, verbose):
                                               match_indxs, smooth_kwargs,
                                               verbose=verbose)
 
-    fout = paths.overlap(nsim0, nsimx, smoothed=True)
+    fout = paths.overlap(simname, nsim0, nsimx, smoothed=True)
     if verbose:
         print(f"{datetime.now()}: saving to ... `{fout}`.", flush=True)
     numpy.savez(fout, smoothed_overlap=smoothed_overlap, sigma=sigma)
diff --git a/scripts_plots/plot_data.py b/scripts_plots/plot_data.py
index 88d7848..d7fa0c5 100644
--- a/scripts_plots/plot_data.py
+++ b/scripts_plots/plot_data.py
@@ -178,6 +178,7 @@ def plot_hmf(pdf=False):
     csiborg5511[x > 3e15] = numpy.nan
 
     with plt.style.context(plt_utils.mplstyle):
+        cols = plt.rcParams["axes.prop_cycle"].by_key()["color"]
         fig, ax = plt.subplots(nrows=2, sharex=True,
                                figsize=(3.5, 2.625 * 1.25),
                                gridspec_kw={"height_ratios": [1, 0.45]})
@@ -186,15 +187,15 @@ def plot_hmf(pdf=False):
         # Upper panel data
         mean_csiborg = numpy.mean(csiborg_counts, axis=0)
         std_csiborg = numpy.std(csiborg_counts, axis=0)
-        ax[0].plot(x, mean_csiborg, label="CSiBORG")
+        ax[0].plot(x, mean_csiborg, label="CSiBORG", c=cols[0])
         ax[0].fill_between(x, mean_csiborg - std_csiborg,
-                           mean_csiborg + std_csiborg, alpha=0.5)
+                           mean_csiborg + std_csiborg, alpha=0.5, color=cols[0])
 
         mean_quijote = numpy.mean(quijote_counts, axis=0)
         std_quijote = numpy.std(quijote_counts, axis=0)
-        ax[0].plot(x, mean_quijote, label="Quijote")
+        ax[0].plot(x, mean_quijote, label="Quijote", c=cols[1])
         ax[0].fill_between(x, mean_quijote - std_quijote,
-                           mean_quijote + std_quijote, alpha=0.5)
+                           mean_quijote + std_quijote, alpha=0.5, color=cols[1])
 
         ax[0].plot(x, csiborg5511, label="CSiBORG 5511", c="k", ls="--")
         std5511 = numpy.sqrt(csiborg5511)
@@ -204,9 +205,11 @@ def plot_hmf(pdf=False):
         log_y = numpy.log10(mean_csiborg / mean_quijote)
         err = numpy.sqrt((std_csiborg / mean_csiborg / numpy.log(10))**2
                          + (std_quijote / mean_quijote / numpy.log(10))**2)
-        ax[1].plot(x, 10**log_y, c="gray")
+        ax[1].plot(x, 10**log_y, c=cols[0])
         ax[1].fill_between(x, 10**(log_y - err), 10**(log_y + err), alpha=0.5,
-                           color="gray")
+                           color=col[0])
+
+        ax[1].plot(x, csiborg5511 / mean_quijote, c="k", ls="--")
 
         # Labels and accesories
         ax[1].axhline(1, color="k", ls="--",
diff --git a/scripts_plots/plot_match.py b/scripts_plots/plot_match.py
index 8cb4ffd..0c602c1 100644
--- a/scripts_plots/plot_match.py
+++ b/scripts_plots/plot_match.py
@@ -143,13 +143,15 @@ def plot_mass_vs_maxpairoverlap(nsim0, nsimx):
 
 
 @cache_to_disk(7)
-def get_overlap(nsim0):
+def get_overlap(simname, nsim0):
     """
     Calculate the summed overlap and probability of no match for a single
     reference simulation.
 
     Parameters
     ----------
+    simname : str
+        Simulation name.
     nsim0 : int
         Simulation index.
 
@@ -167,7 +169,8 @@ def get_overlap(nsim0):
         Probability of no match for each halo in the reference simulation.
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsimxs = csiborgtools.read.get_cross_sims(nsim0, paths, smoothed=True)
+    nsimxs = csiborgtools.read.get_cross_sims(simname, nsim0, paths,
+                                              smoothed=True)
     cat0 = open_cat(nsim0)
 
     catxs = []
@@ -195,7 +198,7 @@ def plot_mass_vsmedmaxoverlap(nsim0):
     nsim0 : int
         Reference simulation index.
     """
-    x, __, max_overlap, __, __ = get_overlap(nsim0)
+    x, __, max_overlap, __, __ = get_overlap("csiborg", nsim0)
 
     for i in trange(max_overlap.shape[0]):
         if numpy.sum(numpy.isnan(max_overlap[i, :])) > 0:
@@ -255,7 +258,7 @@ def plot_summed_overlap_vs_mass(nsim0):
     -------
     None
     """
-    x, __, __, summed_overlap, prob_nomatch = get_overlap(nsim0)
+    x, __, __, summed_overlap, prob_nomatch = get_overlap("csiborg", nsim0)
     del __
     collect()
 
@@ -389,12 +392,14 @@ def plot_mass_vs_separation(nsim0, nsimx, plot_std=False, min_overlap=0.0):
 
 
 @cache_to_disk(7)
-def get_max_key(nsim0, key):
+def get_max_key(simname, nsim0, key):
     """
     Get the value of a maximum overlap halo's property.
 
     Parameters
     ----------
+    simname : str
+        Simulation name.
     nsim0 : int
         Reference simulation index.
     key : str
@@ -412,7 +417,8 @@ def get_max_key(nsim0, key):
         Value of the property of the maximum overlap halo.
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsimxs = csiborgtools.read.get_cross_sims(nsim0, paths, smoothed=True)
+    nsimxs = csiborgtools.read.get_cross_sims(simname, nsim0, paths,
+                                              smoothed=True)
     nsimxs = nsimxs
     cat0 = open_cat(nsim0)
 
@@ -440,7 +446,7 @@ def plot_maxoverlap_mass(nsim0):
     nsim0 : int
         Reference simulation index.
     """
-    mass0, __, __, stat = get_max_key(nsim0, "totpartmass")
+    mass0, __, __, stat = get_max_key("csiborg", nsim0, "totpartmass")
 
     mu = numpy.mean(stat, axis=1)
     std = numpy.std(numpy.log10(stat), axis=1)
@@ -463,8 +469,10 @@ def plot_maxoverlap_mass(nsim0):
 
         axs[0].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
         axs[1].set_xlabel(r"$\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[0].set_ylabel(r"Max. overlap mean of $\log M_{\rm tot} ~ [M_\odot / h]$")
-        axs[1].set_ylabel(r"Max. overlap std. of $\log M_{\rm tot} ~ [M_\odot / h]$")
+        axs[0].set_ylabel(
+            r"Max. overlap mean of $\log M_{\rm tot} ~ [M_\odot / h]$")
+        axs[1].set_ylabel(
+            r"Max. overlap std. of $\log M_{\rm tot} ~ [M_\odot / h]$")
 
         ims = [im0, im1]
         for i in range(2):
@@ -498,7 +506,7 @@ def plot_maxoverlapstat(nsim0, key):
         Property to get.
     """
     assert key != "totpartmass"
-    mass0, key_val, __, stat = get_max_key(nsim0, key)
+    mass0, key_val, __, stat = get_max_key("csiborg", nsim0, key)
 
     xlabels = {"lambda200c": r"\log \lambda_{\rm 200c}"}
     key_label = xlabels.get(key, key)
@@ -546,13 +554,15 @@ def plot_maxoverlapstat(nsim0, key):
 
 
 @cache_to_disk(7)
-def get_expected_mass(nsim0, min_overlap):
+def get_expected_mass(simname, nsim0, min_overlap):
     """
     Get the expected mass of a reference halo given its overlap with halos
     from other simulations.
 
     Parameters
     ----------
+    simname : str
+        Simulation name.
     nsim0 : int
         Reference simulation index.
     min_overlap : float
@@ -570,7 +580,8 @@ def get_expected_mass(nsim0, min_overlap):
         Probability of not matching the reference halo.
     """
     paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsimxs = csiborgtools.read.get_cross_sims(nsim0, paths, smoothed=True)
+    nsimxs = csiborgtools.read.get_cross_sims(simname, nsim0, paths,
+                                              smoothed=True)
     nsimxs = nsimxs
     cat0 = open_cat(nsim0)
 
@@ -602,7 +613,8 @@ def plot_mass_vs_expected_mass(nsim0, min_overlap=0, max_prob_nomatch=1):
     max_prob_nomatch : float, optional
         Maximum probability of no match to consider.
     """
-    mass, mu, std, prob_nomatch = get_expected_mass(nsim0, min_overlap)
+    mass, mu, std, prob_nomatch = get_expected_mass("csiborg", nsim0,
+                                                    min_overlap)
 
     std = std / mu / numpy.log(10)
     mass = numpy.log10(mass)
@@ -1343,7 +1355,7 @@ def plot_kl_vs_overlap(runs, nsim, kwargs, runs_to_mass, plot_std=True,
     for run in runs:
         nn_data = nn_reader.read_single("csiborg", run, nsim, nobs=None)
         nn_hindxs = nn_data["ref_hindxs"]
-        mass, overlap_hindxs, __, summed_overlap, prob_nomatch = get_overlap(nsim)  # noqa
+        mass, overlap_hindxs, __, summed_overlap, prob_nomatch = get_overlap("csiborg", nsim)  # noqa
 
         # We need to match the hindxs between the two.
         hind2overlap_array = {hind: i for i, hind in enumerate(overlap_hindxs)}
diff --git a/setup.py b/setup.py
index cd7cc45..e65b32e 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,24 @@
 from setuptools import find_packages, setup
 
+# List of dependencies:
+#   - Corrfunc  -> To be moved to a separate package.
+#   - NumPy
+#   - SciPy
+#   - Numba
+#   - Pylians
+#   - tqdm
+#   - healpy
+#   - astropy
+#   - scikit-learn
+#   - joblib
+#   - h5py
+#   - MPI
+#   - pyyaml
+#   - taskmaster
+#   - matplotlib
+#   - scienceplots
+#   - cache_to_disk
+
 
 BUILD_REQ = ["numpy", "scipy"]
 INSTALL_REQ = BUILD_REQ