add star parameters to archive

lace/archive/base_archive.py

@@ -74,6 +74,7 @@ def _average_over_samples(self, average="both", drop_axis=None):
             "ind_axis",
             "ind_rescaling",
             "cosmo_params",
+            "star_params",
         ]
 
         keys_spe = ["p1d_Mpc"]
@@ -163,6 +164,7 @@ def _average_over_samples(self, average="both", drop_axis=None):
             dict_av["ind_rescaling"] = ind_rescaling
             dict_av["ind_snap"] = ind_snap
             dict_av["cosmo_params"] = self.data[ind_merge[0]]["cosmo_params"]
+            dict_av["star_params"] = self.data[ind_merge[0]]["star_params"]
 
             # list of available keys to merger
             key_list = self.data[ind_merge[0]].keys()
@@ -377,13 +379,17 @@ def get_training_data(
                 if mask:
                     if all(
                         x in list_keys
-                        or (x in ["A_lya","n_lya","omega_m","H_0","A_UVB"] and "cosmo_params" in list_keys)
+                        or (
+                            x in ["A_lya", "n_lya", "omega_m", "H_0", "A_UVB"]
+                            and "cosmo_params" in list_keys
+                        )
                         for x in emu_params
                     ):
                         if any(
                             np.isnan(arch_av[ii][x])
                             for x in emu_params
-                            if x not in ["A_lya","n_lya","omega_m","H_0","A_UVB"] 
+                            if x
+                            not in ["A_lya", "n_lya", "omega_m", "H_0", "A_UVB"]
                         ) | any(
                             np.any(np.isnan(arch_av[ii][x])) for x in key_power
                         ):
@@ -502,9 +508,16 @@ def get_testing_data(
             if mask:
                 if emu_params is None:
                     testing_data.append(arch_av[ii])
-                #elif all(x in list_keys for x in emu_params):
-                elif all(x in list_keys or (isinstance(list_keys, dict) and x in list_keys.get('cosmo', [])) for x in emu_params.keys()):
-                    print('activated')
+                # elif all(x in list_keys for x in emu_params):
+                elif all(
+                    x in list_keys
+                    or (
+                        isinstance(list_keys, dict)
+                        and x in list_keys.get("cosmo", [])
+                    )
+                    for x in emu_params.keys()
+                ):
+                    print("activated")
                     if any(np.isnan(arch_av[ii][x]) for x in emu_params) | any(
                         np.any(np.isnan(arch_av[ii][x])) for x in key_power
                     ):

lace/archive/gadget_archive.py

@@ -30,6 +30,8 @@ def __init__(
         kp_Mpc=None,
         force_recompute_linP_params=False,
         verbose=False,
+        z_star=3,
+        kp_kms=0.009,
     ):
         """
         Initialize the archive object.
@@ -67,6 +69,8 @@ def __init__(
             if isinstance(kp_Mpc, (int, float, type(None))) == False:
                 raise TypeError("kp_Mpc must be a number or None")
         self.kp_Mpc = kp_Mpc
+        self.z_star = z_star
+        self.kp_kms = kp_kms
 
         if isinstance(verbose, bool) == False:
             raise TypeError("verbose must be boolean")
@@ -327,6 +331,7 @@ def _get_emu_cosmo(self, sim_label, force_recompute_linP_params=False):
                     else:
                         cosmo_params = file_cosmo[ii]["cosmo_params"]
                         linP_params = file_cosmo[ii]["linP_params"]
+                        star_params = file_cosmo[ii]["star_params"]
                     break
 
         if compute_linP_params == True:
@@ -350,6 +355,11 @@ def _get_emu_cosmo(self, sim_label, force_recompute_linP_params=False):
 
             # compute linear power parameters at each z (in Mpc units)
             linP_zs = fit_linP.get_linP_Mpc_zs(sim_cosmo, zs, self.kp_Mpc)
+
+            # compute linear power parameters (in kms units)
+            star_params = fit_linP.parameterize_cosmology_kms(
+                sim_cosmo, None, self.z_star, self.kp_kms
+            )
             # compute conversion from Mpc to km/s using cosmology
             dkms_dMpc_zs = camb_cosmo.dkms_dMpc(sim_cosmo, z=np.array(zs))
 
@@ -366,7 +376,7 @@ def _get_emu_cosmo(self, sim_label, force_recompute_linP_params=False):
                     else:
                         linP_params[lab][ii] = linP_zs[ii][lab]
 
-        return cosmo_params, linP_params
+        return cosmo_params, linP_params, star_params
 
     def _get_file_names(self, sim_label, ind_phase, ind_z, ind_axis):
         """
@@ -539,7 +549,7 @@ def _load_data(self, force_recompute_linP_params):
         ## read info from all sims, all snapshots, all rescalings
         # iterate over simulations
         for sim_label in self.list_sim:
-            cosmo_params, linP_params = self._get_emu_cosmo(
+            cosmo_params, linP_params, star_params = self._get_emu_cosmo(
                 sim_label, force_recompute_linP_params
             )
 
@@ -548,6 +558,7 @@ def _load_data(self, force_recompute_linP_params):
                 # set linear power parameters describing snapshot
                 snap_data = {}
                 snap_data["cosmo_params"] = cosmo_params
+                snap_data["star_params"] = star_params
                 for lab in linP_params.keys():
                     if lab == "kp_Mpc":
                         snap_data[lab] = linP_params[lab]

lace/archive/nyx_archive.py

@@ -20,6 +20,8 @@ def __init__(
         force_recompute_linP_params=False,
         verbose=False,
         nfiles=18,
+        z_star=3,
+        kp_kms=0.009,
     ):
         ## check input
         if isinstance(nyx_version, str) == False:
@@ -41,6 +43,8 @@ def __init__(
             if isinstance(kp_Mpc, (int, float, type(None))) == False:
                 raise TypeError("kp_Mpc must be a number or None")
         self.kp_Mpc = kp_Mpc
+        self.z_star = z_star
+        self.kp_kms = kp_kms
 
         if isinstance(verbose, bool) == False:
             raise TypeError("verbose must be boolean")
@@ -49,7 +53,7 @@ def __init__(
 
         ## sets list simulations available for this suite
         # list of especial simulations
-        self.list_sim_test = ["nyx_central", "nyx_seed", "nyx_wdm"]
+        self.list_sim_test = ["nyx_central", "nyx_seed", "nyx_3_ic", "nyx_wdm"]
         # list of hypercube simulations
         self.list_sim_cube = []
         # simulation 14 was identified as problematic by the Nyx team
@@ -103,7 +107,8 @@ def _set_info_sim(self, nfiles):
         # file and lace bookkeeping
         self.sim_conv = {
             "fiducial": "nyx_central",
-            "bar_ic_grid_3": "nyx_seed",
+            "bar_ic_grid_3": "nyx_3_ic",
+            "CGAN_4096_base": "nyx_seed",
             "wdm_3.5kev_grid_1": "nyx_wdm",
         }
         for ii in range(nfiles):
@@ -218,11 +223,12 @@ def _get_emu_cosmo(
                     else:
                         cosmo_params = file_cosmo[ii]["cosmo_params"]
                         linP_params = file_cosmo[ii]["linP_params"]
+                        star_params = file_cosmo[ii]["star_params"]
                     break
             if sim_in_file == False:
                 file_error = (
                     "The file "
-                    + file_cosmo
+                    + self.file_cosmo
                     + " does not contain "
                     + isim
                     + ". To speed up calculations, "
@@ -244,13 +250,20 @@ def _get_emu_cosmo(
                 cosmo_params["A_s"] = 2.10e-9
                 cosmo_params["n_s"] = 0.966
                 cosmo_params["h"] = cosmo_params["H_0"] / 100
+            elif isim == "nyx_seed":
+                cosmo_params["A_s"] = 2.15865e-9
+                cosmo_params["n_s"] = 0.96
+                cosmo_params["h"] = cosmo_params["H_0"] / 100
             # setup CAMB object
             sim_cosmo = camb_cosmo.get_Nyx_cosmology(cosmo_params)
 
             # compute linear power parameters at each z (in Mpc units)
             linP_zs = fit_linP.get_linP_Mpc_zs(
                 sim_cosmo, self.list_sim_redshifts, self.kp_Mpc
             )
+            star_params = fit_linP.parameterize_cosmology_kms(
+                sim_cosmo, None, self.z_star, self.kp_kms
+            )
             zs = np.array(self.list_sim_redshifts)
             # compute conversion from Mpc to km/s using cosmology
             dkms_dMpc_zs = camb_cosmo.dkms_dMpc(sim_cosmo, z=zs)
@@ -268,7 +281,7 @@ def _get_emu_cosmo(
                     else:
                         linP_params[lab][ii] = linP_zs[ii][lab]
 
-        return cosmo_params, linP_params
+        return cosmo_params, linP_params, star_params
 
     def _load_data(self, nyx_file=None, force_recompute_linP_params=False):
         # set nyx_file if not provided
@@ -312,13 +325,16 @@ def _load_data(self, nyx_file=None, force_recompute_linP_params=False):
             # this simulation seems to have issues
             if isim == "cosmo_grid_14":
                 continue
+            # contained in CGAN_4096_base
+            elif isim == "CGAN_4096_val":
+                continue
 
             if self.verbose:
                 print("read Nyx sim", isim)
 
             # read cosmo information and linear power parameters
             # (will only need CAMB if pivot point kp_Mpc is changed)
-            cosmo_params, linP_params = self._get_emu_cosmo(
+            cosmo_params, linP_params, star_params = self._get_emu_cosmo(
                 ff, isim, force_recompute_linP_params
             )
 
@@ -328,9 +344,13 @@ def _load_data(self, nyx_file=None, force_recompute_linP_params=False):
                 zval = float(split_string(iz)[1])
 
                 # find redshift index to read linear power parameters
-                ind_z = np.where(
-                    np.isclose(self.list_sim_redshifts, zval, 1e-10)
-                )[0][0]
+                try:
+                    ind_z = np.where(
+                        np.isclose(self.list_sim_redshifts, zval, 1e-10)
+                    )[0][0]
+                except IndexError:
+                    print("Could not find redshift", zval, "in ", isim)
+                    continue
 
                 scalings_avail = list(ff[isim][iz].keys())
                 # loop over scalings
@@ -345,6 +365,7 @@ def _load_data(self, nyx_file=None, force_recompute_linP_params=False):
                         # dictionary containing measurements and relevant info
                         _arch = {}
                         _arch["cosmo_params"] = cosmo_params
+                        _arch["star_params"] = star_params
                         for lab in linP_params.keys():
                             if lab == "kp_Mpc":
                                 _arch[lab] = linP_params[lab]
@@ -387,17 +408,18 @@ def _load_data(self, nyx_file=None, force_recompute_linP_params=False):
                         _arch["k_Mpc"] = p1d["k"]
                         _arch["p1d_Mpc"] = p1d["Pk1d"]
                         if "3d power kmu" in _data:
-                            p3d = np.array(
-                                _data["3d power kmu"]["fine binning"]
-                            )
-                            _arch["k3d_Mpc"] = p3d["k"].reshape(
-                                self.kbins, self.mubins
-                            )
-                            _arch["mu3d"] = p3d["mu"].reshape(
-                                self.kbins, self.mubins
-                            )
-                            _arch["p3d_Mpc"] = p3d["Pkmu"].reshape(
-                                self.kbins, self.mubins
-                            )
+                            if "fine binning" in _data["3d power kmu"]:
+                                p3d = np.array(
+                                    _data["3d power kmu"]["fine binning"]
+                                )
+                                _arch["k3d_Mpc"] = p3d["k"].reshape(
+                                    self.kbins, self.mubins
+                                )
+                                _arch["mu3d"] = p3d["mu"].reshape(
+                                    self.kbins, self.mubins
+                                )
+                                _arch["p3d_Mpc"] = p3d["Pkmu"].reshape(
+                                    self.kbins, self.mubins
+                                )
                         self.data.append(_arch)
         ff.close()

scripts/save_mpg_emu_cosmo.py

@@ -25,6 +25,7 @@ def main():
             )[0, 0]
             if archive.data[ind]["ind_snap"] == list_snap[0]:
                 sim_dict["cosmo_params"] = archive.data[ind]["cosmo_params"]
+                sim_dict["star_params"] = archive.data[ind]["star_params"]
                 sim_dict["linP_params"] = {}
                 sim_dict["linP_params"]["kp_Mpc"] = archive.data[ind]["kp_Mpc"]
             for lab in labels:

scripts/save_nyx_emu_cosmo.py

@@ -1,6 +1,6 @@
 import numpy as np
 import h5py
-from lace.archive.nyx_archive import NyxArchive, get_attrs
+from lace.archive.nyx_archive import NyxArchive
 from lace.cosmo import camb_cosmo, fit_linP
 import os
 
@@ -13,6 +13,7 @@ def main():
 
     # file containing Nyx data
     nyx_version = "Oct2023"
+    nyx_version = "Jul2024"
     nyx_fname = os.environ["NYX_PATH"] + "/models_Nyx_" + nyx_version + ".hdf5"
     # file that will be written containing cosmo data for Nyx file
     cosmo_fname = (
@@ -31,6 +32,13 @@ def main():
     all_dict = []
     # iterate over simulations
     for sim_label in sim_avail:
+        # this simulation seems to have issues
+        if sim_label == "cosmo_grid_14":
+            continue
+        # contained in CGAN_4096_base
+        elif sim_label == "CGAN_4096_val":
+            continue
+
         isim = nyx_archive.sim_conv[sim_label]
         sim_dict = {}
         sim_dict["sim_label"] = isim
@@ -50,6 +58,7 @@ def main():
 
             if first:
                 sim_dict["cosmo_params"] = nyx_archive.data[ind]["cosmo_params"]
+                sim_dict["star_params"] = nyx_archive.data[ind]["star_params"]
                 sim_dict["linP_params"] = {}
                 sim_dict["linP_params"]["kp_Mpc"] = nyx_archive.data[ind][
                     "kp_Mpc"