widget_bench.py

 
"""
Widget to use on a bench
 
Usage:
  widget_bench.py [<parameters_filename>] [options]
 
with 'parameters_filename' the path to the parameters file
 
Options:
  -h --help          Show this help message and exit
  --brahma            Distribute data with brahma
  -d, --devices devices      Specify the devices
  -i, --interactive  keep the script interactive
"""
 
import os, sys
import time
 
try:
    from PyQt5 import QtWidgets
    from PyQt5.QtCore import QThread, Qt
except ModuleNotFoundError as e:
    try:    
        from PySide2 import QtWidgets
        from PySide2.QtCore import QThread, Qt
    except ModuleNotFoundError as e:
        raise ModuleNotFoundError("No module named 'PyQt5' or PySide2', please install one of them\nException raised: "+e.msg)
 
from typing import Any
 
from docopt import docopt
 
from shesha.widgets.widget_base import WidgetBase, uiLoader
BenchWindowTemplate, BenchClassTemplate = uiLoader('widget_bench')
 
import matplotlib.pyplot as plt
 
from shesha.supervisor.benchSupervisor import WFSType, BenchSupervisor as Supervisor
 
# For debug
# from IPython.core.debugger import Pdb
# then add this line to create a breakpoint
# Pdb().set_trace()
 
 
class widgetBenchWindow(BenchClassTemplate, WidgetBase):
 
    def __init__(self, config_file: Any = None, brahma: bool = False,
                 devices: str = None) -> None:
        WidgetBase.__init__(self)
        BenchClassTemplate.__init__(self)
 
        self.brahmabrahma = brahma
        self.devicesdevices = devices
 
        self.uiBenchuiBench = BenchWindowTemplate()
        self.uiBenchuiBench.setupUi(self)
 
        
 
        self.supervisorsupervisor = None
        self.configconfig = None
        self.stopstop = False  # type: bool  # Request quit
 
        self.uiBenchuiBench.wao_nbiters.setValue(1000)  # Default GUI nIter box value
        self.nbiternbiter = self.uiBenchuiBench.wao_nbiters.value()
        self.refreshTimerefreshTime = 0  # type: float  # System time at last display refresh
        self.loopThreadloopThread = None  # type: QThread
        self.assistantassistant = None  # type: Any
 
        
        self.defaultParPathdefaultParPathdefaultParPath = os.environ["SHESHA_ROOT"] + "/data/par/bench-config"
        self.defaultAreaPathdefaultAreaPathdefaultAreaPath = os.environ["SHESHA_ROOT"] + "/data/layouts"
        self.loadDefaultConfigloadDefaultConfig()
 
        self.uiBenchuiBench.wao_run.setCheckable(True)
        self.uiBenchuiBench.wao_run.clicked[bool].connect(self.aoLoopClickedaoLoopClicked)
        self.uiBenchuiBench.wao_open_loop.setCheckable(True)
        self.uiBenchuiBench.wao_open_loop.clicked[bool].connect(self.aoLoopOpenaoLoopOpen)
        self.uiBenchuiBench.wao_next.clicked.connect(self.loop_onceloop_once)
 
        self.uiBenchuiBench.wao_forever.stateChanged.connect(self.updateForeverupdateForever)
 
        self.dispStatsInTerminaldispStatsInTerminal = False
 
        self.uiBenchuiBench.wao_run.setDisabled(True)
        self.uiBenchuiBench.wao_next.setDisabled(True)
        self.uiBenchuiBench.wao_unzoom.setDisabled(True)
 
        self.addDockWidget(Qt.DockWidgetArea(1), self.uiBaseuiBase.wao_ConfigDock)
        self.addDockWidget(Qt.DockWidgetArea(1), self.uiBaseuiBase.wao_DisplayDock)
        self.uiBaseuiBase.wao_ConfigDock.setFloating(False)
        self.uiBaseuiBase.wao_DisplayDock.setFloating(False)
 
        self.adjustSize()
 
        if config_file is not None:
            self.uiBaseuiBase.wao_selectConfig.clear()
            self.uiBaseuiBase.wao_selectConfig.addItem(config_file)
            self.load_configload_configload_config()
            self.init_configinit_configinit_config()
 
    
 
    # def updateStatsInTerminal(self, state):
    #     self.dispStatsInTerminal = state
 
    def updateForever(self, state):
        self.uiBenchuiBench.wao_nbiters.setDisabled(state)
 
    def add_dispDock(self, name: str, parent, type: str = "pg_image") -> None:
        d = WidgetBase.add_dispDock(self, name, parent, type)
        if type == "SR":
            d.addWidget(self.uiBenchuiBench.wao_Strehl)
 
    def load_config(self) -> None:
        '''
            Callback when 'LOAD' button is hit
        '''
        WidgetBase.load_config(self)
        config_file = str(self.uiBaseuiBase.wao_selectConfig.currentText())
        sys.path.insert(0, self.defaultParPathdefaultParPathdefaultParPath)
 
        self.supervisorsupervisor = Supervisor(config_file, self.brahmabrahma)
 
        self.configconfig = self.supervisorsupervisor.get_config()
 
        # if self.devices:
        #     self.config.p_loop.set_devices([
        #             int(device) for device in self.devices.split(",")
        #     ])
 
        try:
            sys.path.remove(self.defaultParPathdefaultParPathdefaultParPath)
        except:
            pass
 
        try:
            self.nwfsnwfs = len(self.configconfig.p_wfss)
        except:
            self.nwfsnwfs = 1  # Default. config may very well not define p_wfss
        for wfs in range(self.nwfsnwfs):
            name = 'slpComp_%d' % wfs
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_graphgroup_cbwao_graphgroup_cb, "MPL")
            name = 'wfs_%d' % wfs
            self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
 
        self.uiBenchuiBench.wao_run.setDisabled(True)
        self.uiBenchuiBench.wao_next.setDisabled(True)
        self.uiBenchuiBench.wao_unzoom.setDisabled(True)
 
        self.uiBaseuiBase.wao_init.setDisabled(False)
 
        if (hasattr(self.configconfig, "layout")):
            area_filename = self.defaultAreaPathdefaultAreaPathdefaultAreaPath + "/" + self.configconfig.layout + ".area"
            self.loadArealoadArea(filename=area_filename)
 
        self.adjustSize()
 
    def aoLoopClicked(self, pressed: bool) -> None:
        if pressed:
            self.stopstop = False
            self.refreshTimerefreshTime = time.time()
            self.nbiternbiter = self.uiBenchuiBench.wao_nbiters.value()
            if self.dispStatsInTerminaldispStatsInTerminal:
                if self.uiBenchuiBench.wao_forever.isChecked():
                    print("LOOP STARTED")
                else:
                    print("LOOP STARTED FOR %d iterations" % self.nbiternbiter)
            self.runrunrun()
        else:
            self.stopstop = True
 
    def aoLoopOpen(self, pressed: bool) -> None:
        if (pressed):
            self.supervisorsupervisor.open_loop()
            self.uiAO.wao_open_loop.setText("Close Loop")
        else:
            self.supervisorsupervisor.close_loop()
            self.uiAO.wao_open_loop.setText("Open Loop")
 
    def init_config(self) -> None:
        WidgetBase.init_config(self)
 
    def init_configThread(self) -> None:
        # self.uiBench.wao_deviceNumber.setDisabled(True)
        # self.config.p_loop.devices = self.uiBench.wao_deviceNumber.value()  # using GUI value
        # gpudevice = "ALL"  # using all GPU avalaible
        # gpudevice = np.array([2, 3], dtype=np.int32)
        # gpudevice = np.arange(4, dtype=np.int32) # using 4 GPUs: 0-3
        # gpudevice = 0  # using 1 GPU : 0
        self.supervisorsupervisor.init_config()
 
    def init_configFinished(self) -> None:
        # Thread carmaWrap context reload:
        try:
            self.supervisorsupervisor.force_context()
        except:
            print('Warning: could not call supervisor.force_context().')
 
        for i in range(self.nwfsnwfs):
            if self.configconfig.p_wfss[i].type == WFSType.SH:
                key = "wfs_%d" % i
                self.addSHGridaddSHGrid(self.docksdocks[key].widgets[0],
                               self.configconfig.p_wfss[i].get_validsub(),
                               self.configconfig.p_wfss[i].npix, self.configconfig.p_wfss[i].npix)
 
        self.updateDisplayupdateDisplayupdateDisplay()
 
        self.uiBenchuiBench.wao_run.setDisabled(False)
        self.uiBenchuiBench.wao_next.setDisabled(False)
        self.uiBenchuiBench.wao_open_loop.setDisabled(False)
        self.uiBenchuiBench.wao_unzoom.setDisabled(False)
 
        WidgetBase.init_configFinished(self)
 
    def updateDisplay(self) -> None:
        if (self.supervisorsupervisor is None) or (not self.supervisorsupervisor.is_init()) or (
                not self.uiBaseuiBase.wao_Display.isChecked()):
            # print("Widget not fully initialized")
            return
        if not self.loopLockloopLock.acquire(False):
            return
        else:
            try:
                for key, dock in self.docksdocks.items():
                    if key == "Strehl":
                        continue
                    elif dock.isVisible():
                        index = int(key.split("_")[-1])
                        data = None
                        if "wfs" in key:
                            data = self.supervisorsupervisor.wfs.get_wfs_image(index)
                            if (data is not None):
                                autoscale = True  # self.uiBench.actionAuto_Scale.isChecked()
                                # if (autoscale):
                                #     # inits levels
                                #     self.hist.setLevels(data.min(), data.max())
                                self.imgsimgs[key].setImage(data, autoLevels=autoscale)
                                # self.p1.autoRange()
 
                        elif "slp" in key:  # Slope display
                            if (self.configconfig.p_wfss[index].type == WFSType.PYRHR or
                                        self.configconfig.p_wfss[index].type == WFSType.PYRLR):
                                raise RuntimeError("PYRHR not usable")
                            self.imgsimgs[key].canvas.axes.clear()
                            x, y = self.supervisorsupervisor.config.p_wfss[index].get_validsub()
 
                            nssp = x.size
                            centroids = self.supervisorsupervisor.rtc.get_slopes()
                            vx = centroids[:nssp]
                            vy = centroids[nssp:]
 
                            offset = (self.supervisorsupervisor.config.p_wfss[index].npix - 1) / 2
                            self.imgsimgs[key].canvas.axes.quiver(
                                    x + offset, y + offset, vy, vx, angles='xy',
                                    scale_units='xy',
                                    scale=1)  # wao.supervisor.config.p_wfss[0].pixsize)
                            self.imgsimgs[key].canvas.draw()
 
            finally:
                self.loopLockloopLock.release()
 
    def loop_once(self) -> None:
        if not self.loopLockloopLock.acquire(False):
            print("Display locked")
            return
        else:
            try:
                start = time.time()
                self.supervisorsupervisor.single_next()
                loopTime = time.time() - start
                refreshDisplayTime = 1. / self.uiBaseuiBase.wao_frameRate.value()
 
                if (time.time() - self.refreshTimerefreshTime > refreshDisplayTime):
                    currentFreq = 1 / loopTime
                    self.uiBenchuiBench.wao_currentFreq.setValue(currentFreq)
                    self.refreshTimerefreshTime = start
            except:
                pass
            finally:
                self.loopLockloopLock.release()
 
    def run(self):
        WidgetBase.run(self)
        if not self.uiBenchuiBench.wao_forever.isChecked():
            self.nbiternbiter -= 1
 
        if self.nbiternbiter <= 0:
            self.stopstop = True
            self.uiBenchuiBench.wao_run.setChecked(False)
 
 
if __name__ == '__main__':
    arguments = docopt(__doc__)
    app = QtWidgets.QApplication(sys.argv)
    app.setStyle('cleanlooks')
    wao = widgetBenchWindow(arguments["<parameters_filename>"],
                            brahma=arguments["--brahma"], devices=arguments["--devices"])
    wao.show()
    if arguments["--interactive"]:
        from shesha.util.ipython_embed import embed
        from os.path import basename
        embed(basename(__file__), locals())