Make a scatter colorbar display only a subset of t

2020-07-26 17:33发布

I have a situation where I want to create a colorbar whose colors (associated with a scatter plot) span a particular range, but only display a subset of that range on the colorbar itself. I can do it with contourf, because I can set vmin and vmax independently of the contour levels, but I can't figure out how to do it with scatter. See the following:

import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

x = np.linspace(0, 2*np.pi, 101)
x_arr = np.sin(x)
y_arr = np.cos(x)
arr = y_arr[:,None] * x_arr[None,:]

arr = np.where(arr < 0, arr*4, arr)
ptslist = np.array([-4, -3, -2, -1, 0, 1], dtype=np.float32)

fig, axs = plt.subplots(figsize=(11,5), nrows=1, ncols=2)

# I can achieve my desired behavior with contourf
cont = axs[0].contourf(x, x, arr, levels=np.linspace(-4,1,11),
                       cmap='BrBG_r', vmin=-4, vmax=4)
div0 = make_axes_locatable(axs[0])
cax0 = div0.append_axes('right', '5%', '2%')
plt.colorbar(cont, cax=cax0)

# I can't make this colorbar do a similar thing
sc = axs[1].scatter(np.arange(-4, 2), np.arange(-4, 2), c=ptslist, cmap='BrBG_r',
                    marker='o', s=144, edgecolors='black', vmin=-4, vmax=4)
div1 = make_axes_locatable(axs[1])
cax1 = div1.append_axes('right', '5%', '2%')
cb = plt.colorbar(sc, cax=cax1)

This yields this figure: colorbar -4 to 4

I want the diverging colormap to be centered with white at zero, and the color values displayed linearly on both sides of zero. Both plots do this fine. However, I don't want the extra values from 1 to 4 to display on the right colorbar (see how the left colorbar stops at 1).

My first thought was ylim, but this line:

cb.ax.set_ylim(-4, 1)

causes this strange thing to happen:

colorbar is squished

If I use set_ticks, it just removes absent ticks, and doesn't change the limits. Is there any way to make this happen nicely?

I'm using matplotlib 1.5.0.

p.s. I've also tried a mid-point-centered subclass of Normalize that I found on SO, but it scales the positive and negative values independently, which I don't want (it makes the values of +1.0 dark brown, and I want it to still be light brown, unless I set vmax=4, at which point I have exactly the same problem).

2条回答
Ridiculous、
2楼-- · 2020-07-26 17:59

You can do this in a few different ways, however it sounds like what you're really wanting is a custom colormap created from a portion of another colormap:

For example:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

# Let's create a new colormap from a region of another colormap.
cmap = plt.get_cmap('BrBG_r')
colors = cmap(np.linspace(0, 5 / 8.0, 20))
cmap = LinearSegmentedColormap.from_list('my cmap', colors)

# And now let's plot some data...
x, y, z = np.random.random((3, 10))
z = 5 * z - 4

fig, ax = plt.subplots()
scat = ax.scatter(x, y, c=z, s=200, cmap=cmap, vmin=-4, vmax=1)
cbar = fig.colorbar(scat)

plt.show()

enter image description here

Alternatively, if you'd prefer a discrete colormap, you can do something similar to:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import from_levels_and_colors

# Let's create a new set of colors from a region of another colormap.
ncolors = 10
cmap = plt.get_cmap('BrBG_r')
colors = cmap(np.linspace(0, 5 / 8.0, ncolors - 1))

# We'll then create a new colormap from that set of colors and tie it
# to specific values
levels = np.linspace(-4, 1, ncolors)
cmap, norm = from_levels_and_colors(levels, colors)

x, y, z = np.random.random((3, 10))
z = 5 * z - 4

fig, ax = plt.subplots()
scat = ax.scatter(x, y, c=z, s=200, cmap=cmap, norm=norm)
cbar = fig.colorbar(scat)
cbar.set_ticks(range(-4, 2))

plt.show()

enter image description here

查看更多
相关推荐>>
3楼-- · 2020-07-26 18:26

You can pass the boundaries argument to colorbar:

plt.colorbar(sc, cax=cax1, boundaries=sc.get_array())

I don't know whether sc.get_array() is always the right choice here, but get_array is the ScalarMappable method that is supposed to get the data to be mapped onto colors, so it seems like a reasonable choice. (For contour sets, colorbar automatically grabs the contour levels and uses them as the boundaries.)

查看更多
登录 后发表回答