Python code to convert ligth wave length to RGB color

This is a small python function to convert a wave length of light to its corresponding RGB color. Note that this is a crude approximation. Calculating the actual RGB representation of a specific wave length is rather tricky. You could also say it’s impossible. That’s because your computer screen can only show a limited range of light colors. So, below code should be good enough for most presentational purposes.

def wave2rgb(wave):
    # This is a port of javascript code from  http://stackoverflow.com/a/14917481
    gamma = 0.8
    intensity_max = 1

    if wave < 380:
        red, green, blue = 0, 0, 0
    elif wave < 440:
        red = -(wave - 440) / (440 - 380)
        green, blue = 0, 1
    elif wave < 490:
        red = 0
        green = (wave - 440) / (490 - 440)
        blue = 1
    elif wave < 510:
        red, green = 0, 1
        blue = -(wave - 510) / (510 - 490)
    elif wave < 580:
        red = (wave - 510) / (580 - 510)
        green, blue = 1, 0
    elif wave < 645:
        red = 1
        green = -(wave - 645) / (645 - 580)
        blue = 0
    elif wave <= 780:
        red, green, blue = 1, 0, 0
    else:
        red, green, blue = 0, 0, 0

    # let the intensity fall of near the vision limits
    if wave < 380:
        factor = 0
    elif wave < 420:
        factor = 0.3 + 0.7 * (wave - 380) / (420 - 380)
    elif wave < 700:
        factor = 1
    elif wave <= 780:
        factor = 0.3 + 0.7 * (780 - wave) / (780 - 700)
    else:
        factor = 0

    def f(c):
        if c == 0:
            return 0
        else:
            return intensity_max * pow (c * factor, gamma)

    return f(red), f(green), f(blue)

Below is an example code that uses this function:

import matplotlib.pyplot as p
import numpy as np

N = 100
image = np.zeros((5,N,3))

for i in range(0, 5):
    for j in range(0, N):
        start = 380
        end = 780
        wave = j * (end-start)/N + start
        image[i][j] = wave2rgb(wave)

ax = p.axes()
ax.get_yaxis().set_visible(False)
p.imshow(image)
p.show()

This should produce an output similar to below: