adding better automation

examples/example_wdm.config

@@ -0,0 +1,87 @@
+# Config file for Test run
+
+[run]
+name           : WDM_7.0_MINIHALOS
+
+# output_quantities
+[lightcone_quantities]
+brightness_temp : True
+density         : True
+xH_box          : True
+
+[global_quantities]
+brightness_temp : True
+xH_box          : True
+x_e_box         : True
+Ts_box          : True
+Tk_box          : True
+density         : True
+
+# extra_parameters
+[extra_params]
+redshift   : 3.5
+max_redshift   : 35
+
+# UserParams
+[user_params]
+HII_DIM : 128
+BOX_LEN : 256
+USE_FFTW_WISDOM            : True
+PERTURB_ON_HIGH_RES        : True
+USE_INTERPOLATION_TABLES   : True
+PS_FILTER                  : SHARPK
+POWER_SPECTRUM             : CLASS
+USE_CLASS_TABLES           : False
+USE_SIGMA_8_NORM           : False
+CLASS_FLUID_APPROX_WDM     : 3
+CLASS_FLUID_APPROX_NU      : 2
+NEUTRINO_MASS_HIERARCHY    : DEGENERATE
+USE_RELATIVE_VELOCITIES    : True
+
+# FlagOptions
+[flag_options]
+USE_MINI_HALOS             : True
+USE_MASS_DEPENDENT_ZETA    : True
+SUBCELL_RSD                : True
+INHOMO_RECO                : True
+USE_TS_FLUCT               : True
+
+# AstroParams
+[astro_params]
+F_STAR10            : -1.3
+F_STAR7_MINI        : -2.0
+ALPHA_STAR          : 0.5
+ALPHA_STAR_MINI     : 0.5
+t_STAR              : 0.5
+F_ESC10             : -1.0
+F_ESC7_MINI         : -2.0
+ALPHA_ESC           : -0.5
+L_X                 : 40.0
+L_X_MINI            : 40.0
+NU_X_THRESH         : 500.0 
+
+[cosmo_params]
+FRAC_WDM              : 1.0
+M_WDM                 : 7.0
+NEUTRINO_MASS_1       : 0.02
+VOLUME_FACTOR_SHARP_K : 2.5
+SHETH_q               : 1.0
+SHETH_p               : 0.3
+SHETH_A               : 0.322
+
+[astro_params_vary]
+F_STAR10            : 3.0
+F_STAR7_MINI        : 3.0
+ALPHA_STAR          : 3.0
+ALPHA_STAR_MINI     : 3.0
+t_STAR              : 3.0
+F_ESC10             : 3.0
+F_ESC7_MINI         : 3.0
+ALPHA_ESC           : 3.0
+L_X                 : 0.1
+L_X_MINI            : 0.1
+NU_X_THRESH         : 3.0 
+
+
+[cosmo_params_vary]
+M_WDM  : [ 1.        ,  1.0277316 ,  1.05623224,  1.08552325,  1.11562654, 1.14656465,  1.17836072,  1.21103855,  1.24462259,  1.27913796, 1.3146105 ,  1.35106675,  1.38853399,  1.42704026,  1.46661437, 1.50728593,  1.54908538,  1.59204399,  1.63619392,  1.68156819, 1.72820077,  1.77612654,  1.82538137,  1.87600211,  1.92802665, 1.98149391,  2.03644391,  2.09291775,  2.15095771,  2.2106072 , 2.27191088,  2.3349146 ,  2.39966551,  2.46621207,  2.53460408, 2.6048927 ,  2.67713054,  2.75137165,  2.82767159,  2.90608744, 2.98667789,  3.06950325,  3.15462548,  3.24210829,  3.33201714, 3.4244193 ,  3.51938392,  3.61698207,  3.71728676,  3.82037307, 3.92631812,  4.0352012 ,  4.14710378,  4.2621096 ,  4.38030471, 4.50177757,  4.62661906,  4.7549226 ,  4.88678421,  5.02230255, 5.16157903,  5.30471787,  5.45182618,  5.60301403,  5.75839457, 5.91808406,  6.082202  ,  6.25087118,  6.42421783,  6.60237167, 6.78546599,  6.97363781,  7.16702794,  7.36578108,  7.57004597, 7.77997545,  7.99572661,  8.21746089,  8.44534422,  8.67954712, 8.92024484,  9.16761749,  9.42185019,  9.68313316,  9.95166192, 10.22763742, 10.51126616, 10.80276038, 11.10233819, 11.41022379, 11.72664754, 12.05184622, 12.38606319, 12.72954852, 13.08255926, 13.44535955, 13.81822086, 14.20142222, 14.59525037, 15.        ]

examples/init_fisher.py

@@ -0,0 +1,31 @@
+import py21cmcast as p21c
+import argparse
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("config_file", type = str, help="Path to config file")
+parser.add_argument("-q", "--q", type = float, help="default percentage of variation")
+parser.add_argument("-f", "--force_overwrite", help="overwrite existing config directory", action = 'store_true')
+parser.add_argument("-v", "--verbose", help="the first random seed of the runs", action = 'store_true')
+parser.add_argument("-o", "--output", type = str, help = "output directory (if not given in config file)")
+parser.add_argument("-c", "--cache", type = str, help = "cache directory (if not given in config file)")
+parser.add_argument("-gs", "--generate_script", help = "whether to generate a script", action = 'store_true')
+parser.add_argument("-rs", "--random_seed", type = int, help = "random seed for the runs")
+parser.add_argument("-em", "--email_address", type = str, help = "email address for the cluster script")
+
+
+args = parser.parse_args()
+
+config_file = args.config_file
+
+q           = args.q
+overwrite   = args.force_overwrite
+verbose     = args.verbose
+output      = args.output
+cache       = args.cache
+gen_script  = args.generate_script
+random_seed = args.random_seed
+email       = args.email_address
+
+p21c.init_runs(config_file, q, overwrite, verbose, output, cache, gen_script, random_seed=random_seed, email_address = email)
+

examples/run_batch_lightcones.py

@@ -32,7 +32,7 @@
         lightcone, run_id, output_dir, params = p21c.run_lightcone_from_config(os.path.join(config_dir, "Config_" + str(i) + ".config"), n_omp, random_seed)
 
         # test if was able to pass the reionization test (if it applies)
-        if params.get('Q_max', 1.0) < 0.5:
+        if params.get('Q_max', 1.0) < 0.7 or params.get('Q_mid', 1.0) < 0.4:
             p21c.make_directory(os.path.join(output_dir, "cache"), clean_existing_dir=False)        
             with open(os.path.join(output_dir, 'cache/LateReionization_Run_' + run_id + '.pkl'), 'wb') as file:
                 pickle.dump(params, file)

examples/run_lightcone.py

@@ -22,7 +22,7 @@
 lightcone, run_id, output_dir, params = p21c.run_lightcone_from_config(config_file, n_omp, random_seed)
 
 # test if was able to pass the reionization test (if it applies)
-if params.get('Q_max', 1.0) < 0.5:
+if params.get('Q_max', 1.0) < 0.7 or params.get('Q_mid', 1.0) < 0.4:
     p21c.make_directory(os.path.join(output_dir, "cache"), clean_existing_dir=False)        
     with open(os.path.join(output_dir, 'cache/LateReionization_Run_' + run_id + '.pkl'), 'wb') as file:
         pickle.dump(params, file)

src/py21cmcast/core.py

@@ -389,7 +389,7 @@ def _load(self):
         try:
 
             with open(self._filename_data, 'rb') as file: 
-                data = np.load(file)
+                data = np.load(file, allow_pickle=True)
 
                 self._power_spectrum    = data['power_spectrum']
                 self._ps_poisson_noise  = data['ps_poisson_noise']

src/py21cmcast/runs.py

@@ -44,14 +44,14 @@
 ##################################################################################
 
 import configparser
-from copy import deepcopy
 import os
 
 from py21cmcast import tools as p21c_tools
 import warnings
 import numpy as np
 import random
-import gc, importlib
+
+from scipy import interpolate
 
 PY21CMFAST = True
 try:
@@ -63,31 +63,45 @@
 
 
 
-def init_runs(config_file: str, q_scale: float = 3., clean_existing_dir: bool = False, verbose = False) -> None :
+def init_runs(config_file: str, q_scale: float | None = None, clean_existing_dir: bool = False, 
+              verbose: bool = False, output_dir: str | None = None, cache_dir: str | None = None,
+              generate_script: bool = False, **kwargs) -> None :
 
     """ 
-    Initialise the runs for a fisher analysis according to 
-    a fiducial model defined in config_file
+    Initialise the runs for a fisher analysis according to a fiducial model defined in config_file
     
-    Params : 
-    --------
-    config_file : str
+    Params
+    ------
+    - config_file : str
         Path to the config file representing the fiducial
-    q_scale : float, optional
+    - q_scale : float, optional
         Gives the points where to compute the derivative in pourcentage of the fiducial parameters
-    erase_dir : bool, optional
+    - clean_existing_dir : bool, optional
         If True forces the creation of a folder 
+    - verbose : bool, optional
+    - output_dir : str, optional
+    - cache_dir : str, optional
+    **kwargs: 
+    ---------
+    passed to generate_slurm_script
     """
 
+    if q_scale is None:
+        q_scale = 3.0
+
     config = configparser.ConfigParser(delimiters=':')
     config.optionxform = str
 
     config.read(config_file)
+    name = config.get('run', 'name')
 
-    name            = config.get('run', 'name')
-    output_dir      = config.get('run', 'output_dir')
-    cache_dir       = config.get('run', 'cache_dir')
-
+    try:
+        output_dir = config.get('run', 'output_dir') if output_dir is None else output_dir
+        cache_dir  = config.get('run', 'cache_dir')  if cache_dir  is None else cache_dir
+    except configparser.NoSectionError:
+        output_dir = os.path.join(os.getcwd(), "output_" + name.upper())
+        cache_dir = os.path.join(os.getcwd(), "cache_" + name.upper())
+
     extra_params = {}
 
     try:
@@ -226,9 +240,57 @@ def init_runs(config_file: str, q_scale: float = 3., clean_existing_dir: bool =
                                        lightcone_q, global_q, extra_params, user_params, flag_options, astro_params_fid, cosmo_params, key)
         irun = irun + 1
 
+    if generate_script : 
+        generate_slurm_script(name, irun, output_dir, **kwargs)
+
     return 
 
 
+
+def generate_slurm_script(job_name, nruns, output_dir, random_seed: int | None = None, email_address: str | None = None):
+
+    """
+    Generate a bash script automatically for the job
+    """
+
+    if random_seed is None:
+        random_seed = 1993
+
+    current_dir = os.getcwd()
+
+    with open(os.path.join(current_dir, "submit_" + job_name.upper() + '.sh'), 'w') as file:
+
+        print("#!/bin/bash", file = file)
+        print("#", file = file)
+        print("#SBATCH --job-name=" + job_name.upper(), file = file)
+        print("#SBATCH --output=" + job_name.upper() + "log_%a.txt", file = file)
+        print("#", file = file)
+        print("#SBATCH --ntasks=1", file = file)
+        print("#SBATCH --cpus-per-task=8", file = file)
+        print("#SBATCH --time=02:00:00", file = file)
+        print("#SBATCH --mem=15000", file = file)
+        print("#", file = file)
+
+        if email_address is not None: 
+            print("#SBATCH --mail-type=ALL", file = file)
+            print("#SBATCH --mail-user=" + email_address, file = file)
+
+        print("#SBATCH --array=0-" + str(nruns-2), file = file)
+
+        print("", file = file)
+        print("source ~/.bash_modules", file = file)
+        print("source " + os.path.join(current_dir, "../bin/activate"), file = file)
+
+        print("", file = file)
+        print("export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK", file = file)
+        print("FILES=(" + os.path.join(output_dir, job_name.upper() + "/*.config") + ")", file = file)
+
+        print("", file = file)
+        print("srun python " + os.path.join(current_dir, "../21cmCAST/examples/run_lightcone.py") + " ${FILES[$SLURM_ARRAY_TASK_ID]} -nomp $SLURM_CPUS_PER_TASK -rs " + str(random_seed), file = file)
+
+
+
+
 def init_grid_runs(config_file: str, clean_existing_dir: bool = False, verbose:bool = False) -> None:
 
     config = configparser.ConfigParser(delimiters=':')
@@ -330,17 +392,20 @@ def indices_combinations(*arrays, index_combinations=None):
 
 
 
-def init_random_runs(config_file: str, *, n_sample = 1000, random_seed = None, clean_existing_dir: bool = False, verbose:bool = False) -> None:
+def init_random_runs(config_file: str, *, n_sample = 1000, random_seed = None, clean_existing_dir: bool = False, verbose:bool = False, n_folder:int = None) -> None:
 
     config = configparser.ConfigParser(delimiters=':')
     config.optionxform = str
 
     config.read(config_file)
 
-    name            = config.get('run', 'name')
+    name            = config.get('run', 'name') 
     output_dir      = config.get('run', 'output_dir')
     cache_dir       = config.get('run', 'cache_dir')
 
+    if n_folder is not None:
+        name = name + "_" + str(n_folder)
+
     output_run_dir = output_dir + "/" + name.upper() + "/"
     existing_dir = p21c_tools.make_directory(output_run_dir, clean_existing_dir = clean_existing_dir)
 
@@ -599,10 +664,25 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
                                                                  flag_options=flag_options, 
                                                                  **global_kwargs)
 
-                Q_max = data.get('Q_analytic', np.array([-1.0]))[0]
-                z_min = data.get('z_analytic', np.array([-1.0]))[0]
+                Q = data.get('Q_analytic', np.array([-1.0]))
+                z = data.get('z_analytic', np.array([-1.0]))
+
+                Q_max = Q[0]
+                z_min = z[0]
+                z_max = z[-1]
 
-                if Q_max < 0.5:
+                Q_mid = 1.0
+
+                if z_min <= 5.9 and z_max >= 5.9:
+                    # try block in order to avoid critical errors as, anyway computing Q_mid is
+                    # not crucial and should not make everything crash
+                    try:
+                        # value of Q at redshift 5.9 (where we have strong constraints)
+                        Q_mid = interpolate.interp1d(z, Q)(5.9).item()
+                    except Exception as e:
+                        print("Could not compute Q_mid because: ", e, flush=True)
+
+                if Q_max < 0.7 or Q_mid < 0.4:
 
                     # get all the input parameters
                     (user_params, cosmo_params, astro_params, flag_options) = p21f._setup_inputs({ "user_params": user_params, "cosmo_params": cosmo_params, "astro_params" : astro_params, "flag_options" : flag_options})
@@ -613,15 +693,15 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
                                 'flag_options': dict(flag_options.self),
                                 'cosmo_params' : dict(cosmo_params.self),
                                 'z_min' : z_min,
-                                'Q_max' : Q_max}
+                                'Q_max' : Q_max, 
+                                'Q_mid' : Q_mid,}
 
                     return None, run_id, output_dir, params_out
 
         except Exception as e:
             p21f.free_C_memory()
             print("Impossible to get the analytic calibration value of Q_HII because :", e, flush=True)
 
-
 
     # manually set the number of threads
     if n_omp is not None: 
@@ -653,6 +733,9 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
     if len(global_quantities) == 0:
         global_quantities = ('brightness_temp', 'xH_box',)
 
+
+    error = None
+
     try: 
 
         lightcone = p21f.run_lightcone(
@@ -664,7 +747,7 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
                 global_quantities    = global_quantities,
                 direc                = cache_path, 
                 random_seed          = random_seed,
-                write                = False,
+                write                = True,
                 **extra_params,
                 **kwargs,
             )
@@ -674,9 +757,11 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
         # free the memory of the C code
         p21f.free_C_memory()
         print("The ligthcone could not run because :", e, flush=True)
-
+    
         lightcone = None
 
+        error = e
+
     ####################### Clearing the cache (if exists) ############################
 
     try:
@@ -686,6 +771,7 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
     except FileNotFoundError:
         pass
 
-
+    if error is not None:
+        raise error
 
     return lightcone, run_id, output_dir, {}