Refitting PyStan (3.0+) models with ArviZ#

ArviZ is backend agnostic and therefore does not sample directly. In order to take advantage of algorithms that require refitting models several times, ArviZ uses SamplingWrapper to convert the API of the sampling backend to a common set of functions. Hence, functions like Leave Future Out Cross Validation can be used in ArviZ independently of the sampling backend used.

Below there is one example of SamplingWrapper usage for PyStan exteding arviz.PyStanSamplingWrapper which already implements some default methods targeted to PyStan.

Before starting, it is important to note that PyStan cannot call the C++ functions it uses. Therefore, the code of the model must be slightly modified in order to be compatible with the cross validation refitting functions.

import arviz as az
import stan
import numpy as np
import matplotlib.pyplot as plt
# enable PyStan on Jupyter IDE
import nest_asyncio
nest_asyncio.apply()

For the example we will use a linear regression.

np.random.seed(26)

xdata = np.linspace(0, 50, 100)
b0, b1, sigma = -2, 1, 3
ydata = np.random.normal(loc=b1 * xdata + b0, scale=sigma)
plt.plot(xdata, ydata)
[<matplotlib.lines.Line2D at 0x7fe192db2040>]
../_images/fd96b950c73be9b4be4c29f5515a21affa93b2732095a21e7f26fd91ccbca78e.png

Now we will write the Stan code, keeping in mind that it must be able to compute the pointwise log likelihood on excluded data, that is, data which is not used to fit the model. Thus, the backbone of the code must look like:

data {
    data_for_fitting
    excluded_data
    ...
}
model {
    // fit against data_for_fitting
   ...
}
generated quantities {
    ....
    log_lik for data_for_fitting
    log_lik_excluded for excluded_data
}
refit_lr_code = """
data {
  // Define data for fitting
  int<lower=0> N;
  vector[N] x;
  vector[N] y;
  // Define excluded data. It will not be used when fitting.
  int<lower=0> N_ex;
  vector[N_ex] x_ex;
  vector[N_ex] y_ex;
}

parameters {
  real b0;
  real b1;
  real<lower=0> sigma_e;
}

model {
  b0 ~ normal(0, 10);
  b1 ~ normal(0, 10);
  sigma_e ~ normal(0, 10);
  for (i in 1:N) {
    y[i] ~ normal(b0 + b1 * x[i], sigma_e);  // use only data for fitting
  }
  
}

generated quantities {
    vector[N] log_lik;
    vector[N_ex] log_lik_ex;
    vector[N] y_hat;
    
    for (i in 1:N) {
        // calculate log likelihood and posterior predictive, there are 
        // no restrictions on adding more generated quantities
        log_lik[i] = normal_lpdf(y[i] | b0 + b1 * x[i], sigma_e);
        y_hat[i] = normal_rng(b0 + b1 * x[i], sigma_e);
    }
    for (j in 1:N_ex) {
        // calculate the log likelihood of the excluded data given data_for_fitting
        log_lik_ex[j] = normal_lpdf(y_ex[j] | b0 + b1 * x_ex[j], sigma_e);
    }
}
"""
data_dict = {
    "N": len(ydata),
    "y": ydata,
    "x": xdata,
    # No excluded data in initial fit
    "N_ex": 0,
    "x_ex": [],
    "y_ex": [],
}
sm = stan.build(program_code=refit_lr_code, data=data_dict)
sample_kwargs = {"num_samples": 1000, "num_chains": 4}
fit = sm.sample(**sample_kwargs)
Building... This may take some time.
Done.
Sampling...
    0/8000 [>---------------------------]   0%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0   Messages received during sampling:
  Gradient evaluation took 4.2e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.42 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 4.2e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.42 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 0.00014 seconds
  1000 transitions using 10 leapfrog steps per transition would take 1.4 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 3.9e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.39 seconds.
  Adjust your expectations accordingly!

 8000/8000 [============================] 100%  1 sec/0     
Done.

We have defined a dictionary sample_kwargs that will be passed to the SamplingWrapper in order to make sure that all refits use the same sampler parameters. We follow the same pattern with az.from_pystan.

dims = {"y": ["time"], "x": ["time"], "log_likelihood": ["time"], "y_hat": ["time"]}
idata_kwargs = {
    "posterior_predictive": ["y_hat"],
    "observed_data": "y",
    "constant_data": "x",
    "log_likelihood": ["log_lik", "log_lik_ex"],
    "dims": dims,
}
idata = az.from_pystan(posterior=fit, posterior_model=sm, **idata_kwargs)

We will create a subclass of PyStanSamplingWrapper. Therefore, instead of having to implement all functions required by reloo() we only have to implement sel_observations. As explained in its docs, it takes one argument which are the indices of the data to be excluded and returns modified_observed_data which is passed as data to sampling function of PyStan model and excluded_observed_data which is used to retrieve the log likelihood of the excluded data (as passing the excluded data would make no sense).

class LinearRegressionWrapper(az.PyStanSamplingWrapper):
    def sel_observations(self, idx):
        xdata = self.idata_orig.constant_data.x.values
        ydata = self.idata_orig.observed_data.y.values
        mask = np.full_like(xdata, True, dtype=bool)
        mask[idx] = False
        N_obs = len(mask)
        N_ex = np.sum(~mask)
        observations = {
            "N": int(N_obs - N_ex),
            "x": xdata[mask],
            "y": ydata[mask],
            "N_ex": int(N_ex),
            "x_ex": xdata[~mask],
            "y_ex": ydata[~mask],
        }
        return observations, "log_lik_ex"
loo_orig = az.loo(idata, pointwise=True)
loo_orig
Computed from 4000 by 100 log-likelihood matrix

         Estimate       SE
elpd_loo  -250.85     7.20
p_loo        3.05        -
------

Pareto k diagnostic values:
                         Count   Pct.
(-Inf, 0.5]   (good)      100  100.0%
 (0.5, 0.7]   (ok)          0    0.0%
   (0.7, 1]   (bad)         0    0.0%
   (1, Inf)   (very bad)    0    0.0%

In this case, the Leave-One-Out Cross Validation (LOO-CV) approximation using Pareto Smoothed Importance Sampling (PSIS) works for all observations, so we will use modify loo_orig in order to make reloo() believe that PSIS failed for some observations. This will also serve as a validation of our wrapper, as the PSIS LOO-CV already returned the correct value.

loo_orig.pareto_k[[13, 42, 56, 73]] = np.array([0.8, 1.2, 2.6, 0.9])

We initialize our sampling wrapper

pystan_wrapper = LinearRegressionWrapper(
    refit_lr_code, idata_orig=idata, sample_kwargs=sample_kwargs, idata_kwargs=idata_kwargs
)

And eventually, we can use this wrapper to call az.reloo, and compare the results with the PSIS LOO-CV results.

loo_relooed = az.reloo(pystan_wrapper, loo_orig=loo_orig)
Building...
/home/ahartikainen/github_ubuntu/arviz/arviz/stats/stats_refitting.py:99: UserWarning: reloo is an experimental and untested feature
  warnings.warn("reloo is an experimental and untested feature", UserWarning)
Building...
Found model in cache. Done.
Sampling...
    0/8000 [>---------------------------]   0%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0   Messages received during sampling:
  Gradient evaluation took 4e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.4 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 5.9e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.59 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 4.4e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.44 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 3.8e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.38 seconds.
  Adjust your expectations accordingly!

 8000/8000 [============================] 100%  1 sec/0     
Done.
Building...
Found model in cache. Done.
Sampling...
    0/8000 [>---------------------------]   0%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0   Messages received during sampling:
  Gradient evaluation took 2e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.2 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.9e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.19 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 2.2e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.22 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 2.4e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.24 seconds.
  Adjust your expectations accordingly!

 8000/8000 [============================] 100%  1 sec/0     
Done.
Building...
Found model in cache. Done.
Sampling...
    0/8000 [>---------------------------]   0%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0   Messages received during sampling:
  Gradient evaluation took 1.9e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.19 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.6e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.16 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.8e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.18 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.3e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.13 seconds.
  Adjust your expectations accordingly!

 8000/8000 [============================] 100%  1 sec/0     
Done.
Building...
Found model in cache. Done.
Sampling...
    0/8000 [>---------------------------]   0%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0     
 8000/8000 [============================] 100%  1 sec/0   Messages received during sampling:
  Gradient evaluation took 1.7e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.17 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.8e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.18 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 2.2e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.22 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 2.6e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.26 seconds.
  Adjust your expectations accordingly!

 8000/8000 [============================] 100%  1 sec/0     
Done.
loo_relooed
Computed from 4000 by 100 log-likelihood matrix

         Estimate       SE
elpd_loo  -250.85     7.20
p_loo        3.05        -
------

Pareto k diagnostic values:
                         Count   Pct.
(-Inf, 0.5]   (good)      100  100.0%
 (0.5, 0.7]   (ok)          0    0.0%
   (0.7, 1]   (bad)         0    0.0%
   (1, Inf)   (very bad)    0    0.0%
loo_orig
Computed from 4000 by 100 log-likelihood matrix

         Estimate       SE
elpd_loo  -250.85     7.20
p_loo        3.05        -
------

Pareto k diagnostic values:
                         Count   Pct.
(-Inf, 0.5]   (good)       96   96.0%
 (0.5, 0.7]   (ok)          0    0.0%
   (0.7, 1]   (bad)         2    2.0%
   (1, Inf)   (very bad)    2    2.0%