Plot a 4D hyperspy dataset with the signal in polar coordinates

plot_polar.py

import hyperspy.api as hs
import numpy as np
from ipywidgets.widgets import HBox, VBox, Label, IntSlider, Output, IntRangeSlider
import matplotlib.pyplot as plt
from IPython.display import display

def plot_polar(s, clim=(10000, 100000), clim_max=(0, 250000)):
    """
    Plot a 4D hyperspy dataset with the signal in polar coordinates
    Arguments:
    s: hyperspy signal of shape <X, Y, W| R>
    Notes:
    The theta-axis used is the signal axes' axis, which should be in degree units
    (check s.axes_manager[-1].axis)
    
    How to use:
    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>"))
    """
    
    o = Output()
    # Flip the signal and navigation with .T, then sum over the former signal dimension to produce a X Y W dataset. 
    # This becomes the figure we plot on the left.
    nav = s.T.sum() 

    # Set up the initial values for the rectangular region of interest marker
    # x coordinates
    value2index0 = nav.axes_manager[0].value2index
    index2value0 = nav.axes_manager[0].index2value

    low = nav.axes_manager[0].low_index
    high = nav.axes_manager[0].high_index
    left = index2value0(int(low + 1/4*(high-low)))
    right = index2value0(round(low + 3/4*(high-low)))

    # y coordinates
    value2index1 = nav.axes_manager[1].value2index
    index2value1 = nav.axes_manager[1].index2value
    
    low = nav.axes_manager[1].low_index
    high = nav.axes_manager[1].high_index
    top = index2value1(int(low + 1/4*(high-low)))
    bottom = index2value1(round(low + 3/4*(high-low)))

    # Since the navigator has three dimensions and we can only display two,
    # we add a slider to control the wavelength dimension

    extra_nav_axes = nav.axes_manager.signal_dimension - 2
    nav_sliders = []
    if extra_nav_axes > 0:
        for i in range(extra_nav_axes):
            slider = IntSlider(value=0, min=0, max=nav.axes_manager[i+2].high_index)
            nav_sliders.append(slider)
    nav_sliders_rev = nav_sliders[::-1] # we reverse this here so that it matches the reversed numpy indexing
    
####
    if clim:
        clim_slider = IntRangeSlider(value=clim,
            min=clim_max[0],
            max=clim_max[1],
            step=(clim_max[1] - clim_max[0]) / 25, # 25 steps of 10000 by default
            description='clim:',
            disabled=False,
            continuous_update=True,
            orientation='vertical',
            readout=True,
            readout_format='d',
        )
####
    # create a dataset with only X and Y, slicing at the wavelength position
    slices = (..., ) + tuple(slider.value for slider in nav_sliders_rev)
    nav_first_two = nav.isig[slices] # ... means "leave every axis until the end", equivalent of :,: here.
    
    # We use normal hyperspy plotting to display the first figure, because we want to use hyperspy ROIs.
    # Note that hs plotting really is matplotlib plotting with more on top.
    # Prevent the matplotlib figure from being created immediately by using an Output ipywidget. 
    # This is a separate area that we will display later
    nav_output = Output()
    vmin, vmax = (clim[0], clim[1]) if clim else (None, None)

    with nav_output:
        nav_first_two.plot(vmin=vmin, vmax=vmax)
    # get ahold of the image displayed in this figure
    # we will be changing the array displayed in `im` with `im.set_data()` later.   
    im = nav_first_two._plot.signal_plot.ax.images[0] 
    
    # Create the roi and place it on the left figure, then create labels that hold the name and position of the 
    # coordinates of the ROI and the wavelength index
    roi = hs.roi.RectangularROI(left, top, right, bottom)
    r = roi.interactive(nav_first_two)
    # The labels use f-strings, which are a convenient way to write out strings with objects in them. The :.2f operator means "leave two decimals and show as a float"
    xlabel = Label(value=f"{nav.axes_manager[0].name}: {roi.left:.2f} : {roi.right:.2f}")
    ylabel = Label(value=f"{nav.axes_manager[1].name}: {roi.top:.2f} : {roi.bottom:.2f}")
    labels = []
    if extra_nav_axes:
        for i in range(extra_nav_axes):
            label = Label(value=f"{nav.axes_manager[i+2].name}: {nav.axes_manager[i+2].index2value(nav_sliders[i].value):.2f}")
            labels.append(label)
    # Plot right figure using an Output widget to prevent it being shown immediately
    sig_output = Output()
    with sig_output:
        fig = plt.figure()
    ax = plt.subplot(111, projection='polar')
    ax.set_xticks([0,np.pi/3,2*np.pi/3,np.pi,4*np.pi/3,5*np.pi/3])
    ax.autoscale(enable=True, axis='y')
    thetaaxis = s.axes_manager[-1].axis * np.pi / 180 # convert x axis to radians
    # Plot the polar plot. We return the line plotted into `lns` (which is a list of all lines)
    
    slices = tuple(slider.value for slider in nav_sliders_rev) + (slice(value2index1(roi.top), value2index1(roi.bottom)), slice(value2index0(roi.left), value2index0(roi.right)))
    lns = ax.plot(thetaaxis, s.data[slices].mean((0,1)))
    ln = lns[0] # lns only contains one line, which we will need later to update it

    def update(roi):
        "Update the polar plot and label values from a moving ROI"
        # Here we index the numpy array instead of the hyperspy .inav, because the former is quite a bit faster.
        # the numpy array is indexed in reverse order of hyperspy for some reason
        # The specific "roi" object name is imporant, as the `connect` method we use later expects it.

        # Update plot with set_ydata
        with o:

            # HyperSpy raises an error when we try to value2index a value that is higher than the axis
            # Unfortunately, when choosing the final right and bottom pixels on the nav image,
            # this occurs, since the right/bottom side of the final pixel is axis[-1] + a bit.
            # We solve this by trying, and if it fails we except it and assume the index should be high_index + 1.
            left, top = value2index0(roi.left), value2index1(roi.top)
            try:
                right = value2index0(roi.right)
            except ValueError:
                print('hit the error')
                right = nav.axes_manager[0].high_index + 1
            try:
                bottom = value2index1(roi.bottom)
            except ValueError:
                print('hit the error')
                bottom = nav.axes_manager[1].high_index + 1
            slices = tuple(slider.value for slider in nav_sliders_rev) + (slice(top,bottom), slice(left,right))
            data = s.data[slices]
            ydata = data.mean((0,1)) # 0 and 1 are the X and Y axes here, because we eliminate the wavelength axis by indexing it not using a range
            ln.set_ydata(ydata)
    ########ax.set_ylim(0, ydata.max()) # Polar plots look funny if we index with ydata.min() instead of 0.
            ax.set_ylim(0, ydata.max())


            xlabel.value=f"{nav.axes_manager[0].name}: {roi.left:.2f} : {roi.right:.2f}"
            ylabel.value=f"{nav.axes_manager[1].name}: {roi.top:.2f} : {roi.bottom:.2f}"
            # Don't update other labels, because they can't move when moving the ROI around
        
    def update_from_slider(change):
        "Update the polar plot, label values AND left hand plot by changing the wavelength slider"
        # Here we index the numpy array instead of the hyperspy .inav, because the former is quite a bit faster.
        # the numpy array is indexed in reverse order of hyperspy for some reason
        # The `change` object name is important, ipywidget `observe` methods expect it.
        # While I do use the change['new'] object here, I could equally use `wavelength_slider.value`
        #wavelength_index = change['new']
        #orientation_index = change['new']
        # Update image with set_data
        with o:
            nav_slider_slices = tuple(slider.value for slider in nav_sliders_rev)
            im.set_data(s.data[nav_slider_slices].sum(-1)) # update the navigator, first index here is the wavelength, : or ... after is unnecessary.
            if clim:
                im.set_clim(clim_slider.value)    #.  This sets the min and max of the color scale in the navigation image
            else:
                im.autoscale() # autoscale the nav image  

            # Update plot with set_ydata
            left, top = value2index0(roi.left), value2index1(roi.top)
            try:
                right = value2index0(roi.right)
            except ValueError:
                print('hit the error')
                right = nav.axes_manager[0].high_index + 1
            try:
                bottom = value2index1(roi.bottom)
            except ValueError:
                print('hit the error')
                bottom = nav.axes_manager[1].high_index + 1
            
            slices = nav_slider_slices + (slice(top,bottom), slice(left,right))
            ydata = s.data[slices].mean((0,1)) # 0 and 1 are the X and Y axes here, because we eliminate the wavelength axis by indexing it not using a range
            ln.set_ydata(ydata)
            ax.set_ylim(0, ydata.max()) # Polar plots look funny if we index with ydata.min() instead of 0.
            #ax.set_ylim(0, 1200)

            # Don't update ROI labels, because they're not changing here
            if extra_nav_axes:
                for i in range(extra_nav_axes):
                    labels[i].value = f"{nav.axes_manager[i+2].name}: {nav.axes_manager[i+2].index2value(nav_sliders[i].value):.2f}"

    roi.events.changed.connect(update) # Call the `update(roi=roi)` function every time the ROI changes, with the latest values of the `roi` object
    # Call the `update_from_slider(change=change)` function every time the slider changes, 
    if clim:
        clim_slider.observe(update_from_slider, names='value')
    for slider in nav_sliders:
        slider.observe(update_from_slider, names='value')
    # With ipympl enabled, the figure.canvas object is an ipywidget itself. 
    # That means can we can modify attributes like margin and other positioning. See the ipywidgets docs for more info
    navfig = nav_first_two._plot.signal_plot.figure.canvas 
    
    sigfig = fig.canvas

    navfig.layout.margin = "auto 0px auto 0px" # This centers the figure, I think. Can't actually remember.
    sigfig.layout.margin = "auto 0px auto 0px" # This centers the figure, I think. Can't actually remember.
    sigfig.header_visible = False # removes unnecssary "Figure 1" info on top
    navfig.header_visible = False
    display(HBox([nav_output, sig_output])) # Here we acutally display the figures, wrapped in a "Horizontal Box" widget so that they are next to each other
    
    label_pairs = []
    for i in range(extra_nav_axes):
        label_pairs.append(HBox([Label(f'{nav.axes_manager[i+2].name}'), nav_sliders[i]]))

    clim_slider_list = [clim_slider] if clim else []
    display(HBox(label_pairs + clim_slider_list + [VBox([xlabel, ylabel, *labels]) ]))