Skip to content

Instantly share code, notes, and snippets.

@jonnyhtw
Forked from jakevdp/discrete_cmap.py
Created July 29, 2020 01:52
Show Gist options
  • Save jonnyhtw/0c965e300577b1abac0e396a7edb9fa4 to your computer and use it in GitHub Desktop.
Save jonnyhtw/0c965e300577b1abac0e396a7edb9fa4 to your computer and use it in GitHub Desktop.
Small utility to create a discrete matplotlib colormap
# By Jake VanderPlas
# License: BSD-style
import matplotlib.pyplot as plt
import numpy as np
def discrete_cmap(N, base_cmap=None):
"""Create an N-bin discrete colormap from the specified input map"""
# Note that if base_cmap is a string or None, you can simply do
# return plt.cm.get_cmap(base_cmap, N)
# The following works for string, None, or a colormap instance:
base = plt.cm.get_cmap(base_cmap)
color_list = base(np.linspace(0, 1, N))
cmap_name = base.name + str(N)
return base.from_list(cmap_name, color_list, N)
if __name__ == '__main__':
N = 5
x = np.random.randn(40)
y = np.random.randn(40)
c = np.random.randint(N, size=40)
# Edit: don't use the default ('jet') because it makes @mwaskom mad...
plt.scatter(x, y, c=c, s=50, cmap=discrete_cmap(N, 'cubehelix'))
plt.colorbar(ticks=range(N))
plt.clim(-0.5, N - 0.5)
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment