COMPASS  5.4.4
End-to-end AO simulation tool using GPU acceleration
compassSupervisor.py
1 
37 
38 from shesha.supervisor.genericSupervisor import GenericSupervisor
39 from shesha.supervisor.components import AtmosCompass, DmCompass, RtcCompass, TargetCompass, TelescopeCompass, WfsCompass, CoronagraphCompass
40 from shesha.supervisor.optimizers import ModalBasis, Calibration, ModalGains
41 import numpy as np
42 import time
43 
44 import shesha.constants as scons
45 
46 from typing import Iterable
47 
48 
50  """ This class implements generic supervisor to handle compass simulation
51 
52  Attributes inherited from GenericSupervisor:
53  context : (CarmaContext) : a CarmaContext instance
54 
55  config : (config) : Parameters structure
56 
57  is_init : (bool) : Flag equals to True if the supervisor has already been initialized
58 
59  iter : (int) : Frame counter
60 
61  Attributes:
62  telescope : (TelescopeComponent) : a TelescopeComponent instance
63 
64  atmos : (AtmosComponent) : An AtmosComponent instance
65 
66  target : (targetComponent) : A TargetComponent instance
67 
68  wfs : (WfsComponent) : A WfsComponent instance
69 
70  dms : (DmComponent) : A DmComponent instance
71 
72  rtc : (RtcComponent) : A Rtc component instance
73 
74  cacao : (bool) : CACAO features enabled in the RTC
75 
76  basis : (ModalBasis) : a ModalBasis instance (optimizer)
77 
78  calibration : (Calibration) : a Calibration instance (optimizer)
79 
80  modalgains : (ModalGains) : a ModalGain instance (optimizer) using CLOSE algorithm
81 
82  close_modal_gains : (list of floats) : list of the previous values of the modal gains
83  """
84 
85  def __init__(self, config, *, cacao: bool = False):
86  """ Instantiates a CompassSupervisor object
87 
88  Args:
89  config: (config module) : Configuration module
90 
91  Kwargs:
92  cacao : (bool) : If True, enables CACAO features in RTC (Default is False)
93  Requires OCTOPUS to be installed
94  """
95  self.cacaocacao = cacao
96  self.teltel = None
97  self.atmosatmos = None
98  self.targettarget = None
99  self.wfswfs = None
100  self.dmsdms = None
101  self.rtcrtc = None
102  self.coronocorono = None
103 
104  GenericSupervisor.__init__(self, config)
105  self.basisbasis = ModalBasis(self.configconfig, self.dmsdms, self.targettarget)
106  self.calibrationcalibration = Calibration(self.configconfig, self.teltel, self.atmosatmos, self.dmsdms,
107  self.targettarget, self.rtcrtc, self.wfswfs)
108  if config.p_controllers is not None:
109  self.modalgainsmodalgains = ModalGains(self.configconfig, self.rtcrtc)
110  self.close_modal_gainsclose_modal_gains = []
111 
112 # ___ _ __ __ _ _ _
113 # / __|___ _ _ ___ _ _(_)__ | \/ |___| |_| |_ ___ __| |___
114 # | (_ / -_) ' \/ -_) '_| / _| | |\/| / -_) _| ' \/ _ \/ _` (_-<
115 # \___\___|_||_\___|_| |_\__| |_| |_\___|\__|_||_\___/\__,_/__/
116 
117  def _init_tel(self):
118  """Initialize the telescope component of the supervisor as a TelescopeCompass
119  """
120  self.teltel = TelescopeCompass(self.contextcontext, self.configconfig)
121 
122  def _init_atmos(self):
123  """Initialize the atmosphere component of the supervisor as a AtmosCompass
124  """
125  self.atmosatmos = AtmosCompass(self.contextcontext, self.configconfig)
126 
127  def _init_dms(self):
128  """Initialize the DM component of the supervisor as a DmCompass
129  """
130  self.dmsdms = DmCompass(self.contextcontext, self.configconfig)
131 
132  def _init_target(self):
133  """Initialize the target component of the supervisor as a TargetCompass
134  """
135  if self.teltel is not None:
136  self.targettarget = TargetCompass(self.contextcontext, self.configconfig, self.teltel)
137  else:
138  raise ValueError("Configuration not loaded or Telescope not initilaized")
139 
140  def _init_wfs(self):
141  """Initialize the wfs component of the supervisor as a WfsCompass
142  """
143  if self.teltel is not None:
144  self.wfswfs = WfsCompass(self.contextcontext, self.configconfig, self.teltel)
145  else:
146  raise ValueError("Configuration not loaded or Telescope not initilaized")
147 
148  def _init_rtc(self):
149  """Initialize the rtc component of the supervisor as a RtcCompass
150  """
151  if self.wfswfs is not None:
152  self.rtcrtc = RtcCompass(self.contextcontext, self.configconfig, self.teltel, self.wfswfs,
153  self.dmsdms, self.atmosatmos, cacao=self.cacaocacao)
154  else:
155  raise ValueError("Configuration not loaded or Telescope not initilaized")
156 
157  def _init_components(self) -> None:
158  """ Initialize all the components
159  """
160 
161  if self.configconfig.p_tel is None or self.configconfig.p_geom is None:
162  raise ValueError("Telescope geometry must be defined (p_geom and p_tel)")
163  self._init_tel_init_tel()
164 
165  if self.configconfig.p_atmos is not None:
166  self._init_atmos_init_atmos()
167  if self.configconfig.p_dms is not None:
168  self._init_dms_init_dms()
169  if self.configconfig.p_targets is not None:
170  self._init_target_init_target()
171  if self.configconfig.p_wfss is not None:
172  self._init_wfs_init_wfs()
173  if self.configconfig.p_controllers is not None or self.configconfig.p_centroiders is not None:
174  self._init_rtc_init_rtc()
175  if self.configconfig.p_coronos is not None:
176  self._init_coronagraph_init_coronagraph()
177 
178  GenericSupervisor._init_components(self)
179 
180  def _init_coronagraph(self):
181  """ Initialize the coronagraph
182  """
183  self.coronocorono = CoronagraphCompass()
184  for p_corono in self.configconfig.p_coronos:
185  self.coronocorono.add_corono(self.contextcontext, p_corono, self.configconfig.p_geom, self.targettarget)
186 
187  def next(self, *, move_atmos: bool = True, nControl: int = 0,
188  tar_trace: Iterable[int] = None, wfs_trace: Iterable[int] = None,
189  do_control: bool = True, apply_control: bool = True,
190  compute_tar_psf: bool = True, compute_corono: bool=True) -> None:
191  """Iterates the AO loop, with optional parameters.
192 
193  Overload the GenericSupervisor next() method to handle the GEO controller
194  specific raytrace order operations
195 
196  Kwargs:
197  move_atmos: (bool): move the atmosphere for this iteration. Default is True
198 
199  nControl: (int): Controller number to use. Default is 0 (single control configuration)
200 
201  tar_trace: (List): list of targets to trace. None is equivalent to all (default)
202 
203  wfs_trace: (List): list of WFS to trace. None is equivalent to all (default)
204 
205  do_control : (bool) : Performs RTC operations if True (Default)
206 
207  apply_control: (bool): if True (default), apply control on DMs
208 
209  compute_tar_psf : (bool) : If True (default), computes the PSF at the end of the iteration
210 
211  compute_corono: (bool): If True (default), computes the coronagraphic image
212  """
213  try:
214  iter(nControl)
215  except TypeError:
216  # nControl is not an iterable creating a list
217  nControl = [nControl]
218 
219  #get the index of the first GEO controller (-1 if there is no GEO controller)
220  geo_index = next(( i for i,c in enumerate(self.configconfig.p_controllers)
221  if c.type== scons.ControllerType.GEO ), -1)
222 
223  if tar_trace is None and self.targettarget is not None:
224  tar_trace = range(len(self.configconfig.p_targets))
225  if wfs_trace is None and self.wfswfs is not None:
226  wfs_trace = range(len(self.configconfig.p_wfss))
227 
228  if move_atmos and self.atmosatmos is not None:
229  self.atmosatmos.move_atmos()
230  # in case there is at least 1 controller GEO in the controller list : use this one only
231  self.teltel.update_input_phase()
232  if ( geo_index > -1):
233  nControl = geo_index
234  if tar_trace is not None:
235  for t in tar_trace:
236  if self.atmosatmos.is_enable:
237  self.targettarget.raytrace(t, tel=self.teltel, atm=self.atmosatmos, ncpa=False)
238  else:
239  self.targettarget.raytrace(t, tel=self.teltel, ncpa=False)
240 
241  if do_control and self.rtcrtc is not None:
242  self.rtcrtc.do_control(nControl, sources=self.targettarget.sources)
243  self.targettarget.raytrace(t, dms=self.dmsdms, ncpa=True, reset=False)
244  if apply_control:
245  self.rtcrtc.apply_control(nControl)
246  if self.cacaocacao:
247  self.rtcrtc.publish()
248  else:
249  if tar_trace is not None: # already checked at line 213?
250  for t in tar_trace:
251  if self.atmosatmos.is_enable:
252  self.targettarget.raytrace(t, tel=self.teltel, atm=self.atmosatmos,
253  dms=self.dmsdms)
254  else:
255  self.targettarget.raytrace(t, tel=self.teltel, dms=self.dmsdms)
256 
257  if wfs_trace is not None: # already checked at line 215?
258  for w in wfs_trace:
259  if self.atmosatmos.is_enable:
260  self.wfswfs.raytrace(w, tel=self.teltel, atm=self.atmosatmos)
261  else:
262  self.wfswfs.raytrace(w, tel=self.teltel)
263 
264  if not self.configconfig.p_wfss[w].open_loop and self.dmsdms is not None:
265  self.wfswfs.raytrace(w, dms=self.dmsdms, ncpa=False, reset=False)
266  self.wfswfs.compute_wfs_image(w)
267  if do_control and self.rtcrtc is not None:
268  for ncontrol in nControl : # range(len(self.config.p_controllers)):
269  self.rtcrtc.do_centroids(ncontrol)
270  self.rtcrtc.do_control(ncontrol)
271  self.rtcrtc.do_clipping(ncontrol)
272 
273  if apply_control:
274  for ncontrol in nControl :
275  self.rtcrtc.apply_control(ncontrol)
276 
277  if self.cacaocacao:
278  self.rtcrtc.publish()
279 
280  if compute_tar_psf:
281  for tar_index in tar_trace:
282  self.targettarget.comp_tar_image(tar_index)
283  self.targettarget.comp_strehl(tar_index)
284 
285  if self.coronocorono is not None and compute_corono:
286  for coro_index in range(len(self.configconfig.p_coronos)):
287  self.coronocorono.compute_image(coro_index)
288 
289  if self.configconfig.p_controllers[0].close_opti and (not self.rtcrtc._rtc.d_control[0].open_loop):
290  self.modalgainsmodalgains.update_mgains()
291  self.close_modal_gainsclose_modal_gains.append(self.modalgainsmodalgains.get_modal_gains())
292 
293  self.iteriter += 1
294 
295  def _print_strehl(self, monitoring_freq: int, iters_time: float, total_iters: int, *,
296  tar_index: int = 0):
297  """ Print the Strehl ratio SE and LE from a target on the terminal, the estimated remaining time and framerate
298 
299  Args:
300  monitoring_freq : (int) : Number of frames between two prints
301 
302  iters_time : (float) : time elapsed between two prints
303 
304  total_iters : (int) : Total number of iterations
305 
306  Kwargs:
307  tar_index : (int) : Index of the target. Default is 0
308  """
309  framerate = monitoring_freq / iters_time
310  strehl = self.targettarget.get_strehl(tar_index)
311  etr = (total_iters - self.iteriter) / framerate
312  print("%d \t %.3f \t %.3f\t %.1f \t %.1f" % (self.iteriter + 1, strehl[0],
313  strehl[1], etr, framerate))
314 
315  def loop(self, number_of_iter: int, *, monitoring_freq: int = 100,
316  compute_tar_psf: bool = True, **kwargs):
317  """ Perform the AO loop for <number_of_iter> iterations
318 
319  Args:
320  number_of_iter: (int) : Number of iteration that will be done
321 
322  Kwargs:
323  monitoring_freq: (int) : Monitoring frequency [frames]. Default is 100
324 
325  compute_tar_psf : (bool) : If True (default), computes the PSF at each iteration
326  Else, only computes it each <monitoring_freq> frames
327  """
328  if not compute_tar_psf:
329  print("WARNING: Target PSF will be computed (& accumulated) only during monitoring"
330  )
331 
332  print("----------------------------------------------------")
333  print("iter# | S.E. SR | L.E. SR | ETR (s) | Framerate (Hz)")
334  print("----------------------------------------------------")
335  # self.next(**kwargs)
336  t0 = time.time()
337  t1 = time.time()
338  if number_of_iter == -1: # Infinite loop
339  while (True):
340  self.nextnext(compute_tar_psf=compute_tar_psf, **kwargs)
341  if ((self.iteriter + 1) % monitoring_freq == 0):
342  if not compute_tar_psf:
343  self.targettarget.comp_tar_image(0)
344  self.targettarget.comp_strehl(0)
345  self._print_strehl_print_strehl(monitoring_freq, time.time() - t1, number_of_iter)
346  t1 = time.time()
347 
348  for _ in range(number_of_iter):
349  self.nextnext(compute_tar_psf=compute_tar_psf, **kwargs)
350  if ((self.iteriter + 1) % monitoring_freq == 0):
351  if not compute_tar_psf:
352  self.targettarget.comp_tar_image(0)
353  self.targettarget.comp_strehl(0)
354  self._print_strehl_print_strehl(monitoring_freq, time.time() - t1, number_of_iter)
355  t1 = time.time()
356  t1 = time.time()
357  print(" loop execution time:", t1 - t0, " (", number_of_iter, "iterations), ",
358  (t1 - t0) / number_of_iter, "(mean) ", number_of_iter / (t1 - t0), "Hz")
359 
360  def reset(self):
361  """ Reset the simulation to return to its original state
362  """
363  self.atmosatmos.reset_turbu()
364  self.wfswfs.reset_noise()
365  for tar_index in range(len(self.configconfig.p_targets)):
366  self.targettarget.reset_strehl(tar_index)
367  self.dmsdms.reset_dm()
368  self.rtcrtc.open_loop()
369  self.rtcrtc.close_loop()
370 
371 
372 # ___ _ __ _ __ __ _ _ _
373 # / __|_ __ ___ __(_)/ _(_)__ | \/ |___| |_| |_ ___ __| |___
374 # \__ \ '_ \/ -_) _| | _| / _| | |\/| / -_) _| ' \/ _ \/ _` (_-<
375 # |___/ .__/\___\__|_|_| |_\__| |_| |_\___|\__|_||_\___/\__,_/__/
376 # |_|
377 
379  self, cb_count: int, sub_sample: int = 1, controller_index: int = 0,
380  tar_index: int = 0, see_atmos: bool = True, cube_data_type: str = None,
381  cube_data_file_path: str = "", ncpa: int = 0, ncpa_wfs: np.ndarray = None,
382  ref_slopes: np.ndarray = None, ditch_strehl: bool = True,
383  projection_matrix: np.ndarray = None):
384  """ Used to record a synchronized circular buffer AO loop data.
385 
386  Args:
387  cb_count: (int) : the number of iterations to record.
388 
389  sub_sample: (int) : sub sampling of the data (default=1, I.e no subsampling)
390 
391  controller_index: (int) :
392 
393  tar_index: (int) : target number
394 
395  see_atmos: (int) : used for the next function to enable or not the Atmos
396 
397  cube_data_type: (int) : if specified ("tarPhase" or "psfse") returns the target phase or short exposure PSF data cube in the output variable
398 
399  cube_data_file_path: (int) : if specified it will also save the target phase cube data (full path on the server)
400 
401  ncpa: (int) : !!experimental!!!: Used only in the context of PYRWFS + NCPA compensation on the fly (with optical gain)
402  defines how many iters the NCPA refslopes are updates with the proper optical gain. Ex: if NCPA=10 refslopes will be updates every 10 iters.
403 
404  ncpa_wfs: (int) : the ncpa phase as seen from the wfs array with dims = size of Mpupil
405 
406  ref_slopes: (int) : the reference slopes to use.
407 
408  ditch_strehl: (int) : resets the long exposure SR computation at the beginning of the Circular buffer (default= True)
409 
410  projection_matrix : (np.ndarray) : projection matrix on modal basis to compute residual coefficients
411 
412  Returns:
413  slopes: (int) : the slopes CB
414 
415  volts: (int) : the volts applied to the DM(s) CB
416 
417  ai: (int) : the modal coefficient of the residual phase projected on the currently used modal Basis
418 
419  psf_le: (int) : Long exposure PSF over the <cb_count> iterations (I.e SR is reset at the begining of the CB if ditch_strehl=True)
420 
421  strehl_se_list: (int) : The SR short exposure evolution during CB recording
422 
423  strehl_le_list: (int) : The SR long exposure evolution during CB recording
424 
425  g_ncpa_list: (int) : the gain applied to the NCPA (PYRWFS CASE) if NCPA is set to True
426 
427  cube_data: (int) : the tarPhase or psfse cube data (see cube_data_type)
428  """
429  slopes_data = None
430  volts_data = None
431  cube_data = None
432  ai_data = None
433  k = 0
434  sthrel_se_list = []
435  sthrel_le_list = []
436  g_ncpa_list = []
437 
438  # Resets the target so that the PSF LE is synchro with the data
439  # Doesn't reset it if Ditch_strehl == False (used for real time gain computation)
440  if ditch_strehl:
441  for i in range(len(self.configconfig.p_targets)):
442  self.targettarget.reset_strehl(i)
443 
444  # Starting CB loop...
445  for j in range(cb_count):
446  print(j, end="\r")
447  if (ncpa):
448  if (j % ncpa == 0):
449  ncpa_diff = ref_slopes[None, :]
450  ncpa_turbu = self.calibrationcalibration.do_imat_phase(
451  controller_index, -ncpa_wfs[None, :, :], noise=False,
452  with_turbu=True)
453  g_ncpa = float(
454  np.sqrt(
455  np.dot(ncpa_diff, ncpa_diff.T) / np.dot(
456  ncpa_turbu, ncpa_turbu.T)))
457  if (g_ncpa > 1e18):
458  g_ncpa = 0
459  print('Warning NCPA ref slopes gain too high!')
460  g_ncpa_list.append(g_ncpa)
461  self.rtcrtc.set_ref_slopes(-ref_slopes * g_ncpa)
462  else:
463  g_ncpa_list.append(g_ncpa)
464  print('NCPA ref slopes gain: %4.3f' % g_ncpa)
465  self.rtcrtc.set_ref_slopes(-ref_slopes / g_ncpa)
466 
467  self.atmosatmos.enable_atmos(see_atmos)
468  self.nextnext()
469  for t in range(len(self.configconfig.p_targets)):
470  self.targettarget.comp_tar_image(t)
471 
472  srse, srle, _, _ = self.targettarget.get_strehl(tar_index)
473  sthrel_se_list.append(srse)
474  sthrel_le_list.append(srle)
475  if (j % sub_sample == 0):
476  if (projection_matrix is not None):
477  ai_vector = self.calibrationcalibration.compute_modal_residuals(
478  projection_matrix, selected_actus=self.basisbasis.selected_actus)
479  if (ai_data is None):
480  ai_data = np.zeros((len(ai_vector), int(cb_count / sub_sample)))
481  ai_data[:, k] = ai_vector
482 
483  slopes_vector = self.rtcrtc.get_slopes(controller_index)
484  if (slopes_data is None):
485  slopes_data = np.zeros((len(slopes_vector),
486  int(cb_count / sub_sample)))
487  slopes_data[:, k] = slopes_vector
488 
489  volts_vector = self.rtcrtc.get_command(
490  controller_index) # get_command or get_voltages ?
491  if (volts_data is None):
492  volts_data = np.zeros((len(volts_vector),
493  int(cb_count / sub_sample)))
494  volts_data[:, k] = volts_vector
495 
496  if (cube_data_type):
497  if (cube_data_type == "tarPhase"):
498  dataArray = self.targettarget.get_tar_phase(tar_index, pupil=True)
499  elif (cube_data_type == "psfse"):
500  dataArray = self.targettarget.get_tar_image(tar_index, expo_type="se")
501  else:
502  raise ValueError("unknown dataData" % cube_data_type)
503  if (cube_data is None):
504  cube_data = np.zeros((*dataArray.shape,
505  int(cb_count / sub_sample)))
506  cube_data[:, :, k] = dataArray
507  k += 1
508  if (cube_data_file_path != ""):
509  print("Saving tarPhase cube at: ", cube_data_file_path)
510  from astropy.io import fits as pf
511  pf.writeto(cube_data_file_path, cube_data, overwrite=True)
512 
513  psf_le = self.targettarget.get_tar_image(tar_index, expo_type="le")
514  return slopes_data, volts_data, ai_data, psf_le, sthrel_se_list, sthrel_le_list, g_ncpa_list, cube_data
515 
516  def export_config(self):
517  """
518  Extract and convert compass supervisor configuration parameters
519  into 2 dictionnaries containing relevant AO parameters
520 
521  Args:
522  root: (object), COMPASS supervisor object to be parsed
523 
524  Returns :
525  2 dictionaries... See F. Vidal :)
526  """
527  return self.configconfig.export_config()
528 
529  def get_s_pupil(self):
530  """
531  Returns the so called S Pupil of COMPASS
532 
533  Return:
534  s_pupil: (np.array) : S Pupil of COMPASS
535  """
536  return self.configconfig.p_geom.get_spupil()
537 
538  def get_i_pupil(self):
539  """
540  Returns the so called I Pupil of COMPASS
541 
542  Return:
543  i_pupil: (np.array) : I Pupil of COMPASS
544  """
545  return self.configconfig.p_geom.get_ipupil()
546 
547  def get_m_pupil(self):
548  """
549  Returns the so called M Pupil of COMPASS
550 
551  Return:
552  m_pupil: (np.array) : M Pupil of COMPASS
553  """
554  return self.configconfig.p_geom.get_mpupil()
This class implements generic supervisor to handle compass simulation.
def __init__(self, config, *bool cacao=False)
Instantiates a CompassSupervisor object.
atmos
(AtmosComponent) : An AtmosComponent instance
def record_ao_circular_buffer(self, int cb_count, int sub_sample=1, int controller_index=0, int tar_index=0, bool see_atmos=True, str cube_data_type=None, str cube_data_file_path="", int ncpa=0, np.ndarray ncpa_wfs=None, np.ndarray ref_slopes=None, bool ditch_strehl=True, np.ndarray projection_matrix=None)
Used to record a synchronized circular buffer AO loop data.
None next(self, *bool move_atmos=True, int nControl=0, Iterable[int] tar_trace=None, Iterable[int] wfs_trace=None, bool do_control=True, bool apply_control=True, bool compute_tar_psf=True, bool compute_corono=True)
Iterates the AO loop, with optional parameters.
def _print_strehl(self, int monitoring_freq, float iters_time, int total_iters, *int tar_index=0)
def get_s_pupil(self)
Returns the so called S Pupil of COMPASS.
basis
(ModalBasis) : a ModalBasis instance (optimizer)
rtc
(RtcComponent) : A Rtc component instance
target
(targetComponent) : A TargetComponent instance
cacao
(bool) : CACAO features enabled in the RTC
def get_m_pupil(self)
Returns the so called M Pupil of COMPASS.
def loop(self, int number_of_iter, *int monitoring_freq=100, bool compute_tar_psf=True, **kwargs)
Perform the AO loop for <number_of_iter> iterations.
calibration
(Calibration) : a Calibration instance (optimizer)
def get_i_pupil(self)
Returns the so called I Pupil of COMPASS.
modalgains
(ModalGains) : a ModalGain instance (optimizer) using CLOSE algorithm
def reset(self)
Reset the simulation to return to its original state.
close_modal_gains
(list of floats) : list of the previous values of the modal gains
wfs
(WfsComponent) : A WfsComponent instance
def export_config(self)
Extract and convert compass supervisor configuration parameters into 2 dictionnaries containing relev...
This class defines generic methods and behavior of a supervisor It is not intended to be instantiated...
context
(CarmaContext) : a CarmaContext instance
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
User layer for optimizing AO supervisor loop.