Hyperspy function to emulate GMS behaviour on spectrum images

spanplot.py

%matplotlib widget
import hyperspy.api as hs
import matplotlib.pyplot as plt
from ipywidgets.widgets import HBox, Label, Dropdown

# Optional, for notebook only, not lab
# from IPython.core.display import display, HTML
# display(HTML("<style>.container { width:100% !important; }</style>"))

def plot(s):
    """
    Plot a spectrum image with a span ROI to sum regions of the spectrum
    
    Arguments:
    
    s: hyperspy signal of shape <X, Y | E>
    
    Notes:
    Must use the ipympl backend. %matplotlib widget
    Best viewed in jupyter lab or in the notebook with a stretched viewport, try:
    from IPython.core.display import display, HTML
    display(HTML("<style>.container { width:100% !important; }</style>"))
    """
    s = s.T
    value2index = s.axes_manager[0].value2index
    
    nav = s.T.sum()
    plt.ioff()

    low = nav.axes_manager[-1].low_value
    high = nav.axes_manager[-1].high_value
    left = low + 1/4*(high-low)
    right = low + 3/4*(high-low)
    
    label = Label(value="Span: {:.2f}:{:.2f}".format(left, right))
    colormap = Dropdown(
        options=['viridis', 'jet', 'Greys'],
        value='viridis',
        description='Colormap:',
        disabled=False,)
    nav.plot()
    roi = hs.roi.SpanROI(left=left, right=right)
    r = roi.interactive(nav)

    fig, ax = plt.subplots()
    im = ax.imshow(s.data[value2index(roi.left):value2index(roi.right)].mean(0), cmap=colormap.value)
    ax.axis('off')

    plt.ion()

    def update(roi):
        im.set_data(s.data[value2index(roi.left):value2index(roi.right)].mean(0))
        im.autoscale()
        label.value = "Span: {:.2f}:{:.2f}".format(roi.left, roi.right)
        
    def update_cmap(change):
        im.set_cmap(change['new'])

    roi.events.changed.connect(update)
    colormap.observe(update_cmap, names='value')

    navfig = nav._plot.signal_plot.figure.canvas
    sigfig = fig.canvas

    navfig.layout.margin = "auto 0px auto 0px"
    sigfig.layout.margin = "auto 0px auto 0px"
    sigfig.header_visible = False
    navfig.header_visible = False
    display(HBox([navfig, sigfig]))
    display(HBox([label, colormap]))
    
    
    
    
### Usage
# s = hs.load("signal.hspy")
# plot(s)