COMPASS  5.4.4
End-to-end AO simulation tool using GPU acceleration
widget_twoStages.py
1 
37 """
38 Widget built to simulate a 2 stage AO loop (1st stage = SH; second stage = pyramid)
39 
40 Usage:
41  widget_twoStages.py <parameters_filename1> <parameters_filename2> <freqratio> [options]
42 
43 with 'parameters_filename1' the path to the parameters file for first stage
44 with 'parameters_filename2' the path to the parameters file for second stage
45 with 'freqratio' the ratio of the frequencies of the two stages
46 
47 Options:
48  -a, --adopt used to connect ADOPT (via pyro + shm cacao)
49 
50 Example:
51  ipython -i widget_twoStages.py ../../data/par/SPHERE+/sphere.py ../../data/par/SPHERE+/sphere+.py
52  ipython -i widget_twoStages.py ../../data/par/SPHERE+/sphere.py ../../data/par/SPHERE+/sphere+.py -- --adopt
53 """
54 
55 import os, sys
56 import numpy as np
57 import time
58 
59 import pyqtgraph as pg
60 from shesha.util.tools import plsh, plpyr
61 from rich.progress import track
62 import astropy.io.fits as pfits
63 from PyQt5 import QtWidgets
64 from shesha.supervisor.twoStagesManager import TwoStagesManager
65 
66 from typing import Any, Dict, Tuple, Callable, List
67 from docopt import docopt
68 
69 from shesha.widgets.widget_base import WidgetBase
70 from shesha.widgets.widget_ao import widgetAOWindow, widgetAOWindow
71 
72 global server
73 server = None
74 
75 
77 
78  def __init__(self, config_file1: Any = None, config_file2: Any = None, freqratio : int = None,
79  cacao: bool = False, expert: bool = False) -> None:
80  self.config1config1 = config_file1
81  self.config2config2 = config_file2
82  self.freqratiofreqratio = freqratio
83 
84  from shesha.config import ParamConfig
85 
86 
87  self.wao2wao2=widgetAOWindow(config_file2, cacao=cacao, hide_histograms=True, twoStages=True)
88  self.wao1wao1=widgetAOWindow(config_file1, cacao=cacao, hide_histograms=True, twoStages=True)
89  pupdiam_first_stage = self.wao1wao1.supervisor.config.p_geom.pupdiam
90  pupdiam_second_stage = self.wao2wao2.supervisor.config.p_geom.pupdiam
91  if(pupdiam_first_stage != pupdiam_second_stage):
92  print("---------------ERROR---------------")
93  print("SECOND STAGE PUPDIAM IS SET TO %d" % pupdiam_second_stage)
94  print("FIRST STAGE PUPDIAM IS SET TO %d" % pupdiam_first_stage)
95  raise Exception('ERROR!!!! FIRST STAGE PUPDIAM MUST BE SET TO %d' % pupdiam_second_stage)
96 
97  #Pyro.core.ObjBase.__init__(self)
98  self.CBCB = {}
99  self.wpyrwpyr = None
100  self.current_buffercurrent_buffer = 1
101  self.managermanager = None
102  self.cacaocacao = cacao
103 
109  self.wao1wao1.uiAO.actionShow_Pyramid_Tools.toggled.connect(self.show_pyr_toolsshow_pyr_tools)
110  self.wao2wao2.uiAO.actionShow_Pyramid_Tools.toggled.connect(self.show_pyr_toolsshow_pyr_tools)
111  self.wpyrNbBufferwpyrNbBuffer = 1
112 
115 
116  self.managermanager = TwoStagesManager(self.wao1wao1.supervisor, self.wao2wao2.supervisor, self.freqratiofreqratio)
117  if(self.cacaocacao):
118  global server
119  server = self.start_pyro_serverstart_pyro_server()
120 
121  def loop_once(self) -> None:
122  self.managermanager.next()
123 
124  for wao in [self.wao1wao1, self.wao2wao2]:
125  start = time.time()
126  for t in range(len(wao.supervisor.config.p_targets)):
127  wao.supervisor.target.comp_tar_image(t)
128  loopTime = time.time() - start
129 
130  refreshDisplayTime = 1. / wao.uiBase.wao_frameRate.value()
131 
132  if (time.time() - wao.refreshTime > refreshDisplayTime):
133  signal_le = ""
134  signal_se = ""
135  for t in range(len(wao.config.p_targets)):
136  SR = wao.supervisor.target.get_strehl(t)
137  # TODO: handle that !
138  if (t == wao.uiAO.wao_dispSR_tar.value()
139  ): # Plot on the wfs selected
140  wao.updateSRDisplay(SR[1], SR[0],
141  wao.supervisor.get_frame_counter())
142  signal_se += "%1.2f " % SR[0]
143  signal_le += "%1.2f " % SR[1]
144 
145  currentFreq = 1 / loopTime
146  refreshFreq = 1 / (time.time() - wao.refreshTime)
147 
148  wao.updateSRSE(signal_se)
149  wao.updateSRLE(signal_le)
150  wao.updateCurrentLoopFrequency(currentFreq)
151 
152  if (wao.dispStatsInTerminal):
153  wao.printInPlace(
154  "iter #%d SR: (L.E, S.E.)= (%s, %s) running at %4.1fHz (real %4.1fHz)"
155  % (wao.supervisor.get_frame_counter(), signal_le,
156  signal_se, refreshFreq, currentFreq))
157  if (self.wao2wao2.uiAO.actionShow_Pyramid_Tools.isChecked()): # PYR only
158  self.wpyrwpyr.Fe = 1 / self.config.p_loop.ittime # needs Fe for PSD...
159  if (self.wpyrwpyr.CBNumber == 1):
160  self.aiai = self.compute_modal_residuals()
161  self.set_pyr_tools_paramsset_pyr_tools_params(self.aiai)
162  else:
163  if (self.current_buffercurrent_buffer == 1): # First iter of the CB
164  aiVect = self.compute_modal_residuals()
165  self.aiai = aiVect[np.newaxis, :]
166  self.current_buffercurrent_buffer += 1 # Keep going
167 
168  else: # Keep filling the CB
169  aiVect = self.compute_modal_residuals()
170  self.aiai = np.concatenate((self.aiai, aiVect[np.newaxis, :]))
171  if (self.current_buffercurrent_buffer < self.wpyrwpyr.CBNumber):
172  self.current_buffercurrent_buffer += 1 # Keep going
173  else:
174  self.current_buffercurrent_buffer = 1 # reset buffer
175  self.set_pyr_tools_paramsset_pyr_tools_params(self.aiai) # display
176 
177  def next(self, nbIters):
178  ''' Move atmos -> get_slopes -> applyControl ; One integrator step '''
179  for i in track(range(nbIters)):
180  self.managermanager.next()
181 
182  def initPyrTools(self):
183  ADOPTPATH = os.getenv("ADOPTPATH")
184  sys.path.append(ADOPTPATH + "/widgets")
185  from pyrStats import widget_pyrStats
186  print("OK Pyramid Tools Widget initialized")
187  self.wpyrwpyr = widget_pyrStats()
188  self.wpyrNbBufferwpyrNbBuffer = self.wpyrwpyr.CBNumber
189  self.wpyrwpyr.show()
190 
191  def set_pyr_tools_params(self, ai):
192  self.wpyrwpyr.pup = self.managermanager.second_stage.config.p_geom._spupil
193  self.wpyrwpyr.phase = self.supervisor.target.get_tar_phase(0, pupil=True)
194  self.wpyrwpyr.updateResiduals(ai)
195  if (self.phase_to_modesphase_to_modes is None):
196  print('computing phase 2 Modes basis')
197  self.phase_to_modesphase_to_modes = self.managermanager.second_stage.basis.compute_phase_to_modes(self.modal_basis)
198  self.wpyrwpyr.phase_to_modes = self.phase_to_modesphase_to_modes
199 
200  def show_pyr_tools(self):
201  if (self.wpyrwpyr is None):
202  try:
203  print("Lauching pyramid widget...")
204  self.initPyrToolsinitPyrTools()
205  print("Done")
206  except:
207  raise ValueError("ERROR: ADOPT not found. Cannot launch Pyramid tools")
208  else:
209  if (self.wao2wao2.uiAO.actionShow_Pyramid_Tools.isChecked()):
210  self.wpyrwpyr.show()
211  else:
212  self.wpyrwpyr.hide()
213 
214  def getAi(self):
215  return self.wpyrwpyr.ai
216 
217  def start_pyro_server(self):
218  try:
219  wao_loop = loopHandler( self.wao1wao1)
220 
221  from subprocess import Popen, PIPE
222  from hraa.server.pyroServer import PyroServer
223  import Pyro4
224  Pyro4.config.REQUIRE_EXPOSE = False
225  p = Popen("whoami", shell=True, stdout=PIPE, stderr=PIPE)
226  out, err = p.communicate()
227  if (err != b''):
228  print(err)
229  raise Exception("ERROR CANNOT RECOGNIZE USER")
230  else:
231  user = out.split(b"\n")[0].decode("utf-8")
232  print("User is " + user)
233  supervisor = self.managermanager
234  supervisor1 = self.managermanager.first_stage
235  supervisor2 = self.managermanager.second_stage
236 
237  if(supervisor1.corono == None):
238  from shesha.util.pyroEmptyClass import PyroEmptyClass
239  coro2pyro1 = PyroEmptyClass()
240  else:
241  coro2pyro1 = supervisor1.corono
242 
243  if(supervisor2.corono == None):
244  from shesha.util.pyroEmptyClass import PyroEmptyClass
245  coro2pyro2 = PyroEmptyClass()
246  else:
247  coro2pyro2 = supervisor2.corono
248 
249  devices1 = [
250  supervisor1, supervisor1.rtc, supervisor1.wfs, supervisor1.target,
251  supervisor1.tel, supervisor1.basis, supervisor1.calibration,
252  supervisor1.atmos, supervisor1.dms, supervisor1.config, supervisor1.modalgains,
253  coro2pyro1
254  ]
255  devices2 = [
256  supervisor2, supervisor2.rtc, supervisor2.wfs, supervisor2.target,
257  supervisor2.tel, supervisor2.basis, supervisor2.calibration,
258  supervisor2.atmos, supervisor2.dms, supervisor2.config, supervisor2.modalgains,
259  coro2pyro2
260  ]
261  names = [
262  "supervisor", "supervisor_rtc", "supervisor_wfs", "supervisor_target",
263  "supervisor_tel", "supervisor_basis", "supervisor_calibration",
264  "supervisor_atmos", "supervisor_dms", "supervisor_config", "supervisor_modalgains",
265  "supervisor_corono"
266  ]
267 
268  label = "firstStage"
269  nname = []
270  for name in names:
271  nname.append(name + "_" + user + "_" +label)
272 
273  label = "secondStage"
274  for name in names:
275  nname.append(name + "_" + user + "_" +label)
276 
277  nname.append('twoStagesManager'+ "_" + user ) # Adding master next dedicated to trigger 2-stages loop
278  nname.append("wao_loop"+ "_" + user)
279  devices = devices1 + devices2 + [supervisor, wao_loop]
280  server = PyroServer(listDevices=devices, listNames=nname)
281  #server.add_device(supervisor, "waoconfig_" + user)
282  server.start()
283 
284 
285  except:
286  raise Exception("Error could not connect to Pyro server.\n It can be:\n - Missing dependencies? (check if Pyro4 is installed)\n - pyro server not running")
287  return server
288 
289 
290 class loopHandler:
291 
292  def __init__(self, wao):
293  self.waowao = wao
294 
295  def start(self):
296  self.waowao.aoLoopClicked(True)
297  self.waowao.uiAO.wao_run.setChecked(True)
298 
299  def stop(self):
300  self.waowao.aoLoopClicked(False)
301  self.waowao.uiAO.wao_run.setChecked(False)
302 
303  def alive(self):
304  return "alive"
305 
306 
307 if __name__ == '__main__':
308  arguments = docopt(__doc__)
309  adopt = arguments["--adopt"]
310  app = QtWidgets.QApplication(sys.argv)
311  app.setStyle('cleanlooks')
312  wao = widgetTwoStagesWindowPyro(arguments["<parameters_filename1>"],
313  arguments["<parameters_filename2>"],
314  arguments["<freqratio>"], cacao=adopt)
315 
316  wao.wao1.show()
317  wao.wao2.show()
318  wao.wao2.uiAO.wao_run.hide()
319  wao.wao2.uiAO.wao_next.hide()
320  wao.wao2.uiAO.wao_atmosphere.hide()
321  wao.wao1.loop_once = wao.loop_once # very dirty (for some reason this does not work during class init...)
322  wao.wao2.loop_once = wao.loop_once # very dirty bis
323 
Class handling both supervisors of first stage and second stage.
None __init__(self, Any config_file1=None, Any config_file2=None, int freqratio=None, bool cacao=False, bool expert=False)
def next(self, nbIters)
Move atmos -> get_slopes -> applyControl ; One integrator step.
Parameter classes for COMPASS.
Initialization and execution of a CANAPASS supervisor.
Imported from CANARY.
Definition: tools.py:1
Widget to simulate a closed loop.
Definition: widget_ao.py:1
Abstract Widget base.
Definition: widget_base.py:1