Matplotlib - color under curve based on spectral color

前端 未结 1 1070
逝去的感伤
逝去的感伤 2020-12-30 12:10

I would like to make a plot of of a spectrum, where the area under the curve will be shaded according the the corresponding color of the light. Much like this plot:

1条回答
  •  太阳男子
    2020-12-30 12:53

    First of all you would need a function that takes the wavelength as input and returns an RGB color. Such a function can be found here. One may adapt it to also return an alpha value, which is smaller 1 outside the range of visible colors.

    This function can be used to make create a colormap. Using a decent normalization allows to have the range of wavelengths mapped to the range between 0 and 1, such that this colormap may be used in an imshow plot.

    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.colors
    
    
    def wavelength_to_rgb(wavelength, gamma=0.8):
        ''' taken from http://www.noah.org/wiki/Wavelength_to_RGB_in_Python
        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
        Additionally alpha value set to 0.5 outside range
        '''
        wavelength = float(wavelength)
        if wavelength >= 380 and wavelength <= 750:
            A = 1.
        else:
            A=0.5
        if wavelength < 380:
            wavelength = 380.
        if wavelength >750:
            wavelength = 750.
        if wavelength >= 380 and wavelength <= 440:
            attenuation = 0.3 + 0.7 * (wavelength - 380) / (440 - 380)
            R = ((-(wavelength - 440) / (440 - 380)) * attenuation) ** gamma
            G = 0.0
            B = (1.0 * attenuation) ** gamma
        elif wavelength >= 440 and wavelength <= 490:
            R = 0.0
            G = ((wavelength - 440) / (490 - 440)) ** gamma
            B = 1.0
        elif wavelength >= 490 and wavelength <= 510:
            R = 0.0
            G = 1.0
            B = (-(wavelength - 510) / (510 - 490)) ** gamma
        elif wavelength >= 510 and wavelength <= 580:
            R = ((wavelength - 510) / (580 - 510)) ** gamma
            G = 1.0
            B = 0.0
        elif wavelength >= 580 and wavelength <= 645:
            R = 1.0
            G = (-(wavelength - 645) / (645 - 580)) ** gamma
            B = 0.0
        elif wavelength >= 645 and wavelength <= 750:
            attenuation = 0.3 + 0.7 * (750 - wavelength) / (750 - 645)
            R = (1.0 * attenuation) ** gamma
            G = 0.0
            B = 0.0
        else:
            R = 0.0
            G = 0.0
            B = 0.0
        return (R,G,B,A)
    
    clim=(350,780)
    norm = plt.Normalize(*clim)
    wl = np.arange(clim[0],clim[1]+1,2)
    colorlist = list(zip(norm(wl),[wavelength_to_rgb(w) for w in wl]))
    spectralmap = matplotlib.colors.LinearSegmentedColormap.from_list("spectrum", colorlist)
    
    fig, axs = plt.subplots(1, 1, figsize=(8,4), tight_layout=True)
    
    wavelengths = np.linspace(200, 1000, 1000)
    spectrum = (5 + np.sin(wavelengths*0.1)**2) * np.exp(-0.00002*(wavelengths-600)**2)
    plt.plot(wavelengths, spectrum, color='darkred')
    
    y = np.linspace(0, 6, 100)
    X,Y = np.meshgrid(wavelengths, y)
    
    extent=(np.min(wavelengths), np.max(wavelengths), np.min(y), np.max(y))
    
    plt.imshow(X, clim=clim,  extent=extent, cmap=spectralmap, aspect='auto')
    plt.xlabel('Wavelength (nm)')
    plt.ylabel('Intensity')
    
    plt.fill_between(wavelengths, spectrum, 8, color='w')
    plt.savefig('WavelengthColors.png', dpi=200)
    
    plt.show()
    

    0 讨论(0)
提交回复
热议问题