Search code examples
pythonmatplotlibarviz

How to swap background colour of stripes in arviz.plot_forest

This might be a simple question but I can't figure it out. In arviz's arviz.plot_forest, how can I swap the order of the shaded backgrounds? For instance, in this example figure from their docs, how can I start with grey and have the odd rows white? Or how do I define and customise this formatting?

enter image description here

Thanks.

Edit after @Onuralp Arslan's answer (6.2.).

The issue I'm facing is that there remains some colour from the original plot, so repainting it to other colours (like red and yellow) works but painting them white for some reason doesn't:

non_centered_data = az.load_arviz_data('non_centered_eight')
centered_data = az.load_arviz_data('centered_eight')
axes = az.plot_forest([non_centered_data, centered_data],
                           model_names = ["non centered eight", "centered eight"],
                           kind='forestplot',
                           var_names=["^the"],
                           filter_vars="regex",
                           combined=True,
                           figsize=(9, 7))
axes[0].set_title('Estimated theta for 8 schools models')


# flip the grey and white
ax=axes[0]
y_ticks = ax.get_yticks()
 
# to calculate where to color / size of plot
y_min, y_max = ax.get_ylim()  
total_height = y_max - y_min  
num_rows = len(y_ticks)  
row_height = total_height / num_rows

# even odd alternate
for i, y in enumerate(y_ticks):
    bottom = y - row_height / 2
    top = y + row_height / 2

    if i % 2 == 0:  
        ax.axhspan(bottom, top, color='white', alpha=1, zorder=0)
    else:  
        ax.axhspan(bottom, top, color='lightgrey', alpha=0.2, zorder=-1)

ax.set_axisbelow(True)

produces: enter image description here

Solution

  • Ok, I found a workaround. The issue was that arviz is producing patches for coloured stripes that stay there and no patches for white stripes. So one first has to remove these patches before proceeding with the solution outlined by Onuralp Arslan, thank you for that idea. It is not possible to just repaint the patches as they do not exist for the white background. Note that the repainting can be done in only two lines if one of the colours is white, such as:

    if i % 2 != 0:
    axes[0].axhspan(bottom, top, color='lightgrey', alpha=0.2, zorder=-1)

    The full code is:

    non_centered_data = az.load_arviz_data('non_centered_eight')
centered_data = az.load_arviz_data('centered_eight')
axes = az.plot_forest([non_centered_data, centered_data],
                           model_names = ["non centered eight", "centered eight"],
                           kind='forestplot',
                           var_names=["^the"],
                           filter_vars="regex",
                           combined=True,
                           figsize=(9, 7))
axes[0].set_title('Estimated theta for 8 schools models')

# remove the grey background
for ax in axes:
    for i, patch in enumerate(ax.patches):
        patch.set_facecolor("none")
        patch.set_edgecolor("none")

# now repaint the stripes as proposed by Onuralp Arslan
y_ticks = axes[0].get_yticks()
 
# to calculate where to color / size of plot
y_min, y_max = ax.get_ylim()  
total_height = y_max - y_min  
num_rows = len(y_ticks)  
row_height = total_height / num_rows

# even odd alternate
for i, y in enumerate(y_ticks):
    bottom = y - row_height / 2
    top = y + row_height / 2

    if i % 2 == 0:
        axes[0].axhspan(bottom, top, color='white', alpha=1, zorder=0)
    else:  
        axes[0].axhspan(bottom, top, color='lightgrey', alpha=0.2, zorder=-1)

    which produces the required: enter image description here