4

I want to do something with plt.hist2d and plt.colorbar and I'm having real trouble working out how to do it. To explain, I've written the following example:

import numpy as np
from matplotlib import pyplot as plt

x = np.random.random(1e6)
y = np.random.random(1e6)

plt.hist2d(x, y)
plt.colorbar()

plt.show()

This code generates a plot that looks something like the image below. 2D histogram with colorbar plotted using pyplot

If I generate a histogram, ideally I would like the colour bar to extend beyond the maximum and minimum range of the data to the next step beyond the maximum and minimum. In the example in this question, this would set the colour bar extent from 9660 to 10260 in increments of 60.

How can I force either plt.hist2d or plt.colorbar to set the colour bar such that ticks are assigned to the start and end of the plotted colour bar?

John Coxon
  • 570
  • 1
  • 3
  • 15

2 Answers2

3

I think this is what you're looking for:

h = plt.hist2d(x, y)
mn, mx = h[-1].get_clim()
mn = 60 * np.floor(mn / 60.)
mx = 60 * np.ceil(mx / 60.)
h[-1].set_clim(mn, mx)
cbar = plt.colorbar(h[-1], ticks=np.arange(mn, mx + 1, 60), )

This gives something like,

enter image description here

It's also often convenient to use tickers from the matplotlib.ticker, and use the tick_values method of tickers, but for this purpose I think the above is most convenient.

Good luck!

farenorth
  • 10,165
  • 2
  • 39
  • 45
  • It didn't occur to me to wait until after calling `plt.hist2d` to create the locator, but this is pretty much perfect. Couple of things about this solution though: Firstly, it fails if any of the bins are empty (which will return NaN in the `h[0]` variable) and secondly, it doesn't give round numbers for maximum and minimum. I'll probably actually do something like `maximum = 60 * np.ceil(np.nanmax(h[0]) / 60.)` (and similar for the minimum) before calling `locator.tick_values`. But thank you! – John Coxon Sep 15 '15 at 13:22
  • Ah, actually, this doesn't appear to work with `MaxNLocator` at all, I'm still getting the colour bar extending beyond the largest tick. I'll post here when I've worked out why that's happening… – John Coxon Sep 15 '15 at 13:26
  • 1
    Thank you for posting this up, it got me thinking that I could update the objects after plotting them, once I knew what the maximum and minimum were. (I learned in IDL so object-oriented thinking still sometimes eludes me.) I've posted a separate solution, which uses the `set_clim` property of `h[3]` to update the limits after the fact and then plot the colour bar with predefined ticks using the method from this answer. Thanks so much for your help! – John Coxon Sep 15 '15 at 13:53
  • Yep, we came to the same conclusion independently! – farenorth Sep 15 '15 at 13:55
3

With huge thanks to farenorth, who got me thinking about this in the right way, I came up with a function, get_colour_bar_ticks:

def get_colour_bar_ticks(colourbar):
    import numpy as np

    # Get the limits and the extent of the colour bar.
    limits = colourbar.get_clim()
    extent = limits[1] - limits[0]

    # Get the yticks of the colour bar as values (ax.get_yticks() returns them as fractions).
    fractions = colourbar.ax.get_yticks()
    yticks = (fractions * extent) + limits[0]
    increment = yticks[1] - yticks[0]

    # Generate the expanded ticks.
    if (fractions[0] == 0) & (fractions[-1] == 1):
        return yticks

    else:
        start = yticks[0] - increment
        end = yticks[-1] + increment

        if fractions[0] == 0:
            newticks = np.concatenate((yticks, [end]))
        elif fractions[1] == 1:
            newticks = np.concatenate(([start], yticks))
        else:
            newticks = np.concatenate(([start], yticks, [end]))

        return newticks

With this function I can then do this:

from matplotlib import pyplot as plt

x = np.random.random(1e6)
y = np.random.random(1e6)

h = plt.hist2d(x, y)
cbar = plt.colorbar()

ticks = get_colour_bar_ticks(cbar)

h[3].set_clim(ticks[0], ticks[-1])
cbar.set_clim(ticks[0], ticks[-1])
cbar.set_ticks(ticks)

plt.show()

Which results in this, which is what I really wanted:

enter image description here

John Coxon
  • 570
  • 1
  • 3
  • 15