EmissionAbsorptionPhysicalModel Tutorial

Trey V. Wenger (c) July 2025

EmissionAbsorptionPhysicalModel is a re-parameterization of EmissionAbsorptionModel that places more physically motivated constraints on the parameter space.

[1]:
# General imports
import time

import matplotlib.pyplot as plt
import arviz as az
import pandas as pd
import numpy as np
import pymc as pm

print("pymc version:", pm.__version__)
print("arviz version:", az.__version__)

import bayes_spec
print("bayes_spec version:", bayes_spec.__version__)

import caribou_hi
print("caribou_hi version:", caribou_hi.__version__)

# Notebook configuration
pd.options.display.max_rows = None
pymc version: 5.22.0
arviz version: 0.22.0dev
bayes_spec version: 1.9.0
caribou_hi version: 4.1.0+3.g94a913e.dirty

Simulating Data

[2]:
from bayes_spec import SpecData

# spectral axes definitions
emission_axis = np.linspace(-60.0, 60.0, 200)  # km s-1
absorption_axis = np.linspace(-30.0, 30.0, 100)  # km s-1

# data noise can either be a scalar (assumed constant noise across the spectrum)
# or an array of the same length as the data
rms_emission = 0.1  # K
rms_absorption = 0.001  # 1 - exp(-tau)

# brightness data. In this case, we just throw in some random data for now
# since we are only doing this in order to simulate some actual data.
emission = rms_emission * np.random.randn(len(emission_axis))
absorption = rms_absorption * np.random.randn(len(absorption_axis))

dummy_data = {
    "emission": SpecData(
        emission_axis,
        emission,
        rms_emission,
        xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
        ylabel=r"$T_B$ (K)",
    ),
    "absorption": SpecData(
        absorption_axis,
        absorption,
        rms_absorption,
        xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
        ylabel=r"1 - exp(-$\tau$)",
    ),
}

# Plot dummy data
fig, axes = plt.subplots(2, layout="constrained", sharex=True)
axes[0].plot(dummy_data["emission"].spectral, dummy_data["emission"].brightness, "k-")
axes[0].plot(dummy_data["emission"].spectral, dummy_data["emission"].noise, "r-")
axes[1].plot(dummy_data["absorption"].spectral, dummy_data["absorption"].brightness, "k-", label="Data")
axes[1].plot(dummy_data["absorption"].spectral, dummy_data["absorption"].noise, "r-", label="Noise")
axes[1].set_xlabel(dummy_data["emission"].xlabel)
axes[0].set_ylabel(dummy_data["emission"].ylabel)
axes[1].set_ylabel(dummy_data["absorption"].ylabel)
_ = axes[1].legend(loc="lower left")
../_images/notebooks_emission_absorption_physical_model_3_0.png
[3]:
from caribou_hi import EmissionAbsorptionPhysicalModel

# Initialize and define the model
n_clouds = 3
baseline_degree = 0
model = EmissionAbsorptionPhysicalModel(
    dummy_data,
    n_clouds=n_clouds,
    baseline_degree=baseline_degree,
    bg_temp = 3.77, # assumed background temperature (K)
    seed=1234,
    verbose=True
)
model.add_priors(
    prior_filling_factor=[2.0, 1.0],  # filling factor prior shape
    prior_ff_NHI=1.0e21,  # filling factor * column density prior width (cm-2)
    prior_fwhm2_thermal_fraction=[2.0, 2.0], # thermal FWHM^2 fraction prior shape
    prior_sigma_log10_NHI=0.5, # log-normal column density distribution width
    prior_fwhm2=200.0, # FWHM^2 prior width (km2 s-2)
    prior_velocity=[-15.0, 15.0],  # lower and upper limit of velocity prior (km/s)
    prior_log10_n_alpha=[-6.0, 1.0],  # log10(n_alpha) prior width (cm-3)
    prior_nth_fwhm_1pc=[1.75, 0.25], # non-thermal FWHM at 1 pc prior mean and width (km s-1)
    prior_depth_nth_fwhm_power=[0.3, 0.1], # non-thermal FWHM vs. depth power prior mean and width
    prior_fwhm_L=None, # Assume Gaussian line profile
    prior_baseline_coeffs=None, # Default baseline priors
)
model.add_likelihood()
[4]:
from caribou_hi import physics

# Simulation parameters
filling_factor = np.array([0.2, 0.8, 1.0])
absorption_weight = np.array([0.9, 0.8, 1.0])
log10_nHI = np.array([1.25, 0.0, -0.5])
tkin = np.array([50.0, 300.0, 5000.0])
log10_n_alpha = np.array([-5.5, -6.0, -6.5])
depth = np.array([5.0, 25.0, 300.0])
nth_fwhm_1pc = np.array([2.0, 1.75, 1.5])
depth_nth_fwhm_power = np.array([0.2, 0.3, 0.4])

tspin = physics.calc_spin_temp(tkin, 10.0**log10_nHI, 10.0**log10_n_alpha).eval()
print("tspin", tspin)

fwhm2_nonthermal = physics.calc_nonthermal_fwhm(depth, nth_fwhm_1pc, depth_nth_fwhm_power)**2.0
print("fwhm2_nonthermal", fwhm2_nonthermal)

fwhm2_thermal = physics.calc_thermal_fwhm2(tkin)
print("fwhm2_thermal", fwhm2_thermal)

fwhm2 = fwhm2_nonthermal + fwhm2_thermal
print("fwhm2", fwhm2)

fwhm2_thermal_fraction = fwhm2_thermal/fwhm2
print("fwhm2_thermal_fraction", fwhm2_thermal_fraction)

log10_Pth = np.log10(tkin) + log10_nHI
print("log10_Pth", log10_Pth)

log10_NHI = log10_nHI + np.log10(depth) + 18.489351
print("log10_NHI", log10_NHI)

tau_total = physics.calc_tau_total(10.0**log10_NHI, tspin)
print("tau_total", tau_total)

wt_ff_tkin = absorption_weight / filling_factor / tkin
print("wt_ff_tkin", wt_ff_tkin)

sim_params = {
    "wt_ff_tkin": wt_ff_tkin,
    "absorption_weight": absorption_weight,
    "filling_factor": filling_factor,
    "log10_NHI": log10_NHI,
    "tau_total": tau_total,
    "tspin": tspin,
    "fwhm2": fwhm2,
    "fwhm2_thermal_fraction": fwhm2_thermal_fraction,
    "velocity": np.array([-5.0, 5.0, 10.0]),
    "log10_n_alpha": log10_n_alpha,
    "nth_fwhm_1pc": nth_fwhm_1pc,
    "depth_nth_fwhm_power": depth_nth_fwhm_power,
    "baseline_absorption_norm": [0.0],
}

# add derived quantities to sim_params
for key in model.cloud_deterministics:
    if key not in sim_params.keys():
        sim_params[key] = model.model[key].eval(sim_params, on_unused_input="ignore")

# Evaluate and save simulated observation
emission = model.model["emission"].eval(sim_params, on_unused_input="ignore")
absorption = model.model["absorption"].eval(sim_params, on_unused_input="ignore")

data = {
    "emission": SpecData(
        emission_axis,
        emission,
        rms_emission,
        xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
        ylabel=r"$T_B$ (K)",
    ),
    "absorption": SpecData(
        absorption_axis,
        absorption,
        rms_absorption,
        xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
        ylabel=r"1 - exp(-$\tau$)",
    ),
}
tspin [  49.9920589   297.73994712 2795.52422645]
fwhm2_nonthermal [  7.61461575  21.12711044 215.71459099]
fwhm2_thermal [  2.2876605  13.725963  228.76605  ]
fwhm2 [  9.90227625  34.85307344 444.48064099]
fwhm2_thermal_fraction [0.2310237 0.3938236 0.5146817]
log10_Pth [2.94897    2.47712125 3.19897   ]
log10_NHI [20.438321   19.88729101 20.46647225]
tau_total [3.01140471 0.14216839 0.05745901]
wt_ff_tkin [0.09       0.00333333 0.0002    ]
[5]:
sim_params
[5]:
{'wt_ff_tkin': array([0.09      , 0.00333333, 0.0002    ]),
 'absorption_weight': array([0.9, 0.8, 1. ]),
 'filling_factor': array([0.2, 0.8, 1. ]),
 'log10_NHI': array([20.438321  , 19.88729101, 20.46647225]),
 'tau_total': array([3.01140471, 0.14216839, 0.05745901]),
 'tspin': array([  49.9920589 ,  297.73994712, 2795.52422645]),
 'fwhm2': array([  9.90227625,  34.85307344, 444.48064099]),
 'fwhm2_thermal_fraction': array([0.2310237, 0.3938236, 0.5146817]),
 'velocity': array([-5.,  5., 10.]),
 'log10_n_alpha': array([-5.5, -6. , -6.5]),
 'nth_fwhm_1pc': array([2.  , 1.75, 1.5 ]),
 'depth_nth_fwhm_power': array([0.2, 0.3, 0.4]),
 'baseline_absorption_norm': [0.0],
 'fwhm2_thermal': array([  2.2876605,  13.725963 , 228.76605  ]),
 'tkin': array([  50.,  300., 5000.]),
 'fwhm2_nonthermal': array([  7.61461575,  21.12711044, 215.71459099]),
 'log10_depth': array([0.69897   , 1.39794001, 2.47712125]),
 'log10_nHI': array([ 1.25,  0.  , -0.5 ]),
 'log10_Pth': array([2.94897   , 2.47712125, 3.19897   ]),
 'log10_Pnth': array([4.95888799, 4.1520801 , 4.66111952]),
 'log10_Ptot': array([4.96311227, 4.16116407, 4.67585106])}
[6]:
# Plot data
fig, axes = plt.subplots(2, layout="constrained", sharex=True)
axes[0].plot(data["emission"].spectral, data["emission"].brightness, "k-")
axes[0].plot(data["emission"].spectral, data["emission"].noise, "r-")
axes[1].plot(data["absorption"].spectral, data["absorption"].brightness, "k-", label="Data")
axes[1].plot(data["absorption"].spectral, data["absorption"].noise, "r-", label="Noise")
axes[1].set_xlabel(data["emission"].xlabel)
axes[0].set_ylabel(data["emission"].ylabel)
axes[1].set_ylabel(data["absorption"].ylabel)
_ = axes[1].legend(loc="upper left")
../_images/notebooks_emission_absorption_physical_model_7_0.png

Model Definition

[8]:
# Initialize and define the model
model = EmissionAbsorptionPhysicalModel(
    data,
    n_clouds=n_clouds,
    baseline_degree=baseline_degree,
    bg_temp = 3.77, # assumed background temperature (K)
    seed=1234,
    verbose=True
)
model.add_priors(
    prior_filling_factor=[2.0, 1.0],  # filling factor prior shape
    prior_ff_NHI=1.0e21,  # filling factor * column density prior width (cm-2)
    prior_fwhm2_thermal_fraction=[2.0, 2.0], # thermal FWHM^2 fraction prior shape
    prior_sigma_log10_NHI=0.5, # log-normal column density distribution width
    prior_fwhm2=200.0, # FWHM^2 prior width (km2 s-2)
    prior_velocity=[-20.0, 20.0],  # lower and upper limit of velocity prior (km/s)
    prior_log10_n_alpha=[-6.0, 1.0],  # log10(n_alpha) prior width (cm-3)
    prior_nth_fwhm_1pc=[1.75, 0.25], # non-thermal FWHM at 1 pc prior mean and width (km s-1)
    prior_depth_nth_fwhm_power=[0.3, 0.1], # non-thermal FWHM vs. depth power prior mean and width
    prior_fwhm_L=None, # Assume Gaussian line profile
    prior_baseline_coeffs=None, # Default baseline priors
)
model.add_likelihood()
[9]:
# Plot model graph
model.graph().render("emission_absorption_physical_model", format="png")
model.graph()
[9]:
../_images/notebooks_emission_absorption_physical_model_10_0.svg
[10]:
# model string representation
print(model.model.str_repr())
  baseline_emission_norm ~ Normal(0, 1)
baseline_absorption_norm ~ Normal(0, 1)
              fwhm2_norm ~ Gamma(0.5, f())
           velocity_norm ~ Beta(2, 2)
      log10_n_alpha_norm ~ Normal(0, 1)
            nth_fwhm_1pc ~ TruncatedNormal(1.75, 0.25, 0, inf)
    depth_nth_fwhm_power ~ InverseGamma(11, 3)
             ff_NHI_norm ~ HalfNormal(0, 1)
  fwhm2_thermal_fraction ~ Beta(2, 2)
          filling_factor ~ Beta(2, 1)
              wt_ff_tkin ~ LogNormal(f(fwhm2_thermal_fraction, fwhm2_norm), 1.15)
                   fwhm2 ~ Deterministic(f(fwhm2_norm))
                velocity ~ Deterministic(f(velocity_norm))
           log10_n_alpha ~ Deterministic(f(log10_n_alpha_norm))
           fwhm2_thermal ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm))
                    tkin ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm))
        fwhm2_nonthermal ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm))
             log10_depth ~ Deterministic(f(depth_nth_fwhm_power, nth_fwhm_1pc, fwhm2_thermal_fraction, fwhm2_norm))
               log10_NHI ~ Deterministic(f(filling_factor, ff_NHI_norm))
               log10_nHI ~ Deterministic(f(filling_factor, depth_nth_fwhm_power, nth_fwhm_1pc, ff_NHI_norm, fwhm2_thermal_fraction, fwhm2_norm))
                   tspin ~ Deterministic(f(fwhm2_thermal_fraction, log10_n_alpha_norm, fwhm2_norm, filling_factor, depth_nth_fwhm_power, nth_fwhm_1pc, ff_NHI_norm))
       absorption_weight ~ Deterministic(f(filling_factor, wt_ff_tkin, fwhm2_thermal_fraction, fwhm2_norm))
               tau_total ~ Deterministic(f(filling_factor, ff_NHI_norm, wt_ff_tkin, fwhm2_thermal_fraction, log10_n_alpha_norm, fwhm2_norm, depth_nth_fwhm_power, nth_fwhm_1pc))
               log10_Pth ~ Deterministic(f(fwhm2_thermal_fraction, filling_factor, depth_nth_fwhm_power, nth_fwhm_1pc, ff_NHI_norm, fwhm2_norm))
              log10_Pnth ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm, filling_factor, depth_nth_fwhm_power, nth_fwhm_1pc, ff_NHI_norm))
              log10_Ptot ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm, filling_factor, depth_nth_fwhm_power, nth_fwhm_1pc, ff_NHI_norm))
              absorption ~ Normal(f(baseline_absorption_norm, filling_factor, wt_ff_tkin, ff_NHI_norm, fwhm2_thermal_fraction, velocity_norm, fwhm2_norm, log10_n_alpha_norm, depth_nth_fwhm_power, nth_fwhm_1pc), <constant>)
                emission ~ Normal(f(filling_factor, baseline_emission_norm, fwhm2_thermal_fraction, ff_NHI_norm, wt_ff_tkin, log10_n_alpha_norm, fwhm2_norm, velocity_norm, depth_nth_fwhm_power, nth_fwhm_1pc), <constant>)
[11]:
from bayes_spec.plots import plot_predictive

# prior predictive check
prior = model.sample_prior_predictive(
    samples=1000,  # prior predictive samples
)
axes = plot_predictive(model.data, prior.prior_predictive.sel(draw=slice(None, None, 20)))
axes.ravel()[1].sharex(axes.ravel()[0])
Sampling: [absorption, baseline_absorption_norm, baseline_emission_norm, depth_nth_fwhm_power, emission, ff_NHI_norm, filling_factor, fwhm2_norm, fwhm2_thermal_fraction, log10_n_alpha_norm, nth_fwhm_1pc, velocity_norm, wt_ff_tkin]
../_images/notebooks_emission_absorption_physical_model_12_1.png
[12]:
print(model.cloud_freeRVs)
print(model.cloud_deterministics)
['fwhm2_norm', 'velocity_norm', 'log10_n_alpha_norm', 'nth_fwhm_1pc', 'depth_nth_fwhm_power', 'ff_NHI_norm', 'fwhm2_thermal_fraction', 'filling_factor', 'wt_ff_tkin']
['fwhm2', 'velocity', 'log10_n_alpha', 'fwhm2_thermal', 'tkin', 'fwhm2_nonthermal', 'log10_depth', 'log10_NHI', 'log10_nHI', 'tspin', 'absorption_weight', 'tau_total', 'log10_Pth', 'log10_Pnth', 'log10_Ptot']
[13]:
from bayes_spec.plots import plot_pair

var_names = model.cloud_deterministics + [p for p in model.cloud_freeRVs if "_norm" not in p]
print(var_names)
_ = plot_pair(
    prior.prior, # samples
    var_names, # var_names to plot
    combine_dims=["cloud"], # concatenate clouds
    labeller=model.labeller, # label manager
    kind="scatter", # plot type
    reference_values=sim_params, # truths
)
['fwhm2', 'velocity', 'log10_n_alpha', 'fwhm2_thermal', 'tkin', 'fwhm2_nonthermal', 'log10_depth', 'log10_NHI', 'log10_nHI', 'tspin', 'absorption_weight', 'tau_total', 'log10_Pth', 'log10_Pnth', 'log10_Ptot', 'nth_fwhm_1pc', 'depth_nth_fwhm_power', 'fwhm2_thermal_fraction', 'filling_factor', 'wt_ff_tkin']
../_images/notebooks_emission_absorption_physical_model_14_1.png

Variational Inference

[14]:
start = time.time()
model.fit(
    n = 1_000_000, # maximum number of VI iterations
    draws = 1_000, # number of posterior samples
    rel_tolerance = 0.005, # VI relative convergence threshold
    abs_tolerance = 0.005, # VI absolute convergence threshold
    learning_rate = 0.001, # VI learning rate
    start = {"velocity_norm": np.linspace(0.1, 0.9, model.n_clouds)},
)
end = time.time()
print(f"Runtime: {(end-start)/60.0:.2f} minutes")
Convergence achieved at 101700
Interrupted at 101,699 [10%]: Average Loss = 2.6359e+05
Adding log-likelihood to trace
Runtime: 2.09 minutes
[15]:
pm.summary(model.trace.posterior)
arviz - WARNING - Shape validation failed: input_shape: (1, 1000), minimum_shape: (chains=2, draws=4)
[15]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
baseline_emission_norm[0] 0.635 0.073 0.503 0.770 0.002 0.002 960.0 1026.0 NaN
baseline_absorption_norm[0] 0.218 0.102 0.038 0.411 0.003 0.002 1023.0 1033.0 NaN
log10_n_alpha_norm[0] 0.986 0.472 0.075 1.813 0.016 0.011 859.0 880.0 NaN
log10_n_alpha_norm[1] -2.554 0.001 -2.556 -2.553 0.000 0.000 922.0 979.0 NaN
log10_n_alpha_norm[2] -2.499 0.017 -2.527 -2.466 0.001 0.000 981.0 983.0 NaN
fwhm2_norm[0] 2.222 0.014 2.191 2.246 0.000 0.000 988.0 1031.0 NaN
fwhm2_norm[1] 0.049 0.000 0.049 0.049 0.000 0.000 832.0 859.0 NaN
fwhm2_norm[2] 0.171 0.002 0.167 0.174 0.000 0.000 865.0 984.0 NaN
velocity_norm[0] 0.750 0.001 0.749 0.752 0.000 0.000 986.0 1014.0 NaN
velocity_norm[1] 0.375 0.000 0.375 0.375 0.000 0.000 971.0 944.0 NaN
velocity_norm[2] 0.626 0.001 0.624 0.627 0.000 0.000 931.0 741.0 NaN
nth_fwhm_1pc[0] 1.745 0.252 1.295 2.223 0.008 0.007 1103.0 752.0 NaN
nth_fwhm_1pc[1] 0.092 0.021 0.057 0.132 0.001 0.001 971.0 845.0 NaN
nth_fwhm_1pc[2] 1.205 0.016 1.177 1.237 0.001 0.000 1018.0 1013.0 NaN
depth_nth_fwhm_power[0] 0.316 0.101 0.148 0.501 0.003 0.003 890.0 944.0 NaN
depth_nth_fwhm_power[1] 0.155 0.013 0.131 0.178 0.000 0.000 877.0 1023.0 NaN
depth_nth_fwhm_power[2] 0.148 0.002 0.145 0.151 0.000 0.000 999.0 983.0 NaN
ff_NHI_norm[0] 0.293 0.001 0.292 0.295 0.000 0.000 846.0 794.0 NaN
ff_NHI_norm[1] 0.053 0.000 0.053 0.053 0.000 0.000 788.0 940.0 NaN
ff_NHI_norm[2] 0.063 0.000 0.062 0.064 0.000 0.000 982.0 810.0 NaN
fwhm2_thermal_fraction[0] 0.444 0.160 0.133 0.717 0.006 0.003 832.0 936.0 NaN
fwhm2_thermal_fraction[1] 0.831 0.001 0.829 0.833 0.000 0.000 1026.0 909.0 NaN
fwhm2_thermal_fraction[2] 0.524 0.028 0.475 0.577 0.001 0.001 1037.0 865.0 NaN
filling_factor[0] 0.702 0.217 0.301 0.995 0.007 0.005 854.0 905.0 NaN
filling_factor[1] 0.113 0.003 0.107 0.118 0.000 0.000 909.0 900.0 NaN
filling_factor[2] 0.907 0.058 0.792 0.988 0.002 0.002 782.0 944.0 NaN
wt_ff_tkin[0] 0.000 0.000 0.000 0.000 0.000 0.000 1033.0 907.0 NaN
wt_ff_tkin[1] 0.039 0.000 0.039 0.039 0.000 0.000 969.0 781.0 NaN
wt_ff_tkin[2] 0.001 0.000 0.001 0.001 0.000 0.000 973.0 952.0 NaN
fwhm2[0] 444.418 2.798 438.110 449.103 0.089 0.071 988.0 1031.0 NaN
fwhm2[1] 9.872 0.013 9.846 9.896 0.000 0.000 832.0 859.0 NaN
fwhm2[2] 34.133 0.384 33.408 34.824 0.013 0.008 865.0 984.0 NaN
velocity[0] 10.009 0.035 9.943 10.071 0.001 0.001 986.0 1014.0 NaN
velocity[1] -5.001 0.002 -5.005 -4.998 0.000 0.000 971.0 944.0 NaN
velocity[2] 5.020 0.024 4.977 5.065 0.001 0.001 931.0 741.0 NaN
log10_n_alpha[0] -5.014 0.472 -5.925 -4.187 0.016 0.011 859.0 880.0 NaN
log10_n_alpha[1] -8.554 0.001 -8.556 -8.553 0.000 0.000 922.0 979.0 NaN
log10_n_alpha[2] -8.499 0.017 -8.527 -8.466 0.001 0.000 981.0 983.0 NaN
fwhm2_thermal[0] 197.143 71.289 59.451 319.743 2.488 1.434 828.0 936.0 NaN
fwhm2_thermal[1] 8.203 0.016 8.175 8.235 0.001 0.000 924.0 717.0 NaN
fwhm2_thermal[2] 17.899 0.971 16.026 19.702 0.031 0.022 953.0 905.0 NaN
tkin[0] 4308.829 1558.129 1299.384 6988.426 54.369 31.347 828.0 936.0 NaN
tkin[1] 179.285 0.345 178.673 179.978 0.011 0.007 924.0 717.0 NaN
tkin[2] 391.200 21.216 350.266 430.624 0.688 0.489 953.0 905.0 NaN
fwhm2_nonthermal[0] 247.275 71.358 127.325 387.729 2.474 1.454 839.0 937.0 NaN
fwhm2_nonthermal[1] 1.669 0.012 1.646 1.691 0.000 0.000 1018.0 1013.0 NaN
fwhm2_nonthermal[2] 16.235 0.958 14.517 18.031 0.028 0.022 1130.0 816.0 NaN
log10_depth[0] 3.312 1.104 1.433 5.167 0.038 0.035 860.0 983.0 NaN
log10_depth[1] 7.527 0.930 5.746 9.187 0.031 0.021 922.0 823.0 NaN
log10_depth[2] 3.550 0.103 3.363 3.742 0.003 0.002 1088.0 1026.0 NaN
log10_NHI[0] 20.651 0.186 20.468 20.988 0.006 0.009 854.0 900.0 NaN
log10_NHI[1] 20.674 0.011 20.653 20.696 0.000 0.000 893.0 898.0 NaN
log10_NHI[2] 19.842 0.030 19.804 19.904 0.001 0.001 803.0 850.0 NaN
log10_nHI[0] -1.150 1.125 -3.245 0.664 0.040 0.035 796.0 944.0 NaN
log10_nHI[1] -5.342 0.930 -6.999 -3.552 0.031 0.021 922.0 791.0 NaN
log10_nHI[2] -2.198 0.107 -2.390 -1.998 0.003 0.002 1088.0 942.0 NaN
tspin[0] 3988.260 1390.017 1321.480 6421.729 50.106 28.833 782.0 855.0 NaN
tspin[1] 75.206 0.139 74.985 75.412 0.005 0.012 836.0 869.0 NaN
tspin[2] 101.201 5.816 91.892 112.518 0.176 0.179 1087.0 890.0 NaN
absorption_weight[0] 0.730 0.363 0.098 1.376 0.011 0.008 1018.0 1023.0 NaN
absorption_weight[1] 0.787 0.021 0.749 0.826 0.001 0.000 899.0 901.0 NaN
absorption_weight[2] 0.308 0.027 0.262 0.364 0.001 0.001 965.0 932.0 NaN
tau_total[0] 0.083 0.103 0.024 0.183 0.003 0.027 1004.0 942.0 NaN
tau_total[1] 3.446 0.091 3.266 3.611 0.003 0.002 893.0 898.0 NaN
tau_total[2] 0.378 0.030 0.326 0.435 0.001 0.001 866.0 984.0 NaN
log10_Pth[0] 2.450 1.170 0.417 4.466 0.043 0.035 721.0 982.0 NaN
log10_Pth[1] -3.089 0.930 -4.745 -1.298 0.031 0.021 922.0 791.0 NaN
log10_Pth[2] 0.394 0.125 0.148 0.610 0.004 0.003 1126.0 850.0 NaN
log10_Pnth[0] 4.050 1.108 1.952 5.757 0.038 0.035 849.0 944.0 NaN
log10_Pnth[1] -2.292 0.930 -3.952 -0.508 0.031 0.021 921.0 791.0 NaN
log10_Pnth[2] 1.839 0.087 1.673 1.992 0.003 0.002 1061.0 945.0 NaN
log10_Ptot[0] 4.063 1.108 1.955 5.781 0.038 0.035 846.0 944.0 NaN
log10_Ptot[1] -2.228 0.930 -3.886 -0.443 0.031 0.021 921.0 791.0 NaN
log10_Ptot[2] 1.854 0.088 1.686 2.009 0.003 0.002 1062.0 945.0 NaN
[16]:
posterior = model.sample_posterior_predictive(
    thin=10,  # keep one in {thin} posterior samples
)
axes = plot_predictive(model.data, posterior.posterior_predictive)
axes.ravel()[1].sharex(axes.ravel()[0])
Sampling: [absorption, emission]
../_images/notebooks_emission_absorption_physical_model_18_3.png

Posterior Sampling: MCMC

We can sample from the posterior distribution using MCMC.

[17]:
start = time.time()
model.sample(
    init="advi+adapt_diag",  # initialization strategy
    tune=1000,  # tuning samples
    draws=1000,  # posterior samples
    chains=8,  # number of independent chains
    cores=8,  # number of parallel chains
    init_kwargs={
        "rel_tolerance": 0.005,
        "abs_tolerance": 0.005,
        "learning_rate": 0.001,
        "start": {"velocity_norm": np.linspace(0.1, 0.9, model.n_clouds)},
    },  # VI initialization arguments
    nuts_kwargs={"target_accept": 0.9},  # NUTS arguments
)
end = time.time()
print(f"Runtime: {(end-start)/60.0:.2f} minutes")
Initializing NUTS using custom advi+adapt_diag strategy
Convergence achieved at 101700
Interrupted at 101,699 [10%]: Average Loss = 2.6359e+05
Multiprocess sampling (8 chains in 8 jobs)
NUTS: [baseline_emission_norm, baseline_absorption_norm, fwhm2_norm, velocity_norm, log10_n_alpha_norm, nth_fwhm_1pc, depth_nth_fwhm_power, ff_NHI_norm, fwhm2_thermal_fraction, filling_factor, wt_ff_tkin]
Sampling 8 chains for 1_000 tune and 1_000 draw iterations (8_000 + 8_000 draws total) took 216 seconds.
Adding log-likelihood to trace
Runtime: 5.99 minutes
[18]:
model.solve(
    init_params="random_from_data", # GMM initialization strategy
    n_init=10, # number of GMM initilizations
    max_iter=1_000, # maximum number of GMM iterations
    kl_div_threshold=0.1, # covergence threshold
)
GMM converged to unique solution

Check that the effective sample sizes are large and the covergence statistic r_hat is close to 1! If not, you may have to increase the number of tuning steps (tune=2000) or the NUTS acceptance rate (target_accept=0.9).

[19]:
print("solutions:", model.solutions)
az.summary(model.trace["solution_0"])
# this also works: az.summary(model.trace.solution_0)
solutions: [0]
[19]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
baseline_emission_norm[0] 0.632 0.093 0.454 0.802 0.001 0.001 10573.0 6120.0 1.0
baseline_absorption_norm[0] 0.206 0.151 -0.066 0.493 0.001 0.002 12419.0 5831.0 1.0
log10_n_alpha_norm[0] -0.014 0.938 -1.653 1.903 0.011 0.010 6555.0 5872.0 1.0
log10_n_alpha_norm[1] 0.002 1.013 -1.810 2.047 0.008 0.013 15903.0 5365.0 1.0
log10_n_alpha_norm[2] -0.005 1.003 -1.802 1.987 0.011 0.013 8332.0 4557.0 1.0
fwhm2_norm[0] 2.227 0.026 2.179 2.277 0.000 0.000 7591.0 6015.0 1.0
fwhm2_norm[1] 0.049 0.000 0.049 0.050 0.000 0.000 13524.0 5739.0 1.0
fwhm2_norm[2] 0.171 0.005 0.161 0.181 0.000 0.000 4797.0 5719.0 1.0
velocity_norm[0] 0.750 0.002 0.746 0.753 0.000 0.000 4193.0 5280.0 1.0
velocity_norm[1] 0.375 0.000 0.375 0.375 0.000 0.000 14306.0 6017.0 1.0
velocity_norm[2] 0.626 0.001 0.624 0.627 0.000 0.000 9169.0 6821.0 1.0
nth_fwhm_1pc[0] 1.749 0.251 1.311 2.245 0.002 0.003 13461.0 4753.0 1.0
nth_fwhm_1pc[1] 1.750 0.256 1.282 2.234 0.002 0.003 14814.0 4738.0 1.0
nth_fwhm_1pc[2] 1.753 0.250 1.295 2.229 0.002 0.003 12785.0 5923.0 1.0
depth_nth_fwhm_power[0] 0.301 0.096 0.151 0.479 0.001 0.002 9668.0 5919.0 1.0
depth_nth_fwhm_power[1] 0.300 0.100 0.140 0.479 0.001 0.002 13446.0 5017.0 1.0
depth_nth_fwhm_power[2] 0.302 0.106 0.156 0.500 0.001 0.004 9994.0 5095.0 1.0
ff_NHI_norm[0] 0.294 0.003 0.289 0.299 0.000 0.000 3409.0 3438.0 1.0
ff_NHI_norm[1] 0.064 0.007 0.052 0.077 0.000 0.000 3241.0 5011.0 1.0
ff_NHI_norm[2] 0.060 0.002 0.057 0.063 0.000 0.000 4061.0 4727.0 1.0
fwhm2_thermal_fraction[0] 0.486 0.214 0.103 0.858 0.002 0.002 10281.0 5371.0 1.0
fwhm2_thermal_fraction[1] 0.099 0.061 0.050 0.195 0.001 0.002 3375.0 4162.0 1.0
fwhm2_thermal_fraction[2] 0.493 0.202 0.155 0.869 0.002 0.002 9290.0 5340.0 1.0
filling_factor[0] 0.670 0.232 0.252 1.000 0.002 0.002 9123.0 4768.0 1.0
filling_factor[1] 0.608 0.242 0.210 0.999 0.004 0.002 3390.0 4377.0 1.0
filling_factor[2] 0.709 0.209 0.331 1.000 0.002 0.002 9885.0 3941.0 1.0
wt_ff_tkin[0] 0.000 0.000 0.000 0.000 0.000 0.000 5141.0 4029.0 1.0
wt_ff_tkin[1] 0.078 0.008 0.063 0.093 0.000 0.000 3239.0 5043.0 1.0
wt_ff_tkin[2] 0.003 0.000 0.003 0.004 0.000 0.000 5549.0 3339.0 1.0
fwhm2[0] 445.337 5.218 435.831 455.311 0.060 0.054 7591.0 6015.0 1.0
fwhm2[1] 9.869 0.026 9.822 9.918 0.000 0.000 13524.0 5739.0 1.0
fwhm2[2] 34.143 1.072 32.152 36.141 0.015 0.010 4797.0 5719.0 1.0
velocity[0] 9.991 0.074 9.854 10.131 0.001 0.001 4193.0 5280.0 1.0
velocity[1] -5.001 0.002 -5.004 -4.997 0.000 0.000 14306.0 6017.0 1.0
velocity[2] 5.021 0.029 4.968 5.075 0.000 0.000 9169.0 6821.0 1.0
log10_n_alpha[0] -6.014 0.938 -7.653 -4.097 0.011 0.010 6555.0 5872.0 1.0
log10_n_alpha[1] -5.998 1.013 -7.810 -3.953 0.008 0.013 15903.0 5365.0 1.0
log10_n_alpha[2] -6.005 1.003 -7.802 -4.013 0.011 0.013 8332.0 4557.0 1.0
fwhm2_thermal[0] 216.621 95.502 45.590 382.470 0.927 0.900 10364.0 5447.0 1.0
fwhm2_thermal[1] 0.975 0.600 0.500 1.934 0.011 0.023 3374.0 4162.0 1.0
fwhm2_thermal[2] 16.843 6.974 5.094 29.648 0.071 0.069 9506.0 5186.0 1.0
tkin[0] 4734.544 2087.323 996.424 8359.422 20.254 19.672 10364.0 5447.0 1.0
tkin[1] 21.320 13.110 10.924 42.272 0.234 0.502 3374.0 4162.0 1.0
tkin[2] 368.126 152.436 111.334 648.009 1.550 1.517 9506.0 5186.0 1.0
fwhm2_nonthermal[0] 228.716 95.238 63.970 399.158 0.926 0.893 10296.0 5535.0 1.0
fwhm2_nonthermal[1] 8.894 0.600 7.936 9.393 0.011 0.023 3387.0 4196.0 1.0
fwhm2_nonthermal[2] 17.300 6.849 4.661 28.882 0.072 0.065 8852.0 5429.0 1.0
log10_depth[0] 3.323 1.102 1.427 5.426 0.011 0.015 9036.0 6257.0 1.0
log10_depth[1] 0.864 0.369 0.225 1.519 0.003 0.006 12546.0 5200.0 1.0
log10_depth[2] 1.294 0.642 0.131 2.539 0.008 0.010 7301.0 5212.0 1.0
log10_NHI[0] 20.681 0.209 20.462 21.068 0.003 0.005 9115.0 4666.0 1.0
log10_NHI[1] 20.069 0.187 19.813 20.429 0.003 0.003 3626.0 4506.0 1.0
log10_NHI[2] 19.955 0.165 19.763 20.269 0.002 0.003 9783.0 4024.0 1.0
log10_nHI[0] -1.132 1.118 -3.290 0.816 0.012 0.015 9310.0 6250.0 1.0
log10_nHI[1] 0.715 0.437 -0.113 1.535 0.005 0.006 6986.0 5607.0 1.0
log10_nHI[2] 0.172 0.674 -1.108 1.412 0.008 0.010 7558.0 5945.0 1.0
tspin[0] 3183.742 1795.610 363.277 6556.526 22.032 19.918 5991.0 4379.0 1.0
tspin[1] 21.309 13.103 11.047 42.405 0.234 0.502 3374.0 4162.0 1.0
tspin[2] 362.892 151.354 109.553 642.948 1.552 1.520 9395.0 5177.0 1.0
absorption_weight[0] 0.551 0.393 0.019 1.265 0.005 0.006 6141.0 3672.0 1.0
absorption_weight[1] 0.809 0.102 0.631 1.002 0.001 0.001 6488.0 6282.0 1.0
absorption_weight[2] 0.900 0.472 0.149 1.767 0.005 0.005 9086.0 5117.0 1.0
tau_total[0] 0.143 0.186 0.019 0.368 0.004 0.025 5725.0 3797.0 1.0
tau_total[1] 3.398 0.411 2.574 4.126 0.005 0.004 6485.0 6241.0 1.0
tau_total[2] 0.181 0.130 0.048 0.414 0.002 0.003 9018.0 5261.0 1.0
log10_Pth[0] 2.489 1.209 0.080 4.563 0.012 0.016 9911.0 6462.0 1.0
log10_Pth[1] 1.996 0.560 0.937 3.081 0.008 0.008 5024.0 4814.0 1.0
log10_Pth[2] 2.694 0.806 1.163 4.184 0.009 0.011 8349.0 5604.0 1.0
log10_Pnth[0] 4.002 1.055 2.052 5.907 0.011 0.014 9498.0 5995.0 1.0
log10_Pnth[1] 4.491 0.421 3.680 5.252 0.005 0.005 7594.0 5722.0 1.0
log10_Pnth[2] 4.188 0.549 3.152 5.173 0.006 0.007 8010.0 6003.0 1.0
log10_Ptot[0] 4.024 1.061 2.036 5.910 0.011 0.014 9435.0 6037.0 1.0
log10_Ptot[1] 4.492 0.421 3.681 5.257 0.005 0.005 7567.0 5722.0 1.0
log10_Ptot[2] 4.211 0.561 3.087 5.159 0.006 0.008 7816.0 6070.0 1.0

We generate posterior predictive checks as well as a trace plot of the individual chains. In the posterior predictive plot, we show each chain as a different color. Each line is one posterior sample.

[20]:
posterior = model.sample_posterior_predictive(
    thin=10,  # keep one in {thin} posterior samples
)
axes = plot_predictive(model.data, posterior.posterior_predictive)
axes.ravel()[0].sharex(axes.ravel()[1])
Sampling: [absorption, emission]
../_images/notebooks_emission_absorption_physical_model_25_3.png
[21]:
from bayes_spec.plots import plot_traces

_ = plot_traces(model.trace.posterior, model.cloud_freeRVs + model.baseline_freeRVs + model.hyper_freeRVs)
../_images/notebooks_emission_absorption_physical_model_26_0.png
[22]:
_ = plot_pair(
    model.trace.solution_0.sel(draw=slice(None, None, 10)), # samples
    model.cloud_freeRVs, # var_names to plot
    combine_dims=["cloud"], # concatenate clouds
    kind="scatter", # plot type
)
../_images/notebooks_emission_absorption_physical_model_27_0.png
[23]:
_ = plot_pair(
    model.trace.solution_0.sel(draw=slice(None, None, 10)), # samples
    ["velocity", "fwhm2"], # var_names to plot
    combine_dims=None, # do not concatenate clouds
    kind="scatter", # plot type
)
../_images/notebooks_emission_absorption_physical_model_28_0.png
[24]:
_ = plot_pair(
    model.trace.solution_0.sel(cloud=0, draw=slice(None, None, 10)), # samples
    model.cloud_freeRVs + model.hyper_freeRVs + model.baseline_freeRVs, # var_names to plot
    kind="scatter", # plot type
)
../_images/notebooks_emission_absorption_physical_model_29_0.png
[25]:
_ = plot_pair(
    model.trace.solution_0.sel(cloud=1, draw=slice(None, None, 10)), # samples
    model.cloud_freeRVs + model.hyper_freeRVs + model.baseline_freeRVs, # var_names to plot
    kind="scatter", # plot type
)
../_images/notebooks_emission_absorption_physical_model_30_0.png
[26]:
_ = plot_pair(
    model.trace.solution_0.sel(cloud=2, draw=slice(None, None, 10)), # samples
    model.cloud_freeRVs + model.hyper_freeRVs + model.baseline_freeRVs, # var_names to plot
    kind="scatter", # plot type
)
../_images/notebooks_emission_absorption_physical_model_31_0.png
[27]:
_ = plot_pair(
    model.trace.solution_0.sel(draw=slice(None, None, 10)), # samples
    var_names, # var_names to plot
    combine_dims=["cloud"], # concatenate clouds
    labeller=model.labeller, # label manager
    kind="scatter", # plot type
    reference_values=sim_params, # truths
)
../_images/notebooks_emission_absorption_physical_model_32_0.png
[28]:
# identify simulation cloud corresponding to each posterior cloud
sim_cloud_map = {}
for i in range(n_clouds):
    posterior_velocity = model.trace.solution_0['velocity'].sel(cloud=i).data.mean()
    match = np.argmin(np.abs(sim_params["velocity"] - posterior_velocity))
    sim_cloud_map[i] = match
sim_cloud_map
[28]:
{0: np.int64(2), 1: np.int64(0), 2: np.int64(1)}
[29]:
cloud = 0

# subset of sim_params
my_sim_params = {}
for var_name in var_names:
    my_sim_params[var_name] = sim_params[var_name][sim_cloud_map[cloud]]

_ = plot_pair(
    model.trace.solution_0.sel(cloud=cloud, draw=slice(None, None, 10)), # samples
    var_names, # var_names to plot
    labeller=model.labeller, # label manager
    kind="scatter", # plot type
    reference_values=my_sim_params, # truths
)
../_images/notebooks_emission_absorption_physical_model_34_0.png
[30]:
cloud = 1

# subset of sim_params
my_sim_params = {}
for var_name in var_names:
    my_sim_params[var_name] = sim_params[var_name][sim_cloud_map[cloud]]

_ = plot_pair(
    model.trace.solution_0.sel(cloud=cloud, draw=slice(None, None, 10)), # samples
    var_names, # var_names to plot
    labeller=model.labeller, # label manager
    kind="scatter", # plot type
    reference_values=my_sim_params, # truths
)
../_images/notebooks_emission_absorption_physical_model_35_0.png
[31]:
cloud = 2

# subset of sim_params
my_sim_params = {}
for var_name in var_names:
    my_sim_params[var_name] = sim_params[var_name][sim_cloud_map[cloud]]

_ = plot_pair(
    model.trace.solution_0.sel(cloud=cloud, draw=slice(None, None, 10)), # samples
    var_names, # var_names to plot
    labeller=model.labeller, # label manager
    kind="scatter", # plot type
    reference_values=my_sim_params, # truths
)
../_images/notebooks_emission_absorption_physical_model_36_0.png
[32]:
point_stats = az.summary(model.trace.solution_0, kind="stats")
print("BIC:", model.bic())
display(point_stats)
BIC: -1276.649801628776
mean sd hdi_3% hdi_97%
baseline_emission_norm[0] 0.632 0.093 0.454 0.802
baseline_absorption_norm[0] 0.206 0.151 -0.066 0.493
log10_n_alpha_norm[0] -0.014 0.938 -1.653 1.903
log10_n_alpha_norm[1] 0.002 1.013 -1.810 2.047
log10_n_alpha_norm[2] -0.005 1.003 -1.802 1.987
fwhm2_norm[0] 2.227 0.026 2.179 2.277
fwhm2_norm[1] 0.049 0.000 0.049 0.050
fwhm2_norm[2] 0.171 0.005 0.161 0.181
velocity_norm[0] 0.750 0.002 0.746 0.753
velocity_norm[1] 0.375 0.000 0.375 0.375
velocity_norm[2] 0.626 0.001 0.624 0.627
nth_fwhm_1pc[0] 1.749 0.251 1.311 2.245
nth_fwhm_1pc[1] 1.750 0.256 1.282 2.234
nth_fwhm_1pc[2] 1.753 0.250 1.295 2.229
depth_nth_fwhm_power[0] 0.301 0.096 0.151 0.479
depth_nth_fwhm_power[1] 0.300 0.100 0.140 0.479
depth_nth_fwhm_power[2] 0.302 0.106 0.156 0.500
ff_NHI_norm[0] 0.294 0.003 0.289 0.299
ff_NHI_norm[1] 0.064 0.007 0.052 0.077
ff_NHI_norm[2] 0.060 0.002 0.057 0.063
fwhm2_thermal_fraction[0] 0.486 0.214 0.103 0.858
fwhm2_thermal_fraction[1] 0.099 0.061 0.050 0.195
fwhm2_thermal_fraction[2] 0.493 0.202 0.155 0.869
filling_factor[0] 0.670 0.232 0.252 1.000
filling_factor[1] 0.608 0.242 0.210 0.999
filling_factor[2] 0.709 0.209 0.331 1.000
wt_ff_tkin[0] 0.000 0.000 0.000 0.000
wt_ff_tkin[1] 0.078 0.008 0.063 0.093
wt_ff_tkin[2] 0.003 0.000 0.003 0.004
fwhm2[0] 445.337 5.218 435.831 455.311
fwhm2[1] 9.869 0.026 9.822 9.918
fwhm2[2] 34.143 1.072 32.152 36.141
velocity[0] 9.991 0.074 9.854 10.131
velocity[1] -5.001 0.002 -5.004 -4.997
velocity[2] 5.021 0.029 4.968 5.075
log10_n_alpha[0] -6.014 0.938 -7.653 -4.097
log10_n_alpha[1] -5.998 1.013 -7.810 -3.953
log10_n_alpha[2] -6.005 1.003 -7.802 -4.013
fwhm2_thermal[0] 216.621 95.502 45.590 382.470
fwhm2_thermal[1] 0.975 0.600 0.500 1.934
fwhm2_thermal[2] 16.843 6.974 5.094 29.648
tkin[0] 4734.544 2087.323 996.424 8359.422
tkin[1] 21.320 13.110 10.924 42.272
tkin[2] 368.126 152.436 111.334 648.009
fwhm2_nonthermal[0] 228.716 95.238 63.970 399.158
fwhm2_nonthermal[1] 8.894 0.600 7.936 9.393
fwhm2_nonthermal[2] 17.300 6.849 4.661 28.882
log10_depth[0] 3.323 1.102 1.427 5.426
log10_depth[1] 0.864 0.369 0.225 1.519
log10_depth[2] 1.294 0.642 0.131 2.539
log10_NHI[0] 20.681 0.209 20.462 21.068
log10_NHI[1] 20.069 0.187 19.813 20.429
log10_NHI[2] 19.955 0.165 19.763 20.269
log10_nHI[0] -1.132 1.118 -3.290 0.816
log10_nHI[1] 0.715 0.437 -0.113 1.535
log10_nHI[2] 0.172 0.674 -1.108 1.412
tspin[0] 3183.742 1795.610 363.277 6556.526
tspin[1] 21.309 13.103 11.047 42.405
tspin[2] 362.892 151.354 109.553 642.948
absorption_weight[0] 0.551 0.393 0.019 1.265
absorption_weight[1] 0.809 0.102 0.631 1.002
absorption_weight[2] 0.900 0.472 0.149 1.767
tau_total[0] 0.143 0.186 0.019 0.368
tau_total[1] 3.398 0.411 2.574 4.126
tau_total[2] 0.181 0.130 0.048 0.414
log10_Pth[0] 2.489 1.209 0.080 4.563
log10_Pth[1] 1.996 0.560 0.937 3.081
log10_Pth[2] 2.694 0.806 1.163 4.184
log10_Pnth[0] 4.002 1.055 2.052 5.907
log10_Pnth[1] 4.491 0.421 3.680 5.252
log10_Pnth[2] 4.188 0.549 3.152 5.173
log10_Ptot[0] 4.024 1.061 2.036 5.910
log10_Ptot[1] 4.492 0.421 3.681 5.257
log10_Ptot[2] 4.211 0.561 3.087 5.159
[ ]: