From f50c1ba9dfcd2bbebaeada2243889866e62d3831 Mon Sep 17 00:00:00 2001 From: Karl Gordon Date: Tue, 13 Aug 2024 09:24:32 -0400 Subject: [PATCH] adding simulate script --- megabeast/tools/simulate_obs.py | 210 ++++++++++++++++++++++++++++++++ setup.cfg | 1 + 2 files changed, 211 insertions(+) create mode 100644 megabeast/tools/simulate_obs.py diff --git a/megabeast/tools/simulate_obs.py b/megabeast/tools/simulate_obs.py new file mode 100644 index 0000000..e6ec8da --- /dev/null +++ b/megabeast/tools/simulate_obs.py @@ -0,0 +1,210 @@ +import numpy as np +import argparse +import asdf + +from beast.physicsmodel.grid import SEDGrid +import beast.observationmodel.noisemodel.generic_noisemodel as noisemodel +from beast.observationmodel.observations import gen_SimObs_from_sedgrid + +from astropy.table import vstack + + +def simulate_obs( + physgrid_list, + noise_model_list, + output_catalog, + beastinfo_list=None, + nsim=100, + compl_filter="max", + complcut=None, + magcut=None, + weight_to_use="weight", + ranseed=None, +): + """ + Simulate photometry based on a physicsmodel grid(s) and observation model(s). + + Parameters + ---------- + physgrid_list : list of strings + Name of the physics model file. If there are multiple physics model + grids (i.e., if there are subgrids), list them all here, and they will + each be sampled nsim/len(physgrid_list) times. + + noise_model_list : list of strings + Name of the noise model file. If there are multiple files for + physgrid_list (because of subgrids), list the noise model files + associated with each physics model file. + + output_catalog : string + Name of the output simulated photometry catalog + + beastinfo_list : list of strings (default=None) + Name of the beast info file. The mass and age prior models are read + from this model to use to compute the number of stars to simulate. If + there are multiple files for physgrid_list (because of subgrids), list + the beast info files associated with each physics model file. + + n_sim : int (default=100) + Number of simulated objects to create if beastinfo_list is not given. If + nsim/len(physgrid_list) isn't an integer, this will be increased so that + each grid has the same number of samples. + + compl_filter : str (default=max) + filter to use for completeness (required for toothpick model) + set to max to use the max value in all filters + + complcut : float (defualt=None) + completeness cut for only including model seds above the given + completeness, where the cut ranges from 0 to 1. + + magcut : float (defualt=None) + faint-end magnitude cut for only including model seds brighter + than the given magnitude in compl_filter. + + weight_to_use : string (default='weight') + Set to either 'weight' (prior+grid), 'prior_weight', 'grid_weight', + or 'uniform' (this option is valid only when nsim is supplied) to + choose the weighting for SED selection. + + ranseed : int + seed for random number generator + """ + # numbers of samples to do + # (ensure there are enough for even sampling of multiple model grids) + n_phys = len(np.atleast_1d(physgrid_list)) + nsim = int(nsim) + samples_per_grid = int(np.ceil(nsim / n_phys)) + + if complcut is not None: + complcut = float(complcut) + + if magcut is not None: + magcut = float(magcut) + + if ranseed is not None: + ranseed = int(ranseed) + + # list to hold all simulation tables + simtable_list = [] + + # make a table for each physics model + noise model + for k, physgrid in enumerate(np.atleast_1d(physgrid_list)): + noise_model = np.atleast_1d(noise_model_list)[k] + + # get the physics model grid - includes priors + modelsedgrid = SEDGrid(str(physgrid)) + + # read in the noise model - includes bias, unc, and completeness + noisegrid = noisemodel.get_noisemodelcat(str(noise_model)) + + if beastinfo_list is not None: + with asdf.open(np.atleast_1d(beastinfo_list)[k]) as af: + binfo = af.tree + age_prior_model = binfo["age_prior_model"] + mass_prior_model = binfo["mass_prior_model"] + else: + age_prior_model = None + mass_prior_model = None + + # generate the table + simtable = gen_SimObs_from_sedgrid( + modelsedgrid, + noisegrid, + age_prior_model=age_prior_model, + mass_prior_model=mass_prior_model, + nsim=samples_per_grid, + compl_filter=compl_filter, + complcut=complcut, + magcut=magcut, + weight_to_use=weight_to_use, + ranseed=ranseed, + ) + + # append to the list + simtable_list.append(simtable) + + # stack all the tables into one and write it out + vstack(simtable_list).write(output_catalog, overwrite=True) + + +if __name__ == "__main__": # pragma: no cover + + # commandline parser + parser = argparse.ArgumentParser() + parser.add_argument( + "--physgrid_list", + "-p", + metavar="PHYS_MODEL", + required=True, + nargs="+", + help="filename(s) of physics grid(s)", + ) + parser.add_argument( + "--noise_model_list", + "-n", + metavar="NOISE_MODEL", + required=True, + nargs="+", + help="filename(s) of observation/noise grid(s)", + ) + parser.add_argument( + "--output_catalog", + "-c", + required=True, + help="filename for simulated observations", + ) + parser.add_argument( + "--beastinfo_list", + metavar="BEAST_INFO", + nargs="+", + help="filename(s) of beast info file(s)", + ) + parser.add_argument( + "--nsim", default=100, type=int, help="number of simulated objects" + ) + parser.add_argument( + "--compl_filter", + default="F475W", + help="filter for completeness, set to max for max of values from all filters", + ) + parser.add_argument( + "--complcut", + default=None, + type=float, + help="completeness cut for selecting seds above the completeness cut" + ) + parser.add_argument( + "--magcut", + default=None, + type=float, + help="magnitdue cut for selecting seds brighter than the magcut in compl_filter" + ) + parser.add_argument( + "--weight_to_use", + default="weight", + type=str, + help="""Set to either 'weight' (prior+grid), 'prior_weight', or + 'grid_weight' to choose the weighting for SED selection.""", + ) + parser.add_argument( + "--ranseed", + default=None, + type=int, + help="seed for random number generator" + ) + args = parser.parse_args() + + # run observation simulator + simulate_obs( + args.physgrid_list, + args.noise_model_lists, + args.output_catalog, + beastinfo_list=args.beastinfo_list, + nsim=args.nsim, + compl_filter=args.compl_filter, + complcut=args.complcut, + magcut=args.magcut, + weight_to_use=args.weight_to_use, + ranseed=args.ranseed, + ) diff --git a/setup.cfg b/setup.cfg index 2c1de7d..ddfd804 100644 --- a/setup.cfg +++ b/setup.cfg @@ -45,6 +45,7 @@ addopts = --doctest-rst [options.entry_points] console_scripts = mb_fit_single = megabeast.fit_single:main + mb_simulate_obs = megabeast.tools.simulate_obs:main [coverage:run] omit =