# 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]