Skip to content

Commit

Permalink
modifying runs and examples
Browse files Browse the repository at this point in the history
  • Loading branch information
Gaetan Facchinetti ([email protected]) committed Jul 5, 2024
1 parent 1aac0fe commit a0863fd
Show file tree
Hide file tree
Showing 4 changed files with 120 additions and 51 deletions.
36 changes: 22 additions & 14 deletions examples/run_batch_lightcones.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import py21cmcast as p21c
from astropy import units
import os, gc, importlib

import pickle

parser = argparse.ArgumentParser()
parser.add_argument("config_dir", type = str, help="Path to config file")
Expand All @@ -29,19 +29,27 @@
try:

# run the lightcone with the given seed
lightcone, run_id, output_dir = p21c.run_lightcone_from_config(os.path.join(config_dir, "Config_" + str(i) + ".config"), n_omp, random_seed)

# define the p21c object to save
run = p21c.Fiducial(output_dir, z_bins, z_centers, k_bins, False, lightcone = lightcone, frac_noise = 0.2,
load = False, save = True, verbose = False, name = "Run_" + str(i), rs = lightcone.random_seed)

# compute the noise associated
run.compute_sensitivity()

# enforce freeing memory
del run
del lightcone
gc.collect()
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:
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)

if lightcone is not None :

# define the p21c object to save
run = p21c.Fiducial(output_dir, z_bins, z_centers, k_bins, False, lightcone = lightcone, frac_noise = 0.2,
load = False, save = True, verbose = False, name = "Run_" + str(i), rs = lightcone.random_seed)

# compute the noise associated
run.compute_sensitivity()

# enforce freeing memory
del run
del lightcone
gc.collect()

except Exception as e:
print(str(s_id) + " : Run " + str(i) + " failed: \n" + str(e))
Expand Down
21 changes: 15 additions & 6 deletions examples/run_lightcone.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import argparse
import py21cmcast as p21c
from astropy import units
import os.path
import pickle

parser = argparse.ArgumentParser()
parser.add_argument("config_file", type = str, help="Path to config file")
Expand All @@ -17,11 +19,18 @@
preprocess = args.pre_process

# Run the lightcone with the given seed
lightcone, run_id, output_dir = p21c.run_lightcone_from_config(config_file, n_omp, random_seed)
lightcone, run_id, output_dir, params = p21c.run_lightcone_from_config(config_file, n_omp, random_seed)

if preprocess is True:
z_bins, z_centers, k_bins = p21c.define_grid_modes_redshifts(6., 8 * units.MHz, z_max = 22, k_min = 0.1 / units.Mpc, k_max = 1 / units.Mpc)
p21c.Run(output_dir, "Lightcone_rs" + str(lightcone.random_seed) + "_" + run_id + ".h5", z_bins, z_centers, k_bins, False, lightcone = lightcone, load = False, save = True, verbose = False)
else:
lightcone.save(fname = "Lightcone_rs" + str(lightcone.random_seed) + "_" + run_id + ".h5", direc = output_dir)
# test if was able to pass the reionization test (if it applies)
if params.get('Q_max', 1.0) < 0.5:
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)

if lightcone is not None:
if preprocess is True:
z_bins, z_centers, k_bins = p21c.define_grid_modes_redshifts(6., 8 * units.MHz, z_max = 22, k_min = 0.1 / units.Mpc, k_max = 1 / units.Mpc)
p21c.Run(output_dir, "Lightcone_rs" + str(lightcone.random_seed) + "_" + run_id + ".h5", z_bins, z_centers, k_bins, False, lightcone = lightcone, load = False, save = True, verbose = False)
else:
lightcone.save(fname = "Lightcone_rs" + str(lightcone.random_seed) + "_" + run_id + ".h5", direc = output_dir)

1 change: 1 addition & 0 deletions src/py21cmcast/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
prepare_triangle_plot,
confidence_ellipse,
ellipse_from_covariance,
make_directory,
)

from .power import (
Expand Down
113 changes: 82 additions & 31 deletions src/py21cmcast/runs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -538,10 +538,17 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
----------
lightcone: Lightcone object (see 21cmFAST)
lightcone of the run
run_id: string
run_id: str
identifier of the run
output_dir: str
where to save the data
reio_test: bool
whether we pass the test for for reionization
"""

if not PY21CMFAST:
ImportError("Need 21cmFAST to run ligthcones")

####################### Getting the data ############################

config = configparser.ConfigParser(delimiters=':')
Expand All @@ -568,6 +575,54 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:
cosmo_params = p21c_tools.read_config_params(config.items('cosmo_params'), int_type=False)


# Special treatment to pass an int as global parameter
recombPhotCons = extra_params.get('RecombPhotonCons', None)
if recombPhotCons is not None:
extra_params['RecombPhotonCons'] = int(recombPhotCons)

keys_to_remove = ['redshift', 'max_redshift']
global_kwargs = {k: v for k, v in extra_params.items() if k not in keys_to_remove}


# If photon conservation used, first check if reionization happens soon enough
# for that run the analytical calibration curve without inhomogeneous recombination
# ask that the HII volume filling factor be higher than 0.5 at the smallest redshift
# this way runs with no complete ionization can be flagged before running the lightcone
# Note that the analytical estimation tends to overestimate Q_HII because it does not
# account for the recombinations (even though it does not accoung also for X-rays (pre)ionization)
if flag_options.get('PHOTON_CONS', False) is True:
try:
with p21f.global_params.use(**global_kwargs):
data = p21f.get_Q_analytic_nonconservation_data(user_params=user_params,
cosmo_params=cosmo_params,
astro_params=astro_params,
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]

if Q_max < 0.5:

# 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})

# returns a list of parameters we can save
params_out = {'astro_params': dict(astro_params.self),
'user_params': dict(user_params.self),
'flag_options': dict(flag_options.self),
'cosmo_params' : dict(cosmo_params.self),
'z_min' : z_min,
'Q_max' : Q_max}

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:
user_params['N_THREADS'] = int(n_omp)
Expand Down Expand Up @@ -600,41 +655,37 @@ def run_lightcone_from_config(config_file: str, n_omp: int = None, random_seed:

try:

if PY21CMFAST:

lightcone = p21f.run_lightcone(
user_params = user_params,
astro_params = astro_params,
cosmo_params = cosmo_params,
flag_options = flag_options,
lightcone_quantities = lightcone_quantities,
global_quantities = global_quantities,
direc = cache_path,
random_seed = random_seed,
write = False,
**extra_params,
**kwargs,
)
else:
ImportError("Need 21cmFAST to run ligthcones")

lightcone = p21f.run_lightcone(
user_params = user_params,
astro_params = astro_params,
cosmo_params = cosmo_params,
flag_options = flag_options,
lightcone_quantities = lightcone_quantities,
global_quantities = global_quantities,
direc = cache_path,
random_seed = random_seed,
write = False,
**extra_params,
**kwargs,
)

except Exception as e :

print(e)

# free the memory of the C code
p21f.free_C_memory()
print("The ligthcone could not run because :", e, flush=True)

lightcone = None
run_id = None

try:
# at the end, we clear the cache
if PY21CMFAST:
p21f.cache_tools.clear_cache(direc=cache_path)

else:
ImportError("Need 21cmFAST to run ligthcones")
####################### Clearing the cache (if exists) ############################

try:
p21f.cache_tools.clear_cache(direc=cache_path)
# delete the directory once it has been emptied
os.rmdir(cache_path)
os.rmdir(cache_path)
except FileNotFoundError:
pass

return lightcone, run_id, output_dir


return lightcone, run_id, output_dir, {}

0 comments on commit a0863fd

Please sign in to comment.