import numpy as np from sedpy.observate import load_filters import h5py import prospect.io.read_results as pread from prospect.models import priors, transforms,sedmodel from prospect.sources import FastStepBasis from scipy.stats import truncnorm from prospect.io import write_results as writer from prospect.fitting import fit_model import sys, os #from astropy.cosmology import Planck13 from astropy.cosmology import FlatLambdaCDM from astropy import units as u from astropy import constants import fsps import sedpy #------------------------ # Convienence Functions #------------------------ def get_best(res, **kwargs): imax = np.argmax(res['lnprobability']) theta_best = res['chain'][imax, :].copy() return theta_best def find_nearest(array,value): idx = (np.abs(np.array(array)-value)).argmin() return idx def zfrac_to_masses_log(logmass=None, z_fraction=None, agebins=None, **extras): sfr_fraction = np.zeros(len(z_fraction) + 1) sfr_fraction[0] = 1.0 - z_fraction[0] for i in range(1, len(z_fraction)): sfr_fraction[i] = np.prod(z_fraction[:i]) * (1.0 - z_fraction[i]) sfr_fraction[-1] = 1 - np.sum(sfr_fraction[:-1]) # convert to mass fractions time_per_bin = np.diff(10**agebins, axis=-1)[:, 0] mass_fraction = sfr_fraction * np.array(time_per_bin) mass_fraction /= mass_fraction.sum() if (mass_fraction < 0).any(): idx = mass_fraction < 0 if np.isclose(mass_fraction[idx],0,rtol=1e-8): mass_fraction[idx] = 0.0 else: raise ValueError('The input z_fractions are returning negative masses!') masses = 10**logmass * mass_fraction return masses #---------------------- # SSP function #----------------------- def build_sps(**kwargs): """ This is our stellar population model which generates the spectra for stars of a given age and mass. Because we are using a non parametric SFH model, we do have to use a different SPS model than before """ from prospect.sources import FastStepBasis sps = FastStepBasis(zcontinuous=1) return sps def build_model(z,**kwargs): print('z:',z) cosmo = FlatLambdaCDM(Om0=0.3,Tcmb0 = 2.725,H0 = 67.11) if(z<10**-4): dl = (10*u.Mpc) else: dl = cosmo.luminosity_distance(z).to(u.Mpc) print('building model') model_params = [] #basics model_params.append({'name': "lumdist", "N": 1, "isfree": False,"init": dl.value,"units": "Mpc"}) model_params.append({'name':'zred','N':1,'isfree':False,'init':z}) model_params.append({'name': 'imf_type', 'N': 1,'isfree': False,'init': 1}) model_params.append({'name': 'dust_type', 'N': 1,'isfree': False,'init': 2,'prior': None}) model_params.append({'name': 'dust2', 'N': 1,'isfree': True, 'init': 0.1,'prior': priors.ClippedNormal(mini=0.0, maxi=2.0, mean=0.0, sigma=0.3)}) model_params.append({'name': 'add_dust_emission', 'N': 1,'isfree': False,'init': 1,'prior': None}) model_params.append({'name': 'duste_gamma', 'N': 1,'isfree': True,'init': 0.01,'prior': priors.TopHat(mini=0.0, maxi=1.0)}) model_params.append({'name': 'duste_umin', 'N': 1,'isfree': True,'init': 1.0,'prior': priors.TopHat(mini=0.1, maxi=25.0)}) model_params.append({'name': 'duste_qpah', 'N': 1,'isfree': True,'init': 3.0,'prior': priors.TopHat(mini=0.0, maxi=10.0)}) model_params.append({'name': 'add_agb_dust_model', 'N': 1,'isfree': False,'init': 0}) #M-Z model_params.append({'name': 'logmass', 'N': 1,'isfree': True,'init': 8.0,'prior': priors.Uniform(mini=7., maxi=14.)}) model_params.append({'name': 'logzsol', 'N': 1,'isfree': True,'init': -0.5,'prior': priors.Uniform(mini=-1.5, maxi=0.5)}) model_params.append({'name': "sfh", "N": 1, "isfree": False, "init": 3}) model_params.append({'name': "mass", 'N': 3, 'isfree': False, 'init': 1., 'depends_on':zfrac_to_masses_log}) model_params.append({'name': "agebins", 'N': 1, 'isfree': False,'init': []}) model_params.append({'name': "z_fraction", "N": 2, 'isfree': True, 'init': [0, 0],'prior': priors.Beta(alpha=1.0, beta=1.0, mini=0.0, maxi=1.0)}) #here we set the number and location of the timebins, and edit the other SFH parameters to match in size n = [p['name'] for p in model_params] tuniv = np.round(cosmo.age(z).to('Gyr').value,decimals=1) #Gyr, age at z=2 agelims = np.linspace(8,np.log10(tuniv)+9,9) agelims = [0]+agelims.tolist() print(agelims) nbins = len(agelims)-1 agebins = np.array([agelims[:-1], agelims[1:]]) zinit = np.array([(i-1)/float(i) for i in range(nbins, 1, -1)]) # Set up the prior in `z` variables that corresponds to a dirichlet in sfr # fraction. alpha = np.arange(nbins-1, 0, -1) zprior = priors.Beta(alpha=alpha, beta=np.ones_like(alpha), mini=0.0, maxi=1.0) model_params[n.index('mass')]['N'] = nbins model_params[n.index('agebins')]['N'] = nbins model_params[n.index('agebins')]['init'] = agebins.T model_params[n.index('z_fraction')]['N'] = nbins-1 model_params[n.index('z_fraction')]['init'] = zinit model_params[n.index('z_fraction')]['prior'] = zprior model = sedmodel.SedModel(model_params) return model #------------------ # Build Observations #------------------- def build_obs(spec_file,z,snr,all_filt=True,**kwargs): obs = {} cosmo = FlatLambdaCDM(H0=67.11, Om0=0.3, Tcmb0=2.725) t_H = cosmo.age(z).to('Gyr').value t_0 = cosmo.age(0).to('Gyr').value print() info = np.load('galaxy_sed.npz') # radiative transfer saves wavelength, nuL_nu (also equivalent to lambda L_lambda) # convert to observed quantities wav = info['wav']*u.micron flux = info['lum']*u.erg/u.s wav = wav.to(u.AA) dl = cosmo.luminosity_distance(z) flux = flux/(4.*3.14*dl**2.) nu = constants.c.cgs/(wav.to(u.cm)) nu = nu.to(u.Hz) fnu = flux*(1+z)/nu fnu = fnu.to(u.Jy) maggies = fnu/3631. flux_aa = flux/wav/(1+z) flux_aa = flux_aa.to(u.erg/u.s/(u.cm**2)/u.AA) other_filts = ['galex_FUV', 'galex_NUV', 'wfc3_uvis_f275w', 'wfc3_uvis_f336w', 'sdss_u0', 'sdss_g0', 'wfc3_uvis_f475w', 'wfc3_uvis_f555w', 'wfc3_uvis_f606w', 'sdss_r0', 'jwst_f070w', 'sdss_i0', 'wfc3_uvis_f814w', 'sdss_z0', 'jwst_f090w', 'wfc3_ir_f105w', 'wfc3_ir_f110w', 'jwst_f115w', 'wfc3_ir_f125w', 'wfc3_ir_f140w', 'jwst_f150w', 'wfc3_ir_f160w', 'jwst_f200w', 'jwst_f277w', 'wise_w1', 'jwst_f356w', 'jwst_f444w', 'wise_w2', 'jwst_f560w', 'jwst_f770w', 'jwst_f1000w', 'wise_w3', 'jwst_f1280w', 'jwst_f1500w'] reddest_filts = ['jwst_f1800w', 'jwst_f2100w', 'wise_w4', 'spitzer_mips_24', 'herschel_pacs_70', 'herschel_pacs_100', 'herschel_pacs_160', 'herschel_spire_250', 'herschel_spire_350', 'herschel_spire_500'] if(all_filt): all_filters = other_filts+reddest_filts else: all_filters = reddest_filts filters = load_filters(all_filters) # redshift wavelengths wav = wav*(1.+z) gal_phot = sedpy.observate.getSED(wav.value,flux_aa.value,filterlist = filters,linear_flux = True) flux_mag = np.asarray(gal_phot) unc_mag = gal_phot/snr obs['filters'] = filters obs['maggies'] = flux_mag obs['maggies_unc'] = unc_mag obs['phot_mask'] = np.isfinite(flux_mag) obs['wavelength'] = None obs['spectrum'] = None obs['rest_sed'] = maggies obs['wav'] = wav obs['z_val'] = z return obs def build_all(phot_file,snr,z,all_filt,**kwargs): return (build_obs(phot_file,z,snr,all_filt,**kwargs), build_model(z,**kwargs), build_sps(**kwargs)) run_params = {'verbose':False, 'debug':False, 'output_pickles': True, 'nested_bound': 'multi', # bounding method 'nested_sample': 'auto', # sampling method 'nested_nlive_init': 400, 'nested_nlive_batch': 200, 'nested_bootstrap': 0, 'nested_dlogz_init': 0.05, 'nested_weight_kwargs': {"pfrac": 1.0}, } if __name__ == '__main__': phot_file = f'galaxy_sed.npz' z = 2 SNR = 10 all_filt = True obs, model, sps = build_all(phot_file,SNR,z,all_filt,**run_params) outfile = 'test_prosp_rt_all_filt.h5' run_params["sps_libraries"] = sps.ssp.libraries run_params["param_file"] = __file__ print('Running fits') output = fit_model(obs, model, sps, [None,None],**run_params) print('Done. Writing now') print(model) print(obs) writer.write_hdf5(outfile, run_params, model, obs, output["sampling"][0], output["optimization"][0], tsample=output["sampling"][1], toptimize=output["optimization"][1])