import sv
from provabgs import util as UT
from provabgs import infer as Infer
from provabgs import models as Models
from provabgs import flux_calib as FluxCalib
i0 = int(sys.argv[1])
i1 = int(sys.argv[2])
sample = sys.argv[3]
niter = int(sys.argv[4])
n_cpu = int(sys.argv[5])
# declare SPS model
m_nmf = Models.NMF(burst=True, emulator=True)
# declare flux calibration
m_fluxcalib = FluxCalib.constant_flux_factor
def run_mcmc(igal):
# get observations
zred_i, photo_flux_i, photo_ivar_i, w_obs, f_obs, i_obs, f_fiber, sigma_f_fiber\
= sv.get_spectrophotometry(igal, sample=sample)
# set prior
prior = Infer.load_priors([
Infer.UniformPrior(7., 12.5, label='sed'),
Infer.FlatDirichletPrior(4, label='sed'), # flat dirichilet priors
Infer.UniformPrior(0., 1., label='sed'), # burst fraction
Infer.UniformPrior(1e-2, 13.27, label='sed'), # tburst
def fsps_prior_samples(name, ibatch, ncpu=1):
''' run FSPS SPS model and get composite stellar population luminosity.
'''
if ibatch == 'test':
np.random.seed(123456)
nspec = 100000 # batch size
else:
np.random.seed(ibatch)
nspec = 10000 # batch size
# load priors
priors = prior_nmf(4)
dat_dir = '/global/cscratch1/sd/chahah/provabgs/emulator' # hardcoded to NERSC directory
if ibatch == 'test':
ftheta = os.path.join(dat_dir, 'fsps.%s.theta.test.npy' % name)
fspectrum = os.path.join(dat_dir, 'fsps.%s.lnspectrum.test.npy' % name)
else:
ftheta = os.path.join(dat_dir,
'fsps.%s.theta.seed%i.npy' % (name, ibatch))
fspectrum = os.path.join(
dat_dir, 'fsps.%s.lnspectrum.seed%i.npy' % (name, ibatch))
if os.path.isfile(ftheta) and os.path.isfile(fspectrum):
print()
print('--- batch %s already exists ---' % str(ibatch))
print('--- do not overwrite ---')
print()
return None
# load SPS model
Msps = Models.NMF(burst=False, emulator=False)
# sample prior and transform
_thetas = np.array(
[priors.transform(priors.sample()) for i in range(nspec)])
zred = _thetas[:, -1]
tage = Msps.cosmo.age(zred).value # age in Gyr
thetas = np.zeros((nspec, 11))
thetas[:, 1:-1] = _thetas[:, :-1]
thetas[:, -1] = tage
w_fsps, _ = Msps._fsps(thetas[0, :-1], thetas[0, -1])
# wavelength range set to cover the DESI spectra and photometric filter (up
# to W2
wmin, wmax = 2300., 60000.
wlim = (w_fsps >= wmin) & (w_fsps <= wmax)
fwave = os.path.join(dat_dir, 'wave_fsps.npy')
if not os.path.isfile(fwave): # save FSPS wavelength if not saved
np.save(fwave, w_fsps[wlim])
print()
print('--- batch %s ---' % str(ibatch))
# save parameters sampled from prior
print(' saving thetas to %s' % ftheta)
np.save(ftheta, _thetas)
if (ncpu == 1): # run on serial
logspectra = []
for _theta in thetas:
_, _spectrum = Msps._fsps(_theta[:-1], _theta[-1])
logspectra.append(np.log(_spectrum[wlim]))
else:
def _fsps_model_wrapper(theta):
_, _spectrum = Msps._fsps(theta[:-1], theta[-1])
return np.log(_spectrum[wlim])
pewl = mp.Pool(ncpu)
logspectra = pewl.map(_fsps_model_wrapper, thetas)
print(' saving ln(spectra) to %s' % fspectrum)
np.save(fspectrum, np.array(logspectra))
print()
return None
if ibatch == 'test':
np.random.seed(123456)
nspec = 100000 # batch size
ftheta = os.path.join(dat_dir, 'fsps.%s.v%s.theta.test.npy' % (name, version))
ftheta_unt = os.path.join(dat_dir, 'fsps.%s.v%s.theta_unt.test.npy' % (name, version))
fspectrum = os.path.join(dat_dir, 'fsps.%s.v%s.lnspectrum.test.npy' % (name, version))
else:
np.random.seed(ibatch)
nspec = 10000 # batch size
ftheta = os.path.join(dat_dir, 'fsps.%s.v%s.theta.seed%i.npy' % (name, version, ibatch))
ftheta_unt = os.path.join(dat_dir, 'fsps.%s.v%s.theta_unt.seed%i.npy' % (name, version, ibatch))
fspectrum = os.path.join(dat_dir, 'fsps.%s.v%s.lnspectrum.seed%i.npy' % (name, version, ibatch))
# load SPS model
Msps = Models.NMF(burst=False, emulator=False)
Msps._ssp_initiate()
# sample prior and transform
_thetas = np.array([priors.transform(priors.sample()) for i in range(nspec)])
zred = _thetas[:,-1]
tage = Msps.cosmo.age(zred).value # age in Gyr
thetas = np.zeros((nspec, 11))
thetas[:,1:-1] = _thetas[:,:-1]
thetas[:,-1] = tage
_thetas_unt = priors.untransform(_thetas)
thetas_unt = np.zeros((nspec, 10))
thetas_unt[:,0] = thetas[:,0]
thetas_unt[:,1:-1] = _thetas_unt[:,:-1]
thetas_unt[:,-1] = tage
def fsps_burst_prior_samples(ibatch, ncpu=1):
''' run FSPS SPS model and get stellar population luminosity for a single
burst. This is to build a separate emulator for the burst component
'''
if ibatch == 'test':
np.random.seed(123456)
nspec = 100000 # batch size
else:
np.random.seed(ibatch)
nspec = 10000 # batch size
# load priors
name = 'burst'
priors = prior_burst()
#if os.environ['machine'] in ['della', 'tiger']:
# dat_dir='/tigress/chhahn/provabgs/'
#else:
dat_dir = '/global/cscratch1/sd/chahah/provabgs/emulator' # hardcoded to NERSC directory
if ibatch == 'test':
ftheta = os.path.join(dat_dir, 'fsps.%s.theta.test.npy' % name)
fspectrum = os.path.join(dat_dir, 'fsps.%s.lnspectrum.test.npy' % name)
else:
ftheta = os.path.join(dat_dir,
'fsps.%s.theta.seed%i.npy' % (name, ibatch))
fspectrum = os.path.join(
dat_dir, 'fsps.%s.lnspectrum.seed%i.npy' % (name, ibatch))
#if os.path.isfile(ftheta) and os.path.isfile(fspectrum):
# print()
# print('--- batch %s already exists ---' % str(ibatch))
# print('--- do not overwrite ---')
# print()
# return None
# sample prior for burst
thetas = np.array([priors.sample() for i in range(nspec)])
# load SPS model
Msps = Models.NMF(burst=True, emulator=False)
Msps._ssp_initiate()
def lssp_burst(_theta):
tburst, zburst, dust1, dust2, dust_index = _theta
# luminosity of SSP at tburst
Msps._ssp.params['logzsol'] = np.log10(zburst / 0.0190) # log(Z/Zsun)
Msps._ssp.params['dust1'] = dust1
Msps._ssp.params['dust2'] = dust2
Msps._ssp.params['dust_index'] = dust_index
return Msps._ssp.get_spectrum(tage=np.clip(tburst, 1e-8, None),
peraa=True) # in units of Lsun/AA
w_fsps, _ = lssp_burst(thetas[0])
# wavelength range set to cover the DESI spectra and photometric filter
# wavelength range
wmin, wmax = 2300., 60000.
wlim = (w_fsps >= wmin) & (w_fsps <= wmax)
fwave = os.path.join(dat_dir, 'wave_fsps.npy')
if not os.path.isfile(fwave): # save FSPS wavelength if not saved
np.save(fwave, w_fsps[wlim])
print()
print('--- batch %s ---' % str(ibatch))
# save parameters sampled from prior
print(' saving thetas to %s' % ftheta)
np.save(ftheta, thetas)
if (ncpu == 1): # run on serial
logspectra = []
for _theta in thetas:
_, _spectrum = lssp_burst(_theta)
logspectra.append(np.log(_spectrum[wlim]))
else:
def _fsps_model_wrapper(theta):
_, _spectrum = lssp_burst(theta)
return np.log(_spectrum[wlim])
pewl = mp.Pool(ncpu)
logspectra = pewl.map(_fsps_model_wrapper, thetas)
print(' saving ln(spectra) to %s' % fspectrum)
np.save(fspectrum, np.array(logspectra))
print()
return None