COMPASS  5.4.4
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 time
54 
55 try:
56  from PyQt5 import QtWidgets
57  from PyQt5.QtCore import QThread, Qt
58 except ModuleNotFoundError as e:
59  try:
60  from PySide2 import QtWidgets
61  from PySide2.QtCore import QThread, Qt
62  except ModuleNotFoundError as e:
63  raise ModuleNotFoundError("No module named 'PyQt5' or PySide2', please install one of them\nException raised: "+e.msg)
64 
65 from typing import Any
66 
67 from docopt import docopt
68 
69 from shesha.widgets.widget_base import WidgetBase, uiLoader
70 BenchWindowTemplate, BenchClassTemplate = uiLoader('widget_bench')
71 
72 import matplotlib.pyplot as plt
73 
74 from shesha.supervisor.benchSupervisor import WFSType, BenchSupervisor as Supervisor
75 
76 # For debug
77 # from IPython.core.debugger import Pdb
78 # then add this line to create a breakpoint
79 # Pdb().set_trace()
80 
81 
83 
84  def __init__(self, config_file: Any = None, brahma: bool = False,
85  devices: str = None) -> None:
86  WidgetBase.__init__(self)
87  BenchClassTemplate.__init__(self)
88 
89  self.brahmabrahma = brahma
90  self.devicesdevices = devices
91 
92  self.uiBenchuiBench = BenchWindowTemplate()
93  self.uiBenchuiBench.setupUi(self)
94 
95 
98 
99  self.supervisorsupervisor = None
100  self.configconfig = None
101  self.stopstop = False # type: bool # Request quit
102 
103  self.uiBenchuiBench.wao_nbiters.setValue(1000) # Default GUI nIter box value
104  self.nbiternbiter = self.uiBenchuiBench.wao_nbiters.value()
105  self.refreshTimerefreshTime = 0 # type: float # System time at last display refresh
106  self.loopThreadloopThread = None # type: QThread
107  self.assistantassistant = None # type: Any
108 
109 
113  self.defaultParPathdefaultParPathdefaultParPath = os.environ["SHESHA_ROOT"] + "/data/par/bench-config"
114  self.defaultAreaPathdefaultAreaPathdefaultAreaPath = os.environ["SHESHA_ROOT"] + "/data/layouts"
115  self.loadDefaultConfigloadDefaultConfig()
116 
117  self.uiBenchuiBench.wao_run.setCheckable(True)
118  self.uiBenchuiBench.wao_run.clicked[bool].connect(self.aoLoopClickedaoLoopClicked)
119  self.uiBenchuiBench.wao_open_loop.setCheckable(True)
120  self.uiBenchuiBench.wao_open_loop.clicked[bool].connect(self.aoLoopOpenaoLoopOpen)
121  self.uiBenchuiBench.wao_next.clicked.connect(self.loop_onceloop_once)
122 
123  self.uiBenchuiBench.wao_forever.stateChanged.connect(self.updateForeverupdateForever)
124 
125  self.dispStatsInTerminaldispStatsInTerminal = False
126 
127  self.uiBenchuiBench.wao_run.setDisabled(True)
128  self.uiBenchuiBench.wao_next.setDisabled(True)
129  self.uiBenchuiBench.wao_unzoom.setDisabled(True)
130 
131  self.addDockWidget(Qt.DockWidgetArea(1), self.uiBaseuiBase.wao_ConfigDock)
132  self.addDockWidget(Qt.DockWidgetArea(1), self.uiBaseuiBase.wao_DisplayDock)
133  self.uiBaseuiBase.wao_ConfigDock.setFloating(False)
134  self.uiBaseuiBase.wao_DisplayDock.setFloating(False)
135 
136  self.adjustSize()
137 
138  if config_file is not None:
139  self.uiBaseuiBase.wao_selectConfig.clear()
140  self.uiBaseuiBase.wao_selectConfig.addItem(config_file)
141  self.load_configload_configload_config()
142  self.init_configinit_configinit_config()
143 
144 
147 
148  # def updateStatsInTerminal(self, state):
149  # self.dispStatsInTerminal = state
150 
151  def updateForever(self, state):
152  self.uiBenchuiBench.wao_nbiters.setDisabled(state)
153 
154  def add_dispDock(self, name: str, parent, type: str = "pg_image") -> None:
155  d = WidgetBase.add_dispDock(self, name, parent, type)
156  if type == "SR":
157  d.addWidget(self.uiBenchuiBench.wao_Strehl)
158 
159  def load_config(self) -> None:
160  '''
161  Callback when 'LOAD' button is hit
162  '''
163  WidgetBase.load_config(self)
164  config_file = str(self.uiBaseuiBase.wao_selectConfig.currentText())
165  sys.path.insert(0, self.defaultParPathdefaultParPathdefaultParPath)
166 
167  self.supervisorsupervisor = Supervisor(config_file, self.brahmabrahma)
168 
169  self.configconfig = self.supervisorsupervisor.get_config()
170 
171  # if self.devices:
172  # self.config.p_loop.set_devices([
173  # int(device) for device in self.devices.split(",")
174  # ])
175 
176  try:
177  sys.path.remove(self.defaultParPathdefaultParPathdefaultParPath)
178  except:
179  pass
180 
181  try:
182  self.nwfsnwfs = len(self.configconfig.p_wfss)
183  except:
184  self.nwfsnwfs = 1 # Default. config may very well not define p_wfss
185  for wfs in range(self.nwfsnwfs):
186  name = 'slpComp_%d' % wfs
187  self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_graphgroup_cbwao_graphgroup_cb, "MPL")
188  name = 'wfs_%d' % wfs
189  self.add_dispDockadd_dispDockadd_dispDock(name, self.wao_imagesgroup_cbwao_imagesgroup_cb)
190 
191  self.uiBenchuiBench.wao_run.setDisabled(True)
192  self.uiBenchuiBench.wao_next.setDisabled(True)
193  self.uiBenchuiBench.wao_unzoom.setDisabled(True)
194 
195  self.uiBaseuiBase.wao_init.setDisabled(False)
196 
197  if (hasattr(self.configconfig, "layout")):
198  area_filename = self.defaultAreaPathdefaultAreaPathdefaultAreaPath + "/" + self.configconfig.layout + ".area"
199  self.loadArealoadArea(filename=area_filename)
200 
201  self.adjustSize()
202 
203  def aoLoopClicked(self, pressed: bool) -> None:
204  if pressed:
205  self.stopstop = False
206  self.refreshTimerefreshTime = time.time()
207  self.nbiternbiter = self.uiBenchuiBench.wao_nbiters.value()
208  if self.dispStatsInTerminaldispStatsInTerminal:
209  if self.uiBenchuiBench.wao_forever.isChecked():
210  print("LOOP STARTED")
211  else:
212  print("LOOP STARTED FOR %d iterations" % self.nbiternbiter)
213  self.runrunrun()
214  else:
215  self.stopstop = True
216 
217  def aoLoopOpen(self, pressed: bool) -> None:
218  if (pressed):
219  self.supervisorsupervisor.open_loop()
220  self.uiAO.wao_open_loop.setText("Close Loop")
221  else:
222  self.supervisorsupervisor.close_loop()
223  self.uiAO.wao_open_loop.setText("Open Loop")
224 
225  def init_config(self) -> None:
226  WidgetBase.init_config(self)
227 
228  def init_configThread(self) -> None:
229  # self.uiBench.wao_deviceNumber.setDisabled(True)
230  # self.config.p_loop.devices = self.uiBench.wao_deviceNumber.value() # using GUI value
231  # gpudevice = "ALL" # using all GPU avalaible
232  # gpudevice = np.array([2, 3], dtype=np.int32)
233  # gpudevice = np.arange(4, dtype=np.int32) # using 4 GPUs: 0-3
234  # gpudevice = 0 # using 1 GPU : 0
235  self.supervisorsupervisor.init_config()
236 
237  def init_configFinished(self) -> None:
238  # Thread carmaWrap context reload:
239  try:
240  self.supervisorsupervisor.force_context()
241  except:
242  print('Warning: could not call supervisor.force_context().')
243 
244  for i in range(self.nwfsnwfs):
245  if self.configconfig.p_wfss[i].type == WFSType.SH:
246  key = "wfs_%d" % i
247  self.addSHGridaddSHGrid(self.docksdocks[key].widgets[0],
248  self.configconfig.p_wfss[i].get_validsub(),
249  self.configconfig.p_wfss[i].npix, self.configconfig.p_wfss[i].npix)
250 
251  self.updateDisplayupdateDisplayupdateDisplay()
252 
253  self.uiBenchuiBench.wao_run.setDisabled(False)
254  self.uiBenchuiBench.wao_next.setDisabled(False)
255  self.uiBenchuiBench.wao_open_loop.setDisabled(False)
256  self.uiBenchuiBench.wao_unzoom.setDisabled(False)
257 
258  WidgetBase.init_configFinished(self)
259 
260  def updateDisplay(self) -> None:
261  if (self.supervisorsupervisor is None) or (not self.supervisorsupervisor.is_init()) or (
262  not self.uiBaseuiBase.wao_Display.isChecked()):
263  # print("Widget not fully initialized")
264  return
265  if not self.loopLockloopLock.acquire(False):
266  return
267  else:
268  try:
269  for key, dock in self.docksdocks.items():
270  if key == "Strehl":
271  continue
272  elif dock.isVisible():
273  index = int(key.split("_")[-1])
274  data = None
275  if "wfs" in key:
276  data = self.supervisorsupervisor.wfs.get_wfs_image(index)
277  if (data is not None):
278  autoscale = True # self.uiBench.actionAuto_Scale.isChecked()
279  # if (autoscale):
280  # # inits levels
281  # self.hist.setLevels(data.min(), data.max())
282  self.imgsimgs[key].setImage(data, autoLevels=autoscale)
283  # self.p1.autoRange()
284 
285  elif "slp" in key: # Slope display
286  if (self.configconfig.p_wfss[index].type == WFSType.PYRHR or
287  self.configconfig.p_wfss[index].type == WFSType.PYRLR):
288  raise RuntimeError("PYRHR not usable")
289  self.imgsimgs[key].canvas.axes.clear()
290  x, y = self.supervisorsupervisor.config.p_wfss[index].get_validsub()
291 
292  nssp = x.size
293  centroids = self.supervisorsupervisor.rtc.get_slopes()
294  vx = centroids[:nssp]
295  vy = centroids[nssp:]
296 
297  offset = (self.supervisorsupervisor.config.p_wfss[index].npix - 1) / 2
298  self.imgsimgs[key].canvas.axes.quiver(
299  x + offset, y + offset, vy, vx, angles='xy',
300  scale_units='xy',
301  scale=1) # wao.supervisor.config.p_wfss[0].pixsize)
302  self.imgsimgs[key].canvas.draw()
303 
304  finally:
305  self.loopLockloopLock.release()
306 
307  def loop_once(self) -> None:
308  if not self.loopLockloopLock.acquire(False):
309  print("Display locked")
310  return
311  else:
312  try:
313  start = time.time()
314  self.supervisorsupervisor.single_next()
315  loopTime = time.time() - start
316  refreshDisplayTime = 1. / self.uiBaseuiBase.wao_frameRate.value()
317 
318  if (time.time() - self.refreshTimerefreshTime > refreshDisplayTime):
319  currentFreq = 1 / loopTime
320  self.uiBenchuiBench.wao_currentFreq.setValue(currentFreq)
321  self.refreshTimerefreshTime = start
322  except:
323  pass
324  finally:
325  self.loopLockloopLock.release()
326 
327  def run(self):
328  WidgetBase.run(self)
329  if not self.uiBenchuiBench.wao_forever.isChecked():
330  self.nbiternbiter -= 1
331 
332  if self.nbiternbiter <= 0:
333  self.stopstop = True
334  self.uiBenchuiBench.wao_run.setChecked(False)
335 
336 
337 if __name__ == '__main__':
338  arguments = docopt(__doc__)
339  app = QtWidgets.QApplication(sys.argv)
340  app.setStyle('cleanlooks')
341  wao = widgetBenchWindow(arguments["<parameters_filename>"],
342  brahma=arguments["--brahma"], devices=arguments["--devices"])
343  wao.show()
344  if arguments["--interactive"]:
345  from shesha.util.ipython_embed import embed
346  from os.path import basename
347  embed(basename(__file__), locals())
def loadArea(self, widget=None, filename=None)
Definition: widget_base.py:207
None load_config(self, *args, **kwargs)
Callback when 'LOAD' button is hit.
Definition: widget_base.py:311
def addSHGrid(self, pg_image, valid_sub, sspsize, pitch)
Definition: widget_base.py:388
Dock add_dispDock(self, str name, parent, str type="pg_image")
Definition: widget_base.py:262
None load_config(self)
Callback when 'LOAD' button is hit.
None __init__(self, Any config_file=None, bool brahma=False, str devices=None)
Definition: widget_bench.py:85
None add_dispDock(self, str name, parent, str type="pg_image")
def updateForever(self, state)
METHODS #.
Initialization and execution of a Bench supervisor.
Abstract Widget base.
Definition: widget_base.py:1