Simple python function to get an approximate RGB value for a given wavelength

wavelength_to_rgb.py

#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt

def wavelength_to_rgb(wavelengths, gamma=0.8):
    '''This converts a given wavelength of light to an
    approximate RGB color value. The wavelength must be given
    in nanometers in the range from 380 nm through 750 nm
    (789 THz through 400 THz).
    Based on code by Dan Bruton
    http://www.physics.sfasu.edu/astro/color/spectra.html
    '''
    wavelengths = np.array(wavelengths).astype(float)
    rgbs = np.zeros(wavelengths.shape + (3,))

    attenuations = np.full(wavelengths.shape, np.nan)

    idx = ((wavelengths >= 380) & (wavelengths <= 440))
    attenuations[idx] = 0.3 + 0.7 * (wavelengths[idx] - 380) / (440 - 380)
    rgbs[idx,0] = ((-(wavelengths[idx] - 440) / (440 - 380)) * attenuations[idx]) ** gamma
    rgbs[idx,2] = (1.0 * attenuations[idx]) ** gamma

    idx = ((wavelengths >= 440) & (wavelengths <= 490))
    rgbs[idx,1] = ((wavelengths[idx] - 440) / (490 - 440)) ** gamma
    rgbs[idx,2] = 1.0

    idx = ((wavelengths >= 490) & (wavelengths <= 510))
    rgbs[idx,1] = 1.0
    rgbs[idx,2] = (-(wavelengths[idx] - 510) / (510 - 490)) ** gamma

    idx = ((wavelengths >= 510) & (wavelengths <= 580))
    rgbs[idx,0] = ((wavelengths[idx] - 510) / (580 - 510)) ** gamma
    rgbs[idx,1] = 1.0

    idx = ((wavelengths >= 580) & (wavelengths <= 645))
    rgbs[idx,0] = 1.0
    rgbs[idx,1] = (-(wavelengths[idx] - 645) / (645 - 580)) ** gamma

    idx = ((wavelengths >= 645) & (wavelengths <= 750))
    attenuations[idx] = 0.3 + 0.7 * (750 - wavelengths[idx]) / (750 - 645)
    rgbs[idx,0] = (1.0 * attenuations[idx]) ** gamma

    rgbs = (255*rgbs).astype(int)
    return rgbs

# Display a simple rainbow
#rgbs = wavelength_to_rgb(np.arange(380,750,30))
#plt.imshow(rgbs[np.newaxis,:,:])
#plt.show()

# Display a 3 by 3 grid with a choice of wavelengths
gridlambdas = np.zeros((3,3))
gridlambdas[0,0] = 380
gridlambdas[0,1] = 420
gridlambdas[0,2] = 460
gridlambdas[1,0] = 500
gridlambdas[1,1] = 540
gridlambdas[1,2] = 580
gridlambdas[2,0] = 620
gridlambdas[2,1] = 660
gridlambdas[2,2] = 700

gridrgbs = wavelength_to_rgb(gridlambdas)

plt.figure()
plt.imshow(gridrgbs)
plt.title("Example wavelength grid")
ax = plt.gca();
ax.grid(color='k', linestyle='-', linewidth=2)
ax.set_xticks(np.arange(-.5, 3, 1));
ax.set_yticks(np.arange(-.5, 3, 1));
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.tick_params(length=0)
for i in range(3):
  for j in range(3):
    plt.text(j, i, f"{gridlambdas[i,j]:.0f} nm",
             horizontalalignment='center',
             verticalalignment='center',
             backgroundcolor=(1,1,1,0.5),
             color='black')
plt.show()