[matplotlib] colormap for visible wavelength

visible.py

import numpy as np
from spectrumRGB import rgb

ww=np.arange(380.,781.)
rgb=rgb()

from matplotlib import pyplot as plt
from matplotlib import mpl

plt.rc(['axes','lines'], lw=2)
plt.rc('lines',markeredgewidth=2)

fig=plt.figure()
ax1 = fig.add_axes([.1, .1, .8, .7])
ax1.plot(ww, rgb[:,0], "r")
ax1.plot(ww, rgb[:,1], "g")
ax1.plot(ww, rgb[:,2], "b")
ax1.set_xlabel('Wavelength (nm)')
ax1.set_ylim(-0.05,1.05)

ax2 = fig.add_axes([.1, .8, .8, .1])

def update_ax2(ax1):
    wmin, wmax = ax1.get_xlim()
    l1,b1,w1,h1 = ax1.get_position().bounds
    ax2box = ax2.get_position()
    [[x1,y1],[x2,y2]] = ax2box.get_points()
    ax2bottom = y1
    ax2width  = x2-x1
    ax2height = y2-y1
    
    ax2.set_position(
        [ l1+w1*( ww[0]-wmin) /(wmax-wmin), ax2bottom,
             w1*(ww[-1]-ww[0])/(wmax-wmin), ax2height])
    ax2.figure.canvas.draw()

ax1.callbacks.connect("xlim_changed", update_ax2)

cmap = mpl.colors.ListedColormap(rgb)
norm = mpl.colors.Normalize(vmin=380., vmax=780.)
cb2 = mpl.colorbar.ColorbarBase(ax2, cmap=cmap,
                                     norm=norm,
                                     orientation='horizontal')
ax2.get_xaxis().set_ticks_position('top')

plt.draw()
update_ax2(ax1)
plt.show()

"""
Python 2.5.4 |EPD 5.1.1| (r254:67916, Sep 25 2009, 12:11:02) [MSC v.1310 32 bit (Intel)] on win32
matplotlib 0.99.1.1
"""

spectrumRGB.py

"""
Based on 
  http://www.physics.sfasu.edu/astro/color/spectra.html
  RGB VALUES FOR VISIBLE WAVELENGTHS   by Dan Bruton (astro@tamu.edu)
"""

import numpy as np
select=np.select
power=np.power
transpose=np.transpose
arange=np.arange

def factor(wl):
    return select(
        [ wl > 700.,
          wl < 420.,
          True ],
        [ .3+.7*(780.-wl)/(780.-700.),
          .3+.7*(wl-380.)/(420.-380.),
          1.0 ] )

def raw_r(wl):
    return select(
        [ wl >= 580.,
          wl >= 510.,
          wl >= 440.,
          wl >= 380.,
          True ],
        [ 1.0,
          (wl-510.)/(580.-510.),
          0.0,
          (wl-440.)/(380.-440.),
          0.0 ] )

def raw_g(wl):
    return select(
        [ wl >= 645.,
          wl >= 580.,
          wl >= 490.,
          wl >= 440.,
          True ],
        [ 0.0,
          (wl-645.)/(580.-645.),
          1.0,
          (wl-440.)/(490.-440.),
          0.0 ] )

def raw_b(wl):
    return select(
        [ wl >= 510.,
          wl >= 490.,
          wl >= 380.,
          True ],
        [ 0.0,
          (wl-510.)/(490.-510.),
          1.0,
          0.0 ] )

gamma = 0.80
def correct_r(wl):
    return power(factor(wl)*raw_r(wl),gamma)
def correct_g(wl):
    return power(factor(wl)*raw_g(wl),gamma)
def correct_b(wl):
    return power(factor(wl)*raw_b(wl),gamma)

ww=arange(380.,781.)
def rgb():
    return transpose([correct_r(ww),correct_g(ww),correct_b(ww)])