def split_create_physicsmodel(nsubs=1, nprocs=1):
"""
Making the physics model grid takes a while for production runs. This
creates scripts to run each subgrid as a separate job.
Parameters
----------
nsubs : int (default=1)
number of subgrids to split the physics model into
nprocs : int (default=1)
Number of parallel processes to use
(currently only implemented for subgrids)
"""
# before doing ANYTHING, force datamodel to re-import (otherwise, any
# changes within this python session will not be loaded!)
importlib.reload(datamodel)
# check input parameters
verify_params.verify_input_format(datamodel)
# make sure the project directory exists
create_project_dir(datamodel.project)
# directory for scripts
job_path = "./{0}/model_batch_jobs/".format(datamodel.project)
if not os.path.isdir(job_path):
os.mkdir(job_path)
log_path = job_path + "logs/"
if not os.path.isdir(log_path):
os.mkdir(log_path)
for i in range(nsubs):
joblist_file = job_path + "create_physicsmodel_" + str(i) + ".job"
with open(joblist_file, "w") as jf:
jf.write("python -m beast.tools.run.create_physicsmodel " +
" --nsubs " + str(nsubs) + " --nprocs " + str(nprocs) +
" --subset " + str(i) + " " + str(i + 1) + " >> " +
log_path + "create_physicsmodel_" + str(i) + ".log\n")
# slurm needs it to be executable
os.chmod(joblist_file, stat.S_IRWXU | stat.S_IRGRP | stat.S_IROTH)
def gen_subobsmodel(modelsedgridfile):
modelsedgrid = FileSEDGrid(modelsedgridfile)
# generate the AST noise model
noisefile = modelsedgridfile.replace("seds", "noisemodel")
astfile = settings.astfile
# If we are treating regions with different
# backgrounds/source densities separately, pick one of the
# split ast files, and put the results in a subfolder.
if args.dens_bin is not None:
noisefile = os.path.join(bin_subfolder, noisefile)
create_project_dir(os.path.dirname(noisefile))
astfile = subcatalog_fname(astfile, args.dens_bin)
outname = noisemodel.make_toothpick_noise_model(
noisefile,
astfile,
modelsedgrid,
absflux_a_matrix=settings.absflux_a_matrix,
)
return outname
def create_physicsmodel(nsubs=1, nprocs=1, subset=[None, None]):
"""
Create the physics model grid. If nsubs > 1, this will make sub-grids.
Parameters
----------
nsubs : int (default=1)
number of subgrids to split the physics model into
nprocs : int (default=1)
Number of parallel processes to use
(currently only implemented for subgrids)
subset : list of two ints (default=[None,None])
Only process subgrids in the range [start,stop].
(only relevant if nsubs > 1)
"""
# before doing ANYTHING, force datamodel to re-import (otherwise, any
# changes within this python session will not be loaded!)
importlib.reload(datamodel)
# check input parameters
verify_params.verify_input_format(datamodel)
# filename for the SED grid
modelsedgrid_filename = "%s/%s_seds.grid.hd5" % (
datamodel.project,
datamodel.project,
)
# grab the current subgrid slice
subset_slice = slice(subset[0], subset[1])
# make sure the project directory exists
create_project_dir(datamodel.project)
# download and load the isochrones
(iso_fname, oiso) = make_iso_table(
datamodel.project,
oiso=datamodel.oiso,
logtmin=datamodel.logt[0],
logtmax=datamodel.logt[1],
dlogt=datamodel.logt[2],
z=datamodel.z,
)
# remove the isochrone points with logL=-9.999
oiso = ezIsoch(oiso.selectWhere("*", "logL > -9"))
if hasattr(datamodel, "add_spectral_properties_kwargs"):
extra_kwargs = datamodel.add_spectral_properties_kwargs
else:
extra_kwargs = None
if hasattr(datamodel, "velocity"):
redshift = (datamodel.velocity / const.c).decompose().value
else:
redshift = 0
# generate the spectral library (no dust extinction)
(spec_fname, g_spec) = make_spectral_grid(
datamodel.project,
oiso,
osl=datamodel.osl,
redshift=redshift,
distance=datamodel.distances,
distance_unit=datamodel.distance_unit,
extLaw=datamodel.extLaw,
add_spectral_properties_kwargs=extra_kwargs,
)
# add the stellar priors as weights
# also computes the grid weights for the stellar part
(pspec_fname, g_pspec) = add_stellar_priors(
datamodel.project,
g_spec,
age_prior_model=datamodel.age_prior_model,
mass_prior_model=datamodel.mass_prior_model,
met_prior_model=datamodel.met_prior_model,
)
# --------------------
# no subgrids
# --------------------
if nsubs == 1:
# generate the SED grid by integrating the filter response functions
# effect of dust extinction applied before filter integration
# also computes the dust priors as weights
make_extinguished_sed_grid(
datamodel.project,
g_pspec,
datamodel.filters,
extLaw=datamodel.extLaw,
av=datamodel.avs,
rv=datamodel.rvs,
fA=datamodel.fAs,
rv_prior_model=datamodel.rv_prior_model,
av_prior_model=datamodel.av_prior_model,
fA_prior_model=datamodel.fA_prior_model,
spec_fname=modelsedgrid_filename,
add_spectral_properties_kwargs=extra_kwargs,
)
# --------------------
# use subgrids
# --------------------
if nsubs > 1:
# Work with the whole grid up to there (otherwise, priors need a
# rework - they don't like having only a subset of the parameter
# space, especially when there's only one age for example)
# Make subgrids, by splitting the spectral grid into equal sized pieces
custom_sub_pspec = subgridding_tools.split_grid(pspec_fname, nsubs)
file_prefix = "{0}/{0}_".format(datamodel.project)
# function to process the subgrids individually
def gen_subgrid(i, sub_name):
sub_g_pspec = FileSEDGrid(sub_name)
sub_seds_fname = "{}seds.gridsub{}.hd5".format(file_prefix, i)
# generate the SED grid by integrating the filter response functions
# effect of dust extinction applied before filter integration
# also computes the dust priors as weights
(sub_seds_fname, sub_g_seds) = make_extinguished_sed_grid(
datamodel.project,
sub_g_pspec,
datamodel.filters,
extLaw=datamodel.extLaw,
av=datamodel.avs,
rv=datamodel.rvs,
fA=datamodel.fAs,
rv_prior_model=datamodel.rv_prior_model,
av_prior_model=datamodel.av_prior_model,
fA_prior_model=datamodel.fA_prior_model,
add_spectral_properties_kwargs=extra_kwargs,
seds_fname=sub_seds_fname,
)
return sub_seds_fname
# run the above function
par_tuples = [
(i, sub_name) for i, sub_name in enumerate(custom_sub_pspec)
][subset_slice]
parallel_wrapper(gen_subgrid, par_tuples, nprocs=nprocs)
# Save a list of subgrid names that we expect to see
required_names = [
"{}seds.gridsub{}.hd5".format(file_prefix, i) for i in range(nsubs)
]
outdir = os.path.join(".", datamodel.project)
subgrid_names_file = os.path.join(outdir, "subgrid_fnames.txt")
with open(subgrid_names_file, "w") as fname_file:
for fname in required_names:
fname_file.write(fname + "\n")
)
parser.add_argument(
"--dens_bin",
type=int,
default=None,
help="""Run for a certain source/background
density bin. Use the split_catalog_using_map `
tool to generate the right files for this
option.""",
)
args = parser.parse_args()
if args.dens_bin is not None:
bin_subfolder = "bin{}".format(args.dens_bin)
pdir = create_project_dir(bin_subfolder)
def subcatalog_fname(full_cat_fname, dens_bin):
return full_cat_fname.replace(".fits", "_bin{}.fits".format(dens_bin))
def parallel_wrapper(function, argument):
parallel = args.nprocs > 1
if parallel:
p = Pool(args.nprocs)
for r in p.imap_unordered(function, argument, chunksize=1):
print(r)
else:
for a in argument:
r = function(a)
print(r)
def run_beast_production(basename,
physicsmodel=False,
ast=False,
observationmodel=False,
trim=False,
fitting=False,
resume=False,
source_density='',
sub_source_density=''):
"""
Turns the original command-line version of run_beast_production.py into
something callable from within a function
Parameters
----------
basename : string
name of the gst file (assuming it's located in ./data/)
For the info related to the other inputs, see the argparse info at the bottom
"""
# before doing ANYTHING, force datamodel to re-import (otherwise, any
# changes within this python session will not be loaded!)
importlib.reload(datamodel)
# check input parameters, print what is the problem, stop run_beast
verify_params.verify_input_format(datamodel)
# update the filenames as needed for production
# - photometry sub-file
datamodel.obsfile = basename.replace(
'.fits',
'_with_sourceden' + '_SD_' + source_density.replace('_', '-') +
'_sub' + sub_source_density + '.fits')
# - stats files
stats_filebase = "%s/%s"%(datamodel.project,datamodel.project) \
+ '_sd' + source_density.replace('_','-') \
+ '_sub' + sub_source_density
sed_trimname = stats_filebase + '_sed_trim.grid.hd5'
# - trimmed noise model
noisemodel_trimname = stats_filebase + '_noisemodel_trim.hd5'
# - SED grid
#modelsedgrid_filename = "%s/%s_seds.grid.hd5"%(datamodel.project,
# datamodel.project)
modelsedgrid_filename = "METAL_seds.grid.hd5"
print("***run information***")
print(" project = " + datamodel.project)
print(" obsfile = " + datamodel.obsfile)
print(" astfile = " + datamodel.astfile)
print(" noisefile = " + datamodel.noisefile)
print(" trimmed sedfile = " + sed_trimname)
print("trimmed noisefiles = " + noisemodel_trimname)
print(" stats filebase = " + stats_filebase)
# make sure the project directory exists
pdir = create_project_dir(datamodel.project)
if physicsmodel:
# download and load the isochrones
(iso_fname, oiso) = make_iso_table(datamodel.project,
oiso=datamodel.oiso,
logtmin=datamodel.logt[0],
logtmax=datamodel.logt[1],
dlogt=datamodel.logt[2],
z=datamodel.z)
if hasattr(datamodel, 'add_spectral_properties_kwargs'):
extra_kwargs = datamodel.add_spectral_properties_kwargs
else:
extra_kwargs = None
if hasattr(datamodel, 'velocity'):
redshift = (datamodel.velocity / const.c).decompose().value
else:
redshift = 0
# generate the spectral library (no dust extinction)
(spec_fname, g_spec) = make_spectral_grid(
datamodel.project,
oiso,
osl=datamodel.osl,
redshift=redshift,
distance=datamodel.distances,
distance_unit=datamodel.distance_unit,
add_spectral_properties_kwargs=extra_kwargs)
# add the stellar priors as weights
# also computes the grid weights for the stellar part
(pspec_fname, g_pspec) = add_stellar_priors(datamodel.project, g_spec)
# generate the SED grid by integrating the filter response functions
# effect of dust extinction applied before filter integration
# also computes the dust priors as weights
(seds_fname, g_seds) = make_extinguished_sed_grid(
datamodel.project,
g_pspec,
datamodel.filters,
extLaw=datamodel.extLaw,
av=datamodel.avs,
rv=datamodel.rvs,
fA=datamodel.fAs,
rv_prior_model=datamodel.rv_prior_model,
av_prior_model=datamodel.av_prior_model,
fA_prior_model=datamodel.fA_prior_model,
spec_fname=modelsedgrid_filename,
add_spectral_properties_kwargs=extra_kwargs)
if ast:
N_models = datamodel.ast_models_selected_per_age
Nfilters = datamodel.ast_bands_above_maglimit
Nrealize = datamodel.ast_realization_per_model
mag_cuts = datamodel.ast_maglimit
obsdata = datamodel.get_obscat(basename, datamodel.filters)
if len(mag_cuts) == 1:
tmp_cuts = mag_cuts
min_mags = np.zeros(len(datamodel.filters))
for k, filtername in enumerate(obsdata.filters):
sfiltername = obsdata.data.resolve_alias(filtername)
sfiltername = sfiltername.replace('rate', 'vega')
sfiltername = sfiltername.replace('RATE', 'VEGA')
keep, = np.where(obsdata[sfiltername] < 99.)
min_mags[k] = np.percentile(obsdata[keep][sfiltername], 90.)
# max. mags from the gst observation cat.
mag_cuts = min_mags + tmp_cuts
outfile = './' + datamodel.project + '/' + datamodel.project + '_inputAST.txt'
outfile_params = './' + datamodel.project + '/' + datamodel.project + '_ASTparams.fits'
chosen_seds = pick_models(modelsedgrid_filename,
datamodel.filters,
mag_cuts,
Nfilter=Nfilters,
N_stars=N_models,
Nrealize=Nrealize,
outfile=outfile,
outfile_params=outfile_params)
if datamodel.ast_with_positions == True:
separation = datamodel.ast_pixel_distribution
filename = datamodel.project + '/' + datamodel.project + '_inputAST.txt'
if datamodel.ast_reference_image is not None:
# With reference image, use the background or source density map if available
if datamodel.ast_density_table is not None:
pick_positions_from_map(
obsdata,
chosen_seds,
datamodel.ast_density_table,
datamodel.ast_N_bins,
datamodel.ast_realization_per_model,
outfile=filename,
refimage=datamodel.ast_reference_image,
refimage_hdu=0,
Nrealize=1,
set_coord_boundary=datamodel.ast_coord_boundary)
else:
pick_positions(obsdata,
filename,
separation,
refimage=datamodel.ast_reference_image)
else:
# Without reference image, we can only use this function
if datamodel.ast_density_table is None:
pick_positions(obsdata, filename, separation)
else:
print(
"To use ast_density_table, ast_reference_image must be specified."
)
if observationmodel:
print('Generating noise model from ASTs and absflux A matrix')
# get the modesedgrid on which to generate the noisemodel
modelsedgrid = FileSEDGrid(modelsedgrid_filename)
# generate the AST noise model
noisemodel.make_toothpick_noise_model( \
datamodel.noisefile,
datamodel.astfile,
modelsedgrid,
use_rate=True,
absflux_a_matrix=datamodel.absflux_a_matrix)
if trim:
print('Trimming the model and noise grids')
# read in the observed data
obsdata = datamodel.get_obscat(basename, datamodel.filters)
# get the modesedgrid on which to generate the noisemodel
modelsedgrid = FileSEDGrid(modelsedgrid_filename)
# read in the noise model just created
noisemodel_vals = noisemodel.get_noisemodelcat(datamodel.noisefile)
# trim the model sedgrid
trim_grid.trim_models(modelsedgrid,
noisemodel_vals,
obsdata,
sed_trimname,
noisemodel_trimname,
sigma_fac=3.)
if fitting:
start_time = time.clock()
# read in the the AST noise model
noisemodel_vals = noisemodel.get_noisemodelcat(noisemodel_trimname)
# read in the observed data
obsdata = datamodel.get_obscat(datamodel.obsfile, datamodel.filters)
# output files
statsfile = stats_filebase + '_stats.fits'
pdf1dfile = statsfile.replace('stats.fits', 'pdf1d.fits')
lnpfile = statsfile.replace('stats.fits', 'lnp.hd5')
fit.summary_table_memory(obsdata,
noisemodel_vals,
sed_trimname,
resume=resume,
threshold=-10.,
save_every_npts=100,
lnp_npts=500,
stats_outname=statsfile,
pdf1d_outname=pdf1dfile,
lnp_outname=lnpfile,
surveyname=datamodel.surveyname)
new_time = time.clock()
print('time to fit: ', (new_time - start_time) / 60., ' min')
"--fit",
help="Fit the observed data",
action="store_true")
parser.add_argument("-r",
"--resume",
help="Resume a fitting run",
action="store_true")
args = parser.parse_args()
# check input parameters, print what is the problem, stop run_beast
verify_params.verify_input_format(datamodel)
if args.physicsmodel:
# make sure the project directory exists
pdir = create_project_dir(datamodel.project)
# download and load the isochrones
(iso_fname, oiso) = make_iso_table(datamodel.project,
oiso=datamodel.oiso,
logtmin=datamodel.logt[0],
logtmax=datamodel.logt[1],
dlogt=datamodel.logt[2],
z=datamodel.z)
# calculate the distance in pc
if datamodel.distanceModulus.unit == units.mag:
dmod = datamodel.distanceModulus.value
distance = 10**((dmod / 5.) + 1) * units.pc
else:
raise ValueError("distance modulus does not have mag units")
def create_physicsmodel(beast_settings_info, nsubs=1, nprocs=1, subset=[None, None]):
"""
Create the physics model grid. If nsubs > 1, this will make sub-grids.
Parameters
----------
beast_settings_info : string or beast.tools.beast_settings.beast_settings instance
if string: file name with beast settings
if class: beast.tools.beast_settings.beast_settings instance
nsubs : int (default=1)
number of subgrids to split the physics model into
nprocs : int (default=1)
Number of parallel processes to use
(currently only implemented for subgrids)
subset : list of two ints (default=[None,None])
Only process subgrids in the range [start,stop].
(only relevant if nsubs > 1)
"""
# process beast settings info
if isinstance(beast_settings_info, str):
settings = beast_settings.beast_settings(beast_settings_info)
elif isinstance(beast_settings_info, beast_settings.beast_settings):
settings = beast_settings_info
else:
raise TypeError(
"beast_settings_info must be string or beast.tools.beast_settings.beast_settings instance"
)
# filename for the SED grid
modelsedgrid_filename = "%s/%s_seds.grid.hd5" % (
settings.project,
settings.project,
)
# grab the current subgrid slice
subset_slice = slice(subset[0], subset[1])
# make sure the project directory exists
create_project_dir(settings.project)
# download and load the isochrones
(iso_fname, oiso) = make_iso_table(
settings.project,
oiso=settings.oiso,
logtmin=settings.logt[0],
logtmax=settings.logt[1],
dlogt=settings.logt[2],
z=settings.z,
)
if hasattr(settings, "add_spectral_properties_kwargs"):
extra_kwargs = settings.add_spectral_properties_kwargs
else:
extra_kwargs = None
if hasattr(settings, "velocity"):
redshift = (settings.velocity / const.c).decompose().value
else:
redshift = 0
# generate the spectral library (no dust extinction)
(spec_fname, g_spec) = make_spectral_grid(
settings.project,
oiso,
osl=settings.osl,
redshift=redshift,
distance=settings.distances,
distance_unit=settings.distance_unit,
extLaw=settings.extLaw,
add_spectral_properties_kwargs=extra_kwargs,
)
# add the stellar priors as weights
# also computes the grid weights for the stellar part
(pspec_fname, g_pspec) = add_stellar_priors(
settings.project,
g_spec,
age_prior_model=settings.age_prior_model,
mass_prior_model=settings.mass_prior_model,
met_prior_model=settings.met_prior_model,
distance_prior_model=settings.distance_prior_model,
)
# --------------------
# no subgrids
# --------------------
if nsubs == 1:
# generate the SED grid by integrating the filter response functions
# effect of dust extinction applied before filter integration
# also computes the dust priors as weights
make_extinguished_sed_grid(
settings.project,
g_pspec,
settings.filters,
extLaw=settings.extLaw,
av=settings.avs,
rv=settings.rvs,
fA=settings.fAs,
rv_prior_model=settings.rv_prior_model,
av_prior_model=settings.av_prior_model,
fA_prior_model=settings.fA_prior_model,
spec_fname=modelsedgrid_filename,
add_spectral_properties_kwargs=extra_kwargs,
)
# --------------------
# use subgrids
# --------------------
if nsubs > 1:
# Work with the whole grid up to there (otherwise, priors need a
# rework - they don't like having only a subset of the parameter
# space, especially when there's only one age for example)
# Make subgrids, by splitting the spectral grid into equal sized pieces
custom_sub_pspec = subgridding_tools.split_grid(pspec_fname, nsubs)
file_prefix = "{0}/{0}_".format(settings.project)
# function to process the subgrids individually
def gen_subgrid(i, sub_name):
sub_g_pspec = SpectralGrid(sub_name)
sub_seds_fname = "{}seds.gridsub{}.hd5".format(file_prefix, i)
# generate the SED grid by integrating the filter response functions
# effect of dust extinction applied before filter integration
# also computes the dust priors as weights
(sub_seds_fname, sub_g_seds) = make_extinguished_sed_grid(
settings.project,
sub_g_pspec,
settings.filters,
extLaw=settings.extLaw,
av=settings.avs,
rv=settings.rvs,
fA=settings.fAs,
rv_prior_model=settings.rv_prior_model,
av_prior_model=settings.av_prior_model,
fA_prior_model=settings.fA_prior_model,
add_spectral_properties_kwargs=extra_kwargs,
seds_fname=sub_seds_fname,
)
return sub_seds_fname
# run the above function
par_tuples = [(i, sub_name) for i, sub_name in enumerate(custom_sub_pspec)][
subset_slice
]
parallel_wrapper(gen_subgrid, par_tuples, nprocs=nprocs)
# Save a list of subgrid names that we expect to see
required_names = [
"{}seds.gridsub{}.hd5".format(file_prefix, i) for i in range(nsubs)
]
outdir = os.path.join(".", settings.project)
subgrid_names_file = os.path.join(outdir, "subgrid_fnames.txt")
with open(subgrid_names_file, "w") as fname_file:
for fname in required_names:
fname_file.write(fname + "\n")
def split_create_physicsmodel(beast_settings_info, nsubs=1, nprocs=1):
"""
Making the physics model grid takes a while for production runs. This
creates scripts to run each subgrid as a separate job.
Parameters
----------
beast_settings_info : string or beast.tools.beast_settings.beast_settings instance
if string: file name with beast settings
if class: beast.tools.beast_settings.beast_settings instance
nsubs : int (default=1)
number of subgrids to split the physics model into
nprocs : int (default=1)
Number of parallel processes to use
(currently only implemented for subgrids)
"""
# process beast settings info
if isinstance(beast_settings_info, str):
settings = beast_settings.beast_settings(beast_settings_info)
elif isinstance(beast_settings_info, beast_settings.beast_settings):
settings = beast_settings_info
else:
raise TypeError(
"beast_settings_info must be string or beast.tools.beast_settings.beast_settings instance"
)
# make sure the project directory exists
create_project_dir(settings.project)
# directory for scripts
job_path = "./{0}/model_batch_jobs/".format(settings.project)
if not os.path.isdir(job_path):
os.mkdir(job_path)
log_path = job_path + "logs/"
if not os.path.isdir(log_path):
os.mkdir(log_path)
for i in range(nsubs):
joblist_file = job_path + "create_physicsmodel_" + str(i) + ".job"
with open(joblist_file, "w") as jf:
jf.write(
"python -m beast.tools.run.create_physicsmodel "
+ " {0} ".format(settings.settings_file)
+ " --nsubs "
+ str(nsubs)
+ " --nprocs "
+ str(nprocs)
+ " --subset "
+ str(i)
+ " "
+ str(i + 1)
+ " >> "
+ log_path
+ "create_physicsmodel_"
+ str(i)
+ ".log\n"
)
# slurm needs it to be executable
os.chmod(joblist_file, stat.S_IRWXU | stat.S_IRGRP | stat.S_IROTH)
)
parser.add_argument(
"--dens_bin",
type=int,
default=None,
help="""Run for a certain source/background
density bin. Use the split_catalog_using_map `
tool to generate the right files for this
option.""",
)
args = parser.parse_args()
if args.dens_bin is not None:
bin_subfolder = "bin{}".format(args.dens_bin)
pdir = create_project_dir(bin_subfolder)
def subcatalog_fname(full_cat_fname, dens_bin):
return full_cat_fname.replace(".fits", "_bin{}.fits".format(dens_bin))
# check input parameters, print what is the problem, stop run_beast
verify_params.verify_input_format(datamodel)
def parallel_wrapper(function, argument):
parallel = args.nprocs > 1
if parallel:
p = Pool(args.nprocs)
for r in p.imap_unordered(function, argument, chunksize=1):
print(r)
else:
