CSiBORG FoF switch (#75)

* Add moving FoF membership files

* add FoF membership path

* Add notes where its PHEW

* Add FoF catalogue path

* Correct typo

* Add more functionalities

* Make work with halo IDs from FoF

* Edit print statement

* Fix copy bug

* copy

* Add FoF catalogue reading

* Clean up script

* Fix typo

* Little edits

* Fix naming convention

* Rename key

* Remove loading substructure particles

* Rename CSiBORG Cat

* Rename clumps cat

* Rename cat

* Remove misplaced import

* Switch to halos

* rm import

* structfit of only halos

* Add FoF halo reading

* Add a short comment

* Fix __getitem__ to work with int

* Fix problems

* Improve __getitem__

* Add more conversion

* Fix indexing

* Fix __getitem__ assertion

* Fix numbers

* Rename

* Fix verbosity flags

* Add full Quijote HMF option

* Add plot of Quijote only

* Add quijote full paths

* Fix the fit_init script

* Renam arg

* Update .gitignore

* add default argument name

* Change default verbosity flag

* Modernise script structure

* Fix dictionary

* Fix reading to include m200c

* Modernise script

* Add args

scripts/fit_init.py

@@ -24,6 +24,8 @@ import numpy
 from mpi4py import MPI
 from tqdm import tqdm
 
+from utils import get_nsims
+
 try:
     import csiborgtools
 except ModuleNotFoundError:
@@ -41,16 +43,16 @@ verbose = nproc == 1
 
 # Argument parser
 parser = ArgumentParser()
-parser.add_argument("--ics", type=int, nargs="+", default=None,
+parser.add_argument("--simname", type=str, default="csiborg",
+                    choices=["csiborg", "quijote"],
+                    help="Simulation name")
+parser.add_argument("--nsims", type=int, nargs="+", default=None,
                     help="IC realisations. If `-1` processes all simulations.")
 args = parser.parse_args()
 paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
 partreader = csiborgtools.read.ParticleReader(paths)
 
-if args.ics is None or args.ics[0] == -1:
-    ics = paths.get_ics("csiborg")
-else:
-    ics = args.ics
+nsims = get_nsims(args, paths)
 
 cols_collect = [("index", numpy.int32),
                 ("x", numpy.float32),
@@ -61,8 +63,8 @@ cols_collect = [("index", numpy.int32),
 
 
 # MPI loop over simulations
-jobs = csiborgtools.fits.split_jobs(len(ics), nproc)[rank]
-for nsim in [ics[i] for i in jobs]:
+jobs = csiborgtools.fits.split_jobs(len(nsims), nproc)[rank]
+for nsim in [nsims[i] for i in jobs]:
     nsnap = max(paths.get_snapshots(nsim))
     overlapper = csiborgtools.match.ParticleOverlap()
     print(f"{datetime.now()}: rank {rank} calculating simulation `{nsim}`.",
@@ -70,22 +72,18 @@ for nsim in [ics[i] for i in jobs]:
 
     parts = csiborgtools.read.read_h5(paths.initmatch(nsim, "particles"))
     parts = parts['particles']
-    clump_map = csiborgtools.read.read_h5(paths.particles(nsim))
-    clump_map = clump_map["clumpmap"]
-    clumps_cat = csiborgtools.read.ClumpsCatalogue(nsim, paths, rawdata=True,
-                                                   load_fitted=False)
-    clid2map = {clid: i for i, clid in enumerate(clump_map[:, 0])}
-    ismain = clumps_cat.ismain
+    halo_map = csiborgtools.read.read_h5(paths.particles(nsim))
+    halo_map = halo_map["halomap"]
+    cat = csiborgtools.read.CSiBORGHaloCatalogue(
+        nsim, paths, rawdata=True, load_fitted=False, load_initial=False)
+    hid2map = {hid: i for i, hid in enumerate(halo_map[:, 0])}
 
-    out = csiborgtools.read.cols_to_structured(len(clumps_cat), cols_collect)
-    indxs = clumps_cat["index"]
-    for i, hid in enumerate(tqdm(indxs) if verbose else indxs):
+    out = csiborgtools.read.cols_to_structured(len(cat), cols_collect)
+    for i, hid in enumerate(tqdm(cat["index"]) if verbose else cat["index"]):
         out["index"][i] = hid
-        if not ismain[i]:
-            continue
+        part = csiborgtools.read.load_halo_particles(hid, parts, halo_map,
+                                                     hid2map)
 
-        part = csiborgtools.read.load_parent_particles(hid, parts, clump_map,
-                                                       clid2map, clumps_cat)
         # Skip if the halo is too small.
         if part is None or part.size < 100:
             continue
@@ -101,7 +99,6 @@ for nsim in [ics[i] for i in jobs]:
         delta = overlapper.make_delta(part[:, :3], part[:, 3], subbox=True)
         out["lagpatch_ncells"][i] = csiborgtools.fits.delta2ncells(delta)
 
-    out = out[ismain]
     # Now save it
     fout = paths.initmatch(nsim, "fit")
     print(f"{datetime.now()}: dumping fits to .. `{fout}`.",