Read merger branch (#104)

* Add support to read merger tree main branch

* Add imports

* Add trees path

* Edit processing

csiborgtools/read/__init__.py

@@ -12,7 +12,8 @@
 # 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.
-from .catalogue import CSiBORG1Catalogue, CSiBORG2Catalogue, QuijoteCatalogue   # noqa
+from .catalogue import (CSiBORG1Catalogue, CSiBORG2Catalogue,                   # noqa
+                        CSiBORG2MergerTreeReader, QuijoteCatalogue)             # noqa
 from .snapshot import (CSIBORG1Snapshot, CSIBORG2Snapshot, QuijoteSnapshot,     # noqa
                        CSiBORG1Field, CSiBORG2Field, QuijoteField)              # noqa
 from .obs import (SDSS, MCXCClusters, PlanckClusters, TwoMPPGalaxies,           # noqa

csiborgtools/read/catalogue.py

@@ -741,10 +741,17 @@ class CSiBORG2Catalogue(BaseCatalogue):
             676.6, [338.3, 338.3, 338.3], observer_velocity, cache_maxsize)
 
         self._custom_keys = ["GroupFirstSub", "GroupContamination",
-                             "GroupNsubs"]
+                             "GroupNsubs", "Group_M_Crit200"]
 
     @property
     def kind(self):
+        """
+        Simulation kind.
+
+        Returns
+        -------
+        str
+        """
         return self._simname.split("_")[-1]
 
     def _read_fof_catalogue(self, kind):
@@ -802,12 +809,257 @@ class CSiBORG2Catalogue(BaseCatalogue):
     def GroupNsubs(self):
         return self._read_fof_catalogue("GroupNsubs")
 
+    @property
+    def Group_M_Crit200(self):
+        return self._read_fof_catalogue("Group_M_Crit200")
+
     @property
     def GroupContamination(self):
         mass_type = self._read_fof_catalogue("GroupMassType")
         return mass_type[:, 5] / (mass_type[:, 1] + mass_type[:, 5])
 
 
+###############################################################################
+#                      CSiBORG2 merger tree reader                            #
+###############################################################################
+
+
+class CSiBORG2MergerTreeReader:
+    """
+    Merger tree reader for CSiBORG2. Currently supports reading the main branch
+    of the most massive FoF member at `z = 99`. Documentation of the merger
+    trees can be found at:
+
+        https://www.mtng-project.org/03_output/#descendant-and-progenitor-information       # noqa
+
+    Parameters
+    ----------
+    nsim : int
+        IC realisation index.
+    kind : str
+        Simulation kind. Must be one of 'main', 'varysmall', or 'random'.
+    paths : py:class`csiborgtools.read.Paths`, optional
+        Paths object.
+    """
+    _simname = None
+    _nsim = None
+
+    def __init__(self, nsim, kind, paths=None):
+        self.simname = f"csiborg2_{kind}"
+        self.nsim = nsim
+        self._paths = paths
+
+        self._group2treeid = self.make_group_to_treeid()
+
+    @property
+    def simname(self):
+        """
+        Simulation name.
+
+        Returns
+        -------
+        str
+        """
+        if self._simname is None:
+            raise RuntimeError("`simname` is not set!")
+        return self._simname
+
+    @simname.setter
+    def simname(self, simname):
+        if not isinstance(simname, str):
+            raise TypeError("`simname` must be a string.")
+        self._simname = simname
+
+    @property
+    def nsim(self):
+        """
+        Simulation IC realisation index.
+
+        Returns
+        -------
+        int
+        """
+        if self._nsim is None:
+            raise RuntimeError("`nsim` is not set!")
+        return self._nsim
+
+    @nsim.setter
+    def nsim(self, nsim):
+        if not isinstance(nsim, (int, numpy.integer)):
+            raise TypeError("`nsim` must be an integer.")
+        self._nsim = int(nsim)
+
+    @property
+    def kind(self):
+        """
+        Simulation kind.
+
+        Returns
+        -------
+        str
+        """
+        return self._simname.split("_")[-1]
+
+    @property
+    def paths(self):
+        """
+        Paths manager.
+
+        Returns
+        -------
+        py:class:`csiborgtools.read.Paths`
+        """
+        if self._paths is None:
+            return Paths(**paths_glamdring)
+        return self._paths
+
+    def make_group_to_treeid(self):
+        """
+        Create a mapping from group number at snapshot 99 to its tree ID.
+
+        Returns
+        -------
+        group2treeid : dict
+            Dictionary with group number as key and tree ID as value.
+        """
+        print("Creating group to tree ID mapping...")
+        with File(self.paths.trees(self.nsim, self.simname), 'r') as f:
+            groupnr = f["TreeHalos/GroupNr"][:]
+            snapnum = f["TreeHalos/SnapNum"][:]
+            treeid = f["TreeHalos/TreeID"][:]
+
+            # Select only groups in the last snapshot
+            mask = snapnum == 99
+            groupnr = groupnr[mask]
+            treeid = treeid[mask]
+
+            group2treeid = {g: t for g, t in zip(groupnr, treeid)}
+
+        return group2treeid
+
+    def get_tree(self, tree_id):
+        """
+        Extract a tree from the merger tree file.
+
+        Parameters
+        ----------
+        tree_id : int
+            Tree ID.
+
+        Returns
+        -------
+        dict
+        """
+        keys_extract = ["GroupNr",
+                        "SnapNum",
+                        "TreeFirstHaloInFOFgroup",
+                        "TreeMainProgenitor",
+                        "TreeNextProgenitor",
+                        "TreeProgenitor",
+                        "SubhaloMass",
+                        ]
+
+        with File(self.paths.trees(self.nsim, self.simname), 'r') as f:
+            max_treeid = f["TreeTable/TreeID"][-1]
+            if not (0 <= tree_id <= max_treeid):
+                raise ValueError(f"Tree ID must be between 0 and {max_treeid}.")  # noqa
+
+            i = f["TreeTable/StartOffset"][tree_id]
+            j = i + f["TreeTable/Length"][tree_id]
+
+            tree = {}
+            for key in keys_extract:
+                tree[key] = f[f"TreeHalos/{key}"][i:j]
+
+            tree["Redshift"] = f["TreeTimes/Redshift"][:]
+            tree["Time"] = f["TreeTimes/Time"][:]
+
+        return tree
+
+    def find_first_fof_member(self, group_nr, tree):
+        """
+        Find the first FOF member of a group at snapshot 99.
+
+        Parameters
+        ----------
+        group_nr : int
+            Group number at snapshot 99.
+        tree : dict
+            Merger tree.
+
+        Returns
+        -------
+        int
+        """
+        fsnap_mask = tree["SnapNum"] == 99
+        group_mask = tree["GroupNr"][fsnap_mask] == group_nr
+        first_fof = tree["TreeFirstHaloInFOFgroup"][fsnap_mask][group_mask]
+
+        if not numpy.all(first_fof == first_fof[0]):
+            raise RuntimeError("We get non-unique first FOF group.")
+
+        return first_fof[0]
+
+    def main_progenitor(self, group_nr):
+        """
+        Return the mass and redshift of the main progenitor of a group.
+
+        Parameters
+        ----------
+        group_nr : int
+            Group number at snapshot 99.
+
+        Returns
+        -------
+        dict
+        """
+        tree_id = self._group2treeid[group_nr]
+        tree = self.get_tree(tree_id)
+        first_fof = self.find_first_fof_member(group_nr, tree)
+
+        n = first_fof  # Index of the current main progenitor
+
+        time, redshift, main_progenitor_mass = [], [], []
+        max_next_progenitor_mass, total_next_progenitor_mass = [], []
+        while True:
+            # First off attempt to find the next progenitors of the current
+            # halo. Deal with the main progenitor later.
+            next_prog = tree["TreeNextProgenitor"][n]
+            if next_prog != -1:
+                minors = []
+                while True:
+                    minors.append(tree["SubhaloMass"][next_prog])
+
+                    next_prog = tree["TreeNextProgenitor"][next_prog]
+
+                    if next_prog == -1:
+                        break
+            else:
+                # Fiducially set it to zero.
+                minors = [0]
+
+            # Update data with information from the current main progenitor.
+            major = tree["SubhaloMass"][n]
+            main_progenitor_mass.append(major)
+            max_next_progenitor_mass.append(max(minors))
+            total_next_progenitor_mass.append(sum(minors))
+            redshift.append(tree["Redshift"][tree["SnapNum"][n]])
+            time.append(tree["Time"][tree["SnapNum"][n]])
+
+            # Update n to the next main progenitor.
+            n = tree["TreeMainProgenitor"][n]
+
+            if n == -1:
+                break
+
+        return {"Time": numpy.array(time),
+                "Redshift": numpy.array(redshift),
+                "MainProgenitorMass": numpy.array(main_progenitor_mass) * 1e10,
+                "MaxNextProgenitorMass": numpy.array(max_next_progenitor_mass) * 1e10,  # noqa
+                "TotalNextProgenitorMass": numpy.array(total_next_progenitor_mass) * 1e10,  # noqa
+                }
+
+
 ###############################################################################
 #                         Quijote halo catalogue                              #
 ###############################################################################

csiborgtools/read/paths.py

@@ -231,6 +231,38 @@ class Paths:
         else:
             raise ValueError(f"Unknown simulation name `{simname}`.")
 
+    def trees(self, nsim, simname):
+        """
+        Path to the halo trees of a simulation snapshot.
+
+        Parameters
+        ----------
+        nsim : int
+            IC realisation index.
+        simname : str
+            Simulation name.
+
+        Returns
+        -------
+        str
+        """
+        if simname == "csiborg1":
+            raise ValueError("Trees not available for CSiBORG1.")
+        elif simname == "csiborg2_main":
+            return join(self.csiborg2_main_srcdir, f"chain_{nsim}", "output",
+                        "trees.hdf5")
+        elif simname == "csiborg2_random":
+            return join(self.csiborg2_ranodm_srcdir, f"chain_{nsim}", "output",
+                        "trees.hdf5")
+        elif simname == "csiborg2_varysmall":
+            return join(self.csiborg2_varysmall_srcdir,
+                        f"chain_16417_{str(nsim).zfill(3)}", "output",
+                        "trees.hdf5")
+        elif simname == "quijote":
+            raise ValueError("Trees not available for Quijote.")
+        else:
+            raise ValueError(f"Unknown simulation name `{simname}`.")
+
     def overlap(self, simname, nsim0, nsimx, min_logmass, smoothed):
         """
         Path to the overlap files between two CSiBORG simulations.

scripts_independent/run_process_snapshot.py

@@ -20,8 +20,12 @@ if __name__ == "__main__":
     # simname = "csiborg2_main"
     # mode = 1
 
+    # chains = [1] + [25 + n * 25 for n in range(19)]
+    # simname = "csiborg2_varysmall"
+    # mode = 1
+
     chains = [1] + [25 + n * 25 for n in range(19)]
-    simname = "csiborg2_varysmall"
+    simname = "csiborg2_random"
     mode = 1
 
     # chains = [7444 + n * 24 for n in range(1, 101)]