Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

don't save HaloField by default in high level functions #436

Closed
wants to merge 2 commits into from
Closed

don't save HaloField by default in high level functions #436

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
617ba38
don't save HaloField by default in high level functions
daviesje Oct 14, 2024
083a9d4
fix some earlier merge errors
daviesje Oct 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 4 additions & 9 deletions src/py21cmfast/outputs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,18 +129,13 @@ def prepare_for_perturb(self, flag_options: FlagOptions, force: bool = False):

self.prepare(keep=keep, force=force)

def prepare_for_halos(self, flag_options: FlagOptions, force: bool = False):
"""Ensure ICs have all boxes required for the halos, and no more."""
keep = ["hires_density"] # for dexm
if flag_options.HALO_STOCHASTICITY:
keep.append("lowres_density") # for the sampler
if self.user_params.USE_RELATIVE_VELOCITIES:
keep.append("lowres_vcb")
self.prepare(keep=keep, force=force)

def prepare_for_spin_temp(self, flag_options: FlagOptions, force: bool = False):
"""Ensure ICs have all boxes required for spin_temp, and no more."""
keep = []
# below the sampler resoltion we use Lagrangian density to get CHMF integrals
if flag_options.HALO_STOCHASTICITY:
keep.append("lowres_density")

if self.user_params.USE_RELATIVE_VELOCITIES:
keep.append("lowres_vcb")
self.prepare(keep=keep, force=force)
Expand Down
100 changes: 46 additions & 54 deletions src/py21cmfast/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,10 +367,14 @@ def _call_c_simple(fnc, *args):
return result[0]


def _get_config_direc(direc):
return str(os.path.expanduser(config["direc"] if direc is None else direc))


def _get_config_options(
direc, regenerate, write, hooks
) -> tuple[str, bool, dict[Callable, dict[str, Any]]]:
direc = str(os.path.expanduser(config["direc"] if direc is None else direc))
direc = _get_config_direc(direc)

if hooks is None or len(hooks) > 0:
hooks = hooks or {}
Expand Down Expand Up @@ -2692,8 +2696,7 @@ def run_coeval(
if redshift is None and perturb is None:
raise ValueError("Either redshift or perturb must be given")

direc, regenerate, hooks = _get_config_options(direc, regenerate, write, hooks)

direc = _get_config_direc(direc)
singleton = False
# Ensure perturb is a list of boxes, not just one.
if perturb is None:
Expand Down Expand Up @@ -2721,7 +2724,12 @@ def run_coeval(
if use_interp_perturb_field and flag_options.USE_MINI_HALOS:
raise ValueError("Cannot use an interpolated perturb field with minihalos!")

iokw = {"regenerate": regenerate, "hooks": hooks, "direc": direc}
iokw = {
"regenerate": regenerate,
"hooks": hooks,
"direc": direc,
"write": write,
}

if init_box is None:
init_box = initial_conditions(
Expand Down Expand Up @@ -2791,18 +2799,14 @@ def run_coeval(

perturb = perturb_

# Now we can purge init_box further.
try:
init_box.prepare_for_halos(flag_options=flag_options, force=always_purge)
except OSError:
pass

# get the halos (reverse redshift order)
pt_halos = []
# since the HaloField can be quite large, by default only save PerturbHaloField
kw_hf = {**kw, "write": write if write else False}
if flag_options.USE_HALO_FIELD and not flag_options.FIXED_HALO_GRIDS:
halos_desc = None
for i, z in enumerate(redshifts[::-1]):
halos = determine_halo_list(redshift=z, halos_desc=halos_desc, **kw)
halos = determine_halo_list(redshift=z, halos_desc=halos_desc, **kw_hf)
pt_halos += [perturb_halo_list(redshift=z, halo_field=halos, **kw)]

# we never want to store every halofield
Expand Down Expand Up @@ -3169,7 +3173,7 @@ def run_lightcone(
regenerate, write, direc, random_seed
See docs of :func:`initial_conditions` for more information.
"""
direc, regenerate, hooks = _get_config_options(direc, regenerate, write, hooks)
direc = _get_config_direc(direc)

if cosmo_params is None and lightconer is not None:
cosmo_params = CosmoParams.from_astropy(lightconer.cosmo)
Expand Down Expand Up @@ -3276,7 +3280,12 @@ def run_lightcone(
scrollz, coeval_callback, coeval_callback_redshifts
)

iokw = {"hooks": hooks, "regenerate": regenerate, "direc": direc}
iokw = {
"hooks": hooks,
"regenerate": regenerate,
"direc": direc,
"write": write,
}

if init_box is None: # no need to get cosmo, user params out of it.
init_box = initial_conditions(
Expand Down Expand Up @@ -3364,14 +3373,6 @@ def run_lightcone(
perturb = perturb_
perturb_min = perturb[np.argmin(scrollz)]

# Now that we've got all the perturb fields, we can purge init more.
try:
init_box.prepare_for_spin_temp(
flag_options=flag_options, force=always_purge
)
except OSError:
pass

kw = {
**{
"init_boxes": init_box,
Expand All @@ -3381,6 +3382,29 @@ def run_lightcone(
**iokw,
}

# since the HaloField can be quite large, by default only save PerturbHaloField
pt_halos = []
kw_hf = {**kw, "write": write if write else False}
if flag_options.USE_HALO_FIELD and not flag_options.FIXED_HALO_GRIDS:
halos_desc = None
for iz, z in enumerate(scrollz[::-1]):
halo_field = determine_halo_list(
redshift=z,
halos_desc=halos_desc,
**kw_hf,
)
halos_desc = halo_field
pt_halos += [perturb_halo_list(redshift=z, halo_field=halo_field, **kw)]

# we never want to store every halofield
try:
pt_halos[iz].purge(force=always_purge)
except OSError:
pass

# reverse the halo lists to be in line with the redshift lists
pt_halos = pt_halos[::-1]

photon_nonconservation_data = None
if flag_options.PHOTON_CONS_TYPE != 0:
setup_photon_cons(**kw)
Expand Down Expand Up @@ -3420,37 +3444,6 @@ def run_lightcone(
st, ib, prev_perturb = None, None, None
pf = None

pf = None

# Now we can purge init_box further.
try:
init_box.prepare_for_halos(flag_options=flag_options, force=always_purge)
except OSError:
pass

# we explicitly pass the descendant halos here since we have a redshift list prior
# this will generate the extra fields if STOC_MINIMUM_Z is given
pt_halos = []
if flag_options.USE_HALO_FIELD and not flag_options.FIXED_HALO_GRIDS:
halos_desc = None
for iz, z in enumerate(scrollz[::-1]):
halo_field = determine_halo_list(
redshift=z,
halos_desc=halos_desc,
**kw,
)
halos_desc = halo_field
pt_halos += [perturb_halo_list(redshift=z, halo_field=halo_field, **kw)]

# we never want to store every halofield
try:
pt_halos[iz].purge(force=always_purge)
except OSError:
pass

# reverse the halo lists to be in line with the redshift lists
pt_halos = pt_halos[::-1]

# Now that we've got all the perturb fields, we can purge init more.
try:
init_box.prepare_for_spin_temp(
Expand All @@ -3459,8 +3452,6 @@ def run_lightcone(
except OSError:
pass

ph = None

perturb_files = []
spin_temp_files = []
ionize_files = []
Expand All @@ -3473,6 +3464,7 @@ def run_lightcone(
z_halos = []
coeval = None
prev_coeval = None
ph = None
st2 = None
hbox2 = None
hbox = None
Expand Down
Loading