# How to forestplot in arviz with only a subset of trace variables and samples

**URL:** <https://discourse.pymc.io/t/how-to-forestplot-in-arviz-with-only-a-subset-of-trace-variables-and-samples/3656>\
**Category:** Questions\
**Created:** [July 30, 2019, 7:43pm UTC](https://discourse.pymc.io/t/how-to-forestplot-in-arviz-with-only-a-subset-of-trace-variables-and-samples/3656 "2019-07-30T19:43:02Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![clausherther](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/clausherther/32/414_2.png) [@clausherther](https://discourse.pymc.io/u/clausherther)\
**Post date:** [July 30, 2019, 7:43pm UTC](https://discourse.pymc.io/t/how-to-forestplot-in-arviz-with-only-a-subset-of-trace-variables-and-samples/3656/1 "2019-07-30T19:43:02Z")

</div>

I think I’m answering my own question here, so this is to share and get feedback (if any).

I’ve been wanting to use the `arviz.plot_forest()` function to visualize a trace. However, the trace contains a number of variables, and for each variable, hundreds if not thousands of parameters.

For example, I have a trace that contains posteriors for a parameter for 156 products by 52 weeks. I wanted to create a forestplot for only one of those products.  
My data is arranged row-wise, so I need to index into the correct product.

Here’s what I’ve come up with:

```python
import numpy as np
import arviz as az

```

```python
weeks = [...array of 52 weeks...]
products = [....array of 156 products...]
product_week_idx = [...index of 8,112 product/week combinations]
trace = [...PyMC3 trace with 6000 samples from 3 chains, 8112 parameters per variable]

# e.g. 
# trace["my_trace_variable].shape
# > (6000, 8112)

```

```python
# find index of selected product in products array
product_name = "The Product I Want to Plot"
selected_product_index = np.where(products == product_name)[0][0]

```

```python
# find index of selected product/week combos in products x weeks array
selected_product_week_index = np.where(product_week_idx == selected_product_index)[0]
# selected_product_week_index.shape
# > (52,)

```

The key is to convert the relevant subset of the trace to an `xarray.Dataset` before passing it to `forestplot`:

```python
trace_data = az.from_pymc3(trace=trace)

plot_data = (product_week_trace_data
             .posterior["my_trace_variable"][:,:, selected_product_week_index]
             .to_dataset()
            )

```

```python
_, axes = az.plot_forest(plot_data, 
               credible_interval=0.95,
               combined=True)
# reverse the week names for the plot
axes[0].set_yticklabels(weeks[::-1]);
axes[0].set_title(product_name)

```

---

<div class="post-metadata">

**Author:** ![aseyboldt](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/aseyboldt/32/5795_2.png) [@aseyboldt](https://discourse.pymc.io/u/aseyboldt)\
**Post date:** [July 30, 2019, 9:04pm UTC](https://discourse.pymc.io/t/how-to-forestplot-in-arviz-with-only-a-subset-of-trace-variables-and-samples/3656/2 "2019-07-30T21:04:19Z")

</div>

If you use the `dims` and `coords` kwargs in `az.from_pymc3` you could get rid of all that index juggling. You can then plot using something like

```python
plot_data = (product_week_trace_data
             .posterior["my_trace_variable"].sel(product=list_of_product_names)
             .to_dataset()
            )

```

---

<div class="post-metadata">

**Author:** ![colcarroll](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/colcarroll/32/9_2.png) [@colcarroll](https://discourse.pymc.io/u/colcarroll)\
**Post date:** [July 30, 2019, 9:06pm UTC](https://discourse.pymc.io/t/how-to-forestplot-in-arviz-with-only-a-subset-of-trace-variables-and-samples/3656/3 "2019-07-30T21:06:02Z")

</div>

The `.to_dataset` is a nice method I hadn’t used.

I was hoping you could use named coordinates and dimensions in a nice way here: it is difficult because I don’t _think_ xarray supports compound indices, and it seems like that is what you want with `product_to_index`.

---

<div class="post-metadata">

**Author:** ![aseyboldt](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/aseyboldt/32/5795_2.png) [@aseyboldt](https://discourse.pymc.io/u/aseyboldt)\
**Post date:** [July 30, 2019, 9:28pm UTC](https://discourse.pymc.io/t/how-to-forestplot-in-arviz-with-only-a-subset-of-trace-variables-and-samples/3656/4 "2019-07-30T21:28:29Z")

</div>

Hm, seems I didn’t read the question properly 🙂  
How about something like this?

```auto
ds = xr.Dataset({
    'item': ('item', ['item1', 'item2', 'item3']),
    'week': ['week1', 'week2', 'week3'],
    'product': ('product', ['bar', 'foo']),

    'item_product': ('item', ['bar', 'foo', 'bar']),
    'item_week': ('item', ['week1', 'week2', 'week1']),
    'item_data': ('item', [1, 2, 3]),
})
ds

```

```auto
<xarray.Dataset>
Dimensions: (item: 3, product: 2, week: 3)
Coordinates:
  * item (item) <U5 'item1' 'item2' 'item3'
  * week (week) <U5 'week1' 'week2' 'week3'
  * product (product) <U3 'bar' 'foo'
Data variables:
    item_product (item) <U3 'bar' 'foo' 'bar'
    item_week (item) <U5 'week1' 'week2' 'week1'
    item_data (item) int64 1 2 3

```

```auto
ds.sel(item=ds.item_product.isin(['bar']))

```

```auto
<xarray.Dataset>
Dimensions: (item: 2, product: 2, week: 3)
Coordinates:
  * item (item) <U5 'item1' 'item3'
  * week (week) <U5 'week1' 'week2' 'week3'
  * product (product) <U3 'bar' 'foo'
Data variables:
    item_product (item) <U3 'bar' 'bar'
    item_week (item) <U5 'week1' 'week1'
    item_data (item) int64 1 3

```

## Update

You can also tell xarray that item is a multiindex thing:

```auto
ds.set_index(item=['item_product', 'item_week'])

```

I’m not too convinced by the current Multiindex support in xarray though, so maybe it is better to do this manually for now.

---

<div class="post-metadata">

**Author:** ![clausherther](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/clausherther/32/414_2.png) [@clausherther](https://discourse.pymc.io/u/clausherther)\
**Post date:** [July 30, 2019, 9:29pm UTC](https://discourse.pymc.io/t/how-to-forestplot-in-arviz-with-only-a-subset-of-trace-variables-and-samples/3656/5 "2019-07-30T21:29:56Z")

</div>

Nice! I’ll give that a shot. Thx!