COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
widget_bench.py
1 
37 """
38 Widget to use on a bench
39 
40 Usage:
41  widget_bench.py [<parameters_filename>] [options]
42 
43 with 'parameters_filename' the path to the parameters file
44 
45 Options:
46  -h --help Show this help message and exit
47  --brahma Distribute data with brahma
48  -d, --devices devices Specify the devices
49  -i, --interactive keep the script interactive
50 """
51 
52 import os, sys
53 import numpy as np
54 import time
55 
56 import pyqtgraph as pg
57 from pyqtgraph.dockarea import Dock, DockArea
58 from shesha.util.tools import plsh, plpyr
59 
60 import warnings
61 
62 from PyQt5 import QtGui, QtWidgets
63 from PyQt5.uic import loadUiType
64 from PyQt5.QtCore import QThread, QTimer, Qt
65 
66 from subprocess import Popen, PIPE
67 
68 from typing import Any, Dict, Tuple, Callable, List
69 
70 from docopt import docopt
71 from collections import deque
72 
73 from shesha.widgets.widget_base import WidgetBase, uiLoader
74 BenchWindowTemplate, BenchClassTemplate = uiLoader('widget_bench')
75 
76 import matplotlib.pyplot as plt
77 
78 from shesha.supervisor.benchSupervisor import WFSType, BenchSupervisor as Supervisor
79 
80 # For debug
81 # from IPython.core.debugger import Pdb
82 # then add this line to create a breakpoint
83 # Pdb().set_trace()
84 
85 
87 
88  def __init__(self, config_file: Any = None, brahma: bool = False,
89  devices: str = None) -> None:
90  WidgetBase.__init__(self)
91  BenchClassTemplate.__init__(self)
92 
93  self.brahma = brahma
94  self.devices = devices
95 
97  self.uiBench.setupUi(self)
98 
99 
102 
103  self.supervisor = None
104  self.config = None
105  self.stop = False # type: bool # Request quit
106 
107  self.uiBench.wao_nbiters.setValue(1000) # Default GUI nIter box value
108  self.nbiter = self.uiBench.wao_nbiters.value()
109  self.refreshTime = 0 # type: float # System time at last display refresh
110  self.loopThread = None # type: QThread
111  self.assistant = None # type: Any
112 
113 
117  self.defaultParPath = os.environ["SHESHA_ROOT"] + "/data/par/bench-config"
118  self.defaultAreaPath = os.environ["SHESHA_ROOT"] + "/data/layouts"
119  self.loadDefaultConfig()
120 
121  self.uiBench.wao_run.setCheckable(True)
122  self.uiBench.wao_run.clicked[bool].connect(self.aoLoopClicked)
123  self.uiBench.wao_open_loop.setCheckable(True)
124  self.uiBench.wao_open_loop.clicked[bool].connect(self.aoLoopOpen)
125  self.uiBench.wao_next.clicked.connect(self.loop_once)
126 
127  self.uiBench.wao_forever.stateChanged.connect(self.updateForever)
128 
129  self.dispStatsInTerminal = False
130 
131  self.uiBench.wao_run.setDisabled(True)
132  self.uiBench.wao_next.setDisabled(True)
133  self.uiBench.wao_unzoom.setDisabled(True)
134 
135  self.addDockWidget(Qt.DockWidgetArea(1), self.uiBase.wao_ConfigDock)
136  self.addDockWidget(Qt.DockWidgetArea(1), self.uiBase.wao_DisplayDock)
137  self.uiBase.wao_ConfigDock.setFloating(False)
138  self.uiBase.wao_DisplayDock.setFloating(False)
139 
140  self.adjustSize()
141 
142  if config_file is not None:
143  self.uiBase.wao_selectConfig.clear()
144  self.uiBase.wao_selectConfig.addItem(config_file)
145  self.load_config()
146  self.init_config()
147 
148 
151 
152  # def updateStatsInTerminal(self, state):
153  # self.dispStatsInTerminal = state
154 
155  def updateForever(self, state):
156  self.uiBench.wao_nbiters.setDisabled(state)
157 
158  def add_dispDock(self, name: str, parent, type: str = "pg_image") -> None:
159  d = WidgetBase.add_dispDock(self, name, parent, type)
160  if type == "SR":
161  d.addWidget(self.uiBench.wao_Strehl)
162 
163  def load_config(self) -> None:
164  '''
165  Callback when 'LOAD' button is hit
166  '''
167  WidgetBase.load_config(self)
168  config_file = str(self.uiBase.wao_selectConfig.currentText())
169  sys.path.insert(0, self.defaultParPath)
170 
171  self.supervisor = Supervisor(config_file, self.brahma)
172 
173  self.config = self.supervisor.get_config()
174 
175  # if self.devices:
176  # self.config.p_loop.set_devices([
177  # int(device) for device in self.devices.split(",")
178  # ])
179 
180  try:
181  sys.path.remove(self.defaultParPath)
182  except:
183  pass
184 
185  try:
186  self.nwfs = len(self.config.p_wfss)
187  except:
188  self.nwfs = 1 # Default. config may very well not define p_wfss
189  for wfs in range(self.nwfs):
190  name = 'slpComp_%d' % wfs
191  self.add_dispDock(name, self.wao_graphgroup_cb, "MPL")
192  name = 'wfs_%d' % wfs
193  self.add_dispDock(name, self.wao_imagesgroup_cb)
194 
195  self.uiBench.wao_run.setDisabled(True)
196  self.uiBench.wao_next.setDisabled(True)
197  self.uiBench.wao_unzoom.setDisabled(True)
198 
199  self.uiBase.wao_init.setDisabled(False)
200 
201  if (hasattr(self.config, "layout")):
202  area_filename = self.defaultAreaPath + "/" + self.config.layout + ".area"
203  self.loadArea(filename=area_filename)
204 
205  self.adjustSize()
206 
207  def aoLoopClicked(self, pressed: bool) -> None:
208  if pressed:
209  self.stop = False
210  self.refreshTime = time.time()
211  self.nbiter = self.uiBench.wao_nbiters.value()
212  if self.dispStatsInTerminal:
213  if self.uiBench.wao_forever.isChecked():
214  print("LOOP STARTED")
215  else:
216  print("LOOP STARTED FOR %d iterations" % self.nbiter)
217  self.run()
218  else:
219  self.stop = True
220 
221  def aoLoopOpen(self, pressed: bool) -> None:
222  if (pressed):
223  self.supervisor.open_loop()
224  self.uiAO.wao_open_loop.setText("Close Loop")
225  else:
226  self.supervisor.close_loop()
227  self.uiAO.wao_open_loop.setText("Open Loop")
228 
229  def init_config(self) -> None:
230  WidgetBase.init_config(self)
231 
232  def init_configThread(self) -> None:
233  # self.uiBench.wao_deviceNumber.setDisabled(True)
234  # self.config.p_loop.devices = self.uiBench.wao_deviceNumber.value() # using GUI value
235  # gpudevice = "ALL" # using all GPU avalaible
236  # gpudevice = np.array([2, 3], dtype=np.int32)
237  # gpudevice = np.arange(4, dtype=np.int32) # using 4 GPUs: 0-3
238  # gpudevice = 0 # using 1 GPU : 0
239  self.supervisor.init_config()
240 
241  def init_configFinished(self) -> None:
242  # Thread carmaWrap context reload:
243  try:
244  self.supervisor.force_context()
245  except:
246  print('Warning: could not call supervisor.force_context().')
247 
248  for i in range(self.nwfs):
249  if self.config.p_wfss[i].type == WFSType.SH:
250  key = "wfs_%d" % i
251  self.addSHGrid(self.docks[key].widgets[0],
252  self.config.p_wfss[i].get_validsub(),
253  self.config.p_wfss[i].npix, self.config.p_wfss[i].npix)
254 
255  self.updateDisplay()
256 
257  self.uiBench.wao_run.setDisabled(False)
258  self.uiBench.wao_next.setDisabled(False)
259  self.uiBench.wao_open_loop.setDisabled(False)
260  self.uiBench.wao_unzoom.setDisabled(False)
261 
262  WidgetBase.init_configFinished(self)
263 
264  def updateDisplay(self) -> None:
265  if (self.supervisor is None) or (not self.supervisor.is_init()) or (
266  not self.uiBase.wao_Display.isChecked()):
267  # print("Widget not fully initialized")
268  return
269  if not self.loopLock.acquire(False):
270  return
271  else:
272  try:
273  for key, dock in self.docks.items():
274  if key == "Strehl":
275  continue
276  elif dock.isVisible():
277  index = int(key.split("_")[-1])
278  data = None
279  if "wfs" in key:
280  data = self.supervisor.wfs.get_wfs_image(index)
281  if (data is not None):
282  autoscale = True # self.uiBench.actionAuto_Scale.isChecked()
283  # if (autoscale):
284  # # inits levels
285  # self.hist.setLevels(data.min(), data.max())
286  self.imgs[key].setImage(data, autoLevels=autoscale)
287  # self.p1.autoRange()
288 
289  elif "slp" in key: # Slope display
290  if (self.config.p_wfss[index].type == WFSType.PYRHR or
291  self.config.p_wfss[index].type == WFSType.PYRLR):
292  raise RuntimeError("PYRHR not usable")
293  self.imgs[key].canvas.axes.clear()
294  x, y = self.supervisor.config.p_wfss[index].get_validsub()
295 
296  nssp = x.size
297  centroids = self.supervisor.rtc.get_slopes()
298  vx = centroids[:nssp]
299  vy = centroids[nssp:]
300 
301  offset = (self.supervisor.config.p_wfss[index].npix - 1) / 2
302  self.imgs[key].canvas.axes.quiver(
303  x + offset, y + offset, vy, vx, angles='xy',
304  scale_units='xy',
305  scale=1) # wao.supervisor.config.p_wfss[0].pixsize)
306  self.imgs[key].canvas.draw()
307 
308  finally:
309  self.loopLock.release()
310 
311  def loop_once(self) -> None:
312  if not self.loopLock.acquire(False):
313  print("Display locked")
314  return
315  else:
316  try:
317  start = time.time()
318  self.supervisor.single_next()
319  loopTime = time.time() - start
320  refreshDisplayTime = 1. / self.uiBase.wao_frameRate.value()
321 
322  if (time.time() - self.refreshTime > refreshDisplayTime):
323  currentFreq = 1 / loopTime
324  self.uiBench.wao_currentFreq.setValue(currentFreq)
325  self.refreshTime = start
326  except:
327  pass
328  finally:
329  self.loopLock.release()
330 
331  def run(self):
332  WidgetBase.run(self)
333  if not self.uiBench.wao_forever.isChecked():
334  self.nbiter -= 1
335 
336  if self.nbiter <= 0:
337  self.stop = True
338  self.uiBench.wao_run.setChecked(False)
339 
340 
341 if __name__ == '__main__':
342  arguments = docopt(__doc__)
343  app = QtWidgets.QApplication(sys.argv)
344  app.setStyle('cleanlooks')
345  wao = widgetBenchWindow(arguments["<parameters_filename>"],
346  brahma=arguments["--brahma"], devices=arguments["--devices"])
347  wao.show()
348  if arguments["--interactive"]:
349  from shesha.util.ipython_embed import embed
350  from os.path import basename
351  embed(basename(__file__), locals())
shesha.widgets.widget_bench.widgetBenchWindow.run
def run(self)
Definition: widget_bench.py:331
shesha.widgets.widget_base.WidgetBase.wao_graphgroup_cb
wao_graphgroup_cb
Definition: widget_base.py:322
shesha.widgets.widget_bench.widgetBenchWindow.aoLoopOpen
None aoLoopOpen(self, bool pressed)
Definition: widget_bench.py:221
shesha.widgets.widget_bench.widgetBenchWindow.dispStatsInTerminal
dispStatsInTerminal
Definition: widget_bench.py:128
shesha.util.tools
Imported from CANARY.
Definition: tools.py:1
shesha.widgets.widget_bench.widgetBenchWindow
Definition: widget_bench.py:86
shesha.widgets.widget_bench.widgetBenchWindow.add_dispDock
None add_dispDock(self, str name, parent, str type="pg_image")
Definition: widget_bench.py:158
shesha.widgets.widget_bench.widgetBenchWindow.refreshTime
refreshTime
Definition: widget_bench.py:108
shesha.widgets.widget_bench.widgetBenchWindow.updateForever
def updateForever(self, state)
METHODS #.
Definition: widget_bench.py:155
shesha.widgets.widget_bench.widgetBenchWindow.load_config
None load_config(self)
Callback when 'LOAD' button is hit.
Definition: widget_bench.py:166
shesha.widgets.widget_base.WidgetBase.loadArea
def loadArea(self, widget=None, filename=None)
Definition: widget_base.py:197
shesha.widgets.widget_bench.widgetBenchWindow.stop
stop
Definition: widget_bench.py:104
shesha.widgets.widget_bench.widgetBenchWindow.__init__
None __init__(self, Any config_file=None, bool brahma=False, str devices=None)
Definition: widget_bench.py:88
shesha.widgets.widget_bench.BenchClassTemplate
BenchClassTemplate
Definition: widget_bench.py:74
shesha.widgets.widget_bench.widgetBenchWindow.loop_once
None loop_once(self)
Definition: widget_bench.py:311
shesha.widgets.widget_bench.widgetBenchWindow.aoLoopClicked
None aoLoopClicked(self, bool pressed)
Definition: widget_bench.py:207
shesha.widgets.widget_base.WidgetBase.addSHGrid
def addSHGrid(self, pg_image, valid_sub, sspsize, pitch)
Definition: widget_base.py:373
shesha.supervisor.benchSupervisor
Initialization and execution of a Bench supervisor.
Definition: benchSupervisor.py:1
open_loop
Definition: open_loop.py:1
shesha.widgets.widget_bench.widgetBenchWindow.supervisor
supervisor
ATTRIBUTES #.
Definition: widget_bench.py:102
shesha.widgets.widget_base.WidgetBase.add_dispDock
Dock add_dispDock(self, str name, parent, str type="pg_image")
Definition: widget_base.py:252
shesha.widgets.widget_base.WidgetBase.docks
docks
Definition: widget_base.py:125
shesha.widgets.widget_base
Abstract Widget base.
Definition: widget_base.py:1
shesha.widgets.widget_bench.widgetBenchWindow.brahma
brahma
Definition: widget_bench.py:92
shesha.widgets.widget_base.WidgetBase.loopLock
loopLock
Definition: widget_base.py:95
shesha.widgets.widget_bench.widgetBenchWindow.init_config
None init_config(self)
Definition: widget_bench.py:229
shesha.widgets.widget_base.WidgetBase.load_config
None load_config(self, *args, **kwargs)
Callback when 'LOAD' button is hit.
Definition: widget_base.py:296
shesha.widgets.widget_bench.BenchWindowTemplate
BenchWindowTemplate
Definition: widget_bench.py:74
shesha.widgets.widget_bench.widgetBenchWindow.updateDisplay
None updateDisplay(self)
Definition: widget_bench.py:264
shesha.widgets.widget_base.WidgetBase.defaultParPath
defaultParPath
Definition: widget_base.py:110
shesha.widgets.widget_bench.widgetBenchWindow.init_configThread
None init_configThread(self)
Definition: widget_bench.py:232
shesha.widgets.widget_bench.widgetBenchWindow.nwfs
nwfs
Definition: widget_bench.py:186
shesha.widgets.widget_base.WidgetBase.imgs
imgs
Definition: widget_base.py:127
shesha.widgets.widget_bench.widgetBenchWindow.loopThread
loopThread
Definition: widget_bench.py:109
shesha.widgets.widget_base.WidgetBase.init_config
None init_config(self)
Definition: widget_base.py:348
shesha.widgets.widget_bench.widgetBenchWindow.uiBench
uiBench
Definition: widget_bench.py:95
shesha.util.ipython_embed
Definition: ipython_embed.py:1
shesha.widgets.widget_base.WidgetBase.uiBase
uiBase
Definition: widget_base.py:84
shesha.widgets.widget_base.WidgetBase
Definition: widget_base.py:79
shesha.widgets.widget_base.WidgetBase.defaultAreaPath
defaultAreaPath
Definition: widget_base.py:111
shesha.widgets.widget_base.WidgetBase.wao_imagesgroup_cb
wao_imagesgroup_cb
Definition: widget_base.py:328
shesha.widgets.widget_bench.widgetBenchWindow.nbiter
nbiter
Definition: widget_bench.py:107
shesha.widgets.widget_bench.widgetBenchWindow.devices
devices
Definition: widget_bench.py:93
shesha.widgets.widget_base.WidgetBase.run
def run(self)
Definition: widget_base.py:394
shesha.widgets.widget_bench.widgetBenchWindow.config
config
Definition: widget_bench.py:103
shesha.widgets.widget_bench.widgetBenchWindow.assistant
assistant
Definition: widget_bench.py:110
shesha.widgets.widget_bench.widgetBenchWindow.init_configFinished
None init_configFinished(self)
Definition: widget_bench.py:241
shesha.widgets.widget_base.WidgetBase.loadDefaultConfig
None loadDefaultConfig(self)
Definition: widget_base.py:340
shesha.widgets.widget_base.WidgetBase.updateDisplay
None updateDisplay(self)
Definition: widget_base.py:364