Matching paper plots (#91)

* Fix calculations of expected mass

* Add paper plots

* Edits to pltos

* Add overlap summary

* Add imports

* Add import

* Add binned stat

* Add fit

* Add more plots

* Add basic env

* Add histogram mode

* Edit expected mass

* Improve expected plots

* Clean up plot

* Improve separation plot

* Update plots

* Edit expected calculation

* Update plotting

* Update plots

* Update plots

* Update plots

* Add conc fraction

* Add halo maker sorting

* Renaming

* Add import

* Add NaN treatment

* add import

* Move cosine smi

* Update plots

* Move similarity

* Fix little bugs

* Shorten documentation

* Update plots

scripts/sort_halomaker.py

@@ -0,0 +1,100 @@
+# 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.
+"""
+Script to sort the HaloMaker's `particle_membership` file to match the ordering
+of particles in the simulation snapshot.
+"""
+from argparse import ArgumentParser
+from datetime import datetime
+from glob import iglob
+
+import h5py
+import numpy
+import pynbody
+from mpi4py import MPI
+from taskmaster import work_delegation
+from tqdm import trange
+
+import csiborgtools
+
+
+def sort_particle_membership(nsim, nsnap, method):
+    """
+    Read the FoF particle halo membership and sort the halo IDs to the ordering
+    of particles in the PHEW clump IDs.
+
+    Parameters
+    ----------
+    nsim : int
+        IC realisation index.
+    verbose : bool, optional
+        Verbosity flag.
+    """
+    print(f"{datetime.now()}:   starting simulation {nsim}, snapshot {nsnap} and method {method}.")  # noqa
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+
+    fpath = next(iglob(f"/mnt/extraspace/rstiskalek/CSiBORG/halo_maker/ramses_{nsim}/output_{str(nsnap).zfill(5)}/**/*particle_membership*", recursive=True), None)  # noqa
+    print(f"{datetime.now()}:   loading particle membership `{fpath}`.")
+    # Columns are halo ID, particle ID
+    membership = numpy.genfromtxt(fpath, dtype=int)
+
+    print(f"{datetime.now()}:   loading particle IDs from the snapshot.")
+    sim = pynbody.load(paths.snapshot(nsnap, nsim, "csiborg"))
+    pids = numpy.asanyarray(sim["iord"])
+
+    print(f"{datetime.now()}:   mapping particle IDs to their indices.")
+    pids_idx = {pid: i for i, pid in enumerate(pids)}
+
+    print(f"{datetime.now()}:   mapping HIDs to their array indices.")
+    # Unassigned particle IDs are assigned a halo ID of 0.
+    hids = numpy.zeros(pids.size, dtype=numpy.int32)
+    for i in trange(membership.shape[0]):
+        hid, pid = membership[i]
+        hids[pids_idx[pid]] = hid
+
+    fout = fpath + "_sorted.hdf5"
+    print(f"{datetime.now()}:   saving the sorted data to ... `{fout}`")
+
+    header = """
+    This dataset represents halo indices for each particle.
+        - The particles are ordered as they appear in the simulation snapshot.
+        - Unassigned particles are given a halo index of 0.
+        """
+    with h5py.File(fout, 'w') as hdf:
+        dset = hdf.create_dataset('hids_dataset', data=hids)
+        dset.attrs['header'] = header
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--method", type=str, required=True,
+                        help="HaloMaker method")
+    parser.add_argument("--nsim", type=int, required=False, default=None,
+                        help="IC index. If not set process all.")
+    args = parser.parse_args()
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+
+    if args.nsim is None:
+        ics = paths.get_ics("csiborg")
+    else:
+        ics = [args.nsim]
+
+    snaps = numpy.array([max(paths.get_snapshots(nsim, "csiborg"))
+                         for nsim in ics])
+
+    def main(n):
+        sort_particle_membership(ics[n], snaps[n], args.method)
+
+    work_delegation(main, list(range(len(ics))), MPI.COMM_WORLD)

scripts/sort_initsnap.py

@@ -29,16 +29,9 @@ import numpy
 from mpi4py import MPI
 from taskmaster import work_delegation
 
+import csiborgtools
 from utils import get_nsims
 
-try:
-    import csiborgtools
-except ModuleNotFoundError:
-    import sys
-
-    sys.path.append("../")
-    import csiborgtools
-
 
 def _main(nsim, simname, verbose):
     """
@@ -51,9 +44,7 @@ def _main(nsim, simname, verbose):
     else:
         partreader = csiborgtools.read.QuijoteReader(paths)
 
-    if verbose:
-        print(f"{datetime.now()}: reading and processing simulation `{nsim}`.",
-              flush=True)
+    print(f"{datetime.now()}:   processing simulation `{nsim}`.", flush=True)
     # We first load the particle IDs in the final snapshot.
     pidf = csiborgtools.read.read_h5(paths.particles(nsim, simname))
     pidf = pidf["particle_ids"]