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