Improving halo fits (#76)

* Add periodic distances

* Little corrections

* Fix little bug

* Modernise the script

* Small updates

* Remove clump

* Add new halo routines

* Fix weights

* Modernise the script

* Add check ups on convergence

* More convergence check ups

* Edit bounds

* Add default argument

* Update fit heuristic and NaNs

* Change maxiter

* Switch NFW minimization to log-sapce

* Remove print statement

* Turn convert_from_box abstract property required for all boxes.

* Move files

* Simplify script

* Improve the argument parser

* Remove optinal argument

* Improve argument parser

* Add a minimum concentration limit

scripts/fit_init.py

@@ -22,6 +22,7 @@ from datetime import datetime
 
 import numpy
 from mpi4py import MPI
+from taskmaster import work_delegation
 from tqdm import tqdm
 
 from utils import get_nsims
@@ -35,73 +36,83 @@ except ModuleNotFoundError:
     import csiborgtools
 
 
-# Get MPI things
-comm = MPI.COMM_WORLD
-rank = comm.Get_rank()
-nproc = comm.Get_size()
-verbose = nproc == 1
+def _main(nsim, simname, verbose):
+    """
+    Calculate the Lagrangian halo centre of mass and Lagrangian patch size in
+    the initial snapshot.
 
-# Argument parser
-parser = ArgumentParser()
-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)
+    Parameters
+    ----------
+    nsim : int
+        IC realisation index.
+    simname : str
+        Simulation name.
+    verbose : bool
+        Verbosity flag.
+    """
+    if simname == "quijote":
+        raise NotImplementedError("Quijote not implemented yet.")
 
-nsims = get_nsims(args, paths)
-
-cols_collect = [("index", numpy.int32),
-                ("x", numpy.float32),
-                ("y", numpy.float32),
-                ("z", numpy.float32),
-                ("lagpatch_size", numpy.float32),
-                ("lagpatch_ncells", numpy.int32),]
-
-
-# MPI loop over simulations
-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}`.",
-          flush=True)
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    cols = [("index", numpy.int32),
+            ("x", numpy.float32),
+            ("y", numpy.float32),
+            ("z", numpy.float32),
+            ("lagpatch_size", numpy.float32),
+            ("lagpatch_ncells", numpy.int32),]
 
     parts = csiborgtools.read.read_h5(paths.initmatch(nsim, "particles"))
     parts = parts['particles']
     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(cat), cols_collect)
+    out = csiborgtools.read.cols_to_structured(len(cat), cols)
     for i, hid in enumerate(tqdm(cat["index"]) if verbose else cat["index"]):
         out["index"][i] = hid
         part = csiborgtools.read.load_halo_particles(hid, parts, halo_map,
                                                      hid2map)
 
-        # Skip if the halo is too small.
+        # Skip if the halo has no particles or is too small.
         if part is None or part.size < 100:
             continue
 
+        pos, mass = part[:, :3], part[:, 3]
         # Calculate the centre of mass and the Lagrangian patch size.
-        dist, cm = csiborgtools.fits.dist_centmass(part)
-        # We enforce a maximum patchsize of 0.075 in box coordinates.
-        patchsize = min(numpy.percentile(dist, 99), 0.075)
+        cm = csiborgtools.fits.center_of_mass(pos, mass, boxsize=1.0)
+        distances = csiborgtools.fits.periodic_distance(pos, cm, boxsize=1.0)
         out["x"][i], out["y"][i], out["z"][i] = cm
-        out["lagpatch_size"][i] = patchsize
+        out["lagpatch_size"][i] = numpy.percentile(distances, 99)
 
         # Calculate the number of cells with > 0 density.
-        delta = overlapper.make_delta(part[:, :3], part[:, 3], subbox=True)
+        overlapper = csiborgtools.match.ParticleOverlap()
+        delta = overlapper.make_delta(pos, mass, subbox=True)
         out["lagpatch_ncells"][i] = csiborgtools.fits.delta2ncells(delta)
 
     # Now save it
     fout = paths.initmatch(nsim, "fit")
-    print(f"{datetime.now()}: dumping fits to .. `{fout}`.",
-          flush=True)
+    if verbose:
+        print(f"{datetime.now()}: dumping fits to .. `{fout}`.", flush=True)
     with open(fout, "wb") as f:
         numpy.save(f, out)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    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.")
+    args = parser.parse_args()
+
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = get_nsims(args, paths)
+
+    def main(nsim):
+        _main(nsim, args.simname, MPI.COMM_WORLD.Get_size() == 1)
+
+    work_delegation(main, nsims, MPI.COMM_WORLD)