Plotting spherical polar mapping data

The below code creates and interactive plot, with the mapped data for a set radius on the left, and a plot of all selected vectors on the right. To select a vector simply click anywhere on the left hand 2d plot and the selected point will have its associated vector plotted on the right-hand side 3d axis, along with a star marker on the left-hand side plot. To clear all selected vectors from the right-hand side plot and star markers from the left-hand plot click the 'clear vectors' button displayed as a widget in the cell output.

#import the required modules
#----------------------------
import h5py
import ipywidgets as wg
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

from matplotlib.collections import PathCollection
from scipy.spatial.transform import Rotation as R
import transformations as tf

#define functions needed for plotting
#--------------------------------------
markers=[]
def drawarrow(vec1,vec2,axs,col,style='-',alp=1):
    axs.quiver(vec1[0],vec1[1],vec1[2],vec2[0],vec2[1],vec2[2],color = col,arrow_length_ratio=0.1,ls=style,alpha=alp)

def point2vec(axes,counts,r_int,th_int,phi_int,zoom,click=False):
    global updating
    if updating:
        return
    #ax2.cla()


    markernew=ax1.scatter(phi_int,th_int,marker='*',s=50,c='aqua',ec='black')
    markers.append(markernew)
    vector0=[0,0,0]
    rstart,rstop,rstep=axes[0][1:4]
    thstart,thstop,thstep=axes[1][1:4]
    phistart,phistop,phistep=axes[2][1:4]

    thetarad=thstart+(th_int*thstep)
    phirad=phistart+(phi_int*phistep)
    radius=rstart+(r_int*rstep)
    if click==False:

        title = ax1.get_title().split('=')[0]
        ax1.cla()
        ax1.set_title(title+f"={radius}")
        ax1.imshow(counts[r_int,:,:],norm='log',cmap='jet')
    theta=np.degrees(thetarad)
    phi=np.degrees(phirad)

    vector1=[0,0,radius]
    rotyzmat=R.from_euler('x', theta, degrees=True)
    rotxymat=R.from_euler('z',phi,degrees=True)
    totalrot=rotxymat*rotyzmat
    plotexitvec=totalrot.apply(vector1)
    ax1.set_ylabel('theta: angle from +z axis')
    ax1.set_xlabel('phi: angle from +x axis')
    drawarrow(vector0,plotexitvec,ax2,'orange')
    ax2.set_xlim(-zoom,zoom)
    ax2.set_xlabel('x')
    ax2.set_ylim(-zoom,zoom)
    ax2.set_ylabel('y')
    ax2.set_zlim(-zoom,zoom)
    ax2.set_zlabel('z')

#setup the active figure
#-------------------------
%matplotlib qt
#other options are available for interactive plotting e.g %matplotlib widget
fig = plt.figure(figsize=(20, 10))

fig.canvas.manager.set_window_title("Spherical polar plot and vector selection")
ax1 = fig.add_subplot(121)
ax2= fig.add_subplot(122, projection='3d')

#import data from hdf5 file
#--------------------
filename='/dls/science/users/rpy65944/output/mapped_scan_519528_1.hdf5'
datah5=h5py.File(filename)
ax1.set_title(filename+'\n radius=1')
#'/dls/science/users/rpy65944/output/mapped_scan_519445_1.hdf5')
counts=datah5['binoculars/counts']

#grab axis details from file
axinfo=datah5['binoculars/axes']
r_ax=np.array(axinfo['h'])
th_ax=np.array(axinfo['k'])
phi_ax=np.array(axinfo['l'])

axes=[r_ax,th_ax,phi_ax]

#setup widgets and start the interactive sliders, button and plot
slider_r = wg.IntSlider(value=1, min=0, max=np.shape(counts)[0]-1, step=1,description='select radius',description_width ='300px',layout=wg.Layout(width='350px'))
slider_r.style.description_width='150px'
slider_th = wg.IntSlider(value=1, min=0, max=np.shape(counts)[1]-1, step=1,description='select theta value',description_width ='300px',layout=wg.Layout(width='350px'))
slider_th.style.description_width='150px'
slider_phi = wg.IntSlider(value=1, min=0, max=np.shape(counts)[2]-1, step=1,description='select phi value',description_width ='300px',layout=wg.Layout(width='350px'))
slider_phi.style.description_width='150px'
slider_zoom=wg.FloatSlider(value=1, min=0, max=5.00, step=0.1,description='select zoom value',description_width ='300px',layout=wg.Layout(width='350px'))
slider_zoom.style.description_width='150px'
clear_button=wg.Button(description="clear vectors")
click_sliders=[slider_th,slider_phi]
updating=False

def onclick(event):
    global updating,axes,counts,click_sliders
    if event.inaxes == ax1:
        updating=True
        #ax2.cla()
        newth=int(event.ydata)
        newphi=int(event.xdata)
        click_sliders=[slider_th,slider_phi]
        newslidevals=[newth,newphi]
        for i,slider in enumerate(click_sliders):
            slider.value=newslidevals[i]
        updating=False
        point2vec(axes,counts,r_int=slider_r.value,th_int=slider_th.value,phi_int=slider_phi.value,zoom=slider_zoom.value,click=True)
        fig.canvas.draw_idle()

def clearaxmarkers():
    scatter_markers = [artist for artist in ax1.collections if isinstance(artist, PathCollection)]
    for marker in scatter_markers:
        marker.remove()
    ax2.cla()
clear_button.on_click(lambda b: clearaxmarkers())
display(clear_button)

fig.canvas.mpl_connect('button_press_event', onclick)


wg.interact(point2vec,axes=wg.fixed(axes),counts=wg.fixed(counts),r_int=slider_r,th_int=slider_th,phi_int=slider_phi,zoom=slider_zoom)

<function __main__.point2vec(axes, counts, r_int, th_int, phi_int, zoom, click=False)>

ax1.set_title(filename)

Text(0.5, 1.0, '/dls/science/users/rpy65944/output/mapped_scan_519528_19.hdf5')