COMPASS  5.4.4
End-to-end AO simulation tool using GPU acceleration
benchSupervisor.py
1 
37 
38 import numpy as np
39 
40 from shesha.constants import CentroiderType, WFSType
41 from shesha.init.dm_init import dm_init_standalone
42 from shesha.init.rtc_init import rtc_standalone
43 from shesha.sutra_wrap import carmaWrap_context
44 
45 from shesha.supervisor.components import RtcStandalone
46 
47 from shesha.supervisor.genericSupervisor import GenericSupervisor
48 
49 from shesha.util.utilities import load_config_from_file
50 
51 from typing import Callable
52 
53 
55  """ This class implements generic supervisor to handle compass simulation
56 
57  Attributes inherited from GenericSupervisor:
58  context : (CarmaContext) : a CarmaContext instance
59 
60  config : (config) : Parameters structure
61 
62  is_init : (bool) : Flag equals to True if the supervisor has already been initialized
63 
64  iter : (int) : Frame counter
65 
66  Attributes:
67  rtc : (RtcComponent) : A Rtc component instance
68 
69  cacao : (bool) : CACAO features enabled in the RTC
70 
71  basis : (ModalBasis) : a ModalBasis instance (optimizer)
72 
73  calibration : (Calibration) : a Calibration instance (optimizer)
74  """
75 
76  def __init__(self, config_file: str = None, cacao: bool = False):
77  """ Init the COMPASS wih the config_file
78 
79  Kwargs:
80  config_file : (str) : path to the configuration file. Default is None
81 
82  cacao : (bool) : Flag to use cacao rtc. Default is False
83  """
84  self.pause_looppause_loop = None
85  self.rtcrtc = None
86  self.frameframe = None
87  self.cacaocacao = cacao
88  self.iteriteriter = 0
89  self.slopes_indexslopes_index = None
90  config = load_config_from_file(config_file)
91 
92  GenericSupervisor.__init__(self, config)
93 
94  def _init_rtc(self):
95  """Initialize the rtc component of the supervisor as a RtcCompass
96  """
97  print("->RTC")
98  self.number_of_wfsnumber_of_wfs = len(self.configconfigconfig.p_wfss)
99  print("Configuration of", self.number_of_wfsnumber_of_wfs, "wfs ...")
100 
101  if (hasattr(self.configconfigconfig, 'p_loop') and self.configconfigconfig.p_loop.devices.size > 1):
102  self.contextcontextcontext = carmaWrap_context.get_instance_ngpu(
103  self.configconfigconfig.p_loop.devices.size, self.configconfigconfig.p_loop.devices)
104  else:
105  self.contextcontextcontext = carmaWrap_context.get_instance_1gpu(
106  self.configconfigconfig.p_loop.devices[0])
107  nact = self.configconfigconfig.p_controllers[0].nactu
108 
109  nvalid = []
110  centroider_type = []
111  delay = []
112  offset = []
113  scale = []
114  gain = []
115  cmat_size = []
116  npix = []
117 
118  # Get parameters
119  for wfs in range(self.number_of_wfsnumber_of_wfs):
120 
121  if self.configconfigconfig.p_wfss[wfs].type == WFSType.SH:
122  npix.append(self.configconfigconfig.p_wfss[wfs].npix)
123  if self.configconfigconfig.p_wfss[wfs]._validsubsx is None or \
124  self.configconfigconfig.p_wfss[wfs]._validsubsy is None:
125 
126  from hraa.tools.doit import makessp
127  roiTab = makessp(self.configconfigconfig.p_wfss[wfs].nxsub, obs=0., rmax=0.98)
128  self.configconfigconfig.p_wfss[wfs]._nvalid = roiTab[0].size
129  self.configconfigconfig.p_wfss[
130  wfs]._validsubsx = roiTab[0] * self.configconfigconfig.p_wfss[wfs].npix
131  self.configconfigconfig.p_wfss[
132  wfs]._validsubsy = roiTab[1] * self.configconfigconfig.p_wfss[wfs].npix
133  else:
134  self.configconfigconfig.p_wfss[wfs]._nvalid = self.configconfigconfig.p_wfss[
135  wfs]._validsubsx.size
136 
137  nvalid.append(
138  np.array([self.configconfigconfig.p_wfss[wfs]._nvalid], dtype=np.int32))
139  # print("nvalid : %d" % nvalid[wfs])
140  centroider_type.append(self.configconfigconfig.p_centroiders[wfs].type)
141  delay.append(self.configconfigconfig.p_controllers[0].delay) # ???
142  offset.append((self.configconfigconfig.p_wfss[wfs].npix - 1) / 2)
143  scale.append(1)
144  gain.append(1)
145  cmat_size.append(2 * nvalid[wfs][0])
146 
147  elif self.configconfigconfig.p_wfss[wfs].type == WFSType.PYRHR or self.configconfigconfig.p_wfss[
148  wfs].type == WFSType.PYRLR:
149  nvalid.append(
150  np.array([self.configconfigconfig.p_wfss[wfs]._nvalid],
151  dtype=np.int32)) # Number of valid SUBAPERTURES
152  centroider_type.append(self.configconfigconfig.p_centroiders[wfs].type)
153  delay.append(self.configconfigconfig.p_controllers[0].delay) # ???
154  offset.append(0)
155  scale.append(1)
156  gain.append(1)
157  cmat_size.append(self.configconfigconfig.p_wfss[wfs].nPupils * nvalid[wfs][0])
158  npix.append(0)
159  else:
160  raise ValueError('WFS type not supported')
161 
162  # Create RTC
163  self.rtcrtc = RtcStandalone(self.contextcontextcontext, self.configconfigconfig, self.number_of_wfsnumber_of_wfs, nvalid,
164  nact, centroider_type, delay, offset, scale,
165  cacao=self.cacaocacao)
166 
167  self.slopes_indexslopes_index = np.cumsum([0] +
168  [wfs.nslopes for wfs in self.rtcrtc._rtc.d_centro])
169 
170  # Create centroiders
171  for wfs in range(self.number_of_wfsnumber_of_wfs):
172  self.rtcrtc._rtc.d_centro[wfs].load_validpos(
173  self.configconfigconfig.p_wfss[wfs]._validsubsx,
174  self.configconfigconfig.p_wfss[wfs]._validsubsy,
175  self.configconfigconfig.p_wfss[wfs]._validsubsx.size)
176  if self.configconfigconfig.p_centroiders[wfs].type is CentroiderType.BPCOG:
177  self.rtcrtc._rtc.d_centro[wfs].set_nmax(self.configconfigconfig.p_centroiders[wfs].nmax)
178  self.rtcrtc._rtc.d_centro[wfs].set_npix(npix[wfs])
179  # finally
180  self.configconfigconfig.p_centroiders[wfs]._nslope = self.rtcrtc._rtc.d_centro[wfs].nslopes
181  print("wfs ", wfs, " set as ", centroider_type[wfs])
182  size = sum(cmat_size)
183  cMat = np.zeros((nact, size), dtype=np.float32)
184  print("Size of cMat:", cMat.shape)
185 
186  # Initiate RTC
187  self.rtcrtc._rtc.d_control[0].set_cmat(cMat)
188  self.rtcrtc._rtc.d_control[0].set_decayFactor(
189  np.ones(nact, dtype=np.float32) * (gain[0] - 1))
190  self.rtcrtc._rtc.d_control[0].set_matE(np.identity(nact, dtype=np.float32))
191  self.rtcrtc._rtc.d_control[0].set_modal_gains(
192  np.ones(nact, dtype=np.float32) * -gain[0])
193 
194  print("RTC initialized")
195  self.is_initis_initis_init = True
196 
197  def _init_components(self) -> None:
198  """ Initialize all the components
199  """
200  if self.configconfigconfig.p_controllers is not None or self.configconfigconfig.p_centroiders is not None:
201  self._init_rtc_init_rtc()
202 
203  GenericSupervisor._init_components(self)
204 
205  # _ _ _ _
206  # / \ | |__ ___| |_ _ __ __ _ ___| |_
207  # / _ \ | '_ \/ __| __| '__/ _` |/ __| __|
208  # / ___ \| |_) \__ \ |_| | | (_| | (__| |_
209  # /_/ \_\_.__/|___/\__|_| \__,_|\___|\__|
210  #
211  # __ __ _ _ _
212  # | \/ | ___| |_| |__ ___ __| |___
213  # | |\/| |/ _ \ __| '_ \ / _ \ / _` / __|
214  # | | | | __/ |_| | | | (_) | (_| \__ \
215  # |_| |_|\___|\__|_| |_|\___/ \__,_|___/
216 
217  def next(self) -> None:
218  """ Performs a single loop iteration
219  """
220  self.load_new_wfs_frameload_new_wfs_frame()
221 
222  if (self.pause_looppause_loop is not True):
223  self.compute_wfs_framecompute_wfs_frame()
224  self.set_commandset_command(0, np.array(self.rtcrtc._rtc.d_control[0].d_voltage))
225  if self.cacaocacao:
226  self.rtcrtc.publish()
227  self.iteriteriter += 1
228 
229  def get_tar_image(self, tar_index, expo_type: str = "se") -> np.ndarray:
230  """ NOT IMPLEMENTED
231  """
232  raise NotImplementedError("Not implemented")
233 
234  def set_command(self, nctrl: int, command: np.ndarray) -> None:
235  """ Immediately sets provided command to DMs - does not affect integrator
236 
237  Args:
238  nctrl : (int) : Controller index (unused)
239 
240  command : (np.ndarray) : Command vector to send
241  """
242  # Do stuff
243  self.dm_set_callbackdm_set_callback(command)
244  # Btw, update the RTC state with the information
245  # self.rtc._rtc.d_control[nctrl].set_com(command, command.size)
246 
247  def get_command(self) -> np.ndarray:
248  """ Get command from DM
249 
250  Returns:
251  command : (np.ndarray) : Command vector
252  """
253  # Do something
254  command = self.dm_get_callbackdm_get_callback()
255  # Btw, update the RTC state with the information
256  # self.rtc._rtc.d_control[nControl].set_com(command, command.size)
257 
258  return command
259 
260  # ____ _ _ _ __ __ _ _ _
261  # / ___| _ __ ___ ___(_) |_(_) ___ | \/ | ___| |_| |__ ___ __| |___
262  # \___ \| '_ \ / _ \/ __| | __| |/ __| | |\/| |/ _ \ __| '_ \ / _ \ / _` / __|
263  # ___) | |_) | __/ (__| | |_| | (__ | | | | __/ |_| | | | (_) | (_| \__ \
264  # |____/| .__/ \___|\___|_|\__|_|\___| |_| |_|\___|\__|_| |_|\___/ \__,_|___/
265  # |_|
266 
267  def __repr__(self):
268 
269  s = '--- BenchSupervisor ---\nRTC: ' + repr(self.rtcrtc)
270  if hasattr(self, '_cam'):
271  s += '\nCAM: ' + repr(cam)
272  if hasattr(self, '_dm'):
273  s += '\nDM: ' + repr(dm)
274  return s
275 
276  def load_new_wfs_frame(self, centro_index: int = 0) -> None:
277  """ Acquire a new WFS frame and load it
278 
279  Args:
280  centro_index : (int) : Index of the centroider where to load the frame
281  """
282  self.frameframe = self.cam_callbackcam_callback()
283  if (type(self.frameframe) is tuple):
284  centro_index = len(self.frameframe)
285  for i in range(centro_index):
286  self.rtcrtc._rtc.d_centro[i].load_img(self.frameframe[i], self.frameframe[i].shape[0],
287  self.frameframe[i].shape[1], -1)
288  else:
289  self.rtcrtc._rtc.d_centro[centro_index].load_img(self.frameframe,
290  self.frameframe.shape[0],
291  self.frameframe.shape[1], -1)
292 
293  def compute_wfs_frame(self):
294  """ Compute the WFS frame: calibrate, centroid, commands.
295  """
296  # for index, centro in enumerate(self.rtc._rtc.d_centro):
297  for centro in self.rtcrtc._rtc.d_centro:
298  centro.calibrate_img()
299  self.rtcrtc.do_centroids(0)
300  self.rtcrtc.do_control(0)
301  self.rtcrtc.do_clipping(0)
302  self.rtcrtc._rtc.comp_voltage(0)
303 
304  def set_one_actu(self, nctrl: int, nactu: int, *, ampli: float = 1,
305  reset: bool = True) -> None:
306  """ Push the selected actuator
307 
308  Args:
309  nctrl : (int) : controller index
310 
311  nactu : (int) : actuator index to push
312 
313  Kwargs:
314  ampli : (float) : amplitude to apply. Default is 1 volt
315 
316  reset : (bool) : reset the previous command vector. Default is True
317  """
318  command = self.get_commandget_command()
319  if reset:
320  command *= 0
321  command[nactu] = ampli
322  self.set_commandset_command(nctrl, command)
323 
324  def force_context(self) -> None:
325  """ Active all the GPU devices specified in the parameters file
326  Required for using with widgets, due to multithreaded init
327  and in case GPU 0 is not used by the simu
328  """
329  if self.is_initis_initis_init and self.contextcontextcontext is not None:
330  current_device = self.contextcontextcontext.active_device
331  for device in range(len(self.configconfigconfig.p_loop.devices)):
332  self.contextcontextcontext.set_active_device_force(device)
333  self.contextcontextcontext.set_active_device(current_device)
334 
335  def reset_dm(self) -> None:
336  """ Reset the DM
337  """
338  if hasattr(self, '_dm'):
339  dm.reset_dm()
340 
341  def reset_command(self, nctrl: int = -1) -> None:
342  """ Reset the nctrl Controller command buffer, reset all controllers if nctrl == -1
343 
344  Kwargs:
345  nctrl : (int) : Controller index. If -1 (default), all controllers commands are reset
346  """
347  if (nctrl == -1): # All Dms reset
348  for control in self.rtcrtc._rtc.d_control:
349  control.d_com.reset()
350  else:
351  self.rtcrtc._rtc.d_control[nctrl].d_com.reset()
352 
353  def load_config(self, config_file: str = None) -> None:
354  """ Init the COMPASS with the config_file
355 
356  Args:
357  config_file : (str) : path to the configuration file
358  """
359  from shesha.config import ParamConfig
360  self.configconfigconfig = ParamConfig(config_file)
361 
362  def set_cam_callback(self, cam_callback: Callable):
363  """ Set the externally defined function that allows to grab frames
364 
365  Args:
366  cam_callback : (Callable) : function that allows to grab frames
367  """
368  self.cam_callbackcam_callback = cam_callback
369 
370  def set_dm_callback(self, dm_get_callback: Callable, dm_set_callback: Callable):
371  """ Set the externally defined function that allows to communicate with the DM
372 
373  Args:
374  dm_get_callback : (Callable) : function that allows to retrieve commands
375  dm_set_callback : (Callable) : function that allows to set commands
376  """
377  self.dm_get_callbackdm_get_callback = dm_get_callback
378  self.dm_set_callbackdm_set_callback = dm_set_callback
379 
380  def adaptive_windows(self, init_config=False, centro_index: int = 0):
381  """ Re-centre the centroiding boxes around the spots, and loads
382  the new box coordinates in the slopes computation supervisor
383  pipeline.
384 
385  Args:
386  init_config : (bool): Flag to reset to the default positions of boxes. Default is False
387 
388  centro_index : (int) : centroider index
389  """
390  if init_config:
391  # reset de la configuration initiale
392  ij_subap = self.configconfigconfig.p_wfss[centro_index].get_validsub()
393  nsubap = ij_subap.shape[1]
394  self.rtcrtc._rtc.d_centro[centro_index].load_validpos(
395  ij_subap[0], ij_subap[1], nsubap)
396  else:
397  # acquire slopes first
398  nslopes = 10
399  s = 0.
400  for i in range(nslopes):
401  self.load_new_wfs_frameload_new_wfs_frame() # sinon toutes les slopes sont les memes
402  self.compute_wfs_framecompute_wfs_frame()
403  s = s + self.get_slopes()[self.slopes_indexslopes_index[centro_index]:self.
404  slopes_index[centro_index + 1]]
405  s /= nslopes
406  # get coordinates of valid sub-apertures
407  #ij_subap = self.config.p_wfss[centro_index].get_validsub()
408  i_subap = np.array(self.rtcrtc._rtc.d_centro[centro_index].d_validx)
409  j_subap = np.array(self.rtcrtc._rtc.d_centro[centro_index].d_validy)
410  # get number of subaps
411  nsubap = i_subap.shape[0]
412  # reshape the array <s> to be conformable with <ij_subap>
413  s = np.resize(s, (2, nsubap))
414  # re-centre the boxes around the spots
415  new_i_subap = (i_subap + s[0, :].round()).astype(int)
416  new_j_subap = (j_subap + s[1, :].round()).astype(int)
417  # load the new positions of boxes
418  self.rtcrtc._rtc.d_centro[centro_index].load_validpos(
419  new_i_subap, new_j_subap, nsubap)
420 
421  def get_current_windows_pos(self, centro_index: int = 0):
422  """ Returns the currently used subapertures positions
423 
424  Args:
425  centro_index : (int) : Index of the centroider
426 
427  Returns:
428  current_pos : (tuple) : (i_subap, j_subap)
429  """
430  i_subap = np.array(self.rtcrtc._rtc.d_centro[centro_index].d_validx)
431  j_subap = np.array(self.rtcrtc._rtc.d_centro[centro_index].d_validy)
432  return i_subap, j_subap
433 
434  def get_slopes_index(self):
435  """ Return the index of the first position of each WFS slopes vector
436  inside the global RTC slopes vector
437 
438  Returns:
439  slopes_index : (np.ndarray) : Slopes index
440  """
441  return self.slopes_indexslopes_index
442 
443  def export_config(self):
444  """
445  Extract and convert compass supervisor configuration parameters
446  into 2 dictionnaries containing relevant AO parameters
447 
448  Args:
449  root: (object), COMPASS supervisor object to be parsed
450 
451  Returns : 2 dictionnaries
452  """
453  return self.configconfigconfig.export_config()
#define set_active_device(new_device, silent)
Definition: carma_context.h:72
#define set_active_device_force(new_device, silent)
Definition: carma_context.h:74
This class implements generic supervisor to handle compass simulation.
None load_new_wfs_frame(self, int centro_index=0)
Acquire a new WFS frame and load it.
def get_slopes_index(self)
Return the index of the first position of each WFS slopes vector inside the global RTC slopes vector.
def get_current_windows_pos(self, int centro_index=0)
Returns the currently used subapertures positions.
def __init__(self, str config_file=None, bool cacao=False)
Init the COMPASS wih the config_file.
def set_cam_callback(self, Callable cam_callback)
Set the externally defined function that allows to grab frames.
None next(self)
Performs a single loop iteration.
None force_context(self)
Active all the GPU devices specified in the parameters file Required for using with widgets,...
None set_command(self, int nctrl, np.ndarray command)
Immediately sets provided command to DMs - does not affect integrator.
rtc
(RtcComponent) : A Rtc component instance
np.ndarray get_tar_image(self, tar_index, str expo_type="se")
NOT IMPLEMENTED.
None set_one_actu(self, int nctrl, int nactu, *float ampli=1, bool reset=True)
Push the selected actuator.
None load_config(self, str config_file=None)
Init the COMPASS with the config_file.
None reset_command(self, int nctrl=-1)
Reset the nctrl Controller command buffer, reset all controllers if nctrl == -1.
cacao
(bool) : CACAO features enabled in the RTC
def compute_wfs_frame(self)
Compute the WFS frame: calibrate, centroid, commands.
def set_dm_callback(self, Callable dm_get_callback, Callable dm_set_callback)
Set the externally defined function that allows to communicate with the DM.
np.ndarray get_command(self)
Get command from DM.
def export_config(self)
Extract and convert compass supervisor configuration parameters into 2 dictionnaries containing relev...
def adaptive_windows(self, init_config=False, int centro_index=0)
Re-centre the centroiding boxes around the spots, and loads the new box coordinates in the slopes com...
This class defines generic methods and behavior of a supervisor It is not intended to be instantiated...
context
(CarmaContext) : a CarmaContext instance
is_init
(bool) : Flag equals to True if the supervisor has already been initialized
Parameter classes for COMPASS.
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
Initialization of a Dms object.
Definition: dm_init.py:1
Initialization of a Rtc object.
Definition: rtc_init.py:1
Basic utilities function.
Definition: utilities.py:1