<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 (working_with_InferenceData)=
Working with InferenceData
Here we present a collection of common manipulations you can use while working with InferenceData.
Step1:  display_expand_data=False makes the default view for {class}xarray.DataArray fold the data values to a single line. To explore the values, click on the {fas}database icon on the left of the view, right under the xarray.DataArray text. It has no effect on Dataset objects that already default to folded views.
display_expand_attrs=False folds the attributes in both DataArray and Dataset objects to keep the views shorter. In this page we print DataArrays and Datasets several times and they always have the same attributes.
Step2:  Get the dataset corresponding to a single group
Step3:  
Step4:  Combine chains and draws
Step5:  You can also use {meth}xarray.Dataset.stack if you only want to combine the chain and draw dimensions. {func}arviz.extract_dataset is a convenience function aimed at taking care of the most common subsetting operations with MCMC samples. It can
Step6:  
Step7:  Get the dimension lengths
Let’s check how many groups are in our hierarchical model.
Step8:  Get coordinate values
What are the names of the groups in our hierarchical model? You can access them from the coordinate name school in this case
Step9:  Get a subset of chains
Let’s keep only chain 0 and 2 here. For the subset to take effect on all relevant InferenceData groups
Step10:  Remove the first n draws (burn-in)
Let’s say we want to remove the first 100 samples, from all the chains and all InferenceData groups with draws.
Step11:  If you check the burnin object you will see that the groups posterior, posterior_predictive, prior and sample_stats have 400 draws compared to idata that has 500. The group observed_data has not been affected because it does not have the draw dimension. Alternatively, you can specify which group or groups you want to change.
Step12:  Compute posterior mean values along draw and chain dimensions
To compute the mean value of the posterior samples, do the following
Step13:  This computes the mean along all dimensions. This is probably what you want for mu and tau, which have two dimensions (chain and draw), but maybe not what you expected for theta, which has one more dimension school. 
You can specify along which dimension you want to compute the mean (or other functions).
Step14:  Compute and store posterior pushforward quantities
We use "posterior pushfoward quantities" to refer to quantities that are not variables in the posterior but deterministic computations using posterior variables. 
You can use xarray for these pushforward operations and store them as a new variable in the posterior group. You'll then be able to plot them with ArviZ functions, calculate stats and diagnostics on them (like the {func}~arviz.mcse) or save and share the inferencedata object with the pushforward quantities included. 
Compute the rolling mean of $\log(\tau)$ with {meth}xarray.DataArray.rolling, storing the result in the posterior
Step15:  Using xarray for pusforward calculations has all the advantages of working with xarray. It also inherits the disadvantages of working with xarray, but we believe those to be outweighed by the advantages, and we have already shown how to extract the data as NumPy arrays. Working with InferenceData is working mainly with xarray objects and this is what is shown in this guide.
Some examples of these advantages are specifying operations with named dimensions instead of positional ones (as seen in some previous sections), 
automatic alignment and broadcasting of arrays (as we'll see now),
or integration with Dask (as shown in the {ref}dask_for_arviz guide).
In this cell you will compute pairwise differences between schools on their mean effects (variable theta).
To do so, substract the variable theta after renaming the school dimension to the original variable. 
Xarray then aligns and broadcasts the two variables because they have different dimensions, and
the result is a 4d variable with all the pointwise differences.
Eventually, store the result in the theta_school_diff variable
Step16:  
Step17:  Advanced subsetting
To select the value corresponding to the difference between the Choate and Deerfield schools do
Step18:  For more advanced subsetting (the equivalent to what is sometimes called "fancy indexing" in NumPy) you need to provide the indices as {class}~xarray.DataArray objects
Step19:  Using lists or NumPy arrays instead of DataArrays does colum/row based indexing. As you can see, the result has 9 values of theta_shool_diff instead of the 3 pairs of difference we selected in the previous cell
Step20:  Add new chains using concat
After checking the {func}~arviz.mcse and realizing you need more samples, you rerun the model with two chains
and obtain an idata_rerun object.
Step21:  You can combine the two into a single InferenceData object using {func}arviz.concat
Step22:  Add groups to InferenceData objects
You can also add new groups to InferenceData objects with the {meth}~arviz.InferenceData.extend (if the new groups are already in an InferenceData object) or with {meth}~arviz.InferenceData.add_groups (if the new groups are dictionaries or xarray.Dataset objects).
<ASSISTANT_TASK:>
Python Code:
import arviz as az
import numpy as np
import xarray as xr
xr.set_options(display_expand_data=False, display_expand_attrs=False);
Explanation: (working_with_InferenceData)=
Working with InferenceData
Here we present a collection of common manipulations you can use while working with InferenceData.
End of explanation
idata = az.load_arviz_data("centered_eight")
idata
Explanation: display_expand_data=False makes the default view for {class}xarray.DataArray fold the data values to a single line. To explore the values, click on the {fas}database icon on the left of the view, right under the xarray.DataArray text. It has no effect on Dataset objects that already default to folded views.
display_expand_attrs=False folds the attributes in both DataArray and Dataset objects to keep the views shorter. In this page we print DataArrays and Datasets several times and they always have the same attributes.
End of explanation
post = idata.posterior
post
Explanation: Get the dataset corresponding to a single group
End of explanation
post["log_tau"] = np.log(post["tau"])
idata.posterior
Explanation: :::{tip} 
You'll have noticed we stored the posterior group in a new variable: post. As .copy() was not called, now using idata.posterior or post is equivalent.
Use this to keep your code short yet easy to read. Store the groups you'll need very often as separate variables to use explicitly, but don't delete the InferenceData parent. You'll need it for many ArviZ functions to work properly. For example: {func}~arviz.plot_pair needs data from sample_stats group to show divergences, {func}~arviz.compare needs data from both log_likelihood and posterior groups, {func}~arviz.plot_loo_pit needs not 2 but 3 groups: log_likelihood, posterior_predictive and posterior.
:::
Add a new variable
End of explanation
stacked = az.extract_dataset(idata)
stacked
Explanation: Combine chains and draws
End of explanation
az.extract_dataset(idata, num_samples=100)
Explanation: You can also use {meth}xarray.Dataset.stack if you only want to combine the chain and draw dimensions. {func}arviz.extract_dataset is a convenience function aimed at taking care of the most common subsetting operations with MCMC samples. It can:
- Combine chains and draws
- Return a subset of variables (with optional filtering with regular expressions or string matching)
- Return a subset of samples. Moreover by default it returns a random subset to prevent getting non-representative samples due to bad mixing.
- Acess any group
(idata/random_subset)=
Get a random subset of the samples
End of explanation
stacked.mu.values
Explanation: :::{tip}
Use a random seed to get the same subset from multiple groups: az.extract_dataset(idata, num_samples=100, rng=3) and az.extract_dataset(idata, group="log_likelihood", num_samples=100, rng=3) will continue to have matching samples
:::
Obtain a NumPy array for a given parameter
Let's say we want to get the values for mu as a NumPy array.
End of explanation
len(idata.observed_data.school)
Explanation: Get the dimension lengths
Let’s check how many groups are in our hierarchical model.
End of explanation
idata.observed_data.school
Explanation: Get coordinate values
What are the names of the groups in our hierarchical model? You can access them from the coordinate name school in this case
End of explanation
idata.sel(chain=[0, 2])
Explanation: Get a subset of chains
Let’s keep only chain 0 and 2 here. For the subset to take effect on all relevant InferenceData groups: posterior, sample_stats, log_likelihood, posterior_predictive we will use the {meth}arviz.InferenceData.sel, the method of InferenceData instead of {meth}xarray.Dataset.sel.
End of explanation
idata.sel(draw=slice(100, None))
Explanation: Remove the first n draws (burn-in)
Let’s say we want to remove the first 100 samples, from all the chains and all InferenceData groups with draws.
End of explanation
idata.sel(draw=slice(100, None), groups="posterior")
Explanation: If you check the burnin object you will see that the groups posterior, posterior_predictive, prior and sample_stats have 400 draws compared to idata that has 500. The group observed_data has not been affected because it does not have the draw dimension. Alternatively, you can specify which group or groups you want to change.
End of explanation
post.mean()
Explanation: Compute posterior mean values along draw and chain dimensions
To compute the mean value of the posterior samples, do the following:
End of explanation
post.mean(dim=['chain', 'draw'])
Explanation: This computes the mean along all dimensions. This is probably what you want for mu and tau, which have two dimensions (chain and draw), but maybe not what you expected for theta, which has one more dimension school. 
You can specify along which dimension you want to compute the mean (or other functions).
End of explanation
post["mlogtau"] = post["log_tau"].rolling({'draw': 50}).mean()
Explanation: Compute and store posterior pushforward quantities
We use "posterior pushfoward quantities" to refer to quantities that are not variables in the posterior but deterministic computations using posterior variables. 
You can use xarray for these pushforward operations and store them as a new variable in the posterior group. You'll then be able to plot them with ArviZ functions, calculate stats and diagnostics on them (like the {func}~arviz.mcse) or save and share the inferencedata object with the pushforward quantities included. 
Compute the rolling mean of $\log(\tau)$ with {meth}xarray.DataArray.rolling, storing the result in the posterior
End of explanation
post['theta_school_diff'] = post.theta - post.theta.rename(school="school_bis")
Explanation: Using xarray for pusforward calculations has all the advantages of working with xarray. It also inherits the disadvantages of working with xarray, but we believe those to be outweighed by the advantages, and we have already shown how to extract the data as NumPy arrays. Working with InferenceData is working mainly with xarray objects and this is what is shown in this guide.
Some examples of these advantages are specifying operations with named dimensions instead of positional ones (as seen in some previous sections), 
automatic alignment and broadcasting of arrays (as we'll see now),
or integration with Dask (as shown in the {ref}dask_for_arviz guide).
In this cell you will compute pairwise differences between schools on their mean effects (variable theta).
To do so, substract the variable theta after renaming the school dimension to the original variable. 
Xarray then aligns and broadcasts the two variables because they have different dimensions, and
the result is a 4d variable with all the pointwise differences.
Eventually, store the result in the theta_school_diff variable:
End of explanation
post
Explanation: :::{note}
:class: dropdown
This same operation using NumPy would require manual alignment of the two arrays to make sure they broadcast correctly. The could would be something like:
python
theta_school_diff = theta[:, :, :, None] - theta[:, :, None, :]
:::
The theta_shool_diff variable in the posterior has kept the named dimensions and coordinates:
End of explanation
post['theta_school_diff'].sel(school="Choate", school_bis="Deerfield")
Explanation: Advanced subsetting
To select the value corresponding to the difference between the Choate and Deerfield schools do:
End of explanation
school_idx = xr.DataArray(["Choate", "Hotchkiss", "Mt. Hermon"], dims=["pairwise_school_diff"])
school_bis_idx = xr.DataArray(["Deerfield", "Choate", "Lawrenceville"], dims=["pairwise_school_diff"])
post['theta_school_diff'].sel(school=school_idx, school_bis=school_bis_idx)
Explanation: For more advanced subsetting (the equivalent to what is sometimes called "fancy indexing" in NumPy) you need to provide the indices as {class}~xarray.DataArray objects:
End of explanation
post['theta_school_diff'].sel(
    school=["Choate", "Hotchkiss", "Mt. Hermon"], 
    school_bis=["Deerfield", "Choate", "Lawrenceville"]
)
Explanation: Using lists or NumPy arrays instead of DataArrays does colum/row based indexing. As you can see, the result has 9 values of theta_shool_diff instead of the 3 pairs of difference we selected in the previous cell:
End of explanation
idata_rerun = idata.sel(chain=[0, 1]).copy().assign_coords(coords={"chain":[4,5]},groups="posterior_groups")
Explanation: Add new chains using concat
After checking the {func}~arviz.mcse and realizing you need more samples, you rerun the model with two chains
and obtain an idata_rerun object.
End of explanation
idata_complete = az.concat(idata, idata_rerun, dim="chain")
idata_complete.posterior.dims["chain"]
Explanation: You can combine the two into a single InferenceData object using {func}arviz.concat:
End of explanation
rng = np.random.default_rng(3)
idata.add_groups(
    {"predictions": {"obs": rng.normal(size=(4, 500, 2))}}, 
    dims={"obs": ["new_school"]}, 
    coords={"new_school": ["Essex College", "Moordale"]}
)
idata
Explanation: Add groups to InferenceData objects
You can also add new groups to InferenceData objects with the {meth}~arviz.InferenceData.extend (if the new groups are already in an InferenceData object) or with {meth}~arviz.InferenceData.add_groups (if the new groups are dictionaries or xarray.Dataset objects).
End of explanation
<END_TASK>