Add pynbody and other support (#92)

* Simplify box units * Move old scripts * Add printing * Update readers * Disable boundscheck * Add new ordering * Clean up imports * Enforce dtype and add mass to quijote * Simplify print statements * Fix little typos * Fix key bug * Bug fixing * Delete boring comments * Improve ultimate clumps for PHEW * Delete boring comments * Add basic reading * Remove 0th index HID * Add flipping of X and Z * Updates to halo catalogues * Add ordered caching * Fix flipping * Add new flags * Fix PHEW empty clumps * Stop over-wrriting * Little improvements to angular neighbours * Add catalogue masking * Change if-else statements * Cache only filtered data * Add PHEW cats * Add comments * Sort imports * Get Quijote workign * Docs * Add HMF calculation * Move to old * Fix angular * Add great circle distance * Update imports * Update impotrts * Update docs * Remove unused import * Fix a quick bug * Update compatibility * Rename files * Renaming * Improve compatiblity * Rename snapsht * Fix snapshot bug * Update interface * Finish updating interface * Update all paths * Add old scripts * Add basic halo * Update imports * Improve snapshot processing * Update ordering * Fix how CM positions accessed * Add merger paths * Add imports * Add merger reading * Add making a merger tree * Add a basic merger tree reader * Add imports * Add main branch walking + comments + debuggin * Get tree running * Add working merger tree walking along main branch * Add units conversion for merger data * Add hid_to_array_index * Update merger tree * Add mergertree mass to PHEWcat * Edit comments * Add this to track changes... * Fix a little bug * Add mergertree mass * Add cache clearing * Improve summing substructure code * Littbe bug * Little updates to the merger tree reader * Update .giignore * Add box selection * Add optional deletingf of a group * add to keep track of changes * Update changes * Remove * Add manual tracker * Fix bug * Add m200c_to_r200c * Add manual halo tracking * Remove skipped snapshots * update cosmo params to match csiborg * remove old comments * Add SDSSxALFALFA * Fix bugs * Rename * Edit paths * Updates * Add comments * Add comment * Add hour conversion * Add imports * Add new observation class * Add selection * Add imports * Fix small bug * Add field copying for safety * Add matching to survey without masking * Add P(k) calculation * Add nb * Edit comment * Move files * Remove merger import * Edit setup.yp * Fix typo * Edit import warnigns * update nb * Update README * Update README * Update README * Add skeleton * Add skeleton
2025-05-14 06:31:11 +00:00 · 2023-12-07 14:23:32 +00:00 · 2023-12-07 14:23:32 +00:00 · e972f8e3f2
commit e972f8e3f2
parent 5500fbd2b9
53 changed files with 4627 additions and 1774 deletions
--- a/old/cluster_knn_auto.py
+++ b/old/cluster_knn_auto.py
@ -0,0 +1,155 @@
+# Copyright (C) 2022 Richard Stiskalek
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+"""
+A script to calculate the KNN-CDF for a set of halo catalogues.
+"""
+from argparse import ArgumentParser
+from datetime import datetime
+from distutils.util import strtobool
+
+import joblib
+import numpy
+import yaml
+from mpi4py import MPI
+from sklearn.neighbors import NearestNeighbors
+from taskmaster import work_delegation
+
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+
+    sys.path.append("../")
+    import csiborgtools
+
+from utils import open_catalogues
+
+
+def do_auto(args, config, cats, nsim, paths):
+    """
+    Calculate the kNN-CDF single catalogue auto-correlation.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Command line arguments.
+    config : dict
+        Configuration dictionary.
+    cats : dict
+        Dictionary of halo catalogues. Keys are simulation indices, values are
+        the catalogues.
+    nsim : int
+        Simulation index.
+    paths : csiborgtools.paths.Paths
+        Paths object.
+
+    Returns
+    -------
+    None
+    """
+    cat = cats[nsim]
+    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"],
+        nsamples=int(config["nsamples"]), neval=int(config["neval"]),
+        batch_size=int(config["batch_size"]), random_state=config["seed"])
+    totvol = (4 / 3) * numpy.pi * args.Rmax ** 3
+    fout = paths.knnauto(args.simname, args.run, nsim)
+    if args.verbose:
+        print(f"Saving output to `{fout}`.")
+    joblib.dump({"rs": rs, "cdf": cdf, "ndensity": len(cat) / totvol}, fout)
+
+
+def do_cross_rand(args, config, cats, nsim, paths):
+    """
+    Calculate the kNN-CDF cross catalogue random correlation.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Command line arguments.
+    config : dict
+        Configuration dictionary.
+    cats : dict
+        Dictionary of halo catalogues. Keys are simulation indices, values are
+        the catalogues.
+    nsim : int
+        Simulation index.
+    paths : csiborgtools.paths.Paths
+        Paths object.
+
+    Returns
+    -------
+    None
+    """
+    cat = cats[nsim]
+    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])
+    knn2.fit(pos2)
+
+    knncdf = csiborgtools.clustering.kNN_1DCDF()
+    rs, cdf0, cdf1, joint_cdf = knncdf.joint(
+        knn1, knn2, rvs_gen=rvs_gen, nneighbours=int(config["nneighbours"]),
+        rmin=config["rmin"], rmax=config["rmax"],
+        nsamples=int(config["nsamples"]), neval=int(config["neval"]),
+        batch_size=int(config["batch_size"]), random_state=config["seed"])
+    corr = knncdf.joint_to_corr(cdf0, cdf1, joint_cdf)
+
+    fout = paths.knnauto(args.simname, args.run, nsim)
+    if args.verbose:
+        print(f"Saving output to `{fout}`.", flush=True)
+    joblib.dump({"rs": rs, "corr": corr}, fout)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--run", type=str, help="Run name.")
+    parser.add_argument("--simname", type=str, choices=["csiborg", "quijote"],
+                        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,
+                        help="High-resolution region radius")  # noqa
+    parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
+                        default=False)
+    args = parser.parse_args()
+
+    with open("./cluster_knn_auto.yml", "r") as file:
+        config = yaml.safe_load(file)
+    comm = MPI.COMM_WORLD
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    cats = open_catalogues(args, config, paths, comm)
+
+    if args.verbose and comm.Get_rank() == 0:
+        print(f"{datetime.now()}: starting to calculate the kNN statistic.")
+
+    def do_work(nsim):
+        if "random" in args.run:
+            do_cross_rand(args, config, cats, nsim, paths)
+        else:
+            do_auto(args, config, cats, nsim, paths)
+
+    nsims = list(cats.keys())
+    work_delegation(do_work, nsims, comm, master_verbose=args.verbose)
+
+    comm.Barrier()
+    if comm.Get_rank() == 0:
+        print(f"{datetime.now()}: all finished. Quitting.")