Variational Inference with HSSM¶

Through our PyMC interface, we also gain access to approximate posteriors via Variational Inference. This tutorial illustrates how to perform variational inference via HSSM.

In [1]:

Copied!





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

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

import hssm

Load Data and Specify model¶

In [2]:

Copied!

cav_data = hssm.load_data("cavanagh_theta")
cav_model = hssm.HSSM(data=cav_data, model="angle")
cav_data = hssm.load_data("cavanagh_theta")
cav_model = hssm.HSSM(data=cav_data, model="angle")

Model initialized successfully.

Inference¶

We will run MCMC and VI here to contrast results.

Run MCMC¶

In [3]:

Copied!





# Slight adjustment to initial values
initvals_tmp = cav_model.initvals
initvals_tmp["theta"] = 0.1


mcmc_idata = cav_model.sample(
    chains=2, tune=500, draws=500, sampler="numpyro", initvals=initvals_tmp
)
# Slight adjustment to initial values
initvals_tmp = cav_model.initvals
initvals_tmp["theta"] = 0.1


mcmc_idata = cav_model.sample(
    chains=2, tune=500, draws=500, sampler="numpyro", initvals=initvals_tmp
)

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/gradient.py:1327: FutureWarning: LANLogpOp should implement `pullback` instead of `L_op`/`grad`. Direct `L_op`/`grad` implementations are deprecated and will stop being called in a future version.
  input_grads = node.op.pullback(inputs, node.outputs, new_output_grads)
/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/bambi/backend/pymc.py:224: UserWarning: `init='adapt_diag'` is ignored by `nuts_sampler='numpyro'`; the external sampler uses its own initialization.
  idata = pm.sample(
NUTS[numpyro]: [theta, z, t, a, v]

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pymc/sampling/jax.py:463: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(

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

warmup:   5%|▌         | 50/1000 [00:01<00:25, 37.13it/s, 7 steps of size 5.63e-03. acc. prob=0.74]

warmup:  10%|█         | 100/1000 [00:02<00:23, 38.07it/s, 31 steps of size 9.98e-03. acc. prob=0.76]

warmup:  15%|█▌        | 150/1000 [00:05<00:33, 25.59it/s, 31 steps of size 2.25e-01. acc. prob=0.78]

warmup:  20%|██        | 200/1000 [00:06<00:26, 29.83it/s, 7 steps of size 1.89e-01. acc. prob=0.78]

warmup:  25%|██▌       | 250/1000 [00:07<00:23, 31.95it/s, 3 steps of size 5.62e-01. acc. prob=0.78]

warmup:  30%|███       | 300/1000 [00:09<00:21, 32.52it/s, 7 steps of size 6.84e-01. acc. prob=0.78]

warmup:  35%|███▌      | 350/1000 [00:11<00:20, 30.96it/s, 7 steps of size 4.12e-01. acc. prob=0.78]

warmup:  40%|████      | 400/1000 [00:12<00:18, 33.33it/s, 15 steps of size 1.78e-01. acc. prob=0.78]

warmup:  45%|████▌     | 450/1000 [00:13<00:15, 34.63it/s, 15 steps of size 2.76e-01. acc. prob=0.79]

warmup:  50%|█████     | 500/1000 [00:15<00:14, 34.32it/s, 7 steps of size 2.93e-01. acc. prob=0.79]

sample:  55%|█████▌    | 550/1000 [00:16<00:12, 35.07it/s, 15 steps of size 2.93e-01. acc. prob=0.91]

sample:  60%|██████    | 600/1000 [00:17<00:11, 35.70it/s, 11 steps of size 2.93e-01. acc. prob=0.92]

sample:  65%|██████▌   | 650/1000 [00:19<00:09, 35.80it/s, 31 steps of size 2.93e-01. acc. prob=0.92]

sample:  70%|███████   | 700/1000 [00:20<00:07, 39.23it/s, 7 steps of size 2.93e-01. acc. prob=0.92]

sample:  75%|███████▌  | 750/1000 [00:21<00:06, 37.88it/s, 7 steps of size 2.93e-01. acc. prob=0.93]

sample:  80%|████████  | 800/1000 [00:23<00:05, 37.51it/s, 23 steps of size 2.93e-01. acc. prob=0.93]

sample:  85%|████████▌ | 850/1000 [00:24<00:03, 39.80it/s, 7 steps of size 2.93e-01. acc. prob=0.93]

sample:  90%|█████████ | 900/1000 [00:25<00:02, 39.68it/s, 15 steps of size 2.93e-01. acc. prob=0.93]

sample:  95%|█████████▌| 950/1000 [00:26<00:01, 39.23it/s, 7 steps of size 2.93e-01. acc. prob=0.93]

sample: 100%|██████████| 1000/1000 [00:27<00:00, 40.32it/s, 15 steps of size 2.93e-01. acc. prob=0.93]

sample: 100%|██████████| 1000/1000 [00:27<00:00, 35.75it/s, 15 steps of size 2.93e-01. acc. prob=0.93]

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

warmup:   5%|▌         | 50/1000 [00:01<00:22, 41.47it/s, 7 steps of size 3.24e-03. acc. prob=0.73]

warmup:  10%|█         | 100/1000 [00:02<00:21, 42.76it/s, 31 steps of size 1.06e-02. acc. prob=0.76]

warmup:  15%|█▌        | 150/1000 [00:04<00:29, 29.01it/s, 15 steps of size 2.48e-01. acc. prob=0.78]

warmup:  20%|██        | 200/1000 [00:06<00:26, 30.55it/s, 7 steps of size 2.78e-01. acc. prob=0.78]

warmup:  25%|██▌       | 250/1000 [00:07<00:23, 32.52it/s, 7 steps of size 1.81e-01. acc. prob=0.78]

warmup:  30%|███       | 300/1000 [00:08<00:20, 34.92it/s, 3 steps of size 5.10e-01. acc. prob=0.78]

warmup:  35%|███▌      | 350/1000 [00:10<00:18, 35.92it/s, 3 steps of size 1.75e-01. acc. prob=0.78]

warmup:  40%|████      | 400/1000 [00:11<00:15, 38.32it/s, 7 steps of size 1.88e-01. acc. prob=0.78]

warmup:  45%|████▌     | 450/1000 [00:12<00:13, 40.86it/s, 15 steps of size 4.77e-01. acc. prob=0.79]

warmup:  50%|█████     | 500/1000 [00:13<00:13, 37.71it/s, 7 steps of size 2.41e-01. acc. prob=0.79]

sample:  55%|█████▌    | 550/1000 [00:15<00:12, 35.63it/s, 15 steps of size 2.41e-01. acc. prob=0.93]

sample:  60%|██████    | 600/1000 [00:16<00:11, 36.08it/s, 3 steps of size 2.41e-01. acc. prob=0.93]

sample:  65%|██████▌   | 650/1000 [00:18<00:09, 35.40it/s, 23 steps of size 2.41e-01. acc. prob=0.93]

sample:  70%|███████   | 700/1000 [00:19<00:08, 34.10it/s, 31 steps of size 2.41e-01. acc. prob=0.93]

sample:  75%|███████▌  | 750/1000 [00:21<00:07, 34.07it/s, 7 steps of size 2.41e-01. acc. prob=0.93]

sample:  80%|████████  | 800/1000 [00:22<00:05, 33.68it/s, 23 steps of size 2.41e-01. acc. prob=0.93]

sample:  85%|████████▌ | 850/1000 [00:24<00:04, 32.89it/s, 31 steps of size 2.41e-01. acc. prob=0.93]

sample:  90%|█████████ | 900/1000 [00:26<00:03, 32.44it/s, 7 steps of size 2.41e-01. acc. prob=0.93]

sample:  95%|█████████▌| 950/1000 [00:27<00:01, 30.22it/s, 7 steps of size 2.41e-01. acc. prob=0.94]

sample: 100%|██████████| 1000/1000 [00:29<00:00, 31.34it/s, 15 steps of size 2.41e-01. acc. prob=0.94]

sample: 100%|██████████| 1000/1000 [00:29<00:00, 34.00it/s, 15 steps of size 2.41e-01. acc. prob=0.94]

We recommend running at least 4 chains for robust computation of convergence diagnostics

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

In [4]:

Copied!

az.summary(mcmc_idata)
az.summary(mcmc_idata)

Out[4]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
theta	0.228	0.013	0.21	0.25	464	462	1.01	0.0006	0.00039
z	0.5049	0.0061	0.5	0.51	602	596	1.00	0.00025	0.00019
t	0.2837	0.0091	0.27	0.3	522	516	1.00	0.0004	0.00028
v	0.361	0.021	0.33	0.4	670	556	1.01	0.0008	0.00054
a	1.306	0.021	1.3	1.3	428	484	1.00	0.001	0.0007

Run VI¶

In [5]:

Copied!

vi_idata = cav_model.vi(niter=50000, method="fullrank_advi")
vi_idata = cav_model.vi(niter=50000, method="fullrank_advi")

Using MCMC starting point defaults.

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/gradient.py:1327: FutureWarning: LANLogpOp should implement `pullback` instead of `L_op`/`grad`. Direct `L_op`/`grad` implementations are deprecated and will stop being called in a future version.
  input_grads = node.op.pullback(inputs, node.outputs, new_output_grads)

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

Finished [100%]: Average Loss = 6,038.6

Inspect Outputs¶

From our variational inference runs, we extract two objects.

An xarray.DataTree object stored under cav_model.vi_idata. This stores a slightly cleaned up posterior sample, constructed by sampling from the variational posterior.
A PyMC variational Approximation object stored under cav_model.vi_approx that holds the variational posterior object itself. This is a rich structure and it is beyond the purpose of this tutorial to illustrate all it's details. Amongst other things you will be able to inspect the loss history and take samples such as those stored under cam_model.vi_data.

`.vi_idata`¶

The approximate variational posterior as an xarray.DataTree.

`.vi_approx`¶

The approximate variational posterior as a PyMC variational Approximation object.

We can take draws from the posterior with the .sample() method.

In [7]:

Copied!

with cav_model.pymc_model:
    cav_model.vi_approx.sample(draws=1000)
with cav_model.pymc_model:
    cav_model.vi_approx.sample(draws=1000)

The .hist attribute stores the loss history. We can plot this to see how the loss function converged.

In [8]:

Copied!

plt.plot(cav_model.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.plot(cav_model.vi_approx.hist)
plt.xlabel("Iteration")
plt.ylabel("Loss")

Out[8]:

Text(0, 0.5, 'Loss')

No description has been provided for this image

Contrast outputs between MCMC and VI¶

In [9]:

Copied!





plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)
plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)

Out[9]:

<arviz_plots.plot_matrix.PlotMatrix at 0x1386f8170>

Working directly through PyMC¶

Here we illustrate how to use our attached pymc_model to make use of the object oriented API for variational inference. This allows us a few extra affordances.

Let's define a few helper functions first.

In [11]:

Copied!





import warnings

import xarray as xr
from matplotlib import gridspec
from pymc.blocking import DictToArrayBijection, RaveledVars


def tracker_to_idata(tracker, model):
    """Turn a tracker object into an DataTree object."""
    tracker_groups = list(tracker.whatchdict.keys())
    # n_steps = len(tracker[tracker_groups[0]])
    stacked_results = {
        tracker_group: {
            key: np.stack([d[key] for d in tracker[tracker_group]])
            for key in tracker[tracker_group][0]
        }
        for tracker_group in tracker_groups
    }

    # coords = {"vi_step": np.arange(n_steps)} | {
    #     k: np.array(v) for k, v in model.coords.items()
    # }
    var_to_dims = {
        var.name: ("vi_step", *(model.named_vars_to_dims.get(var.name, ())))
        for var in model.continuous_value_vars
    }
    datasets = {
        key: xr.Dataset(
            {
                var: (var_to_dims[var], stacked_results[key][var])
                for var in stacked_results[key].keys()
            }
        )
        for key in tracker_groups
    }

    with warnings.catch_warnings(action="ignore"):
        return xr.DataTree.from_dict(datasets)


def untransform_params(idata, model):
    """Bring transformed parmater back to their original scale."""
    suffixes = ["_interval__", "_log__"]

    def remove_suffixes(word, suffixes):
        for suffix in suffixes:
            if word.endswith(suffix):
                return word[: -len(suffix)]
        return word

    free_rv_names = [rv_.name for rv_ in model.free_RVs]
    transformed_vars = list(idata["mean"].ds.data_vars.keys())
    collect_untransformed_vars = []
    collect_untransformed_xarray_datasets = []

    for var_ in transformed_vars:
        var_untrans = remove_suffixes(var_, suffixes=suffixes)
        if var_untrans in free_rv_names:
            rv = model.free_RVs[free_rv_names.index(var_untrans)]
            if model.rvs_to_transforms[rv] is not None:
                untransformed_var = (
                    model.rvs_to_transforms[rv]
                    .backward(idata["mean"][var_].values, *rv.owner.inputs)
                    .eval()
                )
                collect_untransformed_vars.append(var_)
                collect_untransformed_xarray_datasets.append(
                    xr.Dataset(
                        data_vars={var_untrans: (("vi_step"), untransformed_var)}
                    )
                )

    return xr.merge(
        [idata["mean"].ds] + collect_untransformed_xarray_datasets
    ).drop_vars(collect_untransformed_vars)


def plot_vi_traces(idata):
    """Plot parameter history of the optimization alogrithm."""
    if not isinstance(idata, xr.DataTree):
        raise ValueError("idata must be a DataTree object")
    if "loss" not in idata:
        raise ValueError("DataTree object must contain a 'loss' group")
    if "mean_untransformed" not in idata:
        print(
            "Using transformed variables because 'mean_untransformed' group not found"
        )
        data_vars = list(idata["mean"].ds.data_vars.keys())
    else:
        data_vars = list(idata["mean_untransformed"].ds.data_vars.keys())

    fig = plt.figure(figsize=(8, 1.5 * len(data_vars)))
    gs = gridspec.GridSpec(
        len(data_vars) // 2 + 2
        if (len(data_vars) % 2) == 0
        else (len(data_vars) // 2) + 3,
        2,
    )

    for i, var_ in enumerate(data_vars):
        ax_tmp = fig.add_subplot(gs[i // 2, i % 2])
        idata["mean_untransformed"][var_].plot(ax=ax_tmp)
        ax_tmp.set_title(var_)

    last_ax = fig.add_subplot(gs[-2:, :])
    idata["loss"].loss.plot(ax=last_ax)
    gs.tight_layout(fig)
    return fig
import warnings

import xarray as xr
from matplotlib import gridspec
from pymc.blocking import DictToArrayBijection, RaveledVars


def tracker_to_idata(tracker, model):
    """Turn a tracker object into an DataTree object."""
    tracker_groups = list(tracker.whatchdict.keys())
    # n_steps = len(tracker[tracker_groups[0]])
    stacked_results = {
        tracker_group: {
            key: np.stack([d[key] for d in tracker[tracker_group]])
            for key in tracker[tracker_group][0]
        }
        for tracker_group in tracker_groups
    }

    # coords = {"vi_step": np.arange(n_steps)} | {
    #     k: np.array(v) for k, v in model.coords.items()
    # }
    var_to_dims = {
        var.name: ("vi_step", *(model.named_vars_to_dims.get(var.name, ())))
        for var in model.continuous_value_vars
    }
    datasets = {
        key: xr.Dataset(
            {
                var: (var_to_dims[var], stacked_results[key][var])
                for var in stacked_results[key].keys()
            }
        )
        for key in tracker_groups
    }

    with warnings.catch_warnings(action="ignore"):
        return xr.DataTree.from_dict(datasets)


def untransform_params(idata, model):
    """Bring transformed parmater back to their original scale."""
    suffixes = ["_interval__", "_log__"]

    def remove_suffixes(word, suffixes):
        for suffix in suffixes:
            if word.endswith(suffix):
                return word[: -len(suffix)]
        return word

    free_rv_names = [rv_.name for rv_ in model.free_RVs]
    transformed_vars = list(idata["mean"].ds.data_vars.keys())
    collect_untransformed_vars = []
    collect_untransformed_xarray_datasets = []

    for var_ in transformed_vars:
        var_untrans = remove_suffixes(var_, suffixes=suffixes)
        if var_untrans in free_rv_names:
            rv = model.free_RVs[free_rv_names.index(var_untrans)]
            if model.rvs_to_transforms[rv] is not None:
                untransformed_var = (
                    model.rvs_to_transforms[rv]
                    .backward(idata["mean"][var_].values, *rv.owner.inputs)
                    .eval()
                )
                collect_untransformed_vars.append(var_)
                collect_untransformed_xarray_datasets.append(
                    xr.Dataset(
                        data_vars={var_untrans: (("vi_step"), untransformed_var)}
                    )
                )

    return xr.merge(
        [idata["mean"].ds] + collect_untransformed_xarray_datasets
    ).drop_vars(collect_untransformed_vars)


def plot_vi_traces(idata):
    """Plot parameter history of the optimization alogrithm."""
    if not isinstance(idata, xr.DataTree):
        raise ValueError("idata must be a DataTree object")
    if "loss" not in idata:
        raise ValueError("DataTree object must contain a 'loss' group")
    if "mean_untransformed" not in idata:
        print(
            "Using transformed variables because 'mean_untransformed' group not found"
        )
        data_vars = list(idata["mean"].ds.data_vars.keys())
    else:
        data_vars = list(idata["mean_untransformed"].ds.data_vars.keys())

    fig = plt.figure(figsize=(8, 1.5 * len(data_vars)))
    gs = gridspec.GridSpec(
        len(data_vars) // 2 + 2
        if (len(data_vars) % 2) == 0
        else (len(data_vars) // 2) + 3,
        2,
    )

    for i, var_ in enumerate(data_vars):
        ax_tmp = fig.add_subplot(gs[i // 2, i % 2])
        idata["mean_untransformed"][var_].plot(ax=ax_tmp)
        ax_tmp.set_title(var_)

    last_ax = fig.add_subplot(gs[-2:, :])
    idata["loss"].loss.plot(ax=last_ax)
    gs.tight_layout(fig)
    return fig

In [12]:

Copied!

# Define the ADVI runner
with cav_model.pymc_model:
    advi = pm.ADVI()
# Define the ADVI runner
with cav_model.pymc_model:
    advi = pm.ADVI()

In [13]:

Copied!





# Set up starting point
start = cav_model.pymc_model.initial_point()
vars_dict = {var.name: var for var in cav_model.pymc_model.continuous_value_vars}
x0 = DictToArrayBijection.map(
    {var_name: value for var_name, value in start.items() if var_name in vars_dict}
)

# Define quantities to track
tracker = pm.variational.callbacks.Tracker(
    mean=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.mean.eval(), x0.point_map_info), start
    ),  # callable that returns mean
    std=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.std.eval(), x0.point_map_info), start
    ),  # callable that returns std
)
# Set up starting point
start = cav_model.pymc_model.initial_point()
vars_dict = {var.name: var for var in cav_model.pymc_model.continuous_value_vars}
x0 = DictToArrayBijection.map(
    {var_name: value for var_name, value in start.items() if var_name in vars_dict}
)

# Define quantities to track
tracker = pm.variational.callbacks.Tracker(
    mean=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.mean.eval(), x0.point_map_info), start
    ),  # callable that returns mean
    std=lambda: DictToArrayBijection.rmap(
        RaveledVars(advi.approx.std.eval(), x0.point_map_info), start
    ),  # callable that returns std
)

In [14]:

Copied!





# Run VI fit

approx = advi.fit(n=30000, callbacks=[tracker])
with cav_model.pymc_model:
    vi_posterior_samples = approx.sample(1000)
    vi_posterior_samples.posterior = vi_posterior_samples.posterior.ds.drop_vars(
        "v_mean"
    )
# Run VI fit

approx = advi.fit(n=30000, callbacks=[tracker])
with cav_model.pymc_model:
    vi_posterior_samples = approx.sample(1000)
    vi_posterior_samples.posterior = vi_posterior_samples.posterior.ds.drop_vars(
        "v_mean"
    )

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/gradient.py:1327: FutureWarning: LANLogpOp should implement `pullback` instead of `L_op`/`grad`. Direct `L_op`/`grad` implementations are deprecated and will stop being called in a future version.
  input_grads = node.op.pullback(inputs, node.outputs, new_output_grads)

/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(
/Users/afengler/Projects/proj_hssmspine/HSSMSpine/repos/HSSM/.venv/lib/python3.12/site-packages/pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(

Finished [100%]: Average Loss = 6,037.4

In [15]:

Copied!





from copy import deepcopy

# Convert tracked quantities to idata
result = tracker_to_idata(tracker, cav_model.pymc_model)

# Add untransformed parameters
result["mean_untransformed"] = untransform_params(
    deepcopy(result), cav_model.pymc_model
)
# Add loss group
result["loss"] = xr.Dataset(data_vars={"loss": ("vi_step", np.array(approx.hist))})
from copy import deepcopy

# Convert tracked quantities to idata
result = tracker_to_idata(tracker, cav_model.pymc_model)

# Add untransformed parameters
result["mean_untransformed"] = untransform_params(
    deepcopy(result), cav_model.pymc_model
)
# Add loss group
result["loss"] = xr.Dataset(data_vars={"loss": ("vi_step", np.array(approx.hist))})

A quick look at our result DataTree object, to understand what happened here.

We now have two additional groups:

mean_untransformed which holds parameter values in the orignal space (instead of the parameters over which the optimization operates, which will always live in an unconstrained space)
loss which holds our training history

Plot Results¶

We can plot the parameter trajectories (histories) over optimization steps, with our little helper function plot_vi_traces().

NOTE:

This is a random run, and we did not thoroughly check if the number of steps we allowed the optimizer were indeed enough to converge.

In [17]:

Copied!

fig = plot_vi_traces(result)
fig = plot_vi_traces(result)

In [18]:

Copied!





# Plot MCMC [nuts]
plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
# Plot VI via .vi() [fullrank_advi]
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)
# Plot VI via pymc interface [advi]
# (We need to make sure the variables are in correct order)
az.plot_pair(
    vi_posterior_samples.posterior[list(cav_model.traces.posterior.data_vars)],
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="green"),
        "dist": dict(alpha=0.8, color="green"),
    },
)
# Plot MCMC [nuts]
plot_matrix = az.plot_pair(
    cav_model.traces,
    visuals={
        "scatter": dict(alpha=0.01, color="blue"),
        "dist": dict(alpha=0.6, color="blue"),
    },
)
# Plot VI via .vi() [fullrank_advi]
az.plot_pair(
    cav_model.vi_idata,
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="red"),
        "dist": dict(alpha=0.8, color="red"),
    },
)
# Plot VI via pymc interface [advi]
# (We need to make sure the variables are in correct order)
az.plot_pair(
    vi_posterior_samples.posterior[list(cav_model.traces.posterior.data_vars)],
    plot_matrix=plot_matrix,
    visuals={
        "scatter": dict(alpha=0.04, color="green"),
        "dist": dict(alpha=0.8, color="green"),
    },
)

Out[18]:

<arviz_plots.plot_matrix.PlotMatrix at 0x139ba3cb0>

NOTE:

It is expected that the posterior of our last run looks a little worse. We chose to run the advi algorithm, which implies only an isotropic Gaussian approximation to the posterior, so we expect to miss the posterior covariances which we pick up via fullrank_advi as well as MCMC.

Caveats¶

Variational Inference is powerful, however it comes with it's own set of sharp edges.

You will not always be in a position to compare VI with MCMC runs (after all if you can run full MCMC, there isn't much benefit to using VI at all) and it can be hard to a priori estimate how many steps you may need to run the algorithm for.
The posteriors will be approximate, if the true posterior includes complex parameter trade-offs, VI might result in inaccurate posterior estimates.
We recommend VI for loglik_kind="approx_differentiable", since the gradients of the analytical log-likelihoods still prove somewhat brittle at this point in time.

Read¶

To learn a bit more about the VI-API in PyMC, we recommend you to read the excellent short tutorial in the main documentation.