HajimeKawahara/exojax

View on GitHub
documents/tutorials/Reverse_modeling.rst

Summary

Maintainability
Test Coverage
Reverse Modelling of an Emission Spectrum (opacity calculator = LPF)
====================================================================

Here, we use the mock spectrum generated by the tutorial of “Foward
modeling”.

.. code:: ipython3

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

.. code:: ipython3

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

We add Gaussian noise to data. nusd is the observing wavenumber grid.

.. code:: ipython3

    wavd=dat["wav"].values[12:-12]
    flux=dat["flux"].values[12:-12]
    nusd=jnp.array(1.e8/wavd[::-1])
    sigmain=0.05
    norm=40000
    nflux=flux/norm+np.random.normal(0,sigmain,len(wavd))
    plt.plot(wavd[::-1],nflux)





.. parsed-literal::

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




.. image:: Reverse_modeling_files/Reverse_modeling_4_1.png


.. code:: ipython3

    from exojax.spec.lpf import xsmatrix
    from exojax.spec.exomol import gamma_exomol
    from exojax.spec.hitran import SijT, doppler_sigma, gamma_natural, gamma_hitran
    from exojax.spec.hitrancia import read_cia, logacia
    from exojax.spec.rtransfer import rtrun, dtauM, dtauCIA, wavenumber_grid
    from exojax.spec import planck, response
    from exojax.spec.lpf import xsvector
    from exojax.spec import molinfo
    from exojax.utils.constants import RJ, pc, Rs, c

The model is almost same as the forward modeling, but we will infer here
Rp, RV, MMR_CO, T0, alpha, and Vsini.

.. code:: ipython3

    from exojax.spec import rtransfer as rt
    NP=100
    Parr, dParr, k=rt.pressure_layer(NP=NP)
    Nx=1500
    nus,wav,res=wavenumber_grid(np.min(wavd)-5.0,np.max(wavd)+5.0,Nx,unit="AA")
    
    R=100000.
    beta=c/(2.0*np.sqrt(2.0*np.log(2.0))*R)
    
    molmassCO=molinfo.molmass("CO")
    mmw=2.33 #mean molecular weight
    mmrH2=0.74
    molmassH2=molinfo.molmass("H2")
    vmrH2=(mmrH2*mmw/molmassH2) #VMR
    
    Mp = 33.2 #fixing mass...


.. parsed-literal::

    xsmode assumes ESLOG in wavenumber space: mode=lpf


.. parsed-literal::

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


Loading the molecular database of CO and the CIA

.. code:: ipython3

    from exojax.spec import api, contdb
    mdbCO=api.MdbExomol('.database/CO/12C-16O/Li2015',nus,crit=1.e-46)
    cdbH2H2=contdb.CdbCIA('.database/H2-H2_2011.cia',nus)


.. parsed-literal::

    Background atmosphere:  H2
    Reading .database/CO/12C-16O/Li2015/12C-16O__Li2015.trans.bz2
    .broad is used.
    Broadening code level= a0
    H2-H2


We have only 39 CO lines.

.. code:: ipython3

    plt.plot(mdbCO.nu_lines,mdbCO.Sij0,".")




.. parsed-literal::

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




.. image:: Reverse_modeling_files/Reverse_modeling_11_1.png


Again, numatrix should be precomputed prior to HMC-NUTS.

.. code:: ipython3

    from exojax.spec import make_numatrix0
    numatrix_CO=make_numatrix0(nus,mdbCO.nu_lines)

.. code:: ipython3

    #Or you can use initspec.init_lpf instead.
    from exojax.spec import initspec
    numatrix_CO=initspec.init_lpf(mdbCO.nu_lines,nus)

.. code:: ipython3

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

.. code:: ipython3

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

Now we write the model, which is used in HMC-NUTS.

.. code:: ipython3

    #response and rotation settings 
    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)
    
    
    def model_c(nu1,y1):
        Rp = numpyro.sample('Rp', dist.Uniform(0.4,1.2))
        RV = numpyro.sample('RV', dist.Uniform(5.0,15.0))
        MMR_CO = numpyro.sample('MMR_CO', dist.Uniform(0.0,0.015))
        T0 = numpyro.sample('T0', dist.Uniform(1000.0,1500.0))
        alpha=numpyro.sample('alpha', dist.Uniform(0.05,0.2))
        vsini = numpyro.sample('vsini', dist.Uniform(15.0,25.0))
        g=2478.57730044555*Mp/Rp**2 #gravity                                        
        u1=0.0
        u2=0.0
        #T-P model//                                                                
        Tarr = T0*(Parr/Pref)**alpha
    
        #line computation CO                                                        
        qt_CO=vmap(mdbCO.qr_interp)(Tarr)
    
        def obyo(y,tag,nusd,nus,numatrix_CO,mdbCO,cdbH2H2):
            #CO                                                                     
            SijM_CO=jit(vmap(SijT,(0,None,None,None,0)))\
                (Tarr,mdbCO.logsij0,mdbCO.dev_nu_lines,mdbCO.elower,qt_CO)
            gammaLMP_CO = jit(vmap(gamma_exomol,(0,0,None,None)))\
                (Parr,Tarr,mdbCO.n_Texp,mdbCO.alpha_ref)
            gammaLMN_CO=gamma_natural(mdbCO.A)
            gammaLM_CO=gammaLMP_CO+gammaLMN_CO[None,:]
            
            sigmaDM_CO=jit(vmap(doppler_sigma,(None,0,None)))\
                (mdbCO.dev_nu_lines,Tarr,molmassCO)
            xsm_CO=xsmatrix(numatrix_CO,sigmaDM_CO,gammaLM_CO,SijM_CO)
            dtaumCO=dtauM(dParr,xsm_CO,MMR_CO*ONEARR,molmassCO,g)
            #CIA                                                                    
            dtaucH2H2=dtauCIA(nus,Tarr,Parr,dParr,vmrH2,vmrH2,\
                              mmw,g,cdbH2H2.nucia,cdbH2H2.tcia,cdbH2H2.logac)
            dtau=dtaumCO+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, RV)
            
            numpyro.sample(tag, dist.Normal(mu, sigmain), obs=y)
        
        obyo(y1,"y1",nu1,nus,numatrix_CO,mdbCO,cdbH2H2)
    


Run a HMC-NUTS. It took ~30min using my gaming laptop (GTX 1080 Max-Q).
Here, the number of warmup is only 300, and that of the sampling is only
600, because the time when the draft on arxiv will be released is very
soon (June 1st 2021 morning in JST!).

.. code:: ipython3

    rng_key = random.PRNGKey(0)
    rng_key, rng_key_ = random.split(rng_key)
    num_warmup, num_samples = 300, 600
    kernel = NUTS(model_c,forward_mode_differentiation=True)
    mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
    mcmc.run(rng_key_, nu1=nusd, y1=nflux)


.. parsed-literal::

    sample: 100%|██████████| 900/900 [49:39<00:00,  3.31s/it, 127 steps of size 2.01e-02. acc. prob=0.94]   


Plotting a prediction and 90% area with the data… looks good.

.. code:: ipython3

    posterior_sample = mcmc.get_samples()
    pred = Predictive(model_c,posterior_sample,return_sites=["y1"])
    predictions = pred(rng_key_,nu1=nusd,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_files/Reverse_modeling_22_0.png


For the above reasons, I haven’t been checking my results properly.
Arviz is useful to visualize the corner plot. Ah, the range of prior
looks too narrow for some parameters. But I have no time to rerun it.
Try to change the priors and run a HMC-NUTS again! The rest is up to
you.

.. code:: ipython3

    import arviz
    pararr=["Rp","T0","alpha","MMR_CO","vsini","RV"]
    arviz.plot_pair(arviz.from_numpyro(mcmc),kind='kde',divergences=False,marginals=True)
    plt.show()



.. image:: Reverse_modeling_files/Reverse_modeling_24_0.png


For fitting to the real spectrum, we may need a more well-considered
model and a better GPU, such as V100 or A100. Read the next section in
detail.