graph.py

import napari
import numpy as np
import matplotlib.pyplot as pl
import h5py
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
#from qtpy.QtWidgets import QApplication, QMainWindow, QPushButton, QMenu, QVBoxLayout, QSizePolicy, QMessageBox, QWidget, QPushButton
import pyqtgraph as pg
from pyqtgraph.Qt import QtGui, QtCore
import sys
import random

class stokes_plot(object):

    def __init__(self):
        self.left = 10
        self.top = 10
        self.title = ''
        self.width = 320
        self.height = 200

        self.plots = [None] * 4
        self.curve_stokes = [None] * 4
        self.lambda_pos = [None] * 4

        self.win = pg.GraphicsWindow(title="Stokes profiles") # creates a window

        self.plots[0] = self.win.addPlot(title="Stokes I", row=0, col=0)
        self.plots[1] = self.win.addPlot(title="Stokes Q", row=0, col=1)
        self.plots[2] = self.win.addPlot(title="Stokes U", row=1, col=0)
        self.plots[3] = self.win.addPlot(title="Stokes V", row=1, col=1)

        for i in range(4):
            self.lambda_pos[i] = pg.InfiniteLine(pos=0, angle=90)
            self.plots[i].addItem(self.lambda_pos[i])

        for i in range(4):
            self.curve_stokes[i] = self.plots[i].plot()
            self.curve_stokes[i].setShadowPen(pg.mkPen((255,255,255), width=2, cosmetic=True))

    def plot(self, stokes, ind_lambda):        

        if (stokes is not None):
            for i in range(4):
                if (i == 0):
                    self.curve_stokes[i].setData(stokes[i,:])
                else:
                    self.curve_stokes[i].setData(stokes[i,:] / stokes[0,0])
            
        for i in range(4):
            self.lambda_pos[i].setValue(ind_lambda)
        

f = h5py.File('ar11967.h5','r')

stokes = np.transpose(f['stokes'],axes=(0,3,1,2)) 
nx, ny = stokes.shape[2:]

with napari.gui_qt():
    viewer = napari.Viewer()

    stokesv = viewer.add_image(stokes[3,:,:,:], name='Stokes V', opacity=0)
    stokesu = viewer.add_image(stokes[2,:,:,:], name='Stokes U', opacity=0)
    stokesq = viewer.add_image(stokes[1,:,:,:], name='Stokes Q', opacity=0)
    stokesi = viewer.add_image(stokes[0,:,:,:], name='Stokes I')

    app = QtGui.QApplication([])
    stokes_1d = stokes_plot()

    def update_slider(event):
        if (event.axis == 0):
            ind_lambda = viewer.dims.indices[0]
            stokes_1d.plot(None, ind_lambda)
    
    def action(viewer, event):
        ind_lambda, ind_x, ind_y = np.round(viewer.coordinates).astype('int')
        if ( (ind_x > 0) and (ind_y > 0) and (ind_x < nx) and (ind_y < ny) ):
            stokes_1d.plot(stokes[:,:,ind_x,ind_y], ind_lambda)

    def check_exit(viewer, event):
        if (event.button == 2):
            sys.exit()

    @viewer.bind_key('q')
    def exit(viewer):
        sys.exit()

    @stokesi.mouse_move_callbacks.append    
    def action_stokesi(viewer, event):        
        action(viewer, event)

    @stokesq.mouse_move_callbacks.append
    def action_stokesq(viewer, event):
        action(viewer, event)

    @stokesu.mouse_move_callbacks.append
    def action_stokesu(viewer, event):
        action(viewer, event)

    @stokesv.mouse_move_callbacks.append
    def action_stokesv(viewer, event):
        action(viewer, event)

    @stokesi.mouse_drag_callbacks.append
    def butt(viewer, event):
        check_exit(viewer, event)

    @stokesq.mouse_drag_callbacks.append
    def butt(viewer, event):
        check_exit(viewer, event)

    @stokesu.mouse_drag_callbacks.append
    def butt(viewer, event):
        check_exit(viewer, event)

    @stokesv.mouse_drag_callbacks.append
    def butt(viewer, event):
        check_exit(viewer, event)

    # viewer.dims.events.connect(update_slider)
    viewer.dims.events.axis.connect(update_slider)
        
    pg.QtGui.QApplication.exec_()