Monday 2 May 2016

A Sunset Colormap for Python

Room Z533 in the Kruyt building in Utrecht can have a wonderful view, especially at sunset. Not too long ago, I took the following photograph.

Having a love-hate relationship with colormaps, and in particular choosing a good one, I INSTANTLY noticed the beauty of the gradient of the sky, and hence, its application as a colormap.

In Python, it is easy to import figures as a matrix, and then extract the RGB values along a line. To make things easier, I first cropped out the part of the sky I wanted:
Then, I had a look at this website.
Let's import the "slice of sky" in Python and have a look:
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
from PIL import Image

im = Image.open("/path/to/sunset.jpg")
pix = im.load()

xs = range(im.size[1])
y = im.size[0]/2
rs = [pix[y,x][0] for x in xs]
gs = [pix[y,x][1] for x in xs]
bs = [pix[y,x][2] for x in xs]

fig = plt.figure(figsize=(5,2))
ax = fig.add_subplot(111)

ax.plot(xs, rs, color='red')
ax.plot(xs, gs, color='green')
ax.plot(xs, bs, color='blue')
ax.set_xlabel("pixel")
ax.set_ylabel('r/g/b value')

plt.savefig('rgb-plot.png', dpi=300, bboxinches='tight')

This results in the following figure:

Now let's make the actual colormap.
Dp = 1 ## determines the number of pixels between "nodes"
xs = np.linspace(0,1,im.size[1]/Dp + 1) ## points between 0 and 1 
idxs = range(0,im.size[1],Dp) ## indices in the original picture matrix

redlist = [rs[idx]/255. for idx in idxs] + [rs[-1]/255., 0]
greenlist = [gs[idx]/255. for idx in idxs] + [gs[-1]/255., 0]
bluelist = [bs[idx]/255. for idx in idxs] + [bs[-1]/255., 0]

## LinearSegmentedColormap wants these weird triples, 
## where some end points are ignored...

redtuples = [(x, redlist[i], redlist[i+1]) for i, x in enumerate(xs)]
greentuples = [(x, greenlist[i], greenlist[i+1]) for i, x in enumerate(xs)]
bluetuples = [(x, bluelist[i], bluelist[i+1]) for i, x in enumerate(xs)]

cdict = {'red' : redtuples, 'green' : greentuples, 'blue' : bluetuples}

cmap = LinearSegmentedColormap('tbz533', cdict) ## choose a name
plt.register_cmap(cmap=cmap)
## now we can use our new colormap!

OK. Let's make some 70s wallpaper to celebrate the new colormap:
xs = np.linspace(-5,5,1000)
ys = np.linspace(-5,5,1000)
zs = np.array([[np.sin(x)*np.cos(y) for x in xs] for y in ys])

fig = plt.figure(figsize=(5.5,5))
ax1 = fig.add_subplot(1,1,1)
C = ax1.pcolormesh(xs, ys, zs, cmap='tbz533')
fig.colorbar(C)

fig.savefig('my70swallpaper.png', dpi=200, bboxinches='tight')

This results in the following figure:

That's it! Please leave your own favorite colormap in the comments.

1 comment: