[BUG-REPORT] Missing colorbar when using slices

[Thalos12]

Description
I am trying the Slices in a 3rd dimension example in a Jupyter notebook, and the result does not have colorbars.

The piece of code I am using:

import vaex
df = vaex.example()
df.viz.heatmap("Lz", "E", z="FeH:-3,-1,8",
               visual=dict(row="z"),
               figsize=(12, 8),
               f="log",
               wrap_columns=3,
               limits='99%');

I expected this to appear, but instead I got

Software information

• Vaex version:

{'vaex-core': '4.16.1',
 'vaex-viz': '0.5.4',
 'vaex-hdf5': '0.14.1',
 'vaex-server': '0.8.1',
 'vaex-astro': '0.9.3',
 'vaex-jupyter': '0.8.1',
 'vaex-ml': '0.18.1'}

• Vaex was installed via: conda-forge
• OS: Ubuntu 20.04

Additional information
I believe I have found a hint towards a solution, but I do not know if this is applicable to any situation. In vaex/viz/mpl.py, on line 802, ax is defined and the colorbar is added. On line 810 it is re-defined.
Changing the code to use the ax created before makes the colorbars appear in the plot.
I replaced

            if facets > 1:
                ax = plt.subplot(gs[row_offset + row * row_scale:row_offset + (row + 1) * row_scale, column * column_scale:(column + 1) * column_scale])
            else:
                ax = plt.gca()

with

            if facets > 1 and colorbar_location!='individual':
                    ax = plt.subplot(gs[row_offset + row * row_scale:row_offset + (row + 1) * row_scale, column * column_scale:(column + 1) * column_scale])
            else:
                ax = plt.gca()

I hope this proves useful to track down the bug, thank you for your time.

[Edited to remove double if]

[JovanVeljanoski]

Hi,

Thanks for the report! Would you like to open a PR?

[Thalos12]

I can certainly open a PR!