Good practices for analysing MCMC results using arviz¶

Assessing the convergence of a set of MCMC chains is not an easy task in general. jaxspec provides a convenient way to analyse the results of a fit using the arviz library. This library provides powerful tool to explore Bayesian inference results, such as trace plots, pair plots, and summary statistics. Let's run some MCMCs!

In [1]:

%%capture 
# Hide the output of this cell
import numpyro

numpyro.enable_x64()
numpyro.set_platform("cpu")
numpyro.set_host_device_count(4)

import numpyro.distributions as dist
from jaxspec.fit import MCMCFitter
from jaxspec.data.util import load_example_obsconf
from jaxspec.model.additive import Blackbodyrad, Powerlaw
from jaxspec.model.multiplicative import Tbabs

spectral_model = Tbabs()*(Powerlaw() + Blackbodyrad())

obsconf = load_example_obsconf("NGC7793_ULX4_PN")

prior = {
    "powerlaw_1_alpha": dist.Uniform(1, 3),
    "powerlaw_1_norm": dist.LogUniform(1e-5, 1e-3),
    "blackbodyrad_1_kT": dist.Uniform(0, 2),
    "blackbodyrad_1_norm": dist.LogUniform(1e-2, 1),
    "tbabs_1_nh": dist.Uniform(0, 1)
}

fitter = MCMCFitter(spectral_model, prior, obsconf)

%%capture # Hide the output of this cell import numpyro numpyro.enable_x64() numpyro.set_platform("cpu") numpyro.set_host_device_count(4) import numpyro.distributions as dist from jaxspec.fit import MCMCFitter from jaxspec.data.util import load_example_obsconf from jaxspec.model.additive import Blackbodyrad, Powerlaw from jaxspec.model.multiplicative import Tbabs spectral_model = Tbabs()*(Powerlaw() + Blackbodyrad()) obsconf = load_example_obsconf("NGC7793_ULX4_PN") prior = { "powerlaw_1_alpha": dist.Uniform(1, 3), "powerlaw_1_norm": dist.LogUniform(1e-5, 1e-3), "blackbodyrad_1_kT": dist.Uniform(0, 2), "blackbodyrad_1_norm": dist.LogUniform(1e-2, 1), "tbabs_1_nh": dist.Uniform(0, 1) } fitter = MCMCFitter(spectral_model, prior, obsconf)

In [2]:

result_nuts = fitter.fit(
    num_chains=4, 
    num_warmup=1000, 
    num_samples=1000,
    sampler="nuts",
    mcmc_kwargs={"progress_bar": True}
)

result_nuts = fitter.fit( num_chains=4, num_warmup=1000, num_samples=1000, sampler="nuts", mcmc_kwargs={"progress_bar": True} )

  0%|          | 0/2000 [00:00<?, ?it/s]

  0%|          | 0/2000 [00:00<?, ?it/s]

  0%|          | 0/2000 [00:00<?, ?it/s]

  0%|          | 0/2000 [00:00<?, ?it/s]

From the result object, you can access the inference_data attribute, which is an arviz.InferenceData object. This leverage the use of every arviz function to analyse the results of the fit.

In [3]:

inference_data = result_nuts.inference_data
inference_data

inference_data = result_nuts.inference_data inference_data

Out[3]:

arviz.InferenceData

• posterior

  <xarray.Dataset> Size: 168kB
  Dimensions:              (chain: 4, draw: 1000)
  Coordinates:
    * chain                (chain) int64 32B 0 1 2 3
    * draw                 (draw) int64 8kB 0 1 2 3 4 5 ... 995 996 997 998 999
  Data variables:
      blackbodyrad_1_kT    (chain, draw) float64 32kB 0.7618 0.7805 ... 0.7203
      blackbodyrad_1_norm  (chain, draw) float64 32kB 0.1917 0.154 ... 0.2119
      powerlaw_1_alpha     (chain, draw) float64 32kB 2.05 2.116 ... 1.957 1.979
      powerlaw_1_norm      (chain, draw) float64 32kB 0.0002669 ... 0.0002556
      tbabs_1_nh           (chain, draw) float64 32kB 0.1076 0.1352 ... 0.08894
  Attributes:
      created_at:     2024-11-22T15:30:08.733726+00:00
      arviz_version:  0.20.0
      model:          SpectralModel(\n  graph=<networkx.classes.digraph.DiGraph...

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
  • Coordinates: (2)
    • chain
      (chain)
      int64
      0 1 2 3

      array([0, 1, 2, 3])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 995 996 997 998 999

      array([  0,   1,   2, ..., 997, 998, 999])

  • Data variables: (5)
    • blackbodyrad_1_kT
      (chain, draw)
      float64
      0.7618 0.7805 ... 0.7312 0.7203

      array([[0.76179887, 0.78054488, 0.73066334, ..., 0.69674459, 0.7130727 ,
              0.71515355],
             [0.69797501, 0.69014638, 0.71991811, ..., 0.67896861, 0.69112968,
              0.71937316],
             [0.72022504, 0.72100237, 0.75737135, ..., 0.81768224, 0.81960399,
              0.76092194],
             [0.6794004 , 0.78573937, 0.80248828, ..., 0.70299306, 0.73121832,
              0.72030423]])

    • blackbodyrad_1_norm
      (chain, draw)
      float64
      0.1917 0.154 ... 0.1956 0.2119

      array([[0.19167592, 0.15402632, 0.21216628, ..., 0.26068175, 0.2130297 ,
              0.20471758],
             [0.21773577, 0.24253916, 0.23952518, ..., 0.2564305 , 0.24351995,
              0.24469645],
             [0.23318756, 0.22914724, 0.18645229, ..., 0.16909257, 0.16605871,
              0.14268358],
             [0.2684988 , 0.15471224, 0.15468832, ..., 0.20698931, 0.19556884,
              0.21192288]])

    • powerlaw_1_alpha
      (chain, draw)
      float64
      2.05 2.116 2.005 ... 1.957 1.979

      array([[2.04984401, 2.1156262 , 2.00474283, ..., 1.97758677, 1.91570946,
              1.94477095],
             [1.85688789, 1.92841279, 1.93586232, ..., 1.85220205, 1.86317785,
              1.99122157],
             [1.99845109, 1.99603348, 1.99696974, ..., 2.25839406, 2.27286476,
              2.06438414],
             [1.86328243, 2.16675767, 2.15913604, ..., 1.99114135, 1.95738891,
              1.97930009]])

    • powerlaw_1_norm
      (chain, draw)
      float64
      0.0002669 0.0002963 ... 0.0002556

      array([[0.00026688, 0.00029633, 0.00024479, ..., 0.00023418, 0.00024809,
              0.00024898],
             [0.00024984, 0.00025153, 0.00022766, ..., 0.00023153, 0.00022839,
              0.0002288 ],
             [0.00025031, 0.00025074, 0.00025158, ..., 0.00026938, 0.00027173,
              0.00031124],
             [0.00022093, 0.00029963, 0.00029205, ..., 0.00026671, 0.00025036,
              0.00025558]])

    • tbabs_1_nh
      (chain, draw)
      float64
      0.1076 0.1352 ... 0.08068 0.08894

      array([[0.10764608, 0.1352217 , 0.08345133, ..., 0.06892088, 0.07571119,
              0.0668106 ],
             [0.07173085, 0.07809168, 0.06654789, ..., 0.06143396, 0.05507334,
              0.06476625],
             [0.08489157, 0.09076794, 0.08323814, ..., 0.12196389, 0.11659585,
              0.13078243],
             [0.04650204, 0.13265337, 0.13143839, ..., 0.08865422, 0.08068027,
              0.08894497]])

  • Indexes: (2)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
            dtype='int64', name='draw', length=1000))

  • Attributes: (3)

    created_at :

    2024-11-22T15:30:08.733726+00:00

    arviz_version :

    0.20.0

    model :

    SpectralModel( graph=<networkx.classes.digraph.DiGraph object at 0x3739b5f40>, modules={'blackbodyrad_1': Blackbodyrad( kT=Param( value=0.5 ), norm=Param( value=1.0 ) ), 'powerlaw_1': Powerlaw( alpha=Param( value=1.7 ), norm=Param( value=0.0001 ) ), 'tbabs_1': Tbabs( energy=Variable( value=Array(shape=(10000,), dtype=float64) ), sigma=Variable( value=Array(shape=(10000,), dtype=float64) ), nh=Param( value=1.0 ) )} )

  
  
• posterior_predictive

  <xarray.Dataset> Size: 3MB
  Dimensions:         (chain: 4, draw: 1000, obs_data_dim_0: 102)
  Coordinates:
    * chain           (chain) int64 32B 0 1 2 3
    * draw            (draw) int64 8kB 0 1 2 3 4 5 6 ... 994 995 996 997 998 999
    * obs_data_dim_0  (obs_data_dim_0) int64 816B 0 1 2 3 4 5 ... 97 98 99 100 101
  Data variables:
      obs_data        (chain, draw, obs_data_dim_0) int64 3MB 206 202 ... 18 18
  Attributes:
      created_at:     2024-11-22T15:30:08.734722+00:00
      arviz_version:  0.20.0
      model:          SpectralModel(\n  graph=<networkx.classes.digraph.DiGraph...

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • obs_data_dim_0: 102
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0 1 2 3

      array([0, 1, 2, 3])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 995 996 997 998 999

      array([  0,   1,   2, ..., 997, 998, 999])

    • obs_data_dim_0
      (obs_data_dim_0)
      int64
      0 1 2 3 4 5 ... 96 97 98 99 100 101

      array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
              14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
              28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
              42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
              56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
              70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
              84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
              98,  99, 100, 101])

  • Data variables: (1)
    • obs_data
      (chain, draw, obs_data_dim_0)
      int64
      206 202 205 216 257 ... 25 14 18 18

      array([[[206, 202, 205, ...,  15,  28,  32],
              [215, 222, 187, ...,  15,  17,  18],
              [226, 224, 239, ...,  13,  14,  24],
              ...,
              [239, 238, 250, ...,  17,  17,  25],
              [234, 215, 230, ...,  21,  24,  24],
              [235, 218, 235, ...,  14,  21,  20]],
      
             [[225, 231, 216, ...,  21,  24,  23],
              [189, 215, 200, ...,  15,  12,  25],
              [218, 203, 221, ...,  14,  21,  34],
              ...,
              [214, 247, 233, ...,   9,  22,  26],
              [238, 212, 230, ...,  19,  23,  31],
              [227, 216, 217, ...,  12,  15,  24]],
      
             [[211, 249, 196, ...,  17,  20,  17],
              [217, 230, 236, ...,   9,  17,  22],
              [225, 235, 233, ...,  17,  23,  27],
              ...,
              [220, 208, 231, ...,  16,  16,  15],
              [209, 222, 249, ...,  11,  16,  13],
              [217, 229, 209, ...,  16,  21,  22]],
      
             [[222, 226, 223, ...,  15,  18,  26],
              [200, 237, 235, ...,  17,  13,  20],
              [227, 197, 236, ...,  17,  16,  13],
              ...,
              [212, 247, 228, ...,  13,  21,  11],
              [220, 231, 220, ...,  17,  19,  34],
              [216, 235, 232, ...,  14,  18,  18]]])

  • Indexes: (3)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
            dtype='int64', name='draw', length=1000))

    • obs_data_dim_0
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
              92,  93,  94,  95,  96,  97,  98,  99, 100, 101],
            dtype='int64', name='obs_data_dim_0', length=102))

  • Attributes: (3)

    created_at :

    2024-11-22T15:30:08.734722+00:00

    arviz_version :

    0.20.0

    model :

    SpectralModel( graph=<networkx.classes.digraph.DiGraph object at 0x3739b5f40>, modules={'blackbodyrad_1': Blackbodyrad( kT=Param( value=0.5 ), norm=Param( value=1.0 ) ), 'powerlaw_1': Powerlaw( alpha=Param( value=1.7 ), norm=Param( value=0.0001 ) ), 'tbabs_1': Tbabs( energy=Variable( value=Array(shape=(10000,), dtype=float64) ), sigma=Variable( value=Array(shape=(10000,), dtype=float64) ), nh=Param( value=1.0 ) )} )

  
  
• log_likelihood

  <xarray.Dataset> Size: 3MB
  Dimensions:         (chain: 4, draw: 1000, obs_data_dim_0: 102)
  Coordinates:
    * chain           (chain) int64 32B 0 1 2 3
    * draw            (draw) int64 8kB 0 1 2 3 4 5 6 ... 994 995 996 997 998 999
    * obs_data_dim_0  (obs_data_dim_0) int64 816B 0 1 2 3 4 5 ... 97 98 99 100 101
  Data variables:
      obs_data        (chain, draw, obs_data_dim_0) float64 3MB -9.419 ... -2.611
  Attributes:
      created_at:     2024-11-22T15:30:08.734448+00:00
      arviz_version:  0.20.0
      model:          SpectralModel(\n  graph=<networkx.classes.digraph.DiGraph...

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 1000
    • obs_data_dim_0: 102
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0 1 2 3

      array([0, 1, 2, 3])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 995 996 997 998 999

      array([  0,   1,   2, ..., 997, 998, 999])

    • obs_data_dim_0
      (obs_data_dim_0)
      int64
      0 1 2 3 4 5 ... 96 97 98 99 100 101

      array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
              14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
              28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
              42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
              56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
              70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
              84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
              98,  99, 100, 101])

  • Data variables: (1)
    • obs_data
      (chain, draw, obs_data_dim_0)
      float64
      -9.419 -4.145 ... -2.422 -2.611

      array([[[ -9.41889465,  -4.14538458,  -4.858972  , ...,  -3.05224585,
                -2.42693309,  -2.49833128],
              [-11.3718389 ,  -4.73163163,  -4.36516953, ...,  -3.1434989 ,
                -2.44108117,  -2.46213489],
              [ -7.6367582 ,  -3.76905023,  -5.48572575, ...,  -3.26271472,
                -2.45726255,  -2.45654967],
              ...,
              [ -6.06523193,  -3.63366378,  -6.61751847, ...,  -3.38912957,
                -2.48259021,  -2.44314617],
              [ -7.28534196,  -3.70115921,  -5.78939882, ...,  -2.68036931,
                -2.50136668,  -2.94829374],
              [ -5.05323131,  -3.78473297,  -7.84011754, ...,  -2.84914644,
                -2.43563072,  -2.72255884]],
      
             [[ -7.35614882,  -3.70133111,  -5.80946961, ...,  -2.4679973 ,
                -2.80761014,  -3.64743028],
              [ -6.80740185,  -3.64839053,  -6.24239978, ...,  -2.7673907 ,
                -2.46195501,  -2.84133412],
              [ -7.434077  ,  -3.73662497,  -5.50998366, ...,  -3.01570762,
                -2.421432  ,  -2.56562135],
      ...
              [ -7.7941727 ,  -3.82045755,  -5.47004548, ...,  -4.07504073,
                -2.89390519,  -2.64448165],
              [ -6.03866191,  -3.63199432,  -6.75172891, ...,  -4.1619456 ,
                -2.94764082,  -2.68679065],
              [ -9.24852494,  -4.04438119,  -5.23119316, ...,  -2.79166254,
                -2.44360812,  -2.72379932]],
      
             [[ -5.47930603,  -3.6837242 ,  -7.056682  , ...,  -2.79298946,
                -2.45629375,  -2.83793644],
              [ -8.58105906,  -3.92621397,  -5.4067758 , ...,  -3.46386331,
                -2.53746648,  -2.42160241],
              [ -9.54959609,  -4.18159293,  -4.91512047, ...,  -3.29140898,
                -2.487379  ,  -2.42295247],
              ...,
              [ -6.23523945,  -3.63221157,  -6.77106341, ...,  -2.98842189,
                -2.4211653 ,  -2.61220065],
              [ -7.38385543,  -3.72166755,  -5.6839176 , ...,  -2.84717898,
                -2.43370289,  -2.69886522],
              [ -7.92619624,  -3.80117762,  -5.47370421, ...,  -2.95552783,
                -2.4216742 ,  -2.61061652]]])

  • Indexes: (3)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
            dtype='int64', name='draw', length=1000))

    • obs_data_dim_0
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
              92,  93,  94,  95,  96,  97,  98,  99, 100, 101],
            dtype='int64', name='obs_data_dim_0', length=102))

  • Attributes: (3)

    created_at :

    2024-11-22T15:30:08.734448+00:00

    arviz_version :

    0.20.0

    model :

    SpectralModel( graph=<networkx.classes.digraph.DiGraph object at 0x3739b5f40>, modules={'blackbodyrad_1': Blackbodyrad( kT=Param( value=0.5 ), norm=Param( value=1.0 ) ), 'powerlaw_1': Powerlaw( alpha=Param( value=1.7 ), norm=Param( value=0.0001 ) ), 'tbabs_1': Tbabs( energy=Variable( value=Array(shape=(10000,), dtype=float64) ), sigma=Variable( value=Array(shape=(10000,), dtype=float64) ), nh=Param( value=1.0 ) )} )

  
  
• prior

  <xarray.Dataset> Size: 192kB
  Dimensions:              (chain: 1, draw: 4000)
  Coordinates:
    * chain                (chain) int64 8B 0
    * draw                 (draw) int64 32kB 0 1 2 3 4 ... 3996 3997 3998 3999
  Data variables:
      blackbodyrad_1_kT    (chain, draw) float64 32kB 0.0196 1.745 ... 1.102
      blackbodyrad_1_norm  (chain, draw) float64 32kB 0.03938 0.01347 ... 0.01858
      powerlaw_1_alpha     (chain, draw) float64 32kB 1.235 1.731 ... 2.349 2.549
      powerlaw_1_norm      (chain, draw) float64 32kB 0.0001388 ... 2.465e-05
      tbabs_1_nh           (chain, draw) float64 32kB 0.3542 0.2841 ... 0.1103
  Attributes:
      created_at:     2024-11-22T15:30:08.735628+00:00
      arviz_version:  0.20.0
      model:          SpectralModel(\n  graph=<networkx.classes.digraph.DiGraph...

  xarray.Dataset

  • Dimensions:
    • chain: 1
    • draw: 4000
  • Coordinates: (2)
    • chain
      (chain)
      int64
      0

      array([0])

    • draw
      (draw)
      int64
      0 1 2 3 4 ... 3996 3997 3998 3999

      array([   0,    1,    2, ..., 3997, 3998, 3999])

  • Data variables: (5)
    • blackbodyrad_1_kT
      (chain, draw)
      float64
      0.0196 1.745 1.904 ... 0.8037 1.102

      array([[0.01959878, 1.74462513, 1.90374797, ..., 0.32562702, 0.80367275,
              1.10164047]])

    • blackbodyrad_1_norm
      (chain, draw)
      float64
      0.03938 0.01347 ... 0.7623 0.01858

      array([[0.03938291, 0.01347397, 0.12420529, ..., 0.46052051, 0.76231501,
              0.01858229]])

    • powerlaw_1_alpha
      (chain, draw)
      float64
      1.235 1.731 2.282 ... 2.349 2.549

      array([[1.23489973, 1.73070401, 2.28230657, ..., 1.71911027, 2.34878376,
              2.54867252]])

    • powerlaw_1_norm
      (chain, draw)
      float64
      0.0001388 2.602e-05 ... 2.465e-05

      array([[1.38755680e-04, 2.60168693e-05, 5.14614600e-04, ...,
              2.93948148e-04, 3.56282855e-05, 2.46498672e-05]])

    • tbabs_1_nh
      (chain, draw)
      float64
      0.3542 0.2841 ... 0.5237 0.1103

      array([[0.35418597, 0.28405092, 0.67734216, ..., 0.95527087, 0.52368377,
              0.1102949 ]])

  • Indexes: (2)
    • chain
      PandasIndex

      PandasIndex(Index([0], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([   0,    1,    2,    3,    4,    5,    6,    7,    8,    9,
             ...
             3990, 3991, 3992, 3993, 3994, 3995, 3996, 3997, 3998, 3999],
            dtype='int64', name='draw', length=4000))

  • Attributes: (3)

    created_at :

    2024-11-22T15:30:08.735628+00:00

    arviz_version :

    0.20.0

    model :

    SpectralModel( graph=<networkx.classes.digraph.DiGraph object at 0x3739b5f40>, modules={'blackbodyrad_1': Blackbodyrad( kT=Param( value=0.5 ), norm=Param( value=1.0 ) ), 'powerlaw_1': Powerlaw( alpha=Param( value=1.7 ), norm=Param( value=0.0001 ) ), 'tbabs_1': Tbabs( energy=Variable( value=Array(shape=(10000,), dtype=float64) ), sigma=Variable( value=Array(shape=(10000,), dtype=float64) ), nh=Param( value=1.0 ) )} )

  
  
• observed_data

  <xarray.Dataset> Size: 2kB
  Dimensions:         (obs_data_dim_0: 102)
  Coordinates:
    * obs_data_dim_0  (obs_data_dim_0) int64 816B 0 1 2 3 4 5 ... 97 98 99 100 101
  Data variables:
      obs_data        (obs_data_dim_0) int64 816B 265 227 196 209 ... 20 20 20 20
  Attributes:
      created_at:     2024-11-22T15:30:08.736478+00:00
      arviz_version:  0.20.0
      model:          SpectralModel(\n  graph=<networkx.classes.digraph.DiGraph...

  xarray.Dataset

  • Dimensions:
    • obs_data_dim_0: 102
  • Coordinates: (1)
    • obs_data_dim_0
      (obs_data_dim_0)
      int64
      0 1 2 3 4 5 ... 96 97 98 99 100 101

      array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
              14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
              28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
              42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
              56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
              70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
              84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
              98,  99, 100, 101])

  • Data variables: (1)
    • obs_data
      (obs_data_dim_0)
      int64
      265 227 196 209 231 ... 20 20 20 20

      array([265, 227, 196, 209, 231, 253, 219, 199, 218, 222, 216, 221, 208,
             247, 254, 229, 216, 203, 214, 194, 212, 196, 190, 190, 221, 163,
             184, 159, 163, 136, 117, 139, 154, 110, 131, 145, 126, 123, 145,
             130,  95,  78,  94,  78,  72,  65,  75,  48,  43,  63,  65,  66,
              69,  68,  47,  43,  44,  42,  43,  62,  44,  46,  43,  74,  40,
              42,  43,  45,  42,  44,  44,  43,  44,  44,  41,  20,  20,  21,
              42,  21,  21,  22,  20,  20,  20,  22,  22,  20,  20,  21,  21,
              20,  21,  21,  21,  20,  21,  20,  20,  20,  20,  20])

  • Indexes: (1)
    • obs_data_dim_0
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
              92,  93,  94,  95,  96,  97,  98,  99, 100, 101],
            dtype='int64', name='obs_data_dim_0', length=102))

  • Attributes: (3)

    created_at :

    2024-11-22T15:30:08.736478+00:00

    arviz_version :

    0.20.0

    model :

    SpectralModel( graph=<networkx.classes.digraph.DiGraph object at 0x3739b5f40>, modules={'blackbodyrad_1': Blackbodyrad( kT=Param( value=0.5 ), norm=Param( value=1.0 ) ), 'powerlaw_1': Powerlaw( alpha=Param( value=1.7 ), norm=Param( value=0.0001 ) ), 'tbabs_1': Tbabs( energy=Variable( value=Array(shape=(10000,), dtype=float64) ), sigma=Variable( value=Array(shape=(10000,), dtype=float64) ), nh=Param( value=1.0 ) )} )

  
  

This object carries all the information we need about our fit : the posterior samples, the associated likelihood, posterior predictives and so on. They can be a bit tricky to manipulate at first, as they are based on the xarray, but they provide a lot of flexibility to analyse the results of the fit.

Trace plot¶

This visualization is useful to see the evolution of the parameters during the sampling process. It can be used to diagnose convergence issues. The ideal situation is when the chains are well mixed and randomly scattered around the target distribution. If instead, chains are stuck in some region of the parameter space, or show some trends, this might indicate that the sampler did not explore the full parameter space.

In [4]:

import arviz as az
import matplotlib.pyplot as plt 

with az.style.context("arviz-darkgrid", after_reset=True):
    az.plot_trace(inference_data, compact=False)

plt.show()

import arviz as az import matplotlib.pyplot as plt with az.style.context("arviz-darkgrid", after_reset=True): az.plot_trace(inference_data, compact=False) plt.show()

However, this kind of plot can become messy pretty quickly if you work with numerous chains/walkers, which what you must do when using Ensemble Samplers such as ESS or AIES. A general trick proposed by Vehtari et al. (2019) is to plot the rank of each sample in the global run instead of its value. This rank should be evenly spread across the steps if the chains are well mixed. This can be done using the kind="rank_vlines" argument in the plot_trace function.

In [5]:

with az.style.context("arviz-darkgrid", after_reset=True):
    az.plot_trace(result_nuts.inference_data, compact=False, kind="rank_vlines")

plt.show()

with az.style.context("arviz-darkgrid", after_reset=True): az.plot_trace(result_nuts.inference_data, compact=False, kind="rank_vlines") plt.show()

A more quantitative way to assess the convergence of the chains is to use the summary function. This function provides a summary of the posterior distribution of the parameters, including the mean, the standard deviation, and the 95% highest posterior density interval.

In [6]:

az.summary(result_nuts.inference_data.posterior)

az.summary(result_nuts.inference_data.posterior)

Out[6]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
blackbodyrad_1_kT	0.743	0.035	0.679	0.808	0.001	0.001	631.0	1056.0	1.01
blackbodyrad_1_norm	0.200	0.037	0.134	0.269	0.002	0.001	596.0	1040.0	1.01
powerlaw_1_alpha	2.029	0.100	1.836	2.205	0.004	0.003	630.0	1050.0	1.01
powerlaw_1_norm	0.000	0.000	0.000	0.000	0.000	0.000	604.0	967.0	1.01
tbabs_1_nh	0.095	0.028	0.043	0.146	0.001	0.001	592.0	796.0	1.01

The r_hat column provides a measure of the splitted Gelman-Rubin statistic with rank normalisartion. The closer this value is to 1, the better. A value larger than 1.01 point to convergence issues. This statistic can be directly computed using the r_hat function, see Vehtari et al. (2019). The ess denotes the Effective Sample Size of the chains, which is a measure of the quality of the samples. The larger the better, in general, we want this value to be larger than 400 for a reliable estimate of the posterior distribution using the NUTS sampler with 1000 warmups and 1000 samples.

In [7]:

rhat = az.rhat(result_nuts.inference_data.posterior)
rhat

rhat = az.rhat(result_nuts.inference_data.posterior) rhat

Out[7]:

<xarray.Dataset> Size: 40B
Dimensions:              ()
Data variables:
    blackbodyrad_1_kT    float64 8B 1.009
    blackbodyrad_1_norm  float64 8B 1.007
    powerlaw_1_alpha     float64 8B 1.008
    powerlaw_1_norm      float64 8B 1.008
    tbabs_1_nh           float64 8B 1.008

xarray.Dataset

• Dimensions:
• Coordinates: (0)
• Data variables: (5)
  • blackbodyrad_1_kT
    ()
    float64
    1.009

    array(1.00851367)

  • blackbodyrad_1_norm
    ()
    float64
    1.007

    array(1.00726069)

  • powerlaw_1_alpha
    ()
    float64
    1.008

    array(1.00781349)

  • powerlaw_1_norm
    ()
    float64
    1.008

    array(1.00772239)

  • tbabs_1_nh
    ()
    float64
    1.008

    array(1.00815119)

• Indexes: (0)
• Attributes: (0)

In [8]:

ess = az.ess(result_nuts.inference_data.posterior)
ess

ess = az.ess(result_nuts.inference_data.posterior) ess

Out[8]:

<xarray.Dataset> Size: 40B
Dimensions:              ()
Data variables:
    blackbodyrad_1_kT    float64 8B 631.1
    blackbodyrad_1_norm  float64 8B 595.9
    powerlaw_1_alpha     float64 8B 630.1
    powerlaw_1_norm      float64 8B 603.5
    tbabs_1_nh           float64 8B 591.8

xarray.Dataset

• Dimensions:
• Coordinates: (0)
• Data variables: (5)
  • blackbodyrad_1_kT
    ()
    float64
    631.1

    array(631.08869478)

  • blackbodyrad_1_norm
    ()
    float64
    595.9

    array(595.85100914)

  • powerlaw_1_alpha
    ()
    float64
    630.1

    array(630.12697476)

  • powerlaw_1_norm
    ()
    float64
    603.5

    array(603.50070759)

  • tbabs_1_nh
    ()
    float64
    591.8

    array(591.82315261)

• Indexes: (0)
• Attributes: (0)

Pair plot¶

This visualization is useful to see the correlation between the parameters. The ideal situation is when the parameters are uncorrelated, which means that the posterior distribution is close to a multivariate Gaussian distribution.

In [9]:

with az.style.context("arviz-darkgrid", after_reset=True):
    az.plot_pair(result_nuts.inference_data)
    
plt.show()

with az.style.context("arviz-darkgrid", after_reset=True): az.plot_pair(result_nuts.inference_data) plt.show()

Take a look at arviz's documentation to see what else you can do with this library.