<jupyter_start><jupyter_text>JWST-ERS Clusters Spectral Population Synthesis FittingThe goal of this notebook is to fit some simple JWST prism spectra with [prospector](https://github.com/bd-j/prospector).Note that we use the mean fluxing vector derived in the [resolution-throughput-and-noise](https://github.com/moustakas/jwstclusters/tree/master/nb/resolution-throughput-and-noise.ipynb) Jupyter notebook.<jupyter_code>import os
import time
import pickle

import numpy as np
import matplotlib.pyplot as plt

import fitsio
import nestle
import emcee
import h5py
from astropy.table import Table, Column
from prospect.utils.smoothing import smoothspec
from prospect import fitting
from prospect.io import write_results
from prospect.sources import CSPSpecBasis
import multiprocessing
#from concurrent.futures import ProcessPoolExecutor

nproc = multiprocessing.cpu_count() // 2 
#pool = ProcessPoolExecutor(nproc)
plt.style.use('seaborn-talk')
%matplotlib inline
maggies2mJy = 10**(0.4 * 16.4)  # maggies --> mJy
maggies2muJy = 10**(0.4 * 23.9) # maggies --> microJy<jupyter_output><empty_output><jupyter_text>Set the prefix, random seed, and top-level path.<jupyter_code>prefix = 'prism'
seed = 123
rand = np.random.RandomState(seed)

jwstdir = os.path.join( os.getenv('IM_PROJECTS_DIR'), 'jwstclusters')<jupyter_output><empty_output><jupyter_text>Define the key prospector functions.<jupyter_code>def load_obs(zred=1.0, wave=None, flux=None, ferr=None, mask=None):
    """Generate the prospector-style "obs" dictionary which contains the input
    photometry, redshift, etc. for a single object.

    Args:
        zred (float): Galaxy redshift.
        wave (np.array, npix): Wavelength array (micron).
        flux (np.array, npix): Galaxy spectrum (units??).
        ferr (np.array, npix): Uncertainty spectrum corresponding to flux (units??).
        mask (np.array, npix): Mask spectrum (1=good, 0=bad).
        
    Returns:
        obs (dict): Dictionary used by prospector.

    """
    obs = {} 

    # Input photometry
    obs['filters'] = None
    obs['maggies'] = None
    obs['maggies_unc'] = None

    # Input spectroscopy
    obs['wavelength'] = wave
    obs['spectrum'] = flux
    obs['unc'] = ferr
    obs['mask'] = mask

    # Store the redshift and any other galaxy properties
    obs['zred'] = zred

    return obs
def logmass2mass(logmass=11.0, **extras):
    return 10**logmass
def logtau2tau(logtau=-1, **extras):
    return 10**logtau
def load_model(zred=1.0, seed=None):
    """Initialize the priors on each free and fixed parameter.

    Args:
      zred (float): input (fixed) galaxy redshift.

    Returns:
      sed (prospect.models.sedmodel.SedModel): SED priors and other stuff.

    Notes:
      FSPS parameters are documented here:
        http://dan.iel.fm/python-fsps/current/stellarpop_api/#api-reference

      Initialization parameters:
        * compute_vega_mags (must be set at initialization)
        * vactoair_flag (must be set at initialization)
        * zcontinuous (must be set at initialization)
    
      Metallicity parameters:
        * zmet (default 1, ignored if zcontinuous>0)
        * logzsol (default 0.0, used if zcontinuous>0)
        * pmetals (default 2.0, only used if zcontinuous=2)

      Dust parameters:
        * add_agb_dust_model (default True)
        * add_dust_emission (default True)
        * cloudy_dust (default False)
        * agb_dust (default 1.0)
        * dust_type (default 0=power law)
        * dust_index, dust1_index
        * dust_tesc
        * dust1 (default 0.0) - extra optical depth towards young stars at 5500A
        * dust2 (default 0.0) - diffuse dust optical depth towards all stars at 5500A
        * dust_clumps, frac_nodust, frac_obrun
        * mwr, uvb, wgp1, wgp2, wgp3, 
        * duste_gamma, duste_umin, duste_qpah

      Star formation history parameters:
        * sfh (default 0=SSP, 1=tau, 4=delayed, 5=truncated delayed tau)
        * tau (default 1)
        * const, sf_start, sf_trunc
        * tage (default 0.0)
        * fburst, tburst, sf_slope
    
      Miscellaneous parameters:
        * add_igm_absorption (default False)
        * igm_factor (default 1.0)
        * smooth_velocity (default True)
        * sigma_smooth, min_wave_smooth, max_wave_smooth
        * redshift_colors (default False, do not use)
        * compute_light_ages (default False, do not use)
       
      Stellar population parameters:
        * add_stellar_remnants (default True)
        * tpagb_norm_type (default 2)
        * dell (default 0.0, do not use)
        * delt (default 0.0, do not use)
        * redgb (default 1.0)
        * fcstar (default 1.0)
        * sbss (default 0.0)
        * fbhb (default 0.0)
        * pagb (default 1.0)
        * imf_type (default 2=Kroupa01)
        * imf1, imf2, imf3, vdmc, mdave, masscut
        * evtype (default 1)
        * tpagb_norm_type

      Emission lines:
        * add_neb_emission (default False)
        * add_neb_continuum (default False)
        * gas_logz (default 0.0)
        * gas_logu (default -2)

      Galaxy properties:
        * zred (default 0.0)

      AGN properties:
        * fagn (default 0.0)
        * agn_tau (default 10)

      Calibration parameters:
        * phot_jitter

    """
    from prospect.models import priors, sedmodel

    model_params = []

    ##################################################
    # Fixed priors

    # Galaxy redshift
    model_params.append({
        'name': 'zred',
        'N': 1,
        'isfree': False,
        'init': zred,
        'units': '',
        'prior': None,       
        })

    model_params.append({ # current mass in stars, not integral of SFH
        'name': 'mass_units',
        'N': 1,
        'isfree': False,
        'init': 'mstar', # 'mformed'
        'prior': None,       
        })

    # IMF (Chabrier)
    model_params.append({
        'name': 'imf_type',
        'N': 1,
        'isfree': False,
        'init':   0, # 0 - Salpeter, 1 - Chabrier
        'units': '',
        'prior': None,       
        })

    # SFH parameterization (delayed-tau)
    model_params.append({
        'name': 'sfh',
        'N': 1,
        'isfree': False,
        'init':   4, # 4 = delayed tau model
        'units': 'type',
        'prior': None,       
        })

    # Do not include dust emission
    model_params.append({
        'name': 'add_dust_emission',
        'N': 1,
        'isfree': False,
        'init':   False, # do not include dust emission
        'units': 'index',
        'prior': None,       
        })

    ##################################################
    # Free priors / parameters

    print('HARD-CODING THE PRIORS!!')
    print('MAKING THE PRIORS VERY NARROW!!')
    
    # Priors on stellar mass and stellar metallicity
    logmass_prior = priors.TopHat(mini=10.9, maxi=11.1)#, seed=seed)
    #logmass_prior = priors.TopHat(mini=10.5, maxi=11.5)#, seed=seed)
    logmass_init = np.diff(logmass_prior.range)/2.0 + logmass_prior.range[0] # logmass_prior.sample()
    model_params.append({
        'name': 'logmass',
        'N': 1,
        'isfree': True,
        'init': 11.0, # logmass_init, # mass, 
        'init_disp': 0.1,     # dex
        'units': r'$M_{\odot}$',
        'prior': logmass_prior,
        })
    
    model_params.append({
        'name': 'mass',
        'N': 1,
        'isfree': False,
        'init': 10**logmass_init,
        'units': '',
        'prior': None,
        'depends_on': logmass2mass,
        })

    logzsol_prior = priors.TopHat(mini=-0.1, maxi=0.1)
    #logzsol_prior = priors.TopHat(mini=np.log10(0.005/0.019), maxi=np.log10(0.04/0.019))#, seed=seed)
    logzsol_init = np.diff(logzsol_prior.range)/2.0 + logzsol_prior.range[0] # logzsol_prior.sample(), # logzsol,
    model_params.append({
        'name': 'logzsol',
        'N': 1,
        'isfree': False,
        'init': 0.0, # logzsol_init,
        'init_disp': 0.1, # logzsol_prior.range[1] * 0.1,
        'units': r'$\log_{10}\, (Z/Z_\odot)$',
        'prior': logzsol_prior, # roughly (0.2-2)*Z_sun
        })

    # Prior(s) on dust content
    #dust2_prior = priors.TopHat(mini=0.0, maxi=3.0)#, seed=seed)
    #dust2_init = np.diff(dust2_prior.range)/2.0 + dust2_prior.range[0] # dust2_prior.sample(), # dust2,
    #model_params.append({
    #    'name': 'dust2',
    #    'N': 1,
    #    'isfree': True,
    #    'init': dust2_init,
    #    'init_disp': 0.5, # dust2_prior.range[1] * 0.1,
    #    'units': '', # optical depth
    #    'prior': dust2_prior,
    #    })
    
    # Prior on tau
    tau_prior = priors.TopHat(mini=0.0, maxi=1.0)#, seed=seed)
    tau_init = np.diff(tau_prior.range)/2.0 + tau_prior.range[0] # tau_prior.sample(), # tau,
    model_params.append({
        'name': 'tau',
        'N': 1,
        'isfree': False,
        'init': 0.01, # 1.0, #tau_init,
        'init_disp': 0.5, # tau_prior.range[1] * 0.1,
        'units': 'Gyr', 
        'prior': tau_prior,
        })
    
    # Prior on (log)tau
    #logtau_prior = priors.TopHat(mini=-3.0, maxi=1.0)#, seed=seed)
    #logtau_init = np.diff(logtau_prior.range)/2.0 + logtau_prior.range[0] # logtau_prior.sample()
    #model_params.append({
    #    'name': 'logtau',
    #    'N': 1,
    #    'isfree': True,
    #    'init': -1, # logtau_init, # tau, 
    #    'init_disp': 0.1, # dex
    #    'units': 'Gyr',
    #    'prior': logtau_prior,
    #    })
    
    #model_params.append({
    #    'name': 'tau',
    #    'N': 1,
    #    'isfree': False,
    #    'init': 10**logtau_init,
    #    'units': '',
    #    'prior': None,
    #    'depends_on': logtau2tau,
    #    })
    
    # Prior on age
    tage_prior = priors.TopHat(mini=2.5, maxi=3.5)#, seed=seed)
    #tage_prior = priors.TopHat(mini=1.0, maxi=5.0)#, seed=seed)
    tage_init = np.diff(tage_prior.range) / 2.0 + tage_prior.range[0] # tage_prior.sample(), # tage,
    model_params.append( {
        'name': 'tage',
        'N': 1,
        'isfree': True,
        'init': 3.0, # tage_init,
        'init_disp': 0.5, # tage_prior.range[1] * 0.1,
        'units': 'Gyr',
        'prior': tage_prior,
        })

    model = sedmodel.SedModel(model_params)
    
    return model
def lnprobfn(theta, spec_noise=None, phot_noise=None, verbose=False):
    """Define the likelihood function.

    Given a parameter vector and a dictionary of observational data and a model
    object, return the ln of the posterior. This requires that an sps object
    (and if using spectra and gaussian processes, a GP object) be instantiated.

    """
    from prospect.likelihood import lnlike_spec, lnlike_phot, chi_spec, chi_phot

    if doemcee:
        lnp_prior = model.prior_product(theta)
    else:
        lnp_prior = model.prior_product(theta, nested=True)
        
    if np.isfinite(lnp_prior):
        # Generate the mean model--
        t1 = time.time()
        try:
            model_spec, model_phot, model_extras = model.mean_model(
                theta, obs, sps=sps, resolution=run_params['resolution'],
                smoothtype=run_params['smoothtype'])
        except(ValueError):
            return -np.infty
        d1 = time.time() - t1

        # Noise modeling--
        if spec_noise:
            spec_noise.update(**model.params)
        if phot_noise:
            phot_noise.update(**model.params)

        vectors = {
            'spec': model_spec,    # model spectrum
            'phot': model_phot,    # model photometry
            'sed': model._spec,    # object spectrum
            'cal': model._speccal, # object calibration spectrum
        }

        # Calculate likelihoods--
        t2 = time.time()
        lnp_spec = lnlike_spec(model_spec, obs=obs, spec_noise=spec_noise, **vectors)
        lnp_phot = lnlike_phot(model_phot, obs=obs, phot_noise=phot_noise, **vectors)
        d2 = time.time() - t2
        if False:
            from prospect.likelihood import write_log
            write_log(theta, lnp_prior, lnp_spec, lnp_phot, d1, d2)

        if False:
            print(model.theta, lnp_spec)
            plt.errorbar(obs['wavelength'], obs['spectrum'], obs['unc'])
            plt.plot(obs['wavelength'], model_spec)
            plt.show()
            import pdb ; pdb.set_trace()
            
        return lnp_prior + lnp_phot + lnp_spec
    else:
        return -np.infty
def load_run_params(prefix, seed=None, resolution=None):
    """Populate the run_params dictionary."""
    run_params = {
        'prefix':  prefix,
        'verbose': True,
        'seed':    seed,
        # Nestle fitting parameters
        'nestle_method': 'multi',
        'nestle_npoints': 100,
        'nestle_maxcall': 100, #int(1e6),
        # emcee fitting parameters
        'nwalkers': 32, # 128
        'nburn': [16, 16], #[32, 32, 64], 
        'niter': 128, # 256, # 512,
        'nthreads': nproc,
        # Smoothing model
        'smoothtype': 'lsf',
        'resolution': resolution,
        # SPS initialization parameters
        'compute_vega_mags': False,
        'vactoair_flag': False, # use wavelengths in air
        'zcontinuous': 1,      # interpolate in metallicity
    }
    return run_params<jupyter_output><empty_output><jupyter_text>Read the instrumental resolution curve and the mean fluxing vector.<jupyter_code>def read_resolution():
    """Read the JWST/prism instrumental resolution function.
    
    """
    pandeia_dir = os.path.join( os.getenv('IM_DATA_DIR'), 'pandeia_data-1.0' )
    
    dispfile = os.path.join(pandeia_dir, 'jwst', 'nirspec', 'dispersion', 
                            'jwst_nirspec_prism_disp.fits')
    print('Reading {}'.format(dispfile))
    disp = Table(fitsio.read(dispfile, ext=1, upper=True))
    
    return disp
def read_fluxvector():
    """Read the mean fluxing vector.
    
    """
    fluxvecfile = os.path.join(jwstdir, 'etc', 'flux_vector.fits')
    print('Reading {}'.format(fluxvecfile))
    fluxvec = Table(fitsio.read(fluxvecfile, ext=1, upper=True))
    
    return fluxvec
disp = read_resolution()
fluxvec = read_fluxvector()<jupyter_output>Reading /Users/ioannis/research/projects/jwstclusters/etc/flux_vector.fits<jupyter_text>Read the ETC object and uncertainty spectra and "flux-calibrate" them.<jupyter_code>cat = Table.read(os.path.join(jwstdir, 'etc', 'input', 'modelspec.fits'))
nspec = len(cat)
npix = 403
print('Hard-coding the number of pixels to {}!'.format(npix))
cat.add_column(Column(name='SUFFIX', dtype='S2', length=nspec))
cat.add_column(Column(name='JWST_WAVELENGTH', shape=(npix, ), dtype='f4', length=nspec))
cat.add_column(Column(name='JWST_FLUX', shape=(npix, ), dtype='f4', length=nspec))
cat.add_column(Column(name='JWST_FERR', shape=(npix, ), dtype='f4', length=nspec))
cat.add_column(Column(name='RESOLUTION', shape=(npix, ), dtype='f4', length=nspec))
for ii in range(nspec):
    suffix = '{0:02d}'.format(ii)
    cat['SUFFIX'][ii] = suffix
    
    snrspecfile = os.path.join(jwstdir, 'etc', 'output', suffix, 
                               'lineplot', 'lineplot_sn.fits')
    simspecfile = os.path.join(jwstdir, 'etc', 'output', suffix, 
                               'lineplot', 'lineplot_extracted_flux.fits')
    simspecerrfile = os.path.join(jwstdir, 'etc', 'output', suffix, 
                                  'lineplot', 'lineplot_extracted_noise.fits')

    snrspec = Table(fitsio.read(snrspecfile, ext=1, upper=True))
    simspec = Table(fitsio.read(simspecfile, ext=1, upper=True))
    simspecerr = Table(fitsio.read(simspecerrfile, ext=1, upper=True))

    obswave = simspec['WAVELENGTH'] # micron
    
    fluxvector = np.interp(obswave, fluxvec['WAVELENGTH'], fluxvec['FLUXVECTOR'])
    simferr = simspecerr['EXTRACTED_NOISE'] / fluxvector
    simflux = simspec['EXTRACTED_FLUX'] / fluxvector + rand.normal(scale=simferr)
        
    modelfile = os.path.join(jwstdir, 'etc', 'input', 'modelspec_{0:02d}.txt'.format(ii))
    modelwave, modelflux = np.loadtxt(modelfile, unpack=True) # [micron and mJy]
    
    resolution = np.interp(modelwave, disp['WAVELENGTH'], disp['DLDS'])    
    smoothflux = smoothspec(modelwave, modelflux, resolution=resolution, 
                            smoothtype='lsf', outwave=obswave, fftsmooth=True)
    
    cat['JWST_WAVELENGTH'][ii] = obswave * 1e4 # [Angstrom]
    cat['JWST_FLUX'][ii] = simflux / maggies2mJy # [maggies]
    cat['JWST_FERR'][ii] = simferr / maggies2mJy # [maggies]
    cat['RESOLUTION'][ii] = np.interp(obswave * 1e4, disp['WAVELENGTH'] * 1e4, disp['DLDS'] * 1e4)

    fig, ax = plt.subplots()
    ax.errorbar(obswave, simflux, simferr, alpha=0.5, label='Simulated Spectrum')
    ax.plot(modelwave, modelflux, label='Input Model Spectrum')
    #ax.plot(obswave, smoothflux)
    ax.set_xlim(0.6, 5.3)
    ax.set_xlabel('Observed-frame Wavelength ($\mu$m)')
    ax.set_ylabel('Flux (mJy)')
    ax.legend(loc='lower right')
catfile = os.path.join( jwstdir, 'spsmodeling', '{}_catalog.fits'.format(prefix) )
print('Writing {}'.format(catfile))
cat.write(catfile, overwrite=True)
cat<jupyter_output>Writing /Users/ioannis/research/projects/jwstclusters/spsmodeling/prism_catalog.fits<jupyter_text>Do the fitting and write out.We write out the nestle results to an HDF5 file and the model as a pickle file.<jupyter_code>def load_sps(resolution=None, resolution_wave=None, zred=None):
    """Read the SSPs into a CSPSpecBasis object.
    
    """
    t0 = time.time()
    print('Initializing CSPSpecBasis...')
    sps = CSPSpecBasis(zcontinuous=1, compute_vega_mags=False, vactoair_flag=False)
    print('...took {:.1f} seconds.'.format(time.time() - t0)) 
    
    return sps
doemcee = True
for ii in range(0, 1, 1):
#for ii in range(nspec):
    run_params = load_run_params(prefix, seed=seed, resolution=cat['RESOLUTION'][ii])
    
    obs = load_obs(zred=cat['ZRED'][ii], wave=cat['JWST_WAVELENGTH'][ii], 
                   flux=cat['JWST_FLUX'][ii], ferr=cat['JWST_FERR'][ii])
    model = load_model(zred=obs['zred'], seed=run_params['seed'])    
    sps = load_sps(cat['RESOLUTION'][ii], cat['JWST_WAVELENGTH'][ii], cat['ZRED'][ii])
    
    hfilename = os.path.join( jwstdir, 'spsmodeling', '{}_prospect_{}_mcmc.h5'.format(
        run_params['prefix'], cat['SUFFIX'][ii]) )
    outroot = hfilename.replace('.h5', '')
    
    print('Free parameters: {}'.format(model.free_params))
    print('Initial parameter values: {}'.format(model.initial_theta))
    initial_theta = model.rectify_theta(model.initial_theta) # make zeros tiny numbers
    initial_prob = lnprobfn(initial_theta)
    
    if doemcee:
        print('Writing to HDF5 output file {}'.format(hfilename))
        hfile = h5py.File(hfilename, 'w')
        write_results.write_h5_header(hfile, run_params, model)
        write_results.write_obs_to_h5(hfile, obs)
        
        print('Starting emcee at {}'.format(time.asctime()))
        tstart = time.time()
        esampler, burn_p0, burn_prob0 = fitting.run_emcee_sampler(lnprobfn, initial_theta, model,
                                                                  nthreads=run_params['nthreads'],
                                                                  nwalkers=run_params['nwalkers'],
                                                                  nburn=run_params['nburn'],
                                                                  niter=run_params['niter'], 
                                                                  prob0=initial_prob, hdf5=hfile)
        dur = time.time() - tstart
        print('emcee took {:.2f} minutes'.format(dur / 60.0))

        write_results.write_pickles(run_params, model, obs, esampler, None,
                                    outroot=outroot, toptimize=0.0, tsample=dur,
                                    sampling_initial_center=initial_theta,
                                    post_burnin_center=burn_p0,
                                    post_burnin_prob=burn_prob0)
        write_results.write_hdf5(hfilename, run_params, model, obs, esampler, 
                                 None, toptimize=0.0, tsample=dur,
                                 sampling_initial_center=initial_theta,
                                 post_burnin_center=burn_p0,
                                 post_burnin_prob=burn_prob0)
        hfile.close()    

    else:
        print('Starting nestle at {}'.format(time.asctime()))
        tstart = time.time()
        nestleout = nestle.sample(lnprobfn, model.prior_transform, model.ndim, #rstate=rand,
                                  callback=None, pool=None, 
                                  method=run_params['nestle_method'],
                                  npoints=run_params['nestle_npoints'],
                                  maxcall=run_params['nestle_maxcall'])
        dur = time.time() - tstart
        print('Nestle took {:.2f} minutes'.format(dur / 60.0))
    
        print('Writing HDF5 output file {}'.format(hfilename))
        hfile = h5py.File(hfilename, 'w')
        write_results.write_h5_header(hfile, run_params, model)
        write_results.write_obs_to_h5(hfile, obs)
        write_results.write_hdf5(hfilename, run_params, model, obs, 
                                 nestleout, None, tsample=dur)
        hfile.close()    

        # Write the model object as a pickle.
        modelfile = os.path.join( jwstdir, 'spsmodeling', '{}_prospect_{}_model'.format(
            run_params['prefix'], cat['SUFFIX'][ii]) )
        print('Writing {}'.format(modelfile))
        partext = write_results.paramfile_string(**run_params)
        write_results.write_model_pickle(modelfile, model, paramfile_text=partext)
    
        # Also write the nestle outputs as a pickle.
        nestlefile = os.path.join( jwstdir, 'spsmodeling', '{}_prospect_{}_nmc.pkl'.format(
            run_params['prefix'], cat['SUFFIX'][ii]) )    
        print('Writing {}'.format(nestlefile))
        with open(nestlefile, 'wb') as f:
            pickle.dump(nestleout, f)
hfile.close()<jupyter_output><empty_output><jupyter_text>Read the results and generate QAplots.<jupyter_code>def _niceparnames(parnames):
    """Replace parameter names with nice names."""

    old = list([
        'tau',
        'logtau',
        'tage',
        'mass',
        'logmass',
        'logzsol',
        'dust2'])
    new = list([
        r'$\tau$ (Gyr)',
        r'$\log_{10}\, \tau$ (Gyr)',
        'Age (Gyr)',
        r'$M / M_{\odot}$',
        r'$\log_{10}\,(M / M_{\odot})$',
        r'$\log_{10}\, (Z / Z_{\odot})$',
        r'$\tau_{diffuse}$'])

    niceparnames = list(parnames).copy()
    for oo, nn in zip( old, new ):
        this = np.where(np.in1d(parnames, oo))[0]
        if len(this) > 0:
            niceparnames[this[0]] = nn
            
    return np.array(niceparnames)
def _galaxysed(obs):
    """Generate the galaxy SED, converting from Angstroms to microns 
    and from maggies to mJy.

    """
    galaxywave = obs['wavelength'] * 1e-4 # Angstrom --> micron
    galaxyspec = obs['spectrum'] * maggies2muJy
    galaxyuncspec = obs['unc'] * maggies2muJy
    
    return galaxywave, galaxyspec, galaxyuncspec
def _modelsed(model, theta, obs, sps):
    """Construct a model SED for a given set of parameters, converting
    from Angstroms to microns and from maggies to mJy.

    """
    modelwave = sps.csp.wavelengths * (1 + obs['zred']) # [observed-frame wavelengths]
    modelwave *= 1e-4 # Angstrom --> micron
    
    modelspec, modelphot, mextra = model.mean_model(theta, obs, sps=sps)
    modelspec *= maggies2muJy
    modelphot *= maggies2muJy
    
    return modelwave, modelspec, modelphot
def bestfit_sed(results=None, theta=None, obs=None, sps=None, 
                model=None, seed=None, nrand=20, start=50):
    """Plot the data and the best-fitting SED model.

    Either pass results (which contains chain, lnprobability, and weights) to 
    visualize samples from the posterior *or* theta and *obs* to visualize a 
    single SED fit.

    """
    rand = np.random.RandomState(seed)

    if results is not None:
        obs = results['obs']
    
    # Build the galaxy spectrum.
    galaxywave, galaxyspec, galaxyuncspec = _galaxysed(obs)
        
    # Build the maximum a posteriori model fit and also grab a random sampling of
    # the chains with weight equal to the posterior probability.    
    if results is not None:
        # Trim the first ~50 samples.
        chain = results['chain'][start:, :]
        lnp = results['lnprobability'][start:]
        weights = results['weights'][start:]

        niter, nparams = chain.shape

        theta = chain[lnp.argmax(), :] # maximum a posteriori values
        print(theta)

        #prob = np.exp(lnp - lnp.max())
        #prob /= prob.sum()
        rand_indx = rand.choice(niter, size=nrand, replace=False)#p=weights / np.sum(weights))
        theta_rand = chain[rand_indx, :]
        
    modelwave, modelspec, _ = _modelsed(model=model, theta=theta, obs=obs, sps=sps)

    # Establish the wavelength and flux limits.
    #minwave, maxwave = 0.1, 6.0
    minwave, maxwave = np.min(galaxywave), np.max(galaxywave)

    inrange = (modelwave > minwave) * (modelwave < maxwave)
    #print(modelwave.min(), modelwave.max(), minwave, maxwave)
    #maxflux = np.hstack( (galaxyspec, modelspec[inrange]) ).max() * 1.05
    maxflux = np.hstack( (galaxyspec, modelspec) ).max() * 1.2
    minflux = -0.05 * maxflux

    fig, ax = plt.subplots(figsize=(12, 8))
    if chain is not None:
        for ii in range(nrand):
            _, r_modelspec, _ = _modelsed(model=model, theta=theta_rand[ii, :], obs=obs, sps=sps)
            ax.plot(galaxywave, r_modelspec, alpha=0.2, color='gray')
    ax.plot(galaxywave, modelspec, alpha=1.0, label='Maximum Likelihood Fit')
    #ax.plot(modelwave, modelspec, alpha=1.0, label='Model spectrum')

    ax.errorbar(galaxywave, galaxyspec, galaxyuncspec, alpha=0.8, 
                label='Observed Spectrum')
                
    ax.set_xlabel(r'Observed-Frame Wavelength (${}$m)'.format('\mu'))
    ax.set_ylabel('Flux Density ($\mu$Jy)')
    ax.set_xlim(minwave, maxwave)
    ax.set_ylim(minflux, maxflux)
    ax.legend(loc='upper right', fontsize=16, frameon=False)
    fig.subplots_adjust(left=0.1, right=0.95, bottom=0.12, top=0.95)

    return fig
def subtriangle(results, showpars=None, start=50, thin=1, truths=None, 
                trim_outliers=None, extents=None, **kwargs):
    """Generate a triangle / corner plot.
    
    """
    import corner as triangle

    # pull out the parameter names and flatten the thinned chains
    parnames = results['theta_labels']
    
    chain = results['chain'][start::thin, :]
    weights = results['weights'][start::thin]

    # logify mass
    if 'mass' in parnames and 'logmass' not in parnames:
        midx = np.where(np.in1d(parnames, 'mass'))[0]
        if len(midx) > 0:
            flatchain[:, midx[0]] = np.log10(flatchain[:, midx[0]])
            parnames[midx[0]] = 'logmass'
    parnames = np.array(parnames)

    # restrict to parameters you want to show
    if showpars is not None:
        ind_show = np.array([p in showpars for p in parnames], dtype=bool)
        flatchain = flatchain[:, ind_show]
        #truths = truths[ind_show]
        parnames = parnames[ind_show]

    # Make nice labels.
    niceparnames = _niceparnames(parnames)
        
    if trim_outliers:
        trim_outliers = len(parnames) * [trim_outliers]

    fig = triangle.corner(chain, weights=weights, labels=niceparnames, truths=truths,  
                          verbose=False, quantiles=[0.25, 0.5, 0.75], range=trim_outliers,
                          color='k', **kwargs)
    
    return fig
def read_nestle(prefix, suffix):
    """Read the nestle output files.
        
    """
    import pickle
    from prospect.io import read_results
    
    #nestlefile = os.path.join( jwstdir, 'spsmodeling', '{}_prospect_{}_nmc.pkl'.format(prefix, suffix) )
    #print('Reading {}'.format(nestlefile))
    #with open(nestlefile, 'rb') as f:
    #    results = pickle.load(f)    
        
    hdf5file = os.path.join( jwstdir, 'spsmodeling', '{}_prospect_{}_mcmc.h5'.format(prefix, suffix) )
    print('Reading {}'.format(hdf5file))
    results, _, model = read_results.results_from(hdf5file)

    return results
stop
for ii in range(0, 1, 1):
    results = read_nestle(run_params['prefix'], cat['SUFFIX'][ii])
    fig = bestfit_sed(results, sps=sps, model=model, seed=results['run_params']['seed'])
    plt.plot(cat['WAVELENGTH'][ii], cat['FLUX'][ii] * 1e3, color='green')
for ii in range(0, 1, 1):
    results = read_nestle(run_params['prefix'], cat['SUFFIX'][ii])
    truths = (cat['LOGMASS'][ii], cat['LOGZSOL'][ii], cat['TAU'][ii], cat['TAGE'][ii])
    fig = subtriangle(results, truths=truths)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13711781279348847, 'fraction_numerical': 0.016209259452502695, 'mean_word_length': 2.580769709801968, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 36, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '433540', 'license': 'MIT', 'n_tokens_mistral': 10269, 'n_tokens_neox': 9626, 'n_words': 2087, 'path': 'nb/sps-modeling.ipynb', 'repo_name': 'moustakas/jwstclusters', 'size': '28749'}