Fix overlap runs (#125)

* Update nb

* Update script

* Update script

* Rename

* Update script

* Update script

* Remove warning

* Ignore minors when extracting MAH

* Fix paths bug

* Move notebooks

* Move files

* Rename and delete things

* Rename file

* Move file

* Rename things

* Remove old print statement

* Add basic MAH plot

* Add random MAH path

* Output snapshot numbers

* Add MAH random extraction

* Fix redshift bug

* Edit script

* Add extracting random MAH

* Little updates

* Add CB2 redshift

* Add some caching

* Add diagnostic plots

* Add caching

* Minor updates

* Update nb

* Update notebook

* Update script

* Add Sorce randoms

* Add CB2 varysmall

* Update nb

* Update nb

* Update nb

* Use catalogue HMF

* Move definition of radec2galactic

* Update nb

* Update import

* Update import

* Add galatic coords to catalogues

* Update nb

csiborgtools/__init__.py

@@ -19,8 +19,9 @@ from .utils import (center_of_mass, delta2ncells, number_counts,
                     binned_statistic, cosine_similarity, fprint,                # noqa
                     hms_to_degrees, dms_to_degrees, great_circle_distance,      # noqa
                     radec_to_cartesian, cartesian_to_radec,                     # noqa
-                    thin_samples_by_acl, numpyro_gof)                           # noqa
+                    thin_samples_by_acl, numpyro_gof, radec_to_galactic)        # noqa
-from .params import paths_glamdring, simname2boxsize, simname2Omega_m           # noqa
+from .params import (paths_glamdring, simname2boxsize, simname2Omega_m,         # noqa
+                     snap2redshift)                                             # noqa
 
 
 ###############################################################################

csiborgtools/flow/flow_model.py

@@ -26,8 +26,6 @@ from warnings import catch_warnings, simplefilter, warn
 import numpy as np
 import numpyro
 import numpyro.distributions as dist
-from astropy import units as u
-from astropy.coordinates import SkyCoord
 from astropy.cosmology import FlatLambdaCDM
 from h5py import File
 from jax import jit
@@ -43,6 +41,7 @@ from sklearn.model_selection import KFold
 from tqdm import trange
 
 from ..params import simname2Omega_m
+from ..utils import radec_to_galactic
 
 SPEED_OF_LIGHT = 299792.458  # km / s
 H0 = 100                     # km / s / Mpc
@@ -53,24 +52,6 @@ def t():
     return datetime.now().strftime("%H:%M:%S")
 
 
-def radec_to_galactic(ra, dec):
-    """
-    Convert right ascension and declination to galactic coordinates (all in
-    degrees.)
-
-    Parameters
-    ----------
-    ra, dec : float or 1-dimensional array
-        Right ascension and declination in degrees.
-
-    Returns
-    -------
-    l, b : float or 1-dimensional array
-    """
-    c = SkyCoord(ra=ra*u.degree, dec=dec*u.degree, frame='icrs')
-    return c.galactic.l.degree, c.galactic.b.degree
-
-
 ###############################################################################
 #                             Data loader                                     #
 ###############################################################################

csiborgtools/params.py

@@ -16,6 +16,62 @@
 Various user parameters for CSiBORGTools.
 """
 
+CB2_REDSHIFT = [69.0000210000063, 40.250007218751264, 28.24050991940438,
+                21.6470609550175, 17.480001404480106, 14.608109099433955,
+                12.508772664512199, 10.90721705951751, 9.64516173673259,
+                8.625000360937513, 7.7832702592057235, 7.0769233254437935,
+                6.475728365821477, 5.95783150553419, 5.50704240932355,
+                5.111111246913583, 4.760598622974984, 4.448113312911626,
+                4.1677853285437605, 3.914893700679041, 3.685598452365574,
+                3.476744253718227, 3.285714346938776, 3.1103203402819117,
+                2.9487179993425383, 2.7993421515051513, 2.6608558268213116,
+                2.5321101306287352, 2.4121122957547967, 2.3000000330000008,
+                2.1950207773798662, 2.096514773533915, 2.003901196522936,
+                1.9166666909722223, 1.8343558508261513, 1.7565632668759008,
+                1.6829268488994646, 1.613122190273029, 1.5468577900064306,
+                1.4838709837669097, 1.4239244641145379, 1.366803292753544,
+                1.3123123255056859, 1.2602739849878026, 1.210526327423823,
+                1.162921359250726, 1.117323566656109, 1.0736086272735772,
+                1.0316622782422846, 0.9913793189283591, 0.9526627299814432,
+                0.9154228931957131, 0.8795768989699038, 0.8450479301016136,
+                0.8117647122768166, 0.7796610229819017, 0.7486752517178681,
+                0.7187500053710938, 0.6898317534223188, 0.6618705083794834,
+                0.6348195374209455, 0.6086351017498701, 0.5832762206018658,
+                0.5587044572276223, 0.5348837244997295, 0.5117801080759505,
+                0.48936170529651424, 0.46759847820604516, 0.4464621192761633,
+                0.42592592856652933, 0.4059647012034677, 0.3865546241790834,
+                0.3676731815824261, 0.34929906746973005, 0.3314121056648591,
+                0.31399317585528075, 0.2970241454144613, 0.28048780643961924,
+                0.2643678175452504, 0.2486486499985392, 0.23331553782343795,
+                0.21835443153641232, 0.20375195520916023, 0.18949536658248856,
+                0.17557251998135315, 0.1619718318042056, 0.14868224838055033,
+                0.13569321600925854, 0.122994653006949, 0.11057692361085425,
+                0.09843081359419292, 0.08654750746436402, 0.0749185671253807,
+                0.06353591189600438, 0.05239179978414388, 0.04147880992632613,
+                0.03078982610853953, 0.020318021291547472,
+                0.010056843069963017, 0.0]
+
+
+def snap2redshift(snapnum, simname):
+    """
+    Convert a snapshot number to redshift.
+
+    Parameters
+    ----------
+    snapnum : int
+        Snapshot number.
+    simname : str
+        Simulation name.
+
+    Returns
+    -------
+    float
+    """
+    if "csiborg2_" in simname:
+        return CB2_REDSHIFT[snapnum]
+    else:
+        raise ValueError(f"Unknown simname: {simname}")
+
 
 def simname2boxsize(simname):
     """
@@ -65,6 +121,7 @@ def simname2Omega_m(simname):
     d = {"csiborg1": 0.307,
          "csiborg2_main": 0.3111,
          "csiborg2_random": 0.3111,
+         "csiborg2_varysmall": 0.3111,
          "borg1": 0.307,
          "Carrick2015": 0.3,
          }

csiborgtools/read/catalogue.py

@@ -23,6 +23,7 @@ from functools import lru_cache
 from gc import collect
 from itertools import product
 from math import floor
+from astropy.cosmology import FlatLambdaCDM
 
 import numpy
 from h5py import File
@@ -30,7 +31,8 @@ from sklearn.neighbors import NearestNeighbors
 
 from ..params import paths_glamdring
 from ..utils import (cartesian_to_radec, great_circle_distance, number_counts,
-                     periodic_distance_two_points, real2redshift)
+                     periodic_distance_two_points, real2redshift,
+                     radec_to_galactic)
 from .paths import Paths
 from .snapshot import is_instance_of_base_snapshot_subclass
 
@@ -46,6 +48,7 @@ class BaseCatalogue(ABC):
     """
     _properties = ["cartesian_pos",
                    "spherical_pos",
+                   "galactic_pos",
                    "dist",
                    "cartesian_redshiftspace_pos",
                    "spherical_redshiftspace_pos",
@@ -596,6 +599,9 @@ class BaseCatalogue(ABC):
             elif key == "spherical_pos":
                 out = cartesian_to_radec(
                     self["__cartesian_pos"] - self.observer_location)
+            elif key == "galactic_pos":
+                out = self["__spherical_pos"]
+                out[:, 1], out[:, 2] = radec_to_galactic(out[:, 1], out[:, 2])
             elif key == "dist":
                 out = numpy.linalg.norm(
                     self["__cartesian_pos"] - self.observer_location, axis=1)
@@ -985,7 +991,6 @@ class CSiBORG2MergerTreeReader:
         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"][:]
@@ -1020,6 +1025,11 @@ class CSiBORG2MergerTreeReader:
                         "TreeNextProgenitor",
                         "TreeProgenitor",
                         "SubhaloMass",
+                        "Group_M_Crit200",
+                        "SubhaloSpin",
+                        "SubhaloVmax",
+                        "SubhaloHalfmassRad",
+                        "SubhaloVmaxRad",
                         ]
 
         with File(self.paths.trees(self.nsim, self.simname), 'r') as f:
@@ -1082,44 +1092,58 @@ class CSiBORG2MergerTreeReader:
 
         n = first_fof  # Index of the current main progenitor
 
-        time, redshift, main_progenitor_mass = [], [], []
+        snap_num, redshift, main_progenitor_mass, group_m200c = [], [], [], []
-        max_next_progenitor_mass, total_next_progenitor_mass = [], []
+        main_progenitor_vmax, main_progenitor_spin = [], []
+        main_progenitor_vmaxrad, main_progenitor_halfmassrad = [], []
         while True:
-            # First off attempt to find the next progenitors of the current
+            # NOTE: 'Minors' are ignored. This is only relevant if we wanted
-            # halo. Deal with the main progenitor later.
+            # to find the other subhaloes in the current FoF group.
-            next_prog = tree["TreeNextProgenitor"][n]
-            if next_prog != -1:
-                minors = []
-                while True:
-                    minors.append(tree["SubhaloMass"][next_prog])
 
-                    next_prog = tree["TreeNextProgenitor"][next_prog]
+            # # 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])
 
-                    if next_prog == -1:
+            #         next_prog = tree["TreeNextProgenitor"][next_prog]
-                        break
+
-            else:
+            #         if next_prog == -1:
-                # Fiducially set it to zero.
+            #             break
-                minors = [0]
+            # 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))
+            group_m200c.append(tree["Group_M_Crit200"][n])
-            total_next_progenitor_mass.append(sum(minors))
+            main_progenitor_vmax.append(tree["SubhaloVmax"][n])
+            main_progenitor_spin.append(tree["SubhaloSpin"][n])
+            main_progenitor_vmaxrad.append(tree["SubhaloVmaxRad"][n])
+            main_progenitor_halfmassrad.append(tree["SubhaloHalfmassRad"][n])
+            snap_num.append(tree["SnapNum"][n])
             redshift.append(tree["Redshift"][tree["SnapNum"][n]])
-            time.append(tree["Time"][tree["SnapNum"][n]])
 
-            # Update n to the next main progenitor.
+            # Update `n` to the next main progenitor.
             n = tree["TreeMainProgenitor"][n]
 
             if n == -1:
                 break
 
-        return {"Time": numpy.array(time),
+        # For calculating age of the Universe at each redshift.
+        cosmo = FlatLambdaCDM(H0=67.66, Om0=0.3111)
+
+        return {"SnapNum": numpy.array(snap_num, dtype=numpy.int32),
+                "Age": numpy.array(cosmo.age(redshift).value),
                 "Redshift": numpy.array(redshift),
+                "Group_M_Crit200": numpy.array(group_m200c) * 1e10,
                 "MainProgenitorMass": numpy.array(main_progenitor_mass) * 1e10,
-                "MaxNextProgenitorMass": numpy.array(max_next_progenitor_mass) * 1e10,  # noqa
+                "MainProgenitorVmax": numpy.array(main_progenitor_vmax),
-                "TotalNextProgenitorMass": numpy.array(total_next_progenitor_mass) * 1e10,  # noqa
+                "MainProgenitorSpin": numpy.array(main_progenitor_spin),
+                "MainProgenitorVmaxRad": numpy.array(main_progenitor_vmaxrad),
+                "MainProgenitorHalfmassRad": numpy.array(main_progenitor_halfmassrad),      # noqa
                 }
 
 
@@ -1148,7 +1172,6 @@ class CSiBORG2SUBFINDCatalogue(BaseCatalogue):
     """
     def __init__(self, nsim, nsnap, kind, paths=None,
                  bounds=None, flip_xz=True, cache_maxsize=64):
-        # TODO: finish all this!
         super().__init__()
         super().init_with_snapshot(
             f"csiborg2_{kind}", nsim, nsnap, paths, None, bounds,
@@ -1156,7 +1179,8 @@ class CSiBORG2SUBFINDCatalogue(BaseCatalogue):
             cache_maxsize)
 
         self._custom_keys = ["SubhaloSpin", "SubhaloVelDisp", "Central",
-                             "ParentMass"]
+                             "SubhaloVmax", "SubhaloVmaxRad",
+                             "SubhaloHalfmassRad", "ParentMass"]
 
     @property
     def kind(self):
@@ -1235,6 +1259,18 @@ class CSiBORG2SUBFINDCatalogue(BaseCatalogue):
     def SubhaloVelDisp(self):
         return self._read_subfind_catalogue("SubhaloVelDisp")
 
+    @property
+    def SubhaloVmax(self):
+        return self._read_subfind_catalogue("SubhaloVmax")
+
+    @property
+    def SubhaloVmaxRad(self):
+        return self._read_subfind_catalogue("SubhaloVmaxRad")
+
+    @property
+    def SubhaloHalfmassRad(self):
+        return self._read_subfind_catalogue("SubhaloHalfmassRad")
+
     @property
     def SubhaloContamination(self):
         mass_type = self._read_subfind_catalogue("SubhaloMassType")

csiborgtools/read/paths.py

@@ -301,7 +301,7 @@ class Paths:
             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",
+            return join(self.csiborg2_random_srcdir, f"chain_{nsim}", "output",
                         "trees.hdf5")
         elif simname == "csiborg2_varysmall":
             return join(self.csiborg2_varysmall_srcdir,
@@ -351,6 +351,26 @@ class Paths:
             fname = fname.replace("overlap", "overlap_smoothed")
         return join(fdir, fname)
 
+    def random_mah(self, simname, nsim):
+        """
+        Path to the files containing MAHs from random simulations.
+
+        Parameters
+        ----------
+        simname : str
+            Simulation name.
+        nsim0 : int
+            IC realisation index of the simulation.
+
+        Returns
+        -------
+        str
+        """
+        fdir = join(self.postdir, "random_mah")
+        try_create_directory(fdir)
+
+        return join(fdir, f"random_mah_{simname}_{nsim}.hdf5")
+
     def match_max(self, simname, nsim0, nsimx, min_logmass, mult):
         """
         Path to the files containing matching based on [1].

csiborgtools/summary/overlap_summary.py

@@ -82,9 +82,9 @@ class PairOverlap:
         self._cat0 = cat0
         self._catx = catx
         self._paths = cat0.paths
-        self.load(cat0, catx, min_logmass, maxdist)
+        self._load(cat0, catx, min_logmass, maxdist)
 
-    def load(self, cat0, catx, paths, min_logmass, maxdist=None):
+    def _load(self, cat0, catx, min_logmass, maxdist=None):
         r"""
         Load overlap calculation results. Matches the results back to the two
         catalogues in question.
@@ -107,14 +107,13 @@ class PairOverlap:
         """
         nsim0 = cat0.nsim
         nsimx = catx.nsim
-        paths = cat0.paths
 
         # We first load in the output files. We need to find the right
         # combination of the reference and cross simulation.
-        fname = paths.overlap(cat0.simname, nsim0, nsimx, min_logmass,
+        fname = self._paths.overlap(cat0.simname, nsim0, nsimx, min_logmass,
-                              smoothed=False)
+                                    smoothed=False)
-        fname_inv = paths.overlap(cat0.simname, nsimx, nsim0, min_logmass,
+        fname_inv = self._paths.overlap(cat0.simname, nsimx, nsim0,
-                                  smoothed=False)
+                                        min_logmass, smoothed=False)
         if isfile(fname):
             data_ngp = numpy.load(fname, allow_pickle=True)
             to_invert = False
@@ -125,8 +124,8 @@ class PairOverlap:
         else:
             raise FileNotFoundError(f"No file found for {nsim0} and {nsimx}.")
 
-        fname_smooth = paths.overlap(cat0.simname, cat0.nsim, catx.nsim,
+        fname_smooth = self._paths.overlap(cat0.simname, cat0.nsim, catx.nsim,
-                                     min_logmass, smoothed=True)
+                                           min_logmass, smoothed=True)
         data_smooth = numpy.load(fname_smooth, allow_pickle=True)
 
         # Create mapping from halo indices to array positions in the catalogue.
@@ -609,6 +608,7 @@ class NPairsOverlap:
 
         self._pairs = pairs
 
+    @lru_cache()
     def max_overlap(self, min_overlap, from_smoothed, verbose=True):
         """
         Calculate maximum overlap of each halo in the reference simulation with
@@ -644,6 +644,7 @@ class NPairsOverlap:
                                        for y_ in pair.overlap(from_smoothed)])
         return numpy.vstack(out).T
 
+    @lru_cache()
     def max_overlap_key(self, key, min_overlap, from_smoothed, verbose=True):
         """
         Calculate maximum overlap mass of each halo in the reference
@@ -674,6 +675,7 @@ class NPairsOverlap:
 
         return numpy.vstack(out).T
 
+    @lru_cache()
     def summed_overlap(self, from_smoothed, verbose=True):
         """
         Calculate summed overlap of each halo in the reference simulation with

csiborgtools/utils.py

@@ -19,6 +19,8 @@ from copy import deepcopy
 from datetime import datetime
 
 import numpy as np
+from astropy import units as u
+from astropy.coordinates import SkyCoord
 from numba import jit
 from numpyro.infer import util
 from scipy.stats import multivariate_normal
@@ -154,6 +156,24 @@ def radec_to_cartesian(X):
         ]).T
 
 
+def radec_to_galactic(ra, dec):
+    """
+    Convert right ascension and declination to galactic coordinates (all in
+    degrees.)
+
+    Parameters
+    ----------
+    ra, dec : float or 1-dimensional array
+        Right ascension and declination in degrees.
+
+    Returns
+    -------
+    l, b : float or 1-dimensional array
+    """
+    c = SkyCoord(ra=ra*u.degree, dec=dec*u.degree, frame='icrs')
+    return c.galactic.l.degree, c.galactic.b.degree
+
+
 @jit(nopython=True, fastmath=True, boundscheck=False)
 def great_circle_distance(x1, x2):
     """

notebooks/MAH/mah.py

@@ -0,0 +1,221 @@
+# Copyright (C) 2024 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 help with `mah.py`."""
+from datetime import datetime
+
+import csiborgtools
+import numpy as np
+from astropy.cosmology import FlatLambdaCDM
+from h5py import File
+from tqdm import tqdm, trange
+from cache_to_disk import cache_to_disk
+from os.path import join
+
+
+RANDOM_MAH_Sorce_Virgo_UPPER = np.array(
+    [[2.18554217, 0.16246594],
+     [2.93253012, 0.17284951],
+     [3.2939759, 0.34169001],
+     [3.75180723, 0.42006683],
+     [4.28192771, 0.44691426],
+     [4.61927711, 0.53819753],
+     [5.34216867, 0.58454257],
+     [5.89638554, 0.68954882],
+     [6.23373494, 0.73361948],
+     [6.45060241, 0.81341823],
+     [7.05301205, 0.92071572],
+     [7.82409639, 0.92071572],
+     [8.28192771, 0.95953933],
+     [8.61927711, 0.97956078],
+     [9.70361446, 1.],
+     [11.17349398, 1.],
+     [13.07710843, 1.],
+     [13.82409639, 1.]]
+    )
+
+RANDOM_MAH_SORCE_Virgo_LOWER = np.array(
+    [[3.36626506e+00, 1.00000000e-02],
+     [3.75180723e+00, 1.10877404e-02],
+     [3.99277108e+00, 1.04216677e-02],
+     [4.30602410e+00, 1.15552746e-02],
+     [4.61927711e+00, 1.67577322e-02],
+     [4.98072289e+00, 2.14703224e-02],
+     [5.39036145e+00, 3.82789169e-02],
+     [5.89638554e+00, 5.00670000e-02],
+     [6.30602410e+00, 5.11116827e-02],
+     [7.29397590e+00, 5.32668971e-02],
+     [7.77590361e+00, 5.55129899e-02],
+     [8.11325301e+00, 6.68516464e-02],
+     [8.57108434e+00, 8.56515893e-02],
+     [9.60722892e+00, 1.32152759e-01],
+     [1.04265060e+01, 1.46527548e-01],
+     [1.07638554e+01, 1.49584947e-01],
+     [1.11493976e+01, 1.72849513e-01],
+     [1.18240964e+01, 2.16931625e-01],
+     [1.21855422e+01, 2.45546942e-01],
+     [1.25951807e+01, 3.48819614e-01],
+     [1.30771084e+01, 5.27197199e-01],
+     [1.36795181e+01, 8.83462949e-01],
+     [1.38000000e+01, 1.00000000e+00]]
+    )
+
+
+def t():
+    return datetime.now()
+
+
+@cache_to_disk(90)
+def load_data(nsim0, simname, min_logmass):
+    """
+    Load the reference catalogue, the cross catalogues, the merger trees and
+    the overlap reader (in this order).
+    """
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = paths.get_ics(simname)
+    if "csiborg2_" in simname:
+        kind = simname.split("_")[-1]
+        print(f"{t()}: loading {len(nsims)} halo catalogues.")
+        cat0 = csiborgtools.read.CSiBORG2Catalogue(nsim0, 99, kind)
+        catxs = [csiborgtools.read.CSiBORG2Catalogue(n, 99, kind)
+                 for n in nsims if n != nsim0]
+
+        print(f"{t()}: loading {len(nsims)} merger trees.")
+        merger_trees = {}
+        for nsim in tqdm(nsims):
+            merger_trees[nsim] = csiborgtools.read.CSiBORG2MergerTreeReader(
+                nsim, kind)
+    else:
+        raise ValueError(f"Unknown simname: {simname}")
+
+    overlaps = csiborgtools.summary.NPairsOverlap(cat0, catxs, min_logmass)
+
+    return cat0, catxs, merger_trees, overlaps
+
+
+def extract_main_progenitor_maxoverlap(group_nr, overlaps, merger_trees):
+    """
+    Follow the main progenitor of a reference group and its maximum overlap
+    group in the cross catalogues.
+    """
+    min_overlap = 0
+
+    # NOTE these can be all cached in the overlap object.
+    max_overlaps = overlaps.max_overlap(0, True)[group_nr]
+    if np.sum(max_overlaps > 0) == 0:
+        raise ValueError(f"No overlaps for group {group_nr}.")
+
+    max_overlap_indxs = overlaps.max_overlap_key(
+        "index", min_overlap, True)[group_nr]
+
+    out = {}
+    for i in trange(len(overlaps), desc="Cross main progenitors"):
+        nsimx = overlaps[i].catx().nsim
+        group_nr_cross = max_overlap_indxs[i]
+
+        if np.isnan(group_nr_cross):
+            continue
+
+        x = merger_trees[nsimx].main_progenitor(int(group_nr_cross))
+        x["Overlap"] = max_overlaps[i]
+
+        out[nsimx] = x
+
+    nsim0 = overlaps.cat0().nsim
+    print(f"Appending main progenitor for {nsim0}.")
+    out[nsim0] = merger_trees[nsim0].main_progenitor(group_nr)
+
+    return out
+
+
+def summarize_extracted_mah(simname, data, nsim0, nsimxs, key,
+                            min_age=0, include_nsim0=True):
+    """
+    Turn the dictionaries of extracted MAHs into a single array.
+    """
+    if "csiborg2_" in simname:
+        nsnap = 100
+    else:
+        raise ValueError(f"Unknown simname: {simname}")
+
+    X = []
+    for nsimx in nsimxs + [nsim0] if include_nsim0 else nsimxs:
+        try:
+            d = data[nsimx]
+        except KeyError:
+            continue
+
+        x = np.full(nsnap, np.nan, dtype=np.float32)
+        x[d["SnapNum"]] = d[key]
+
+        X.append(x)
+
+    cosmo = FlatLambdaCDM(H0=67.76, Om0=csiborgtools.simname2Omega_m(simname))
+    zs = [csiborgtools.snap2redshift(i, simname) for i in range(nsnap)]
+    age = cosmo.age(zs).value
+
+    mask = age > min_age
+    return age[mask], np.vstack(X)[:, mask]
+
+
+def extract_mah(simname, logmass_bounds, key, min_age=0):
+    """
+    Extract the random MAHs for a given simulation and mass range and key.
+    Keys are for example: "MainProgenitorMass" or "GroupMass"
+    """
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = paths.get_ics(simname)
+
+    X = []
+    for i, nsim in enumerate(nsims):
+        with File(paths.random_mah(simname, nsim), 'r') as f:
+            mah = f[key][:]
+            final_mass = mah[:, -1]
+
+            # Select the mass range
+            mask = final_mass >= 10**logmass_bounds[0]
+            mask &= final_mass < 10**logmass_bounds[1]
+
+            X.append(mah[mask])
+
+            if i == 0:
+                redshift = f["Redshift"][:]
+
+    X = np.vstack(X)
+
+    cosmo = FlatLambdaCDM(H0=67.76, Om0=csiborgtools.simname2Omega_m(simname))
+    age = cosmo.age(redshift).value
+
+    mask = age > min_age
+    return age[mask], X[:, mask]
+
+
+def extract_mah_mdpl2(logmass_bounds, min_age=1.5):
+    """
+    MAH extraction for the MDPL2 simulation. Data comes from
+    `https://arxiv.org/abs/2105.05859`
+    """
+    fdir = "/mnt/extraspace/rstiskalek/catalogs/"
+
+    age = np.genfromtxt(join(fdir, "mdpl2_cosmic_time.txt"))
+    with File(join(fdir, "diffmah_mdpl2.h5"), 'r') as f:
+        log_mp = f["logmp_sim"][:]
+        log_mah_sim = f["log_mah_sim"][...]
+
+    xmin, xmax = logmass_bounds
+    ks = np.where((log_mp > xmin) & (log_mp < xmax))[0]
+    X = 10**log_mah_sim[ks]
+
+    mask = age > min_age
+    return age[mask], X[:, mask]

notebooks/diagnostic/hmf.py

@@ -0,0 +1,48 @@
+# Copyright (C) 2024 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 help with `hmf.py`."""
+
+import csiborgtools
+import numpy as np
+from tqdm import tqdm
+
+
+def calculate_hmf(simname, bin_edges, halofinder="FOF", max_distance=135):
+    """
+    Calculate the halo mass function for a given simulation from catalogues.
+    """
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = paths.get_ics(simname)
+    bounds = {"dist": (0, max_distance)}
+
+    hmf = np.full((len(nsims), len(bin_edges) - 1), np.nan)
+    volume = 4 / 3 * np.pi * max_distance**3
+    for i, nsim in enumerate(tqdm(nsims)):
+        if "csiborg2_" in simname:
+            kind = simname.split("_")[-1]
+            if halofinder == "FOF":
+                cat = csiborgtools.read.CSiBORG2Catalogue(
+                    nsim, 99, kind, bounds=bounds)
+            elif halofinder == "SUBFIND":
+                cat = csiborgtools.read.CSiBORG2SUBFINDCatalogue(
+                    nsim, 99, kind, kind, bounds=bounds)
+            else:
+                raise ValueError(f"Unknown halofinder: {halofinder}")
+        else:
+            raise ValueError(f"Unknown simname: {simname}")
+
+        hmf[i] = cat.halo_mass_function(bin_edges, volume, "totmass")[1]
+
+    return hmf

notebooks/field_sample.ipynb

@@ -38,7 +38,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -47,7 +47,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -66,7 +66,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -81,7 +81,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [

notebooks/flow/flow_bulk.py

@@ -60,6 +60,6 @@ def read_enclosed_flow(simname):
 
     for n in range(nsim):
         V_n = csiborgtools.cartesian_to_radec(V[n])
-        l[n], b[n] = csiborgtools.flow.radec_to_galactic(V_n[:, 1], V_n[:, 2])
+        l[n], b[n] = csiborgtools.radec_to_galactic(V_n[:, 1], V_n[:, 2])
 
     return r, Vmag, l, b

notebooks/flow/flow_calibration.py

@@ -122,7 +122,7 @@ def read_samples(catalogue, simname, ksmooth, include_calibration=False,
     # Calculate direction in galactic coordinates of V_ext
     V = np.vstack([Vx, Vy, Vz]).T
     V = csiborgtools.cartesian_to_radec(V)
-    l, b = csiborgtools.flow.radec_to_galactic(V[:, 1], V[:, 2])
+    l, b = csiborgtools.radec_to_galactic(V[:, 1], V[:, 2])
 
     data = [alpha, beta, Vmag, l, b, sigma_v]
     names = ["alpha", "beta", "Vmag", "l", "b", "sigma_v"]

notebooks/utils.py

@@ -0,0 +1,132 @@
+# Copyright (C) 2023 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.
+"""
+Various utility functions.
+"""
+
+import numpy
+from scipy.special import erf
+
+dpi = 600
+fout = "../plots/"
+mplstyle = ["science"]
+
+
+def latex_float(*floats, n=2):
+    """
+    Convert a float or a list of floats to a LaTeX string(s). Taken from [1].
+
+    Parameters
+    ----------
+    floats : float or list of floats
+        The float(s) to be converted.
+    n : int, optional
+        The number of significant figures to be used in the LaTeX string.
+
+    Returns
+    -------
+    latex_floats : str or list of str
+        The LaTeX string(s) representing the float(s).
+
+    References
+    ----------
+    [1] https://stackoverflow.com/questions/13490292/format-number-using-latex-notation-in-python  # noqa
+    """
+    latex_floats = [None] * len(floats)
+    for i, f in enumerate(floats):
+        float_str = "{0:.{1}g}".format(f, n)
+        if "e" in float_str:
+            base, exponent = float_str.split("e")
+            latex_floats[i] = r"{0} \times 10^{{{1}}}".format(base,
+                                                              int(exponent))
+        else:
+            latex_floats[i] = float_str
+
+    if len(floats) == 1:
+        return latex_floats[0]
+    return latex_floats
+
+
+def nan_weighted_average(arr, weights=None, axis=None):
+    if weights is None:
+        weights = numpy.ones_like(arr)
+
+    valid_entries = ~numpy.isnan(arr)
+
+    # Set NaN entries in arr to 0 for computation
+    arr = numpy.where(valid_entries, arr, 0)
+
+    # Set weights of NaN entries to 0
+    weights = numpy.where(valid_entries, weights, 0)
+
+    # Compute the weighted sum and the sum of weights along the axis
+    weighted_sum = numpy.sum(arr * weights, axis=axis)
+    sum_weights = numpy.sum(weights, axis=axis)
+
+    return weighted_sum / sum_weights
+
+
+def nan_weighted_std(arr, weights=None, axis=None, ddof=0):
+    if weights is None:
+        weights = numpy.ones_like(arr)
+
+    valid_entries = ~numpy.isnan(arr)
+
+    # Set NaN entries in arr to 0 for computation
+    arr = numpy.where(valid_entries, arr, 0)
+
+    # Set weights of NaN entries to 0
+    weights = numpy.where(valid_entries, weights, 0)
+
+    # Calculate weighted mean
+    weighted_mean = numpy.sum(
+        arr * weights, axis=axis) / numpy.sum(weights, axis=axis)
+
+    # Calculate the weighted variance
+    variance = numpy.sum(
+        weights * (arr - numpy.expand_dims(weighted_mean, axis))**2, axis=axis)
+    variance /= numpy.sum(weights, axis=axis) - ddof
+
+    return numpy.sqrt(variance)
+
+
+def compute_error_bars(x, y, xbins, sigma):
+    bin_indices = numpy.digitize(x, xbins)
+    y_medians = numpy.array([numpy.median(y[bin_indices == i])
+                             for i in range(1, len(xbins))])
+
+    lower_pct = 100 * 0.5 * (1 - erf(sigma / numpy.sqrt(2)))
+    upper_pct = 100 - lower_pct
+
+    y_lower = numpy.full(len(y_medians), numpy.nan)
+    y_upper = numpy.full(len(y_medians), numpy.nan)
+
+    for i in range(len(y_medians)):
+        if numpy.sum(bin_indices == i + 1) == 0:
+            continue
+
+        y_lower[i] = numpy.percentile(y[bin_indices == i + 1], lower_pct)
+        y_upper[i] = numpy.percentile(y[bin_indices == i + 1], upper_pct)
+
+    yerr = (y_medians - numpy.array(y_lower), numpy.array(y_upper) - y_medians)
+
+    return y_medians, yerr
+
+
+def normalize_hexbin(hb):
+    hexagon_counts = hb.get_array()
+    normalized_counts = hexagon_counts / hexagon_counts.sum()
+    hb.set_array(normalized_counts)
+    hb.set_clim(normalized_counts.min(), normalized_counts.max())

scripts/mah_random.py

@@ -0,0 +1,97 @@
+# Copyright (C) 2024 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 extract the mass accretion histories in random simulations. Follows
+the main progenitor of FoF haloes.
+"""
+from argparse import ArgumentParser
+import csiborgtools
+import numpy as np
+from h5py import File
+from mpi4py import MPI
+from taskmaster import work_delegation  # noqa
+from tqdm import trange
+from cache_to_disk import cache_to_disk
+
+
+@cache_to_disk(30)
+def load_data(nsim, simname, min_logmass):
+    """Load the data for a given simulation."""
+    bnd = {"totmass": (10**min_logmass, None), "dist": (None, 135)}
+    if "csiborg2_" in simname:
+        kind = simname.split("_")[-1]
+        cat = csiborgtools.read.CSiBORG2Catalogue(nsim, 99, kind, bounds=bnd)
+        merger_reader = csiborgtools.read.CSiBORG2MergerTreeReader(nsim, kind)
+    else:
+        raise ValueError(f"Unknown simname: {simname}")
+
+    return cat, merger_reader
+
+
+def main_progenitor_mah(cat, merger_reader, simname, verbose=True):
+    """Follow the main progenitor of each `z = 0` FoF halo."""
+    indxs = cat["index"]
+
+    # Main progenitor information as a function of time
+    shape = (len(cat), cat.nsnap + 1)
+    main_progenitor_mass = np.full(shape, np.nan, dtype=np.float32)
+    group_mass = np.full(shape, np.nan, dtype=np.float32)
+
+    for i in trange(len(cat), disable=not verbose, desc="Haloes"):
+        d = merger_reader.main_progenitor(indxs[i])
+
+        main_progenitor_mass[i, d["SnapNum"]] = d["MainProgenitorMass"]
+        group_mass[i, d["SnapNum"]] = d["Group_M_Crit200"]
+
+    return {"Redshift": [csiborgtools.snap2redshift(i, simname) for i in range(cat.nsnap + 1)],  # noqa
+            "MainProgenitorMass": main_progenitor_mass,
+            "GroupMass": group_mass,
+            "FinalGroupMass": cat["totmass"],
+            }
+
+
+def save_output(data, nsim, simname, verbose=True):
+    """Save the output to a file."""
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+
+    fname = paths.random_mah(simname, nsim)
+    if verbose:
+        print(f"Saving output to `{fname}`")
+
+    with File(fname, "w") as f:
+        for key, value in data.items():
+            f.create_dataset(key, data=value)
+
+
+if "__main__" == __name__:
+    parser = ArgumentParser(description="Extract the mass accretion history in random simulations.")  # noqa
+    parser.add_argument("--simname", help="Name of the simulation.", type=str,
+                        choices=["csiborg2_random"])
+    parser.add_argument("--min_logmass", type=float,
+                        help="Minimum log mass of the haloes.")
+    args = parser.parse_args()
+    COMM = MPI.COMM_WORLD
+
+    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
+    nsims = paths.get_ics(args.simname)
+
+    def main(nsim):
+        verbose = COMM.Get_size() == 1
+        cat, merger_reader = load_data(nsim, args.simname, args.min_logmass)
+        data = main_progenitor_mah(cat, merger_reader, args.simname,
+                                   verbose=verbose)
+        save_output(data, nsim, args.simname, verbose=verbose)
+
+    work_delegation(main, list(nsims), MPI.COMM_WORLD)

scripts/mah_random.sh

@@ -0,0 +1,23 @@
+memory=4
+on_login=${1}
+nthreads=5
+
+queue="berg"
+env="/mnt/users/rstiskalek/csiborgtools/venv_csiborg/bin/python"
+file="mah_random.py"
+
+min_logmass=13.0
+simname="csiborg2_random"
+
+
+pythoncm="$env $file --simname $simname --min_logmass $min_logmass"
+if [ $on_login -eq 1 ]; then
+    echo $pythoncm
+    $pythoncm
+else
+    cm="addqueue -q $queue -n $nthreads -m $memory $pythoncm"
+    echo "Submitting:"
+    echo $cm
+    echo
+    eval $cm
+fi

scripts/match_overlap_all.py

@@ -56,7 +56,8 @@ if __name__ == "__main__":
                         choices=["overlap", "max"], help="Kind of matching.")
     parser.add_argument("--simname", type=str, required=True,
                         help="Simulation name.",
-                        choices=["csiborg", "quijote"])
+                        choices=["csiborg1", "quijote", "csiborg2_main",
+                                 "csiborg2_random", "csiborg2_varysmall"])
     parser.add_argument("--nsim0", type=int, default=None,
                         help="Reference IC for Max's matching method.")
     parser.add_argument("--min_logmass", type=float, required=True,

scripts/match_overlap_all.sh

@@ -1,16 +1,16 @@
 #!/bin/bash
-nthreads=11
+nthreads=41
-memory=4
+memory=12
-queue="cmb"
+queue="berg"
 env="/mnt/zfsusers/rstiskalek/csiborgtools/venv_csiborg/bin/python"
-file="match_all.py"
+file="match_overlap_all.py"
 
-simname="quijote"
+simname=${1}
-min_logmass=13.25
+min_logmass=12.25
-nsim0=0
-kind="max"
-mult=10
 sigma=1
+kind="overlap"
+mult=10         # Only for Max's method
+nsim0=0         # Only for Max's method
 verbose="false"
 
 

scripts/match_overlap_single.sh

@@ -8,8 +8,8 @@ file="match_overlap_single.py"
 
 simname="csiborg2_main"
 kind="overlap"
-min_logmass=13.25
+min_logmass=12
-mult=5
+mult=5  # Only for Max's method
 sigma=1
 
 # sims=(7444 7468)
@@ -37,14 +37,14 @@ do
 
         pythoncm="$env $file --kind $kind --nsim0 $nsim0 --nsimx $nsimx  --simname $simname --min_logmass $min_logmass --sigma $sigma --mult $mult --verbose $verbose"
 
-        # $pythoncm
+        $pythoncm
 
-        cm="addqueue -q $queue -n 1x1 -m $memory $pythoncm"
+        # cm="addqueue -q $queue -n 1x1 -m $memory $pythoncm"
-        echo "Submitting:"
+        # echo "Submitting:"
-        echo $cm
+        # echo $cm
-        echo
+        # echo
-        $cm
+        # $cm
-        sleep 0.05
+        # sleep 0.05
 
     done
 done

scripts_plots/plot_vobs.py

@@ -1,89 +0,0 @@
-# Copyright (C) 2023 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.
-
-from os.path import join
-
-import matplotlib.pyplot as plt
-import numpy
-import scienceplots  # noqa
-from tqdm import tqdm
-
-import csiborgtools
-import plt_utils
-
-
-def observer_peculiar_velocity(MAS, grid, ext="png"):
-    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
-    nsims = paths.get_ics("csiborg")
-
-    for n, nsim in enumerate(nsims):
-        fpath = paths.observer_peculiar_velocity(MAS, grid, nsim)
-        f = numpy.load(fpath)
-
-        if n == 0:
-            data = numpy.full((len(nsims), *f["observer_vp"].shape), numpy.nan)
-            smooth_scales = f["smooth_scales"]
-
-        data[n] = f["observer_vp"]
-
-    for n, smooth_scale in enumerate(tqdm(smooth_scales,
-                                          desc="Plotting smooth scales")):
-        with plt.style.context(plt_utils.mplstyle):
-            fig, axs = plt.subplots(ncols=3, figsize=(3.5 * 3, 2.625),
-                                    sharey=True)
-            fig.subplots_adjust(wspace=0)
-            for i, ax in enumerate(axs):
-                ax.hist(data[:, n, i], bins="auto", histtype="step")
-                mu, sigma = numpy.mean(data[:, n, i]), numpy.std(data[:, n, i])
-                ax.set_title(r"$V_{{\rm obs, i}} = {:.3f} \pm {:.3f} ~ \mathrm{{km}} / \mathrm{{s}}$".format(mu, sigma)) # noqa
-
-            axs[0].set_xlabel(r"$V_{\rm obs, x} ~ [\mathrm{km} / \mathrm{s}]$")
-            axs[1].set_xlabel(r"$V_{\rm obs, y} ~ [\mathrm{km} / \mathrm{s}]$")
-            axs[2].set_xlabel(r"$V_{\rm obs, z} ~ [\mathrm{km} / \mathrm{s}]$")
-            axs[0].set_ylabel(r"Counts")
-
-            fig.suptitle(r"$N_{{\rm grid}} = {}$, $\sigma_{{\rm smooth}} = {:.2f} ~ [\mathrm{{Mpc}} / h]$".format(grid, smooth_scale)) # noqa
-
-            fig.tight_layout(w_pad=0)
-            fout = join(plt_utils.fout,
-                        f"observer_vp_{grid}_{smooth_scale}.{ext}")
-            fig.savefig(fout, dpi=plt_utils.dpi, bbox_inches="tight")
-            plt.close()
-
-    with plt.style.context(plt_utils.mplstyle):
-        fig, axs = plt.subplots(ncols=3, figsize=(3.5 * 3, 2.625))
-        for i, ax in enumerate(axs):
-
-            ymu = numpy.mean(data[..., i], axis=0)
-            ystd = numpy.std(data[..., i], axis=0)
-            ylow, yupp = ymu - ystd, ymu + ystd
-            ax.plot(smooth_scales, ymu, c="k")
-            ax.fill_between(smooth_scales, ylow, yupp, color="k", alpha=0.2)
-
-            ax.set_xlabel(r"$\sigma_{{\rm smooth}} ~ [\mathrm{{Mpc}} / h]$")
-
-        axs[0].set_ylabel(r"$V_{\rm obs, x} ~ [\mathrm{km} / \mathrm{s}]$")
-        axs[1].set_ylabel(r"$V_{\rm obs, y} ~ [\mathrm{km} / \mathrm{s}]$")
-        axs[2].set_ylabel(r"$V_{\rm obs, z} ~ [\mathrm{km} / \mathrm{s}]$")
-        fig.suptitle(r"$N_{{\rm grid}} = {}$".format(grid))
-
-        fig.tight_layout(w_pad=0)
-        fout = join(plt_utils.fout, f"observer_vp_summary_{grid}.{ext}")
-        fig.savefig(fout, dpi=plt_utils.dpi, bbox_inches="tight")
-        plt.close()
-
-
-if __name__ == "__main__":
-    observer_peculiar_velocity("PCS", 512)