COMPASS  5.4.4
End-to-end AO simulation tool using GPU acceleration
twoStagesManager.py
1 
37 """
38 Initialization and execution of an AO 2-stages manager.
39 It instanciates in one process two compass simulations:
40 1 for first stage and 1 for the second stage (as defined in their relative .par files)
41 
42 IMPORTANT:
43 The next method of this manager --superseeds-- the compass next method so that the loop is fully handled by the manager.
44 
45 Usage:
46  twoStagesManager.py <parameters_filename1> <parameters_filename2> <freqratio> [options]
47 
48 with 'parameters_filename1' the path to the parameters file the first stage
49 with 'parameters_filename2' the path to the parameters file the second stage
50 with 'freqratio' the ratio of the frequencies of the two stages
51 Options:
52  -a, --adopt used to connect optional ADOPT client to the manager (via pyro + shm cacao)
53 
54 Example:
55  ipython -i twoStagesManager.py ../../data/par/SPHERE+/sphere.py ../../data/par/SPHERE+/sphere+.py 3
56  ipython -i twoStagesManager.py ../../data/par/SPHERE+/sphere.py ../../data/par/SPHERE+/sphere+.py 3 -- --adopt
57 """
58 
59 import numpy as np
60 import time
61 from typing import Any, Dict, Tuple, Callable, List
62 from shesha.supervisor.stageSupervisor import StageSupervisor
63 
64 class TwoStagesManager(object):
65  """
66  Class handling both supervisors of first stage and second stage.
67 
68  Attributes:
69  first_stage : (StageSupervisor) : first stage StageSupervisor instance
70 
71  second_stage : (StageSupervisor) : second stage StageSupervisor instance
72 
73  iterations : (int) : frame counter
74 
75  second_stage_input : (array) : input phase screen for second stage
76 
77  mpup_offset : (int) : number of padding pixel from first stage s_pupil to second stage m_pupil
78 
79  frequency_ratio : (int) : second stage simulated frequency over first stage simulated frequency
80  """
81  def __init__(self, first_stage : StageSupervisor, second_stage : StageSupervisor, frequency_ratio : int):
82  """
83  Init of the TwoStagesManager object
84 
85  Args:
86  first_stage : (StageSupervisor) : first stage StageSupervisor instance
87 
88  second_stage : (StageSupervisor) : second stage StageSupervisor instance
89 
90  frequency_ratio : (int) : ratio between second stage frequency and first stage frequency. Only integers are accepted.
91  """
92 
93  self.first_stagefirst_stage = first_stage
94  self.second_stagesecond_stage = second_stage
95  self.second_stagesecond_stage.atmos.enable_atmos(False) # second stage atmos is not used
96  self.iterationsiterations = 0
97  mpup_shape = self.second_stagesecond_stage.config.p_geom._mpupil.shape
98  self.second_stage_inputsecond_stage_input = np.zeros((mpup_shape[0], mpup_shape[1], 1))
99  residual_shape = self.first_stagefirst_stage.config.p_geom._spupil.shape
100  self.mpup_offsetmpup_offset = (mpup_shape[0] - residual_shape[0]) // 2
101  self.frequency_ratiofrequency_ratio = int(frequency_ratio)
102 
103  # flags for enabling coronagraphic images computation
104  self.compute_first_stage_coronocompute_first_stage_corono = True
105  self.compute_second_stage_coronocompute_second_stage_corono = True
106 
107  def next(self, *, do_control: bool = True) -> None:
108  """
109  MAIN method that allows to manage properly the 2 AO stages of SAXO+ system.
110  The phase residuals (including turbulence + AO loop residuals) of the first stage simulation is sent to second stage simulation
111  at each iteration of the manager.
112  The manager disable the seconds stage turbulence simulation (as it is propageated through the first stage residals if any).
113  This next method sould ALWAYS be called to perform a regular SAXO+ simulation
114  instead of the individual stage next methods to ensure the correct synchronisation of the 2 systems.
115  """
116  # Iteration time of the first stage is set as the same as the second stage to
117  # allow correct atmosphere movement for second stage integration. Then,
118  # first stage is controlled only once every frequency_ratio times
119 
120  # Turbulence always disabled on 2nd instance of COMPASS
121  self.second_stagesecond_stage.atmos.enable_atmos(False)
122 
123  if do_control:
124  # compute flags to specify which action need to be done in this first stage:
125  # 1. check if go on stacking WFS image
126  first_stage_stack_wfs = bool(self.iterationsiterations % self.frequency_ratiofrequency_ratio)
127  # 2. check if centroids need to be computed (end of WFS exposure)
128  first_stage_centroids = not(bool((self.iterationsiterations + 1) % self.frequency_ratiofrequency_ratio))
129  # 3. Check if a new command is computed (when new centroids appear)
130  first_stage_control = first_stage_centroids
131  self.first_stagefirst_stage.next(do_control = first_stage_control,
132  apply_control = True,
133  do_centroids = first_stage_centroids,
134  compute_tar_psf = True,
135  stack_wfs_image = first_stage_stack_wfs)
136  else:
137  self.first_stagefirst_stage.next(do_control=False,
138  do_centroids=True,
139  apply_control=True,
140  compute_tar_psf = True)
141 
142  # FIRST STAGE IS DONE.
143 
144  # Get residual of first stage to put it into second stage
145  # For now, involves GPU-CPU memory copies, can be improved later if speed is
146  # a limiting factor here...
147  first_stage_residual = self.first_stagefirst_stage.target.get_tar_phase(0)
148  self.second_stage_inputsecond_stage_input[self.mpup_offsetmpup_offset:-self.mpup_offsetmpup_offset,
149  self.mpup_offsetmpup_offset:-self.mpup_offsetmpup_offset,:] = first_stage_residual[:,:,None]
150  self.second_stagesecond_stage.tel.set_input_phase(self.second_stage_inputsecond_stage_input) # 1st stage residuals sent to seconds stage simulation.
151 
152  # SECOND STAGE LOOP STARTS...
153 
154  #"Updates the second stage simulation accordingly".
155  # WFS exposure is always reset (default).
156  self.second_stagesecond_stage.next(move_atmos=False, do_control=do_control)
157  # SECOND STAGE IS DONE.
158  self.iterationsiterations += 1
159 
160  def enable_corono(self, stage=None):
161  """ Enable coronagraphic image computation for both stages.
162 
163  Args:
164  stage: (str, optional): If 'first', enable only first stage coronagrapic image computation.
165  If 'second', enable only second stage coronagraphic image computation.
166  Default = None.
167  """
168  if stage == 'first':
169  self.compute_first_stage_coronocompute_first_stage_corono = True
170  elif stage == 'second':
171  self.compute_second_stage_coronocompute_second_stage_corono = True
172  else:
173  self.compute_first_stage_coronocompute_first_stage_corono = True
174  self.compute_second_stage_coronocompute_second_stage_corono = True
175 
176  def disable_corono(self):
177  """ Disable all coronagraphic image computation
178  """
179  self.compute_first_stage_coronocompute_first_stage_corono = False
180  self.compute_second_stage_coronocompute_second_stage_corono = False
181 
182  def reset_exposure(self):
183  """ Reset long exposure psf and coronagraphic images for both stages
184  """
185  self.first_stagefirst_stage.corono.reset()
186  self.second_stagesecond_stage.corono.reset()
187  self.first_stagefirst_stage.target.reset_strehl(0)
188  self.second_stagesecond_stage.target.reset_strehl(0)
189 
190  def get_frame_counter(self):
191  """ Returns the current iteration number of the manager
192 
193  Returns:
194  iterations : (int) : Number of manager iterations already performed
195  """
196  return self.iterationsiterations
197 
198  def loop(self, number_of_iter: int, *, monitoring_freq: int = 100, **kwargs):
199  """ Perform the AO loop for <number_of_iter> iterations
200 
201  Args:
202  number_of_iter: (int) : Number of iteration that will be done
203 
204  Kwargs:
205  monitoring_freq: (int) : Monitoring frequency [frames]. Default is 100
206  """
207 
208  print("----------------------------------------------------")
209  print("iter# | S.E. SR | L.E. SR | ETR (s) | Framerate (Hz)")
210  print("----------------------------------------------------")
211  # self.next(**kwargs)
212  t0 = time.time()
213  t1 = time.time()
214  if number_of_iter == -1: # Infinite loop
215  while (True):
216  self.nextnext()
217  if ((self.iterationsiterations + 1) % monitoring_freq == 0):
218  self.second_stagesecond_stage._print_strehl(monitoring_freq, time.time() - t1, number_of_iter)
219  t1 = time.time()
220 
221  for _ in range(number_of_iter):
222  self.nextnext()
223  if ((self.iterationsiterations + 1) % monitoring_freq == 0):
224  self.second_stagesecond_stage._print_strehl(monitoring_freq, time.time() - t1, number_of_iter)
225  t1 = time.time()
226  t1 = time.time()
227  print(" loop execution time:", t1 - t0, " (", number_of_iter, "iterations), ",
228  (t1 - t0) / number_of_iter, "(mean) ", number_of_iter / (t1 - t0), "Hz")
229 
230 class loopHandler:
231 
232  def __init__(self):
233  pass
234 
235  def start(self):
236  pass
237 
238  def stop(self):
239  pass
240 
241  def alive(self):
242  return "alive"
243 
244 if __name__ == '__main__':
245  from docopt import docopt
246  from shesha.config import ParamConfig
247  arguments = docopt(__doc__)
248  adopt = arguments["--adopt"]
249 
250  config1 = ParamConfig(arguments["<parameters_filename1>"])
251  config2 = ParamConfig(arguments["<parameters_filename2>"])
252  frequency_ratio = arguments["<freqratio>"]
253 
254  first_stage = StageSupervisor(config1, cacao=adopt)
255  second_stage = StageSupervisor(config2, cacao=adopt)
256  manager = TwoStagesManager(first_stage, second_stage, frequency_ratio)
257 
258  if(adopt):
259 
260  supervisor1 = manager.first_stage
261  supervisor2 = manager.second_stage
262 
263 
264  try:
265  from subprocess import Popen, PIPE
266  from hraa.server.pyroServer import PyroServer
267  import Pyro4
268  Pyro4.config.REQUIRE_EXPOSE = False
269  p = Popen("whoami", shell=True, stdout=PIPE, stderr=PIPE)
270  out, err = p.communicate()
271  if (err != b''):
272  print(err)
273  raise Exception("ERROR CANNOT RECOGNIZE USER")
274  else:
275  user = out.split(b"\n")[0].decode("utf-8")
276  print("User is " + user)
277 
278 
279  if (supervisor1.corono == None):
280  from shesha.util.pyroEmptyClass import PyroEmptyClass
281  coro2pyro1 = PyroEmptyClass()
282  else:
283  coro2pyro1 = supervisor1.corono
284 
285  if (supervisor2.corono == None):
286  from shesha.util.pyroEmptyClass import PyroEmptyClass
287  coro2pyro2 = PyroEmptyClass()
288  else:
289  coro2pyro2 = supervisor2.corono
290 
291 
292  devices1 = [
293  supervisor1, supervisor1.rtc, supervisor1.wfs, supervisor1.target,
294  supervisor1.tel, supervisor1.basis, supervisor1.calibration,
295  supervisor1.atmos, supervisor1.dms, supervisor1.config, supervisor1.modalgains,
296  coro2pyro1
297  ]
298  devices2 = [
299  supervisor2, supervisor2.rtc, supervisor2.wfs, supervisor2.target,
300  supervisor2.tel, supervisor2.basis, supervisor2.calibration,
301  supervisor2.atmos, supervisor2.dms, supervisor2.config, supervisor2.modalgains,
302  coro2pyro2
303  ]
304  names = [
305  "supervisor", "supervisor_rtc", "supervisor_wfs", "supervisor_target",
306  "supervisor_tel", "supervisor_basis", "supervisor_calibration",
307  "supervisor_atmos", "supervisor_dms", "supervisor_config", "supervisor_modalgains",
308  "supervisor_corono"
309  ]
310 
311  label = "firstStage"
312  nname = []
313  for name in names:
314  nname.append(name + "_" + user + "_" +label)
315 
316  label = "secondStage"
317  for name in names:
318  nname.append(name + "_" + user + "_" +label)
319 
320  nname.append('twoStagesManager'+ "_" + user ) # Adding master next 2-stages loop
321  devices = devices1 + devices2 + [manager]
322  server = PyroServer(listDevices=devices, listNames=nname)
323  #server.add_device(supervisor, "waoconfig_" + user)
324  server.start()
325  except:
326  raise EnvironmentError(
327  "Missing dependencies (code HRAA or Pyro4 or Dill Serializer)")
328 
Shesha parameters configuration class.
Definition: pconfig.py:51
This class implements a single stage (e.g.
Class handling both supervisors of first stage and second stage.
None next(self, *bool do_control=True)
MAIN method that allows to manage properly the 2 AO stages of SAXO+ system.
def __init__(self, StageSupervisor first_stage, StageSupervisor second_stage, int frequency_ratio)
def loop(self, int number_of_iter, *int monitoring_freq=100, **kwargs)
Perform the AO loop for <number_of_iter> iterations.
def disable_corono(self)
Disable all coronagraphic image computation.
second_stage
(StageSupervisor) : second stage StageSupervisor instance
frequency_ratio
(int) : second stage simulated frequency over first stage simulated frequency
second_stage_input
(array) : input phase screen for second stage
def enable_corono(self, stage=None)
Enable coronagraphic image computation for both stages.
def get_frame_counter(self)
Returns the current iteration number of the manager.
mpup_offset
(int) : number of padding pixel from first stage s_pupil to second stage m_pupil
def reset_exposure(self)
Reset long exposure psf and coronagraphic images for both stages.
first_stage
(StageSupervisor) : first stage StageSupervisor instance
This file is here only to avoid pyro server crash when a coronograph class is not in the pramam file.
Parameter classes for COMPASS.
Initialization and execution of a single stage supervisor for cascaded AO systems.