Lodge It

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])