<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 (pystan2_refitting)=
Refitting PyStan (2.x) 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 {class}~arviz.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 an example of SamplingWrapper usage for PyStan extending {class}arviz.PyStan2SamplingWrapper that 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.
Step1:  For the example, we will use a linear regression model.
Step3:  Now we will write the Stan code, keeping in mind that it must be able to compute the pointwise log likelihood on excluded data, i.e., data that is not used to fit the model. Thus, the backbone of the code must look like the following
Step4:  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 will follow the same pattern with az.from_pystan.
Step5:  We will create a subclass of {class}~arviz.PyStan2SamplingWrapper. Therefore, instead of having to implement all functions required by {func}~arviz.reloo we only have to implement {func}~arviz.PyStan2SamplingWrapper.sel_observations. As explained in its docs, it takes one argument which is 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).
Step6:  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 {func}~arviz.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.
Step7:  We initialize our sampling wrapper
Step8:  And eventually, we can use this wrapper to call az.reloo, and compare the results with the PSIS LOO-CV results.
<ASSISTANT_TASK:>
Python Code:
import arviz as az
import pystan
import numpy as np
import matplotlib.pyplot as plt
Explanation: (pystan2_refitting)=
Refitting PyStan (2.x) 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 {class}~arviz.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 an example of SamplingWrapper usage for PyStan extending {class}arviz.PyStan2SamplingWrapper that 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.
End of explanation
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)
Explanation: For the example, we will use a linear regression model.
End of explanation
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);
    }
}
sm = pystan.StanModel(model_code=refit_lr_code)
data_dict = {
    "N": len(ydata),
    "y": ydata,
    "x": xdata,
    # No excluded data in initial fit
    "N_ex": 0,
    "x_ex": [],
    "y_ex": [],
}
sample_kwargs = {"iter": 1000, "chains": 4}
fit = sm.sampling(data=data_dict, **sample_kwargs)
Explanation: Now we will write the Stan code, keeping in mind that it must be able to compute the pointwise log likelihood on excluded data, i.e., data that is not used to fit the model. Thus, the backbone of the code must look like the following:
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
}
End of explanation
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, **idata_kwargs)
Explanation: 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 will follow the same pattern with az.from_pystan.
End of explanation
class LinearRegressionWrapper(az.PyStan2SamplingWrapper):
    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": N_obs - N_ex,
            "x": xdata[mask],
            "y": ydata[mask],
            "N_ex": N_ex,
            "x_ex": xdata[~mask],
            "y_ex": ydata[~mask],
        }
        return observations, "log_lik_ex"
loo_orig = az.loo(idata, pointwise=True)
loo_orig
Explanation: We will create a subclass of {class}~arviz.PyStan2SamplingWrapper. Therefore, instead of having to implement all functions required by {func}~arviz.reloo we only have to implement {func}~arviz.PyStan2SamplingWrapper.sel_observations. As explained in its docs, it takes one argument which is 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).
End of explanation
loo_orig.pareto_k[[13, 42, 56, 73]] = np.array([0.8, 1.2, 2.6, 0.9])
Explanation: 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 {func}~arviz.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.
End of explanation
pystan_wrapper = LinearRegressionWrapper(
    sm, idata_orig=idata, sample_kwargs=sample_kwargs, idata_kwargs=idata_kwargs
)
Explanation: We initialize our sampling wrapper
End of explanation
loo_relooed = az.reloo(pystan_wrapper, loo_orig=loo_orig)
loo_relooed
loo_orig
Explanation: And eventually, we can use this wrapper to call az.reloo, and compare the results with the PSIS LOO-CV results.
End of explanation
<END_TASK>