HajimeKawahara/exojax

View on GitHub
documents/tutorials/Reverse_modeling_for_methane_using_PreMODIT.rst

Summary

Maintainability
Test Coverage
Reverse Modeling of an Emission Spectrum using the PreMODIT Cross Section
=========================================================================

November 23rd Hajime Kawahara

We try to fit an emission spectrum model to the mock data in which many
methane lines exist. This situation mocks a T-type brown dwarf.

.. code:: ipython3

    import pandas as pd
    import numpy as np
    import matplotlib.pyplot as plt
    import jax.numpy as jnp

Load the mock data generated by “Forward Modeling of an Emission
Spectrum using the PreMODIT Cross Section”. If you do not have the data,
try “Forward Modeling of an Emission Spectrum using the PreMODIT Cross
Section” first.

.. code:: ipython3

    dat=pd.read_csv("spectrum_ch4_high.txt",delimiter=",",names=("wav","flux"))

Inject an additional Gaussian noise into the spectrum

.. code:: ipython3

    wavd=dat["wav"].values[30:-30]
    flux=dat["flux"].values[30:-30]
    nusd=jnp.array(1.e8/wavd[::-1]) #wavenumber
    sigmain=0.1
    norm=1000
    nflux=flux/norm+np.random.normal(0,sigmain,len(wavd))
    plt.figure(figsize=(10,3))
    plt.plot(wavd[::-1],nflux)





.. parsed-literal::

    [<matplotlib.lines.Line2D at 0x7f7e6c0bf070>]




.. image:: Reverse_modeling_for_methane_using_PreMODIT_files/Reverse_modeling_for_methane_using_PreMODIT_6_1.png


Let’s set a model atmospheric layers, wavenumber range for the model, an
instrument,

.. code:: ipython3

    from exojax.spec.rtransfer import pressure_layer,wavenumber_grid
    from exojax.utils.instfunc import resolution_to_gaussian_std
    
    NP=100
    Parr, dParr, k=pressure_layer(NP=NP)
    Nx=7500
    nus,wav,res=wavenumber_grid(np.min(wavd)-5.0,np.max(wavd)+5.0,Nx,unit="AA",xsmode="premodit")
    
    Rinst=100000. #instrumental spectral resolution
    beta_inst=resolution_to_gaussian_std(Rinst)  #equivalent to beta=c/(2.0*np.sqrt(2.0*np.log(2.0))*R)


.. parsed-literal::

    xsmode assumes ESLOG in wavenumber space: mode=premodit


.. parsed-literal::

    /home/kawahara/exojax/src/exojax/utils/grids.py:124: UserWarning: Resolution may be too small. R=611064.2488992558
      warnings.warn('Resolution may be too small. R=' + str(resolution),


Load molecular database, CIA, and define some values. Note that PreMODIT
can use gpu_transfer = False option, which siginificantly reduce divice
memory use.

.. code:: ipython3

    from exojax.spec import api, contdb
    from exojax.spec import molinfo
    
    mmw=2.33 #mean molecular weight
    
    mdbCH4=api.MdbExomol('.database/CH4/12C-1H4/YT10to10/',nus,gpu_transfer=False)
    molmassCH4=molinfo.molmass("CH4")
    
    cdbH2H2=contdb.CdbCIA('.database/H2-H2_2011.cia',nus)
    molmassH2=molinfo.molmass("H2")
    mmrH2=0.74
    vmrH2=(mmrH2*mmw/molmassH2) #VMR
    Mp = 33.2 


.. parsed-literal::

    Background atmosphere:  H2
    Reading .database/CH4/12C-1H4/YT10to10/12C-1H4__YT10to10__06000-06100.trans.bz2
    Reading .database/CH4/12C-1H4/YT10to10/12C-1H4__YT10to10__06100-06200.trans.bz2
    .broad is used.
    Broadening code level= a1
    default broadening parameters are used for  23  J lower states in  40  states
    H2-H2


.. code:: ipython3

    print(len(mdbCH4.Sij0),"lines")


.. parsed-literal::

    81308640 lines


Define some arrays for the model.

.. code:: ipython3

    #reference pressure for a T-P model                                             
    Pref=1.0 #bar
    ONEARR=np.ones_like(Parr)
    ONEWAV=jnp.ones_like(nflux)


Initialize PreMODIT. In this process, we precompute Line Basis Density
(lbd). It takes a while.

.. code:: ipython3

    from exojax.spec.initspec import init_premodit
    
    interval_contrast = 0.1
    dit_grid_resolution = 0.1
    Ttyp = 700.0
    
    lbd, multi_index_uniqgrid, elower_grid, ngamma_ref_grid, n_Texp_grid, R, pmarray = init_premodit(
        mdbCH4.nu_lines,
        nus,
        mdbCH4.elower,
        mdbCH4.alpha_ref,
        mdbCH4.n_Texp,
        mdbCH4.Sij0,
        Ttyp,
        interval_contrast=interval_contrast,
        dit_grid_resolution=dit_grid_resolution,
        warning=False)



.. parsed-literal::

    uniqidx: 100%|██████████| 2/2 [00:03<00:00,  1.83s/it]


Do not confuse R with Rinst. R is the spectral resolution of the raw
spectral model, which should be higher than Rinst, while Rinst is the
instrumental spectral resolution.

.. code:: ipython3

    Rinst, R




.. parsed-literal::

    (100000.0, 611064.2488992558)



We here use numpyro as a PPL (probabilistic programming language).

.. code:: ipython3

    from jax import random
    import numpyro.distributions as dist
    import numpyro
    from numpyro.infer import MCMC, NUTS
    from numpyro.infer import Predictive
    from numpyro.diagnostics import hpdi

.. code:: ipython3

    from exojax.spec.response import ipgauss_sampling
    from exojax.spec.spin_rotation import convolve_rigid_rotation
    from exojax.utils.grids import velocity_grid
    vsini_max = 100.0
    vr_array = velocity_grid(res, vsini_max)


.. code:: ipython3

    from exojax.spec.rtransfer import dtauM, dtauCIA, rtrun
    from exojax.spec import planck
    from exojax.spec import premodit
    from jax import vmap
    
    
    def frun(Tarr, MMR_CH4, Mp, Rp, u1, u2, RV, vsini):
        g = 2478.57730044555 * Mp / Rp**2
        qtarr = vmap(mdbCH4.qr_interp)(Tarr)
        xsm = premodit.xsmatrix(Tarr, Parr, R, pmarray, lbd, nus, ngamma_ref_grid,
                                n_Texp_grid, multi_index_uniqgrid, elower_grid,
                                molmassCH4, qtarr)
        #xsm[xsm<0.0]=0.0
        xsm=jnp.abs(xsm)
        dtaum = dtauM(dParr, xsm, MMR_CH4 * ONEARR, molmassCH4, g)
        #CIA
        dtaucH2H2 = dtauCIA(nus, Tarr, Parr, dParr, vmrH2, vmrH2, mmw, g,
                            cdbH2H2.nucia, cdbH2H2.tcia, cdbH2H2.logac)
        dtau = dtaum + dtaucH2H2
        sourcef = planck.piBarr(Tarr, nus)
        F0 = rtrun(dtau, sourcef) / norm
        Frot = convolve_rigid_rotation(F0, vr_array, vsini, u1, u2)
        mu = ipgauss_sampling(nusd, nus, Frot, beta_inst, RV)
        return mu

Test plot using frun

.. code:: ipython3

    T0=800.0 #K                                                                                                        
    Tarr = T0*(Parr/Pref)**0.1
    mu=frun(Tarr,MMR_CH4=0.0058,Mp=33.5,Rp=0.88,u1=0.0,u2=0.0,RV=10.0,vsini=20.0)
    plt.plot(wavd[::-1],mu,label="frun")
    plt.plot(wavd[::-1],nflux,alpha=0.3,label="data to be fitted")
    plt.legend()
    plt.show()




.. image:: Reverse_modeling_for_methane_using_PreMODIT_files/Reverse_modeling_for_methane_using_PreMODIT_23_0.png


Let’s define the model for a HMC.

.. code:: ipython3

    Mp=33.2
    Rp=0.88
    #we assume we know gravity here.
    
    def model_c(y1):
        RV = numpyro.sample('RV', dist.Uniform(5.0,15.1))
        MMR_CH4 = numpyro.sample('MMR_CH4', dist.Uniform(0.0,0.01))
        T0 = numpyro.sample('T0', dist.Uniform(700.0,900.0))
        alpha=numpyro.sample('alpha', dist.Uniform(0.05,0.2))
        vsini = numpyro.sample('vsini', dist.Uniform(15.0,25.0)) 
        sigma = numpyro.sample('sigma',dist.Exponential(1.0))
        u1=0.0
        u2=0.0                                                                                                 
        Tarr = T0*(Parr/Pref)**alpha  
        mu=frun(Tarr,MMR_CH4,Mp,Rp,u1,u2,RV,vsini)
        numpyro.sample("y1", dist.Normal(mu, sigma), obs=y1)

.. code:: ipython3

    rng_key = random.PRNGKey(0)
    rng_key, rng_key_ = random.split(rng_key)
    num_warmup, num_samples = 100, 200
    kernel = NUTS(model_c,forward_mode_differentiation=True,max_tree_depth=9) #24min
    #kernel = NUTS(model_c,max_tree_depth=9) #
    mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
    mcmc.run(rng_key_, y1=nflux)



.. parsed-literal::

    sample: 100%|██████████| 300/300 [24:01<00:00,  4.81s/it, 57 steps of size 9.51e-02. acc. prob=0.89]  


.. code:: ipython3

    posterior_sample = mcmc.get_samples()
    pred = Predictive(model_c,posterior_sample,return_sites=["y1"])
    predictions = pred(rng_key_,y1=None)
    median_mu1 = jnp.median(predictions["y1"],axis=0)
    hpdi_mu1 = hpdi(predictions["y1"], 0.9)                                      
    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(20,6.0))
    ax.plot(wavd[::-1],median_mu1,color="C0")
    ax.plot(wavd[::-1],nflux,"+",color="black",label="data")
    ax.fill_between(wavd[::-1], hpdi_mu1[0], hpdi_mu1[1], alpha=0.3, interpolate=True,color="C0",label="90% area")
    plt.xlabel("wavelength ($\AA$)",fontsize=16)
    plt.legend(fontsize=16)
    plt.tick_params(labelsize=16)



.. image:: Reverse_modeling_for_methane_using_PreMODIT_files/Reverse_modeling_for_methane_using_PreMODIT_27_0.png


.. code:: ipython3

    import arviz
    refs={};refs["RV"]=10.0;refs["T0"]=800;refs["MMR_CH4"]=0.0059;refs["alpha"]=0.1;refs["vsini"]=20.0;refs["sigma"]=0.1;
    arviz.plot_pair(arviz.from_numpyro(mcmc),kind='kde',divergences=False,marginals=True,
                   reference_values=refs,reference_values_kwargs={'color':"red", "marker":"o", "markersize":12})
    plt.show()



.. image:: Reverse_modeling_for_methane_using_PreMODIT_files/Reverse_modeling_for_methane_using_PreMODIT_28_0.png