Skip to content

Commit

Permalink
Option to run multiple sim indices per job
Browse files Browse the repository at this point in the history
New `num_sims` option for `xfaster run` and `xfaster submit` that runs the
prerequisite checkpoints once, then repeats the data, bandpowers and likelihood
checkpoints, incrementing the value of `sim_index_default` each time.
This option can be used for simple sets of sims where all of the input data
are determined by the `sim_data_components` and `sim_index_default` options.
More complex sim sets can be constructed using the `XFasterJobGroup` structure.

To run 100 simple sims efficiently, run one sim with `sim_index_default=0` to
pre-compute the prerequisite checkpoints, then split the remaining 99 sims into
multiple jobs by setting the `sim_index_default` and `num_sims` options to loop
over subsets of the full range of sim indices.  For example, running the rest of
the sims in a single job can be done with options `sim_index_default=1` and
`num_sims=99`.
  • Loading branch information
arahlin committed Nov 3, 2023
1 parent 5264afb commit bfe3218
Showing 1 changed file with 41 additions and 23 deletions.
64 changes: 41 additions & 23 deletions xfaster/xfaster_exec.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ def xfaster_run(
sim_index_noise=None,
sim_index_foreground=None,
sim_index_default=0,
num_sims=1,
signal_type_sim=None,
sim_data_r=None,
noise_type_sim=None,
Expand Down Expand Up @@ -337,6 +338,11 @@ def xfaster_run(
sim_index_default : int
Default sim index to use for any component with index < 0 or None
in ``sim_index_<comp>``.
num_sims : int
If > 1, repeat the data, bandpowers and likelihood checkpoints this many
times, incrementing the value of ``sim_index_default`` by 1 each time,
starting from the input value. Only used if ``sim_data`` is True.
All other options remain the same for each iteration.
signal_type_sim : str
The variant of signal sims to use for sim_index data maps.
This enables having a different noise sim ensemble to use for
Expand Down Expand Up @@ -589,6 +595,7 @@ def xfaster_run(
sim=sim_data,
components=None if not sim_data else sim_data_components,
index=None if not sim_data else sim_index,
num_sims=None if not sim_data else num_sims,
signal_type_sim=signal_type_sim if sim_data_r is None else "r",
r=sim_data_r,
noise_type_sim=noise_type_sim,
Expand All @@ -605,6 +612,7 @@ def xfaster_run(
config_vars.remove_option("XFaster General", "sim_index")
config_vars.remove_option("XFaster General", "qb_file_data")
config_vars.remove_option("XFaster General", "save_sim_data")
data_opts.pop("num_sims")

kernel_opts = dict(
pixwin=pixwin,
Expand Down Expand Up @@ -751,34 +759,43 @@ def xfaster_run(
X.log("Loading spectrum shape for bandpowers...", "notice")
X.get_signal_shape(**shape_opts)

X.log(
"Computing masked {} cross-spectra...".format(
"simulated data" if sim_data else "data"
),
"notice",
)
X.get_masked_data(**data_opts)
if sim_data:
idx0 = data_opts["index"]["default"]
else:
num_sims = 1
idx0 = 0

for idx in range(idx0, idx0 + num_sims):
if sim_data:
data_opts["index"]["default"] = idx
X.log(
"Computing masked {} cross-spectra...".format(
"simulated data index {}".format(idx) if sim_data else "data"
),
"notice",
)
X.get_masked_data(**data_opts)

if multi_map:
X.log("Computing multi-map bandpowers...", "notice")
qb, inv_fish = X.get_bandpowers(return_qb=True, **bandpwr_opts)
if multi_map:
X.log("Computing multi-map bandpowers...", "notice")
qb, inv_fish = X.get_bandpowers(return_qb=True, **bandpwr_opts)

if likelihood:
X.log("Computing multi-map likelihood...", "notice")
X.get_likelihood(qb, inv_fish, **like_opts)
if likelihood:
X.log("Computing multi-map likelihood...", "notice")
X.get_likelihood(qb, inv_fish, **like_opts)

else:
for map_tag, map_file in zip(X.map_tags, X.map_files):
X.log("Processing map {}: {}".format(map_tag, map_file), "notice")
else:
for map_tag, map_file in zip(X.map_tags, X.map_files):
X.log("Processing map {}: {}".format(map_tag, map_file), "notice")

X.log("Computing bandpowers for map {}".format(map_tag), "info")
qb, inv_fish = X.get_bandpowers(
map_tag=map_tag, return_qb=True, **bandpwr_opts
)
X.log("Computing bandpowers for map {}".format(map_tag), "info")
qb, inv_fish = X.get_bandpowers(
map_tag=map_tag, return_qb=True, **bandpwr_opts
)

if likelihood:
X.log("Computing likelihoods for map {}".format(map_tag), "info")
X.get_likelihood(qb, inv_fish, map_tag=map_tag, **like_opts)
if likelihood:
X.log("Computing likelihoods for map {}".format(map_tag), "info")
X.get_likelihood(qb, inv_fish, map_tag=map_tag, **like_opts)

cpu_elapsed = cpu_time() - cpu_start
time_elapsed = time.time() - time_start
Expand Down Expand Up @@ -1184,6 +1201,7 @@ def add_arg(
metavar="COMP",
)
add_arg(G, "sim_index_default", argtype=int)
add_arg(G, "num_sims", argtype=int)
add_arg(G, "sim_index_signal", argtype=int)
add_arg(G, "sim_index_noise", argtype=int)
add_arg(G, "sim_index_foreground", argtype=int)
Expand Down

0 comments on commit bfe3218

Please sign in to comment.