arviz-devs / arviz-devs/arviz

select individual levels from the Dimensions to plot pystan

Open
#595 3 comments 0 reactions 0 assignees View on GitHub
Usability User Documentation
Dominant language
TeX
Stars
1.9k
Forks
508
Avg merge
2d 20h
Merged PRs (30d)
3

Description

## Short Description

I want to select individual levels from the Dimensions to plot because plotting all of the levels of a variable is slow and the plot uninterpretable.

## Code Example or link

I am trying to reproduce the [PyStan example here](https://mc-stan.org/users/documentation/case-studies/radon.html) showing the use of multilevel modelling.

The code and extraction are below:
```python
varying_intercept = """
data {
int J; // the number of counties
int N; // the number of observations
int county[N]; // the county for each observation
vector[N] x; // predictor/regressor (floor/basement)
vector[N] y; // the output variable (log radon levels)
}
parameters {
vector[J] a; // the random intercept
real b; // the fixed coefficient (FIXED EFFECT)
real mu_a; // mean of the population of counties (CONSTANT FOR POPULATION)
real sigma_a; // the variance of the counties (CONSTANT FOR POPULATION)
real sigma_y; // the variance of the observations (CONSTANT FOR POPULATION)
}
transformed parameters {
vector[N] y_hat; // estimated log radon level for each datapoint

for (i in 1:N) // for each datapoint
y_hat[i] <- a[county[i]] + x[i] * b; // estimate the mean of the log radon as a simple linear regression
}
model {
sigma_a ~ uniform(0, 100); // variation between the counties
a ~ normal (mu_a, sigma_a); // the intercept varying (RANDOM EFFECT)

b ~ normal (0, 1); // the coefficient

sigma_y ~ uniform(0, 100); // the sampling variation of the log-radon
y ~ normal(y_hat, sigma_y); // model the log radon levels
}
"""

varying_intercept_data = {'N': len(log_radon),
'J': len(n_county),
'county': county+1, # Stan counts starting at 1
'x': floor_measure,
'y': log_radon}

varying_intercept_fit = pystan.stan(model_code=varying_intercept, data=varying_intercept_data, iter=1000, chains=2)
```
I then extract the data to ArViz
```python
fit = varying_intercept_fit

data = az.from_pystan(posterior=fit,
posterior_predictive='y_hat',
observed_data=['y'],
coords={'county': n_county},
dims={'a': ['county']}) #, 'y': ['county'], 'log_lik': ['county'], 'y_hat': ['county'], 'theta_tilde': ['county']})

data

Out[]:
Inference data with groups:
> posterior
> sample_stats
> posterior_predictive
> observed_data
```

I want to make a plot of only a few of the counties (the model levels). The following takes an age to run because it is plotting ALL counties traces, but I want to select them.
```python
az.traceplot(data)
```

I found this [help here](https://discourse.pymc.io/t/best-way-to-plot-and-do-ppc-with-variable-that-has-too-many-levels/2276/3):
```python
az.plot_trace(data, var_names='a', coords={'county': range(0, 5)});
az.plot_forest(data.posterior.sel(county=range(0, 5)), var_names='a');
az.plot_parallel(data, var_names='a', coords={'county': range(0, 5)});
az.plot_posterior(data, var_names='a', coords={'county': range(0, 5)});
```

But I get an error:
```
---------------------------------------------------------------------------
InvalidIndexError Traceback (most recent call last)
in
1 # select only some of the levels to plot!!!
2 # https://discourse.pymc.io/t/best-way-to-plot-and-do-ppc-with-variable-that-has-too-many-levels/2276/3
----> 3 az.plot_trace(data, var_names='a', coords={'county': range(0, 5)});
4 # az.plot_forest(data.posterior.sel(county=range(0, 5)), var_names='a');
5 # az.plot_parallel(data, var_names='a', coords={'county': range(0, 5)});

~/miniconda3/envs/stan/lib/python3.7/site-packages/arviz/plots/traceplot.py in plot_trace(data, var_names, coords, divergences, figsize, textsize, lines, combined, kde_kwargs, hist_kwargs, trace_kwargs)
105 lines = ()
106
--> 107 plotters = list(xarray_var_iter(get_coords(data, coords), var_names=var_names, combined=True))
108
109 if figsize is None:

~/miniconda3/envs/stan/lib/python3.7/site-packages/arviz/plots/plot_utils.py in get_coords(data, coords)
322 """
323 try:
--> 324 return data.sel(**coords)
325
326 except ValueError:

~/miniconda3/envs/stan/lib/python3.7/site-packages/xarray/core/dataset.py in sel(self, indexers, method, tolerance, drop, **indexers_kwargs)
1608 indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'sel')
1609 pos_indexers, new_indexes = remap_label_indexers(
-> 1610 self, indexers=indexers, method=method, tolerance=tolerance)
1611 result = self.isel(indexers=pos_indexers, drop=drop)
1612 return result._replace_indexes(new_indexes)

~/miniconda3/envs/stan/lib/python3.7/site-packages/xarray/core/coordinates.py in remap_label_indexers(obj, indexers, method, tolerance, **indexers_kwargs)
353
354 pos_indexers, new_indexes = indexing.remap_label_indexers(
--> 355 obj, v_indexers, method=method, tolerance=tolerance
356 )
357 # attach indexer's coordinate to pos_indexers

~/miniconda3/envs/stan/lib/python3.7/site-packages/xarray/core/indexing.py in remap_label_indexers(data_obj, indexers, method, tolerance)
256 else:
257 idxr, new_idx = convert_label_indexer(index, label,
--> 258 dim, method, tolerance)
259 pos_indexers[dim] = idxr
260 if new_idx is not None:

~/miniconda3/envs/stan/lib/python3.7/site-packages/xarray/core/indexing.py in convert_label_indexer(index, label, index_name, method, tolerance)
192 raise ValueError('Vectorized selection is not available along '
193 'MultiIndex variable: ' + index_name)
--> 194 indexer = get_indexer_nd(index, label, method, tolerance)
195 if np.any(indexer < 0):
196 raise KeyError('not all values found in index %r'

~/miniconda3/envs/stan/lib/python3.7/site-packages/xarray/core/indexing.py in get_indexer_nd(index, labels, method, tolerance)
120
121 flat_labels = np.ravel(labels)
--> 122 flat_indexer = index.get_indexer(flat_labels, **kwargs)
123 indexer = flat_indexer.reshape(labels.shape)
124 return indexer

~/miniconda3/envs/stan/lib/python3.7/site-packages/pandas/core/indexes/base.py in get_indexer(self, target, method, limit, tolerance)
2737
2738 if not self.is_unique:
-> 2739 raise InvalidIndexError('Reindexing only valid with uniquely'
2740 ' valued Index objects')
2741

InvalidIndexError: Reindexing only valid with uniquely valued Index objects
```

Also include the ArviZ version and version of any other relevant packages.
```
Arviz Version: 0.3.2
numpy Version: 1.15.0
pandas Version: 0.24.1
```
## Relevant documentation or public examples
https://mc-stan.org/users/documentation/case-studies/radon.html

Contributor guide

Open the contributing guide

Research direction

Start with the traceback entry points in arviz/plots/traceplot.py and arviz/plots/plot_utils.py, then reproduce az.plot_trace(data, var_names='a', coords={'county': range(0, 5)}) using the supplied PyStan example. Done means plotting only the requested county levels without InvalidIndexError; compare the related plot_forest, plot_parallel, and plot_posterior calls as needed.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, pandas, python
Domain
data-visualization
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.