Plot low-loss EELS spectra with range 1.5 - 3.4 eV filled with the visible colours, otherwise black

eels_colours.py

import numpy as np
import matplotlib.pyplot as plt

# get electropy from thomasaarholt github
from electropy.units import eV_to_wavelength

# get colour by pip install colour-science
from colour.utilities import first_item, normalise_maximum
from colour.plotting import filter_cmfs, XYZ_to_plotting_colourspace, COLOUR_STYLE_CONSTANTS
from colour.colorimetry import wavelength_to_XYZ, ILLUMINANTS

def get_colour_array(xaxis, out_of_gamut_clipping=False):
    'Uses the colour python module to generate the appropriate colours'
    wavelengths = eV_to_wavelength(xaxis)
    w2 = wavelengths.copy()
    w2[w2<360] = 360
    w2[w2>830] = 830
    cmfs='CIE 1931 2 Degree Standard Observer',
    cmfs = first_item(filter_cmfs(cmfs).values())
    colours = XYZ_to_plotting_colourspace(
            wavelength_to_XYZ(w2, cmfs),
            ILLUMINANTS['CIE 1931 2 Degree Standard Observer']['E'],
            apply_cctf_encoding=False)
    colours = normalise_maximum(colours)
    if not out_of_gamut_clipping:
        colours += np.abs(np.min(colours))
    colours = COLOUR_STYLE_CONSTANTS['colour'].colourspace.cctf_encoding(colours)
    return colours

def plot_EELS_with_wavelength_color(s, xlim=(None, None), wavelength_units_top=True, inf=False):
    """
    EELS plot with fill colour according to the wavelengths.
    
    If wavelength_units_top is True, it also adds an x-axis in nm units on top.
    If inf is True, then any eV values below 0 are displayed as inf. Otherwise they are blank.
    
    """
    fig, ax1 = plt.subplots()
    x = s.axes_manager[-1].axis
    y = s.data

    colours = get_colour_array(x)
    ax1.bar(x=x, height=y, width = s.axes_manager[-1].scale, color=colours)
    
    xl, xr = xlim
    left_lim = xl if xl != None else x.min()
    right_lim = xr if xr != None else x.max()
    plt.xlim(left_lim, right_lim)
    plt.ylim(0, None)
    if wavelength_units_top:
        ax2 = ax1.twiny()

        ax1Ticks = ax1.get_xticks()
        ax2Ticks = ax1Ticks#[t for t in ax1Ticks if t > 0.0 ]

        def tick_function(X, inf=False):
            V = eV_to_wavelength(X)
            if inf:
                return ["{:d}".format(int(z)) if z!= np.inf else "inf" for z in V]
            else:
                return ["{:d}".format(int(z)) if z!= np.inf else "" for z in V]

        ax2.set_xticks(ax2Ticks)
        print(ax2Ticks)
        ax2.set_xbound(ax1.get_xbound())
        #ax2Ticks[ax2Ticks == 0.] += 1e-6 # in case x-axis includes zero
        ax2.set_xticklabels(tick_function(ax2Ticks, inf=inf))
        ax2.set_xlabel('Wavelength (nm)')
        
    ax1.set_axisbelow(True)
    ax1.xaxis.grid(color='gray', linestyle='dashed')
    ax1.plot(x, y, color=COLOUR_STYLE_CONSTANTS.colour.dark)
    ax1.set_xlabel('Energy loss (eV)')
    ax1.set_ylabel('Counts')

    return fig, ax1, ax2