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/mv_fofmembership.py

@@ -0,0 +1,151 @@
+# 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.
+"""
+Short script to move and change format of the CSiBORG FoF membership files
+calculated by Julien. Additionally, also orders the particles in the same way
+as the PHEW halo finder output.
+"""
+from argparse import ArgumentParser
+from datetime import datetime
+from gc import collect
+from os.path import join
+from shutil import copy
+
+import numpy
+from mpi4py import MPI
+from taskmaster import work_delegation
+from tqdm import trange
+
+from utils import get_nsims
+
+try:
+    import csiborgtools
+except ModuleNotFoundError:
+    import sys
+    sys.path.append("../")
+    import csiborgtools
+
+
+def copy_membership(nsim, verbose=True):
+    """
+    Copy the FoF particle halo membership to the CSiBORG directory and write it
+    as a NumPy array instead of a text file.
+
+    Parameters
+    ----------
+    nsim : int
+        IC realisation index.
+    verbose : bool, optional
+        Verbosity flag.
+    """
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    fpath = join("/mnt/extraspace/jeg/greenwhale/Constrained_Sims",
+                 f"sim_{nsim}/particle_membership_{nsim}_FOF.txt")
+    if verbose:
+        print(f"Loading from ... `{fpath}`.")
+    data = numpy.genfromtxt(fpath, dtype=int)
+
+    fout = paths.fof_membership(nsim)
+    if verbose:
+        print(f"Saving to ... `{fout}`.")
+    numpy.save(fout, data)
+
+
+def copy_catalogue(nsim, verbose=True):
+    """
+    Move the FoF catalogue to the CSiBORG directory.
+
+    Parameters
+    ----------
+    nsim : int
+        IC realisation index.
+    verbose : bool, optional
+        Verbosity flag.
+    """
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    source = join("/mnt/extraspace/jeg/greenwhale/Constrained_Sims",
+                  f"sim_{nsim}/halo_catalog_{nsim}_FOF.txt")
+    dest = paths.fof_cat(nsim)
+    if verbose:
+        print("Copying`{}` to `{}`.".format(source, dest))
+    copy(source, dest)
+
+
+def sort_fofid(nsim, verbose=True):
+    """
+    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.
+    """
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsnap = max(paths.get_snapshots(nsim))
+    fpath = paths.fof_membership(nsim)
+    if verbose:
+        print(f"{datetime.now()}: loading from ... `{fpath}`.")
+    # Columns are halo ID, particle ID.
+    fof = numpy.load(fpath)
+
+    reader = csiborgtools.read.ParticleReader(paths)
+    pars_extract = ["x"]  # Dummy variable
+    __, pids = reader.read_particle(nsnap, nsim, pars_extract,
+                                    return_structured=False, verbose=verbose)
+    del __
+    collect()
+
+    # Map the particle IDs in pids to their corresponding PHEW array index
+    if verbose:
+        print(f"{datetime.now()}: mapping particle IDs to their indices.")
+    pids_idx = {pid: i for i, pid in enumerate(pids)}
+
+    if verbose:
+        print(f"{datetime.now()}: mapping FoF HIDs to their array indices.")
+    # Unassigned particle IDs are assigned a halo ID of 0. Same as PHEW.
+    fof_hids = numpy.zeros(pids.size, dtype=numpy.int32)
+    for i in trange(fof.shape[0]) if verbose else range(fof.shape[0]):
+        hid, pid = fof[i]
+        fof_hids[pids_idx[pid]] = hid
+
+    fout = paths.fof_membership(nsim, sorted=True)
+    if verbose:
+        print(f"Saving the sorted data to ... `{fout}`")
+    numpy.save(fout, fof_hids)
+
+
+def main(nsim, verbose=True):
+    copy_membership(nsim, verbose=verbose)
+    copy_catalogue(nsim, verbose=verbose)
+    sort_fofid(nsim, verbose=verbose)
+
+
+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="Indices of simulations to cross. If `-1` processes all simulations.")  # noqa
+    args = parser.parse_args()
+
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = get_nsims(args, paths)
+    comm = MPI.COMM_WORLD
+
+    work_delegation(main, nsims, comm)