EmissionPhysicalModel Tutorial

Trey V. Wenger (c) July 2025

EmissionPhysicalModel is a re-parameterization of EmissionModel 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(-100.0, 100.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.1 # K

# 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 * np.random.randn(len(velo_axis))

dummy_data = {"emission": SpecData(
    velo_axis,
    emission,
    rms,
    xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
    ylabel=r"$T_B$ (K)",
)}

# Plot dummy data
fig, ax = plt.subplots(layout="constrained")
ax.plot(dummy_data["emission"].spectral, dummy_data["emission"].brightness, "k-")
ax.set_xlabel(dummy_data["emission"].xlabel)
_ = ax.set_ylabel(dummy_data["emission"].ylabel)
../_images/notebooks_emission_physical_model_3_0.png
[3]:
from caribou_hi import EmissionPhysicalModel

# Initialize and define the model
n_clouds = 3
baseline_degree = 0
model = EmissionPhysicalModel(
    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=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=[-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()
[4]:
from caribou_hi import physics

# Simulation parameters
filling_factor = np.array([0.2, 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([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 = {
    "filling_factor": filling_factor,
    "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_emission_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")
data = {"emission": SpecData(
    velo_axis,
    emission,
    rms,
    xlabel=r"$V_{\rm LSR}$ (km s$^{-1}$)",
    ylabel=r"$T_B$ (K)",
)}
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]:
{'filling_factor': array([0.2, 0.8, 1. ]),
 '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_emission_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([  5.,  25., 300.]),
 'log10_nHI': array([  -3.05103   ,  -23.60205999, -298.02287875]),
 'tspin': array([  49.99124125,  297.4138673 , 1729.59964942]),
 'tau_total': array([3.01145396, 0.14232426, 0.09287008]),
 'log10_Pth': array([  -1.35205999,  -21.12493874, -294.32390874]),
 'log10_Pnth': array([ 2.49618032e-01, -1.94499799e+01, -2.92861759e+02]),
 'log10_Ptot': array([ 2.60351217e-01, -1.94408959e+01, -2.92847028e+02])}
[6]:
# Plot data
fig, ax = plt.subplots(layout="constrained")
ax.plot(data["emission"].spectral, data["emission"].brightness, "k-")
ax.set_xlabel(data["emission"].xlabel)
_ = ax.set_ylabel(data["emission"].ylabel)
../_images/notebooks_emission_physical_model_7_0.png

Model Definition

[7]:
# Initialize and define the model
model = EmissionPhysicalModel(
    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=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=[-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()
[8]:
# Plot model graph
model.graph().render('emission_physical_model', format='png')
model.graph()
[8]:
../_images/notebooks_emission_physical_model_10_0.svg
[9]:
# model string representation
print(model.model.str_repr())
baseline_emission_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)
                 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(depth_nth_fwhm_power, nth_fwhm_1pc, filling_factor, ff_NHI_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, filling_factor, ff_NHI_norm))
             tau_total ~ Deterministic(f(filling_factor, ff_NHI_norm, fwhm2_thermal_fraction, log10_n_alpha_norm, fwhm2_norm, depth_nth_fwhm_power, nth_fwhm_1pc))
             log10_Pth ~ Deterministic(f(fwhm2_thermal_fraction, depth_nth_fwhm_power, nth_fwhm_1pc, filling_factor, ff_NHI_norm, fwhm2_norm))
            log10_Pnth ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm, depth_nth_fwhm_power, nth_fwhm_1pc, filling_factor, ff_NHI_norm))
            log10_Ptot ~ Deterministic(f(fwhm2_thermal_fraction, fwhm2_norm, depth_nth_fwhm_power, nth_fwhm_1pc, filling_factor, ff_NHI_norm))
              emission ~ Normal(f(filling_factor, baseline_emission_norm, fwhm2_thermal_fraction, log10_n_alpha_norm, fwhm2_norm, ff_NHI_norm, velocity_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: [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]
../_images/notebooks_emission_physical_model_12_1.png
[11]:
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']
['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', 'filling_factor']
../_images/notebooks_emission_physical_model_14_1.png

Variational Inference

[33]:
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.25, 0.75, model.n_clouds)},
)
end = time.time()
print(f"Runtime: {(end-start)/60.0:.2f} minutes")
Convergence achieved at 61300
Interrupted at 61,299 [6%]: Average Loss = 1.0583e+05
Adding log-likelihood to trace
Runtime: 1.09 minutes
[34]:
pm.summary(model.trace.posterior)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/numpy/_core/_methods.py:191: RuntimeWarning: invalid value encountered in subtract
  x = asanyarray(arr - arrmean)
arviz - WARNING - Shape validation failed: input_shape: (1, 1000), minimum_shape: (chains=2, draws=4)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/stats_utils.py:39: RuntimeWarning: invalid value encountered in subtract
  ary = ary - ary.mean(axis, keepdims=True)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/diagnostics.py:979: RuntimeWarning: invalid value encountered in subtract
  sims_c2 = (ary - ary.mean()) ** 2
arviz - WARNING - Array contains NaN-value.
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/stats_utils.py:39: RuntimeWarning: invalid value encountered in subtract
  ary = ary - ary.mean(axis, keepdims=True)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/diagnostics.py:979: RuntimeWarning: invalid value encountered in subtract
  sims_c2 = (ary - ary.mean()) ** 2
[34]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
baseline_emission_norm[0] -0.058 7.200000e-02 -0.191 0.071 0.002 0.002 849.0 791.0 NaN
log10_n_alpha_norm[0] 0.009 1.021000e+00 -1.849 2.015 0.030 0.023 1131.0 900.0 NaN
log10_n_alpha_norm[1] 0.246 7.980000e-01 -1.202 1.735 0.026 0.018 952.0 872.0 NaN
log10_n_alpha_norm[2] 0.484 6.680000e-01 -0.770 1.747 0.021 0.014 1068.0 1017.0 NaN
fwhm2_norm[0] 0.034 1.000000e-03 0.032 0.037 0.000 0.000 1021.0 870.0 NaN
fwhm2_norm[1] 0.180 4.000000e-03 0.171 0.188 0.000 0.000 1088.0 970.0 NaN
fwhm2_norm[2] 2.197 1.800000e-02 2.163 2.232 0.001 0.000 918.0 873.0 NaN
velocity_norm[0] 0.375 0.000000e+00 0.374 0.376 0.000 0.000 953.0 637.0 NaN
velocity_norm[1] 0.625 1.000000e-03 0.624 0.627 0.000 0.000 893.0 929.0 NaN
velocity_norm[2] 0.751 1.000000e-03 0.748 0.753 0.000 0.000 1017.0 908.0 NaN
nth_fwhm_1pc[0] 1.749 2.450000e-01 1.303 2.219 0.008 0.006 1040.0 913.0 NaN
nth_fwhm_1pc[1] 1.738 2.530000e-01 1.284 2.219 0.008 0.006 971.0 943.0 NaN
nth_fwhm_1pc[2] 1.736 2.520000e-01 1.268 2.204 0.008 0.007 1013.0 982.0 NaN
depth_nth_fwhm_power[0] 0.305 9.800000e-02 0.149 0.494 0.003 0.003 898.0 914.0 NaN
depth_nth_fwhm_power[1] 0.303 9.300000e-02 0.135 0.471 0.003 0.002 957.0 825.0 NaN
depth_nth_fwhm_power[2] 0.304 9.200000e-02 0.152 0.480 0.003 0.003 876.0 854.0 NaN
ff_NHI_norm[0] 0.038 0.000000e+00 0.037 0.038 0.000 0.000 1054.0 892.0 NaN
ff_NHI_norm[1] 0.062 1.000000e-03 0.061 0.063 0.000 0.000 1052.0 1013.0 NaN
ff_NHI_norm[2] 0.290 1.000000e-03 0.288 0.292 0.000 0.000 911.0 937.0 NaN
fwhm2_thermal_fraction[0] 0.702 7.500000e-02 0.557 0.827 0.003 0.002 853.0 903.0 NaN
fwhm2_thermal_fraction[1] 0.643 1.270000e-01 0.397 0.860 0.004 0.003 914.0 1021.0 NaN
fwhm2_thermal_fraction[2] 0.507 1.990000e-01 0.162 0.874 0.007 0.004 934.0 936.0 NaN
filling_factor[0] 0.799 1.240000e-01 0.570 0.972 0.004 0.003 1068.0 983.0 NaN
filling_factor[1] 0.729 2.010000e-01 0.335 0.987 0.006 0.005 1010.0 908.0 NaN
filling_factor[2] 0.714 2.220000e-01 0.288 0.995 0.007 0.005 964.0 937.0 NaN
fwhm2[0] 6.883 2.280000e-01 6.455 7.315 0.007 0.006 1021.0 870.0 NaN
fwhm2[1] 35.965 8.640000e-01 34.263 37.516 0.026 0.020 1088.0 970.0 NaN
fwhm2[2] 439.494 3.675000e+00 432.573 446.332 0.123 0.088 918.0 873.0 NaN
velocity[0] -5.004 1.900000e-02 -5.037 -4.968 0.001 0.000 953.0 637.0 NaN
velocity[1] 5.002 3.300000e-02 4.941 5.060 0.001 0.001 893.0 929.0 NaN
velocity[2] 10.026 4.600000e-02 9.939 10.109 0.001 0.001 1017.0 908.0 NaN
log10_n_alpha[0] -5.991 1.021000e+00 -7.849 -3.985 0.030 0.023 1131.0 900.0 NaN
log10_n_alpha[1] -5.754 7.980000e-01 -7.202 -4.265 0.026 0.018 952.0 872.0 NaN
log10_n_alpha[2] -5.516 6.680000e-01 -6.770 -4.253 0.021 0.014 1068.0 1017.0 NaN
fwhm2_thermal[0] 4.831 5.420000e-01 3.870 5.888 0.018 0.014 912.0 904.0 NaN
fwhm2_thermal[1] 23.125 4.616000e+00 13.960 30.570 0.152 0.098 943.0 967.0 NaN
fwhm2_thermal[2] 222.967 8.760700e+01 68.432 381.924 2.879 1.771 932.0 934.0 NaN
tkin[0] 105.578 1.185300e+01 84.577 128.685 0.393 0.301 912.0 904.0 NaN
tkin[1] 505.436 1.008990e+02 305.116 668.148 3.324 2.145 943.0 967.0 NaN
tkin[2] 4873.249 1.914780e+03 1495.667 8347.485 62.934 38.704 932.0 934.0 NaN
fwhm2_nonthermal[0] 2.052 5.200000e-01 1.241 3.162 0.018 0.015 834.0 904.0 NaN
fwhm2_nonthermal[1] 12.840 4.573000e+00 4.872 21.514 0.153 0.098 909.0 978.0 NaN
fwhm2_nonthermal[2] 216.527 8.754400e+01 55.148 368.715 2.870 1.764 937.0 936.0 NaN
log10_depth[0] 0.615 6.260000e-01 0.049 1.368 0.021 0.149 955.0 917.0 NaN
log10_depth[1] 40.973 1.710460e+02 0.784 109.309 5.522 35.347 897.0 1072.0 NaN
log10_depth[2] 572139.007 1.301091e+07 1.633 158501.566 410199.826 6396788.026 922.0 787.0 NaN
log10_NHI[0] 19.680 7.600000e-02 19.583 19.813 0.002 0.003 1051.0 937.0 NaN
log10_NHI[1] 19.952 1.620000e-01 19.790 20.261 0.005 0.007 1014.0 908.0 NaN
log10_NHI[2] 20.640 1.920000e-01 20.462 21.000 0.006 0.009 961.0 902.0 NaN
log10_nHI[0] 0.575 6.290000e-01 -0.193 1.228 0.021 0.147 962.0 819.0 NaN
log10_nHI[1] -39.511 1.710500e+02 -107.912 0.780 5.521 35.348 896.0 1072.0 NaN
log10_nHI[2] -572136.856 1.301091e+07 -158499.471 0.894 410199.826 6396788.026 922.0 787.0 NaN
tspin[0] 105.290 1.194200e+01 84.044 128.171 0.394 0.302 918.0 904.0 NaN
tspin[1] 485.351 1.011250e+02 299.914 660.169 3.187 2.190 1017.0 997.0 NaN
tspin[2] 3634.938 1.675554e+03 852.791 6780.765 53.510 34.898 1009.0 803.0 NaN
tau_total[0] 0.257 6.200000e-02 0.175 0.369 0.002 0.003 1044.0 901.0 NaN
tau_total[1] 0.118 9.500000e-02 0.053 0.230 0.003 0.015 964.0 975.0 NaN
tau_total[2] 0.105 1.580000e-01 0.021 0.248 0.005 0.033 1110.0 1026.0 NaN
log10_Pth[0] 2.596 6.530000e-01 1.734 3.276 0.021 0.144 965.0 884.0 NaN
log10_Pth[1] -36.817 1.710660e+02 -105.338 3.609 5.522 35.346 895.0 1026.0 NaN
log10_Pth[2] -572133.212 1.301091e+07 -158496.236 4.861 410199.827 6396788.028 922.0 787.0 NaN
log10_Pnth[0] 3.701 5.880000e-01 3.017 4.327 0.019 0.152 979.0 851.0 NaN
log10_Pnth[1] -35.605 1.710260e+02 -103.840 4.239 5.521 35.352 897.0 1072.0 NaN
log10_Pnth[2] -572131.741 1.301091e+07 -158494.092 4.854 410199.825 6396788.023 924.0 787.0 NaN
log10_Ptot[0] 3.735 5.910000e-01 3.045 4.364 0.020 0.152 982.0 851.0 NaN
log10_Ptot[1] -inf NaN -103.826 4.331 NaN NaN 896.0 1072.0 NaN
log10_Ptot[2] -inf NaN -inf -58.184 NaN NaN 963.0 944.0 NaN
[35]:
posterior = model.sample_posterior_predictive(
    thin=10, # keep one in {thin} posterior samples
)
_ = plot_predictive(model.data, posterior.posterior_predictive)
Sampling: [emission]
../_images/notebooks_emission_physical_model_18_3.png

Posterior Sampling: MCMC

[36]:
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.25, 0.75, 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 61300
Interrupted at 61,299 [6%]: Average Loss = 1.0583e+05
Multiprocess sampling (8 chains in 8 jobs)
NUTS: [baseline_emission_norm, fwhm2_norm, velocity_norm, log10_n_alpha_norm, nth_fwhm_1pc, depth_nth_fwhm_power, ff_NHI_norm, fwhm2_thermal_fraction, filling_factor]
Sampling 8 chains for 1_000 tune and 1_000 draw iterations (8_000 + 8_000 draws total) took 83 seconds.
Adding log-likelihood to trace
There were 2 divergences in converged chains.
Runtime: 2.71 minutes
[37]:
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
[38]:
print("solutions:", model.solutions)
az.summary(model.trace["solution_0"])
# this also works: az.summary(model.trace.solution_0)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/numpy/_core/_methods.py:191: RuntimeWarning: invalid value encountered in subtract
  x = asanyarray(arr - arrmean)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/stats_utils.py:39: RuntimeWarning: invalid value encountered in subtract
  ary = ary - ary.mean(axis, keepdims=True)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/diagnostics.py:979: RuntimeWarning: invalid value encountered in subtract
  sims_c2 = (ary - ary.mean()) ** 2
arviz - WARNING - Array contains NaN-value.
solutions: [0]
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/diagnostics.py:970: RuntimeWarning: invalid value encountered in subtract
  ary_folded = np.abs(ary - np.median(ary))
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/stats_utils.py:39: RuntimeWarning: invalid value encountered in subtract
  ary = ary - ary.mean(axis, keepdims=True)
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/arviz/stats/diagnostics.py:979: RuntimeWarning: invalid value encountered in subtract
  sims_c2 = (ary - ary.mean()) ** 2
[38]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
baseline_emission_norm[0] -0.064 8.000000e-02 -0.211 0.091 0.001 1.000000e-03 11593.0 6435.0 1.0
log10_n_alpha_norm[0] -0.007 1.015000e+00 -1.941 1.874 0.009 1.200000e-02 13159.0 6098.0 1.0
log10_n_alpha_norm[1] 0.021 9.920000e-01 -1.811 1.866 0.010 1.200000e-02 9129.0 4618.0 1.0
log10_n_alpha_norm[2] 0.062 9.460000e-01 -1.714 1.819 0.016 1.400000e-02 3126.0 1410.0 1.0
fwhm2_norm[0] 0.020 4.000000e-03 0.014 0.027 0.000 0.000000e+00 2909.0 1583.0 1.0
fwhm2_norm[1] 0.170 8.000000e-03 0.155 0.187 0.000 0.000000e+00 4260.0 5376.0 1.0
fwhm2_norm[2] 2.233 3.500000e-02 2.171 2.299 0.001 1.000000e-03 4163.0 3254.0 1.0
velocity_norm[0] 0.376 1.000000e-03 0.375 0.377 0.000 0.000000e+00 5358.0 6295.0 1.0
velocity_norm[1] 0.626 1.000000e-03 0.624 0.628 0.000 0.000000e+00 6879.0 5567.0 1.0
velocity_norm[2] 0.747 3.000000e-03 0.741 0.751 0.000 0.000000e+00 3403.0 2627.0 1.0
nth_fwhm_1pc[0] 1.747 2.490000e-01 1.248 2.191 0.002 3.000000e-03 10741.0 4993.0 1.0
nth_fwhm_1pc[1] 1.750 2.540000e-01 1.297 2.246 0.002 3.000000e-03 11894.0 4897.0 1.0
nth_fwhm_1pc[2] 1.749 2.500000e-01 1.253 2.202 0.002 3.000000e-03 11023.0 4946.0 1.0
depth_nth_fwhm_power[0] 0.298 9.700000e-02 0.141 0.475 0.001 2.000000e-03 12137.0 5816.0 1.0
depth_nth_fwhm_power[1] 0.301 1.030000e-01 0.145 0.485 0.001 3.000000e-03 10915.0 4432.0 1.0
depth_nth_fwhm_power[2] 0.302 1.020000e-01 0.139 0.483 0.001 2.000000e-03 13089.0 5187.0 1.0
ff_NHI_norm[0] 0.094 5.200000e-02 0.050 0.136 0.002 1.500000e-02 2630.0 1664.0 1.0
ff_NHI_norm[1] 0.060 3.000000e-03 0.055 0.066 0.000 0.000000e+00 3694.0 3061.0 1.0
ff_NHI_norm[2] 0.296 6.000000e-03 0.287 0.305 0.000 0.000000e+00 2158.0 1323.0 1.0
fwhm2_thermal_fraction[0] 0.252 1.350000e-01 0.126 0.520 0.003 3.000000e-03 3930.0 3055.0 1.0
fwhm2_thermal_fraction[1] 0.512 2.190000e-01 0.128 0.893 0.003 2.000000e-03 5529.0 3096.0 1.0
fwhm2_thermal_fraction[2] 0.501 2.220000e-01 0.113 0.904 0.002 2.000000e-03 11458.0 5253.0 1.0
filling_factor[0] 0.592 2.530000e-01 0.181 0.998 0.004 2.000000e-03 4187.0 3523.0 1.0
filling_factor[1] 0.681 2.250000e-01 0.274 1.000 0.002 2.000000e-03 12039.0 4528.0 1.0
filling_factor[2] 0.671 2.320000e-01 0.259 1.000 0.002 2.000000e-03 12205.0 5223.0 1.0
fwhm2[0] 4.048 7.280000e-01 2.702 5.379 0.016 2.700000e-02 2909.0 1583.0 1.0
fwhm2[1] 34.038 1.695000e+00 30.920 37.316 0.026 1.800000e-02 4260.0 5376.0 1.0
fwhm2[2] 446.527 6.901000e+00 434.181 459.814 0.110 1.130000e-01 4163.0 3254.0 1.0
velocity[0] -4.967 2.400000e-02 -5.009 -4.922 0.000 0.000000e+00 5358.0 6295.0 1.0
velocity[1] 5.042 4.200000e-02 4.966 5.124 0.001 0.000000e+00 6879.0 5567.0 1.0
velocity[2] 9.864 1.050000e-01 9.658 10.053 0.002 1.000000e-03 3403.0 2627.0 1.0
log10_n_alpha[0] -6.007 1.015000e+00 -7.941 -4.126 0.009 1.200000e-02 13159.0 6098.0 1.0
log10_n_alpha[1] -5.979 9.920000e-01 -7.811 -4.134 0.010 1.200000e-02 9129.0 4618.0 1.0
log10_n_alpha[2] -5.938 9.460000e-01 -7.714 -4.181 0.016 1.400000e-02 3126.0 1410.0 1.0
fwhm2_thermal[0] 1.012 5.840000e-01 0.517 2.126 0.010 1.500000e-02 3668.0 3747.0 1.0
fwhm2_thermal[1] 17.475 7.600000e+00 3.842 30.558 0.095 7.300000e-02 5559.0 3059.0 1.0
fwhm2_thermal[2] 223.744 9.920200e+01 43.297 396.722 0.905 9.740000e-01 11550.0 5184.0 1.0
tkin[0] 22.129 1.276200e+01 11.297 46.462 0.228 3.280000e-01 3668.0 3747.0 1.0
tkin[1] 381.930 1.661070e+02 83.967 667.893 2.079 1.587000e+00 5559.0 3059.0 1.0
tkin[2] 4890.244 2.168206e+03 946.326 8670.921 19.772 2.128500e+01 11550.0 5184.0 1.0
fwhm2_nonthermal[0] 3.035 7.800000e-01 1.555 4.497 0.015 1.900000e-02 3331.0 1776.0 1.0
fwhm2_nonthermal[1] 16.563 7.369000e+00 3.508 28.924 0.090 6.600000e-02 6185.0 4213.0 1.0
fwhm2_nonthermal[2] 222.782 9.908400e+01 42.309 394.782 0.907 9.670000e-01 11527.0 5252.0 1.0
log10_depth[0] 1.331 1.732000e+00 0.000 3.019 0.024 1.540000e-01 3637.0 2088.0 1.0
log10_depth[1] 91.303 1.086404e+03 0.002 201.257 12.301 3.153320e+02 8542.0 5918.0 1.0
log10_depth[2] 613233.938 2.575101e+07 0.385 204446.716 291122.528 1.215172e+07 10969.0 5193.0 1.0
log10_NHI[0] 20.230 1.980000e-01 19.919 20.642 0.004 5.000000e-03 3601.0 1953.0 1.0
log10_NHI[1] 19.981 1.950000e-01 19.757 20.363 0.002 3.000000e-03 11785.0 4972.0 1.0
log10_NHI[2] 20.683 2.090000e-01 20.462 21.062 0.002 4.000000e-03 11974.0 5434.0 1.0
log10_nHI[0] 0.409 1.798000e+00 -1.464 2.367 0.026 1.500000e-01 3302.0 1872.0 1.0
log10_nHI[1] -89.811 1.086407e+03 -199.789 2.043 12.301 3.153330e+02 8576.0 5809.0 1.0
log10_nHI[2] -613231.743 2.575101e+07 -204444.403 1.938 291122.528 1.215172e+07 10967.0 5193.0 1.0
tspin[0] 22.114 1.275000e+01 11.295 46.409 0.228 3.280000e-01 3671.0 3719.0 1.0
tspin[1] 357.515 1.616940e+02 68.057 641.853 2.095 1.548000e+00 5075.0 2912.0 1.0
tspin[2] 2773.682 1.934143e+03 46.635 6175.740 27.437 1.857400e+01 3059.0 1333.0 1.0
tau_total[0] 5.056 2.849000e+00 2.250 7.902 0.127 7.600000e-01 2964.0 1653.0 1.0
tau_total[1] 0.228 2.440000e-01 0.045 0.611 0.004 1.100000e-02 5905.0 3583.0 1.0
tau_total[2] 0.282 5.270000e-01 0.018 1.015 0.018 5.600000e-02 2878.0 1321.0 1.0
log10_Pth[0] 1.706 1.841000e+00 -0.266 4.054 0.026 1.460000e-01 4118.0 2969.0 1.0
log10_Pth[1] -87.282 1.086424e+03 -197.297 4.871 12.301 3.153290e+02 8574.0 5876.0 1.0
log10_Pth[2] -613228.113 2.575101e+07 -204440.952 5.939 291122.528 1.215172e+07 10966.0 5193.0 1.0
log10_Pnth[0] 3.699 1.747000e+00 1.925 5.107 0.024 1.530000e-01 4946.0 5165.0 1.0
log10_Pnth[1] -85.825 1.086390e+03 -195.628 4.914 12.301 3.153350e+02 8741.0 5784.0 1.0
log10_Pnth[2] -613226.630 2.575101e+07 -204439.082 4.767 291122.528 1.215172e+07 10978.0 5179.0 1.0
log10_Ptot[0] 3.705 1.749000e+00 1.929 5.131 0.024 1.530000e-01 4868.0 5077.0 1.0
log10_Ptot[1] -inf NaN -195.619 5.068 NaN NaN 8604.0 5856.0 1.0
log10_Ptot[2] -inf NaN -inf -41.035 NaN NaN 5694.0 5179.0 1.0
[39]:
posterior = model.sample_posterior_predictive(
    thin=10, # keep one in {thin} posterior samples
)
_ = plot_predictive(model.data, posterior.posterior_predictive)
Sampling: [emission]
../_images/notebooks_emission_physical_model_23_3.png
[40]:
from bayes_spec.plots import plot_traces

_ = plot_traces(model.trace.posterior, model.cloud_freeRVs + model.baseline_freeRVs + model.hyper_freeRVs)
../_images/notebooks_emission_physical_model_24_0.png
[41]:
_ = 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_physical_model_25_0.png
[42]:
_ = 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_physical_model_26_0.png
[43]:
_ = 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_physical_model_27_0.png
[44]:
_ = 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_physical_model_28_0.png
[45]:
_ = 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_physical_model_29_0.png
[46]:
_ = 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_physical_model_30_0.png
[47]:
# 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
[47]:
{0: np.int64(0), 1: np.int64(1), 2: np.int64(2)}
[48]:
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', 'ff_NHI_norm', 'fwhm2_thermal_fraction', 'filling_factor']
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 []
[49]:
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_physical_model_33_0.png
[50]:
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_physical_model_34_0.png
[51]:
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_physical_model_35_0.png
[52]:
point_stats = az.summary(model.trace.solution_0, kind='stats')
print("BIC:", model.bic())
display(point_stats)
BIC: -220.72045983450883
/home/twenger/miniforge3/envs/bayes_spec-dev/lib/python3.13/site-packages/numpy/_core/_methods.py:191: RuntimeWarning: invalid value encountered in subtract
  x = asanyarray(arr - arrmean)
mean sd hdi_3% hdi_97%
baseline_emission_norm[0] -0.064 8.000000e-02 -0.211 0.091
log10_n_alpha_norm[0] -0.007 1.015000e+00 -1.941 1.874
log10_n_alpha_norm[1] 0.021 9.920000e-01 -1.811 1.866
log10_n_alpha_norm[2] 0.062 9.460000e-01 -1.714 1.819
fwhm2_norm[0] 0.020 4.000000e-03 0.014 0.027
fwhm2_norm[1] 0.170 8.000000e-03 0.155 0.187
fwhm2_norm[2] 2.233 3.500000e-02 2.171 2.299
velocity_norm[0] 0.376 1.000000e-03 0.375 0.377
velocity_norm[1] 0.626 1.000000e-03 0.624 0.628
velocity_norm[2] 0.747 3.000000e-03 0.741 0.751
nth_fwhm_1pc[0] 1.747 2.490000e-01 1.248 2.191
nth_fwhm_1pc[1] 1.750 2.540000e-01 1.297 2.246
nth_fwhm_1pc[2] 1.749 2.500000e-01 1.253 2.202
depth_nth_fwhm_power[0] 0.298 9.700000e-02 0.141 0.475
depth_nth_fwhm_power[1] 0.301 1.030000e-01 0.145 0.485
depth_nth_fwhm_power[2] 0.302 1.020000e-01 0.139 0.483
ff_NHI_norm[0] 0.094 5.200000e-02 0.050 0.136
ff_NHI_norm[1] 0.060 3.000000e-03 0.055 0.066
ff_NHI_norm[2] 0.296 6.000000e-03 0.287 0.305
fwhm2_thermal_fraction[0] 0.252 1.350000e-01 0.126 0.520
fwhm2_thermal_fraction[1] 0.512 2.190000e-01 0.128 0.893
fwhm2_thermal_fraction[2] 0.501 2.220000e-01 0.113 0.904
filling_factor[0] 0.592 2.530000e-01 0.181 0.998
filling_factor[1] 0.681 2.250000e-01 0.274 1.000
filling_factor[2] 0.671 2.320000e-01 0.259 1.000
fwhm2[0] 4.048 7.280000e-01 2.702 5.379
fwhm2[1] 34.038 1.695000e+00 30.920 37.316
fwhm2[2] 446.527 6.901000e+00 434.181 459.814
velocity[0] -4.967 2.400000e-02 -5.009 -4.922
velocity[1] 5.042 4.200000e-02 4.966 5.124
velocity[2] 9.864 1.050000e-01 9.658 10.053
log10_n_alpha[0] -6.007 1.015000e+00 -7.941 -4.126
log10_n_alpha[1] -5.979 9.920000e-01 -7.811 -4.134
log10_n_alpha[2] -5.938 9.460000e-01 -7.714 -4.181
fwhm2_thermal[0] 1.012 5.840000e-01 0.517 2.126
fwhm2_thermal[1] 17.475 7.600000e+00 3.842 30.558
fwhm2_thermal[2] 223.744 9.920200e+01 43.297 396.722
tkin[0] 22.129 1.276200e+01 11.297 46.462
tkin[1] 381.930 1.661070e+02 83.967 667.893
tkin[2] 4890.244 2.168206e+03 946.326 8670.921
fwhm2_nonthermal[0] 3.035 7.800000e-01 1.555 4.497
fwhm2_nonthermal[1] 16.563 7.369000e+00 3.508 28.924
fwhm2_nonthermal[2] 222.782 9.908400e+01 42.309 394.782
log10_depth[0] 1.331 1.732000e+00 0.000 3.019
log10_depth[1] 91.303 1.086404e+03 0.002 201.257
log10_depth[2] 613233.938 2.575101e+07 0.385 204446.716
log10_NHI[0] 20.230 1.980000e-01 19.919 20.642
log10_NHI[1] 19.981 1.950000e-01 19.757 20.363
log10_NHI[2] 20.683 2.090000e-01 20.462 21.062
log10_nHI[0] 0.409 1.798000e+00 -1.464 2.367
log10_nHI[1] -89.811 1.086407e+03 -199.789 2.043
log10_nHI[2] -613231.743 2.575101e+07 -204444.403 1.938
tspin[0] 22.114 1.275000e+01 11.295 46.409
tspin[1] 357.515 1.616940e+02 68.057 641.853
tspin[2] 2773.682 1.934143e+03 46.635 6175.740
tau_total[0] 5.056 2.849000e+00 2.250 7.902
tau_total[1] 0.228 2.440000e-01 0.045 0.611
tau_total[2] 0.282 5.270000e-01 0.018 1.015
log10_Pth[0] 1.706 1.841000e+00 -0.266 4.054
log10_Pth[1] -87.282 1.086424e+03 -197.297 4.871
log10_Pth[2] -613228.113 2.575101e+07 -204440.952 5.939
log10_Pnth[0] 3.699 1.747000e+00 1.925 5.107
log10_Pnth[1] -85.825 1.086390e+03 -195.628 4.914
log10_Pnth[2] -613226.630 2.575101e+07 -204439.082 4.767
log10_Ptot[0] 3.705 1.749000e+00 1.929 5.131
log10_Ptot[1] -inf NaN -195.619 5.068
log10_Ptot[2] -inf NaN -inf -41.035
[ ]: