From bfe3218e2692492805053e72146dae4fe5c6690b Mon Sep 17 00:00:00 2001 From: Sasha Rahlin Date: Sat, 4 Nov 2023 06:22:12 +1300 Subject: [PATCH] Option to run multiple sim indices per job 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`. --- xfaster/xfaster_exec.py | 64 ++++++++++++++++++++++++++--------------- 1 file changed, 41 insertions(+), 23 deletions(-) diff --git a/xfaster/xfaster_exec.py b/xfaster/xfaster_exec.py index 15330d9e..4d54e545 100644 --- a/xfaster/xfaster_exec.py +++ b/xfaster/xfaster_exec.py @@ -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, @@ -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_``. + 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 @@ -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, @@ -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, @@ -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 @@ -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)