AbsorptionPhysicalModel Tutorial
Trey V. Wenger (c) July 2025
AbsorptionPhysicalModel is a re-parameterization of AbsorptionModel 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 definition
velo_axis = np.linspace(-50.0, 50.0, 200) # 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 = 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.
absorption = rms * np.random.randn(len(velo_axis))
dummy_data = {"absorption": SpecData(
velo_axis,
absorption,
rms,
xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
ylabel=r"1 - exp(-$\tau$)",
)}
# Plot dummy data
fig, ax = plt.subplots(layout="constrained")
ax.plot(dummy_data["absorption"].spectral, dummy_data["absorption"].brightness, "k-")
ax.set_xlabel(dummy_data["absorption"].xlabel)
_ = ax.set_ylabel(dummy_data["absorption"].ylabel)
[3]:
from caribou_hi import AbsorptionPhysicalModel
# Initialize and define the model
n_clouds = 3
baseline_degree = 0
model = AbsorptionPhysicalModel(
dummy_data,
n_clouds=n_clouds,
baseline_degree=baseline_degree,
seed=1234,
verbose=True
)
model.add_priors(
prior_NHI_fwhm2_thermal=1.0e20, # column density / thermal FWHM^2 prior width (cm-2 km-2 s2)
prior_fwhm2=200.0, # FWHM^2 prior width (km2 s-2)
prior_fwhm2_thermal_fraction=[2.0, 2.0], # thermal FWHM^2 fraction prior shape
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
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([1.25, 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)
sim_params = {
"log10_NHI": log10_NHI,
"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
absorption = model.model["absorption"].eval(sim_params, on_unused_input="ignore")
data = {"absorption": SpecData(
velo_axis,
absorption,
rms,
xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
ylabel=r"1 - exp(-$\tau$)",
)}
tspin [ 49.9920589 297.73994712 2795.52422645]
fwhm2_nonthermal [ 2.97445928 21.12711044 215.71459099]
fwhm2_thermal [ 2.2876605 13.725963 228.76605 ]
fwhm2 [ 5.26211978 34.85307344 444.48064099]
fwhm2_thermal_fraction [0.43474124 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]
[5]:
sim_params
[5]:
{'log10_NHI': array([20.438321 , 19.88729101, 20.46647225]),
'fwhm2': array([ 5.26211978, 34.85307344, 444.48064099]),
'fwhm2_thermal_fraction': array([0.43474124, 0.3938236 , 0.5146817 ]),
'velocity': array([-5., 5., 10.]),
'log10_n_alpha': array([-5.5, -6. , -6.5]),
'nth_fwhm_1pc': array([1.25, 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([ 2.97445928, 21.12711044, 215.71459099]),
'log10_depth': array([0.69897 , 1.39794001, 2.47712125]),
'log10_nHI': array([ 1.25, 0. , -0.5 ]),
'tspin': array([ 49.9920589 , 297.73994712, 2795.52422645]),
'tau_total': array([3.01140471, 0.14216839, 0.05745901]),
'log10_Pth': array([2.94897 , 2.47712125, 3.19897 ]),
'log10_Pnth': array([4.55064803, 4.1520801 , 4.66111952]),
'log10_Ptot': array([4.56138121, 4.16116407, 4.67585106])}
[6]:
# Plot data
fig, ax = plt.subplots(layout="constrained")
ax.plot(data["absorption"].spectral, data["absorption"].brightness, "k-")
ax.set_xlabel(data["absorption"].xlabel)
_ = ax.set_ylabel(data["absorption"].ylabel)
Model Definition
[7]:
# Initialize and define the model
n_clouds = 3
model = AbsorptionPhysicalModel(
data,
n_clouds=n_clouds,
baseline_degree=baseline_degree,
seed=1234,
verbose=True
)
model.add_priors(
prior_NHI_fwhm2_thermal=1.0e20, # column density / thermal FWHM^2 prior width (cm-2 km-2 s2)
prior_fwhm2=200.0, # FWHM^2 prior width (km2 s-2)
prior_fwhm2_thermal_fraction=[2.0, 2.0], # thermal FWHM^2 fraction prior shape
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()
[8]:
# Plot model graph
model.graph().render('absorption_physical_model', format='png')
model.graph()
[8]:
[9]:
# model string representation
print(model.model.str_repr())
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)
NHI_fwhm2_thermal_norm ~ HalfNormal(0, 1)
fwhm2_thermal_fraction ~ Beta(2, 2)
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(NHI_fwhm2_thermal_norm, fwhm2_thermal_fraction, fwhm2_norm))
log10_nHI ~ Deterministic(f(depth_nth_fwhm_power, nth_fwhm_1pc, NHI_fwhm2_thermal_norm, fwhm2_thermal_fraction, fwhm2_norm))
tspin ~ Deterministic(f(fwhm2_thermal_fraction, log10_n_alpha_norm, fwhm2_norm, depth_nth_fwhm_power, nth_fwhm_1pc, NHI_fwhm2_thermal_norm))
tau_total ~ Deterministic(f(NHI_fwhm2_thermal_norm, fwhm2_thermal_fraction, fwhm2_norm, log10_n_alpha_norm, depth_nth_fwhm_power, nth_fwhm_1pc))
log10_Pth ~ Deterministic(f(fwhm2_thermal_fraction, depth_nth_fwhm_power, nth_fwhm_1pc, NHI_fwhm2_thermal_norm, fwhm2_norm))
log10_Pnth ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm, depth_nth_fwhm_power, nth_fwhm_1pc, NHI_fwhm2_thermal_norm))
log10_Ptot ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm, depth_nth_fwhm_power, nth_fwhm_1pc, NHI_fwhm2_thermal_norm))
absorption ~ Normal(f(baseline_absorption_norm, NHI_fwhm2_thermal_norm, velocity_norm, fwhm2_norm, fwhm2_thermal_fraction, log10_n_alpha_norm, depth_nth_fwhm_power, nth_fwhm_1pc), <constant>)
[10]:
from bayes_spec.plots import plot_predictive
# prior predictive check
prior = model.sample_prior_predictive(
samples=1000, # prior predictive samples
)
_ = plot_predictive(model.data, prior.prior_predictive.sel(draw=slice(None, None, 20)))
Sampling: [NHI_fwhm2_thermal_norm, absorption, baseline_absorption_norm, depth_nth_fwhm_power, fwhm2_norm, fwhm2_thermal_fraction, log10_n_alpha_norm, nth_fwhm_1pc, velocity_norm]
[11]:
print(model.cloud_freeRVs)
print(model.cloud_deterministics)
['fwhm2_norm', 'velocity_norm', 'log10_n_alpha_norm', 'nth_fwhm_1pc', 'depth_nth_fwhm_power', 'NHI_fwhm2_thermal_norm', 'fwhm2_thermal_fraction']
['fwhm2', 'velocity', 'log10_n_alpha', 'fwhm2_thermal', 'tkin', 'fwhm2_nonthermal', 'log10_depth', 'log10_NHI', 'log10_nHI', 'tspin', 'tau_total', 'log10_Pth', 'log10_Pnth', 'log10_Ptot']
[12]:
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', 'tau_total', 'log10_Pth', 'log10_Pnth', 'log10_Ptot', 'nth_fwhm_1pc', 'depth_nth_fwhm_power', 'fwhm2_thermal_fraction']
Variational Inference
[13]:
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 98200
Interrupted at 98,199 [9%]: Average Loss = 78,973
Adding log-likelihood to trace
Runtime: 1.03 minutes
[14]:
pm.summary(model.trace.posterior)
arviz - WARNING - Shape validation failed: input_shape: (1, 1000), minimum_shape: (chains=2, draws=4)
[14]:
| mean | sd | hdi_3% | hdi_97% | mcse_mean | mcse_sd | ess_bulk | ess_tail | r_hat | |
|---|---|---|---|---|---|---|---|---|---|
| baseline_absorption_norm[0] | -0.032 | 0.072 | -0.171 | 0.090 | 0.002 | 0.002 | 877.0 | 770.0 | NaN |
| log10_n_alpha_norm[0] | 1.153 | 0.390 | 0.475 | 1.948 | 0.013 | 0.009 | 966.0 | 876.0 | NaN |
| log10_n_alpha_norm[1] | 0.084 | 1.015 | -1.796 | 2.062 | 0.032 | 0.024 | 1004.0 | 979.0 | NaN |
| log10_n_alpha_norm[2] | 0.332 | 0.781 | -0.989 | 1.904 | 0.024 | 0.019 | 1060.0 | 983.0 | NaN |
| fwhm2_norm[0] | 2.913 | 0.530 | 1.979 | 3.870 | 0.020 | 0.016 | 726.0 | 802.0 | NaN |
| fwhm2_norm[1] | 0.026 | 0.000 | 0.026 | 0.026 | 0.000 | 0.000 | 980.0 | 1012.0 | NaN |
| fwhm2_norm[2] | 0.182 | 0.007 | 0.170 | 0.194 | 0.000 | 0.000 | 1000.0 | 849.0 | NaN |
| velocity_norm[0] | 0.880 | 0.038 | 0.806 | 0.940 | 0.001 | 0.001 | 941.0 | 908.0 | NaN |
| velocity_norm[1] | 0.333 | 0.000 | 0.333 | 0.333 | 0.000 | 0.000 | 1015.0 | 907.0 | NaN |
| velocity_norm[2] | 0.669 | 0.002 | 0.665 | 0.672 | 0.000 | 0.000 | 971.0 | 788.0 | NaN |
| nth_fwhm_1pc[0] | 1.755 | 0.255 | 1.318 | 2.277 | 0.008 | 0.005 | 948.0 | 983.0 | NaN |
| nth_fwhm_1pc[1] | 1.771 | 0.254 | 1.321 | 2.284 | 0.008 | 0.008 | 943.0 | 835.0 | NaN |
| nth_fwhm_1pc[2] | 1.769 | 0.236 | 1.323 | 2.193 | 0.007 | 0.005 | 999.0 | 1016.0 | NaN |
| depth_nth_fwhm_power[0] | 0.317 | 0.103 | 0.139 | 0.493 | 0.003 | 0.004 | 1043.0 | 983.0 | NaN |
| depth_nth_fwhm_power[1] | 0.296 | 0.094 | 0.152 | 0.475 | 0.003 | 0.003 | 862.0 | 936.0 | NaN |
| depth_nth_fwhm_power[2] | 0.314 | 0.099 | 0.148 | 0.486 | 0.003 | 0.003 | 933.0 | 906.0 | NaN |
| NHI_fwhm2_thermal_norm[0] | 0.021 | 0.002 | 0.018 | 0.024 | 0.000 | 0.000 | 1006.0 | 868.0 | NaN |
| NHI_fwhm2_thermal_norm[1] | 1.198 | 0.001 | 1.196 | 1.201 | 0.000 | 0.000 | 966.0 | 902.0 | NaN |
| NHI_fwhm2_thermal_norm[2] | 0.057 | 0.001 | 0.055 | 0.058 | 0.000 | 0.000 | 868.0 | 936.0 | NaN |
| fwhm2_thermal_fraction[0] | 0.457 | 0.211 | 0.102 | 0.829 | 0.007 | 0.004 | 870.0 | 828.0 | NaN |
| fwhm2_thermal_fraction[1] | 0.511 | 0.230 | 0.116 | 0.902 | 0.008 | 0.004 | 861.0 | 868.0 | NaN |
| fwhm2_thermal_fraction[2] | 0.514 | 0.239 | 0.107 | 0.912 | 0.007 | 0.004 | 1032.0 | 1023.0 | NaN |
| fwhm2[0] | 582.584 | 105.903 | 395.759 | 773.945 | 4.020 | 3.181 | 726.0 | 802.0 | NaN |
| fwhm2[1] | 5.251 | 0.011 | 5.231 | 5.273 | 0.000 | 0.000 | 980.0 | 1012.0 | NaN |
| fwhm2[2] | 36.462 | 1.306 | 34.062 | 38.783 | 0.041 | 0.031 | 1000.0 | 849.0 | NaN |
| velocity[0] | 11.397 | 1.127 | 9.172 | 13.203 | 0.037 | 0.032 | 941.0 | 908.0 | NaN |
| velocity[1] | -4.999 | 0.001 | -5.001 | -4.997 | 0.000 | 0.000 | 1015.0 | 907.0 | NaN |
| velocity[2] | 5.058 | 0.057 | 4.948 | 5.156 | 0.002 | 0.001 | 971.0 | 788.0 | NaN |
| log10_n_alpha[0] | -4.847 | 0.390 | -5.525 | -4.052 | 0.013 | 0.009 | 966.0 | 876.0 | NaN |
| log10_n_alpha[1] | -5.916 | 1.015 | -7.796 | -3.938 | 0.032 | 0.024 | 1004.0 | 979.0 | NaN |
| log10_n_alpha[2] | -5.668 | 0.781 | -6.989 | -4.096 | 0.024 | 0.019 | 1060.0 | 983.0 | NaN |
| fwhm2_thermal[0] | 266.815 | 135.581 | 34.668 | 500.379 | 4.888 | 3.369 | 790.0 | 666.0 | NaN |
| fwhm2_thermal[1] | 2.685 | 1.206 | 0.610 | 4.741 | 0.041 | 0.020 | 857.0 | 868.0 | NaN |
| fwhm2_thermal[2] | 18.758 | 8.762 | 3.360 | 33.041 | 0.273 | 0.152 | 1051.0 | 863.0 | NaN |
| tkin[0] | 5831.623 | 2963.308 | 757.725 | 10936.473 | 106.840 | 73.628 | 790.0 | 666.0 | NaN |
| tkin[1] | 58.678 | 26.370 | 13.327 | 103.626 | 0.907 | 0.439 | 857.0 | 868.0 | NaN |
| tkin[2] | 409.989 | 191.515 | 73.437 | 722.156 | 5.971 | 3.319 | 1051.0 | 863.0 | NaN |
| fwhm2_nonthermal[0] | 315.769 | 135.822 | 85.329 | 566.284 | 4.305 | 2.851 | 982.0 | 941.0 | NaN |
| fwhm2_nonthermal[1] | 2.566 | 1.206 | 0.502 | 4.616 | 0.041 | 0.020 | 862.0 | 868.0 | NaN |
| fwhm2_nonthermal[2] | 17.703 | 8.729 | 3.299 | 32.549 | 0.274 | 0.148 | 1017.0 | 903.0 | NaN |
| log10_depth[0] | 3.432 | 1.209 | 1.389 | 5.612 | 0.037 | 0.042 | 1085.0 | 1016.0 | NaN |
| log10_depth[1] | -0.266 | 0.611 | -1.431 | 0.611 | 0.020 | 0.027 | 966.0 | 798.0 | NaN |
| log10_depth[2] | 1.190 | 0.680 | 0.050 | 2.557 | 0.023 | 0.017 | 891.0 | 980.0 | NaN |
| log10_NHI[0] | 20.674 | 0.270 | 20.165 | 21.109 | 0.009 | 0.008 | 805.0 | 755.0 | NaN |
| log10_NHI[1] | 20.449 | 0.248 | 20.004 | 20.795 | 0.008 | 0.007 | 857.0 | 868.0 | NaN |
| log10_NHI[2] | 19.959 | 0.270 | 19.462 | 20.321 | 0.008 | 0.008 | 1053.0 | 975.0 | NaN |
| log10_nHI[0] | -1.248 | 1.294 | -3.757 | 0.864 | 0.040 | 0.042 | 1064.0 | 1001.0 | NaN |
| log10_nHI[1] | 2.226 | 0.803 | 0.805 | 3.682 | 0.026 | 0.027 | 917.0 | 883.0 | NaN |
| log10_nHI[2] | 0.279 | 0.868 | -1.414 | 1.774 | 0.028 | 0.020 | 932.0 | 983.0 | NaN |
| tspin[0] | 5378.721 | 2605.492 | 755.899 | 9803.908 | 89.931 | 58.992 | 851.0 | 898.0 | NaN |
| tspin[1] | 58.672 | 26.371 | 13.324 | 103.623 | 0.907 | 0.439 | 857.0 | 868.0 | NaN |
| tspin[2] | 407.174 | 190.988 | 73.200 | 721.765 | 5.962 | 3.309 | 1049.0 | 941.0 | NaN |
| tau_total[0] | 0.056 | 0.007 | 0.045 | 0.067 | 0.000 | 0.000 | 971.0 | 861.0 | NaN |
| tau_total[1] | 3.008 | 0.004 | 3.002 | 3.015 | 0.000 | 0.000 | 990.0 | 828.0 | NaN |
| tau_total[2] | 0.143 | 0.003 | 0.138 | 0.148 | 0.000 | 0.000 | 747.0 | 862.0 | NaN |
| log10_Pth[0] | 2.448 | 1.425 | -0.230 | 4.875 | 0.046 | 0.043 | 1033.0 | 1016.0 | NaN |
| log10_Pth[1] | 3.936 | 1.020 | 2.068 | 5.697 | 0.034 | 0.028 | 898.0 | 939.0 | NaN |
| log10_Pth[2] | 2.826 | 1.092 | 0.669 | 4.663 | 0.035 | 0.024 | 965.0 | 982.0 | NaN |
| log10_Pnth[0] | 4.031 | 1.213 | 1.658 | 5.971 | 0.037 | 0.043 | 1034.0 | 979.0 | NaN |
| log10_Pnth[1] | 5.391 | 0.556 | 4.333 | 6.413 | 0.018 | 0.019 | 943.0 | 800.0 | NaN |
| log10_Pnth[2] | 4.279 | 0.644 | 2.957 | 5.250 | 0.021 | 0.015 | 911.0 | 977.0 | NaN |
| log10_Ptot[0] | 4.049 | 1.218 | 1.714 | 6.050 | 0.038 | 0.043 | 1038.0 | 898.0 | NaN |
| log10_Ptot[1] | 5.420 | 0.587 | 4.270 | 6.437 | 0.019 | 0.022 | 947.0 | 800.0 | NaN |
| log10_Ptot[2] | 4.308 | 0.668 | 2.898 | 5.328 | 0.022 | 0.016 | 916.0 | 952.0 | NaN |
[15]:
posterior = model.sample_posterior_predictive(
thin=10, # keep one in {thin} posterior samples
)
_ = plot_predictive(model.data, posterior.posterior_predictive)
Sampling: [absorption]
Posterior Sampling: MCMC
[16]:
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.8}, # 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 98200
Interrupted at 98,199 [9%]: Average Loss = 78,973
Multiprocess sampling (8 chains in 8 jobs)
NUTS: [baseline_absorption_norm, fwhm2_norm, velocity_norm, log10_n_alpha_norm, nth_fwhm_1pc, depth_nth_fwhm_power, NHI_fwhm2_thermal_norm, fwhm2_thermal_fraction]
Sampling 8 chains for 1_000 tune and 1_000 draw iterations (8_000 + 8_000 draws total) took 11 seconds.
Adding log-likelihood to trace
There were 242 divergences in converged chains.
Runtime: 1.43 minutes
[20]:
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
[21]:
print("solutions:", model.solutions)
az.summary(model.trace["solution_0"])
# this also works: az.summary(model.trace.solution_0)
solutions: [0]
[21]:
| mean | sd | hdi_3% | hdi_97% | mcse_mean | mcse_sd | ess_bulk | ess_tail | r_hat | |
|---|---|---|---|---|---|---|---|---|---|
| baseline_absorption_norm[0] | -0.050 | 0.102 | -0.248 | 0.133 | 0.002 | 0.001 | 3574.0 | 3811.0 | 1.00 |
| log10_n_alpha_norm[0] | 0.418 | 0.890 | -1.266 | 2.141 | 0.024 | 0.017 | 1440.0 | 542.0 | 1.01 |
| log10_n_alpha_norm[1] | -0.006 | 0.973 | -1.831 | 1.828 | 0.016 | 0.012 | 3880.0 | 3932.0 | 1.00 |
| log10_n_alpha_norm[2] | 0.016 | 0.994 | -1.912 | 1.806 | 0.023 | 0.021 | 1932.0 | 696.0 | 1.01 |
| fwhm2_norm[0] | 3.063 | 0.787 | 1.709 | 4.583 | 0.013 | 0.012 | 3340.0 | 3468.0 | 1.00 |
| fwhm2_norm[1] | 0.026 | 0.000 | 0.026 | 0.026 | 0.000 | 0.000 | 3573.0 | 4452.0 | 1.00 |
| fwhm2_norm[2] | 0.181 | 0.010 | 0.162 | 0.201 | 0.000 | 0.000 | 1997.0 | 1592.0 | 1.00 |
| velocity_norm[0] | 0.864 | 0.053 | 0.767 | 0.963 | 0.001 | 0.001 | 1916.0 | 2220.0 | 1.00 |
| velocity_norm[1] | 0.333 | 0.000 | 0.333 | 0.333 | 0.000 | 0.000 | 4638.0 | 3384.0 | 1.01 |
| velocity_norm[2] | 0.669 | 0.002 | 0.665 | 0.672 | 0.000 | 0.000 | 3646.0 | 4042.0 | 1.00 |
| nth_fwhm_1pc[0] | 1.764 | 0.247 | 1.313 | 2.229 | 0.004 | 0.004 | 3832.0 | 3376.0 | 1.00 |
| nth_fwhm_1pc[1] | 1.753 | 0.247 | 1.310 | 2.227 | 0.004 | 0.003 | 3441.0 | 3848.0 | 1.00 |
| nth_fwhm_1pc[2] | 1.752 | 0.244 | 1.278 | 2.200 | 0.003 | 0.003 | 4925.0 | 3483.0 | 1.00 |
| depth_nth_fwhm_power[0] | 0.321 | 0.110 | 0.144 | 0.512 | 0.002 | 0.002 | 3460.0 | 3360.0 | 1.01 |
| depth_nth_fwhm_power[1] | 0.299 | 0.101 | 0.139 | 0.477 | 0.002 | 0.002 | 2534.0 | 2412.0 | 1.00 |
| depth_nth_fwhm_power[2] | 0.302 | 0.101 | 0.155 | 0.491 | 0.002 | 0.002 | 4779.0 | 4365.0 | 1.00 |
| NHI_fwhm2_thermal_norm[0] | 0.019 | 0.006 | 0.008 | 0.029 | 0.000 | 0.000 | 581.0 | 157.0 | 1.01 |
| NHI_fwhm2_thermal_norm[1] | 1.198 | 0.002 | 1.195 | 1.201 | 0.000 | 0.000 | 2921.0 | 2840.0 | 1.00 |
| NHI_fwhm2_thermal_norm[2] | 0.056 | 0.003 | 0.051 | 0.061 | 0.000 | 0.000 | 1028.0 | 538.0 | 1.01 |
| fwhm2_thermal_fraction[0] | 0.497 | 0.232 | 0.102 | 0.919 | 0.004 | 0.002 | 3832.0 | 3656.0 | 1.00 |
| fwhm2_thermal_fraction[1] | 0.494 | 0.222 | 0.083 | 0.872 | 0.004 | 0.003 | 3600.0 | 3085.0 | 1.00 |
| fwhm2_thermal_fraction[2] | 0.493 | 0.221 | 0.099 | 0.876 | 0.004 | 0.002 | 3227.0 | 3485.0 | 1.00 |
| fwhm2[0] | 612.522 | 157.429 | 341.866 | 916.591 | 2.694 | 2.333 | 3340.0 | 3468.0 | 1.00 |
| fwhm2[1] | 5.251 | 0.012 | 5.230 | 5.274 | 0.000 | 0.000 | 3573.0 | 4452.0 | 1.00 |
| fwhm2[2] | 36.211 | 2.080 | 32.492 | 40.244 | 0.046 | 0.033 | 1997.0 | 1592.0 | 1.00 |
| velocity[0] | 10.929 | 1.590 | 8.021 | 13.899 | 0.036 | 0.019 | 1916.0 | 2220.0 | 1.00 |
| velocity[1] | -4.999 | 0.001 | -5.001 | -4.997 | 0.000 | 0.000 | 4638.0 | 3384.0 | 1.01 |
| velocity[2] | 5.061 | 0.060 | 4.945 | 5.171 | 0.001 | 0.001 | 3646.0 | 4042.0 | 1.00 |
| log10_n_alpha[0] | -5.582 | 0.890 | -7.266 | -3.859 | 0.024 | 0.017 | 1440.0 | 542.0 | 1.01 |
| log10_n_alpha[1] | -6.006 | 0.973 | -7.831 | -4.172 | 0.016 | 0.012 | 3880.0 | 3932.0 | 1.00 |
| log10_n_alpha[2] | -5.984 | 0.994 | -7.912 | -4.194 | 0.023 | 0.021 | 1932.0 | 696.0 | 1.01 |
| fwhm2_thermal[0] | 304.277 | 166.621 | 19.007 | 588.877 | 2.677 | 2.391 | 3913.0 | 4229.0 | 1.00 |
| fwhm2_thermal[1] | 2.595 | 1.164 | 0.429 | 4.574 | 0.019 | 0.013 | 3599.0 | 3085.0 | 1.00 |
| fwhm2_thermal[2] | 17.854 | 8.099 | 3.570 | 31.889 | 0.149 | 0.085 | 3096.0 | 3506.0 | 1.00 |
| tkin[0] | 6650.389 | 3641.736 | 415.418 | 12870.715 | 58.507 | 52.248 | 3913.0 | 4229.0 | 1.00 |
| tkin[1] | 56.710 | 25.444 | 9.367 | 99.961 | 0.422 | 0.291 | 3599.0 | 3085.0 | 1.00 |
| tkin[2] | 390.227 | 177.020 | 78.027 | 696.988 | 3.256 | 1.867 | 3096.0 | 3506.0 | 1.00 |
| fwhm2_nonthermal[0] | 308.245 | 168.221 | 15.725 | 599.978 | 2.826 | 2.036 | 3514.0 | 4058.0 | 1.00 |
| fwhm2_nonthermal[1] | 2.656 | 1.165 | 0.664 | 4.804 | 0.019 | 0.013 | 3599.0 | 3052.0 | 1.00 |
| fwhm2_nonthermal[2] | 18.357 | 8.073 | 3.881 | 32.512 | 0.147 | 0.086 | 3176.0 | 3284.0 | 1.00 |
| log10_depth[0] | 3.332 | 1.243 | 1.313 | 5.688 | 0.022 | 0.017 | 3217.0 | 4285.0 | 1.01 |
| log10_depth[1] | -0.206 | 0.569 | -1.253 | 0.732 | 0.010 | 0.022 | 2913.0 | 2689.0 | 1.00 |
| log10_depth[2] | 1.322 | 0.656 | 0.183 | 2.649 | 0.011 | 0.010 | 3854.0 | 4372.0 | 1.00 |
| log10_NHI[0] | 20.661 | 0.333 | 20.033 | 21.249 | 0.011 | 0.011 | 1222.0 | 423.0 | 1.01 |
| log10_NHI[1] | 20.431 | 0.262 | 19.966 | 20.794 | 0.005 | 0.006 | 3596.0 | 3085.0 | 1.00 |
| log10_NHI[2] | 19.937 | 0.259 | 19.462 | 20.327 | 0.005 | 0.005 | 2754.0 | 3291.0 | 1.00 |
| log10_nHI[0] | -1.160 | 1.370 | -3.676 | 1.191 | 0.025 | 0.019 | 2995.0 | 3288.0 | 1.00 |
| log10_nHI[1] | 2.148 | 0.766 | 0.758 | 3.591 | 0.014 | 0.020 | 3116.0 | 2660.0 | 1.00 |
| log10_nHI[2] | 0.126 | 0.829 | -1.463 | 1.647 | 0.014 | 0.012 | 3453.0 | 3272.0 | 1.00 |
| tspin[0] | 5208.361 | 3135.437 | 346.218 | 10812.899 | 62.144 | 45.460 | 1293.0 | 460.0 | 1.00 |
| tspin[1] | 56.702 | 25.446 | 9.367 | 99.960 | 0.422 | 0.291 | 3598.0 | 3085.0 | 1.00 |
| tspin[2] | 385.415 | 176.536 | 77.747 | 694.889 | 3.324 | 1.914 | 2884.0 | 3217.0 | 1.00 |
| tau_total[0] | 0.060 | 0.011 | 0.042 | 0.081 | 0.000 | 0.000 | 2170.0 | 3382.0 | 1.00 |
| tau_total[1] | 3.008 | 0.004 | 3.001 | 3.015 | 0.000 | 0.000 | 3342.0 | 4088.0 | 1.00 |
| tau_total[2] | 0.142 | 0.005 | 0.132 | 0.152 | 0.000 | 0.000 | 1816.0 | 1642.0 | 1.00 |
| log10_Pth[0] | 2.581 | 1.536 | -0.168 | 5.314 | 0.027 | 0.020 | 3243.0 | 4339.0 | 1.00 |
| log10_Pth[1] | 3.841 | 0.993 | 1.992 | 5.699 | 0.018 | 0.021 | 3266.0 | 2588.0 | 1.00 |
| log10_Pth[2] | 2.657 | 1.036 | 0.697 | 4.573 | 0.018 | 0.016 | 3338.0 | 2899.0 | 1.00 |
| log10_Pnth[0] | 4.074 | 1.239 | 1.776 | 6.168 | 0.023 | 0.018 | 2904.0 | 2855.0 | 1.00 |
| log10_Pnth[1] | 5.340 | 0.543 | 4.352 | 6.402 | 0.010 | 0.013 | 2789.0 | 3286.0 | 1.00 |
| log10_Pnth[2] | 4.159 | 0.648 | 2.941 | 5.226 | 0.011 | 0.010 | 3675.0 | 3615.0 | 1.00 |
| log10_Ptot[0] | 4.100 | 1.253 | 1.767 | 6.210 | 0.023 | 0.018 | 2903.0 | 2832.0 | 1.00 |
| log10_Ptot[1] | 5.365 | 0.570 | 4.345 | 6.484 | 0.011 | 0.017 | 2841.0 | 3307.0 | 1.00 |
| log10_Ptot[2] | 4.183 | 0.665 | 2.869 | 5.258 | 0.011 | 0.011 | 3661.0 | 3590.0 | 1.00 |
[22]:
posterior = model.sample_posterior_predictive(
thin=10, # keep one in {thin} posterior samples
)
_ = plot_predictive(model.data, posterior.posterior_predictive)
Sampling: [absorption]
[23]:
from bayes_spec.plots import plot_traces
_ = plot_traces(model.trace.solution_0, model.cloud_freeRVs + model.baseline_freeRVs + model.hyper_freeRVs)
[24]:
_ = 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
)
[25]:
_ = 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
)
[26]:
_ = 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
)
[27]:
_ = 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
)
[28]:
_ = 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
)
[29]:
_ = 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
)
[30]:
# 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
[30]:
{0: np.int64(2), 1: np.int64(0), 2: np.int64(1)}
[31]:
print("cloud freeRVs", model.cloud_freeRVs)
print("cloud deterministics", model.cloud_deterministics)
print("hyper freeRVs", model.hyper_freeRVs)
print("hyper deterministics", model.hyper_deterministics)
cloud freeRVs ['fwhm2_norm', 'velocity_norm', 'log10_n_alpha_norm', 'nth_fwhm_1pc', 'depth_nth_fwhm_power', 'NHI_fwhm2_thermal_norm', 'fwhm2_thermal_fraction']
cloud deterministics ['fwhm2', 'velocity', 'log10_n_alpha', 'fwhm2_thermal', 'tkin', 'fwhm2_nonthermal', 'log10_depth', 'log10_NHI', 'log10_nHI', 'tspin', 'tau_total', 'log10_Pth', 'log10_Pnth', 'log10_Ptot']
hyper freeRVs []
hyper deterministics []
[32]:
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
)
[33]:
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
)
[34]:
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
)
[35]:
point_stats = az.summary(model.trace.solution_0, kind='stats')
print("BIC:", model.bic())
display(point_stats)
BIC: -2092.722382242852
| mean | sd | hdi_3% | hdi_97% | |
|---|---|---|---|---|
| baseline_absorption_norm[0] | -0.050 | 0.102 | -0.248 | 0.133 |
| log10_n_alpha_norm[0] | 0.418 | 0.890 | -1.266 | 2.141 |
| log10_n_alpha_norm[1] | -0.006 | 0.973 | -1.831 | 1.828 |
| log10_n_alpha_norm[2] | 0.016 | 0.994 | -1.912 | 1.806 |
| fwhm2_norm[0] | 3.063 | 0.787 | 1.709 | 4.583 |
| fwhm2_norm[1] | 0.026 | 0.000 | 0.026 | 0.026 |
| fwhm2_norm[2] | 0.181 | 0.010 | 0.162 | 0.201 |
| velocity_norm[0] | 0.864 | 0.053 | 0.767 | 0.963 |
| velocity_norm[1] | 0.333 | 0.000 | 0.333 | 0.333 |
| velocity_norm[2] | 0.669 | 0.002 | 0.665 | 0.672 |
| nth_fwhm_1pc[0] | 1.764 | 0.247 | 1.313 | 2.229 |
| nth_fwhm_1pc[1] | 1.753 | 0.247 | 1.310 | 2.227 |
| nth_fwhm_1pc[2] | 1.752 | 0.244 | 1.278 | 2.200 |
| depth_nth_fwhm_power[0] | 0.321 | 0.110 | 0.144 | 0.512 |
| depth_nth_fwhm_power[1] | 0.299 | 0.101 | 0.139 | 0.477 |
| depth_nth_fwhm_power[2] | 0.302 | 0.101 | 0.155 | 0.491 |
| NHI_fwhm2_thermal_norm[0] | 0.019 | 0.006 | 0.008 | 0.029 |
| NHI_fwhm2_thermal_norm[1] | 1.198 | 0.002 | 1.195 | 1.201 |
| NHI_fwhm2_thermal_norm[2] | 0.056 | 0.003 | 0.051 | 0.061 |
| fwhm2_thermal_fraction[0] | 0.497 | 0.232 | 0.102 | 0.919 |
| fwhm2_thermal_fraction[1] | 0.494 | 0.222 | 0.083 | 0.872 |
| fwhm2_thermal_fraction[2] | 0.493 | 0.221 | 0.099 | 0.876 |
| fwhm2[0] | 612.522 | 157.429 | 341.866 | 916.591 |
| fwhm2[1] | 5.251 | 0.012 | 5.230 | 5.274 |
| fwhm2[2] | 36.211 | 2.080 | 32.492 | 40.244 |
| velocity[0] | 10.929 | 1.590 | 8.021 | 13.899 |
| velocity[1] | -4.999 | 0.001 | -5.001 | -4.997 |
| velocity[2] | 5.061 | 0.060 | 4.945 | 5.171 |
| log10_n_alpha[0] | -5.582 | 0.890 | -7.266 | -3.859 |
| log10_n_alpha[1] | -6.006 | 0.973 | -7.831 | -4.172 |
| log10_n_alpha[2] | -5.984 | 0.994 | -7.912 | -4.194 |
| fwhm2_thermal[0] | 304.277 | 166.621 | 19.007 | 588.877 |
| fwhm2_thermal[1] | 2.595 | 1.164 | 0.429 | 4.574 |
| fwhm2_thermal[2] | 17.854 | 8.099 | 3.570 | 31.889 |
| tkin[0] | 6650.389 | 3641.736 | 415.418 | 12870.715 |
| tkin[1] | 56.710 | 25.444 | 9.367 | 99.961 |
| tkin[2] | 390.227 | 177.020 | 78.027 | 696.988 |
| fwhm2_nonthermal[0] | 308.245 | 168.221 | 15.725 | 599.978 |
| fwhm2_nonthermal[1] | 2.656 | 1.165 | 0.664 | 4.804 |
| fwhm2_nonthermal[2] | 18.357 | 8.073 | 3.881 | 32.512 |
| log10_depth[0] | 3.332 | 1.243 | 1.313 | 5.688 |
| log10_depth[1] | -0.206 | 0.569 | -1.253 | 0.732 |
| log10_depth[2] | 1.322 | 0.656 | 0.183 | 2.649 |
| log10_NHI[0] | 20.661 | 0.333 | 20.033 | 21.249 |
| log10_NHI[1] | 20.431 | 0.262 | 19.966 | 20.794 |
| log10_NHI[2] | 19.937 | 0.259 | 19.462 | 20.327 |
| log10_nHI[0] | -1.160 | 1.370 | -3.676 | 1.191 |
| log10_nHI[1] | 2.148 | 0.766 | 0.758 | 3.591 |
| log10_nHI[2] | 0.126 | 0.829 | -1.463 | 1.647 |
| tspin[0] | 5208.361 | 3135.437 | 346.218 | 10812.899 |
| tspin[1] | 56.702 | 25.446 | 9.367 | 99.960 |
| tspin[2] | 385.415 | 176.536 | 77.747 | 694.889 |
| tau_total[0] | 0.060 | 0.011 | 0.042 | 0.081 |
| tau_total[1] | 3.008 | 0.004 | 3.001 | 3.015 |
| tau_total[2] | 0.142 | 0.005 | 0.132 | 0.152 |
| log10_Pth[0] | 2.581 | 1.536 | -0.168 | 5.314 |
| log10_Pth[1] | 3.841 | 0.993 | 1.992 | 5.699 |
| log10_Pth[2] | 2.657 | 1.036 | 0.697 | 4.573 |
| log10_Pnth[0] | 4.074 | 1.239 | 1.776 | 6.168 |
| log10_Pnth[1] | 5.340 | 0.543 | 4.352 | 6.402 |
| log10_Pnth[2] | 4.159 | 0.648 | 2.941 | 5.226 |
| log10_Ptot[0] | 4.100 | 1.253 | 1.767 | 6.210 |
| log10_Ptot[1] | 5.365 | 0.570 | 4.345 | 6.484 |
| log10_Ptot[2] | 4.183 | 0.665 | 2.869 | 5.258 |
[ ]: