COMPASS  5.0.0
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
40 from shesha.supervisor.optimizers import ModalBasis, Calibration
41 import numpy as np
42 
43 import shesha.constants as scons
44 from shesha.constants import CONST
45 import shesha.ao.basis as basis
46 import astropy.io.fits as pfits
47 from tqdm import trange, tqdm
48 import time
49 
50 from typing import List, Iterable
51 
52 
54  """ This class implements generic supervisor to handle compass simulation
55 
56  Attributes:
57  context : (CarmaContext) : a CarmaContext instance
58 
59  config : (config) : Parameters structure
60 
61  telescope : (TelescopeComponent) : a TelescopeComponent instance
62 
63  atmos : (AtmosComponent) : An AtmosComponent instance
64 
65  target : (targetComponent) : A TargetComponent instance
66 
67  wfs : (WfsComponent) : A WfsComponent instance
68 
69  dms : (DmComponent) : A DmComponent instance
70 
71  rtc : (RtcComponent) : A Rtc component instance
72 
73  is_init : (bool) : Flag equals to True if the supervisor has already been initialized
74 
75  iter : (int) : Frame counter
76 
77  cacao : (bool) : CACAO features enabled in the RTC
78  """
79  def __init__(self, config, cacao : bool=False):
80  """ Instantiates a CompassSupervisor object
81 
82  Parameters:
83  config: (config module) : Configuration module
84 
85  cacao : (bool, optional) : If True, enables CACAO features in RTC (Default is False)
86  /!\ Requires OCTOPUS to be installed
87  """
88  self.cacao = cacao
89  GenericSupervisor.__init__(self, config)
90  self.basis = ModalBasis(self.config, self.dms, self.target)
91  self.calibration = Calibration(self.config, self.tel, self.atmos, self.dms,
92  self.target, self.rtc, self.wfs)
93 # ___ _ __ __ _ _ _
94 # / __|___ _ _ ___ _ _(_)__ | \/ |___| |_| |_ ___ __| |___
95 # | (_ / -_) ' \/ -_) '_| / _| | |\/| / -_) _| ' \/ _ \/ _` (_-<
96 # \___\___|_||_\___|_| |_\__| |_| |_\___|\__|_||_\___/\__,_/__/
97 
98  def _init_tel(self):
99  """Initialize the telescope component of the supervisor as a TelescopeCompass
100  """
101  self.tel = TelescopeCompass(self.context, self.config)
102 
103  def _init_atmos(self):
104  """Initialize the atmosphere component of the supervisor as a AtmosCompass
105  """
106  self.atmos = AtmosCompass(self.context, self.config)
107 
108  def _init_dms(self):
109  """Initialize the DM component of the supervisor as a DmCompass
110  """
111  self.dms = DmCompass(self.context, self.config)
112 
113  def _init_target(self):
114  """Initialize the target component of the supervisor as a TargetCompass
115  """
116  if self.tel is not None:
117  self.target = TargetCompass(self.context, self.config, self.tel)
118  else:
119  raise ValueError("Configuration not loaded or Telescope not initilaized")
120 
121  def _init_wfs(self):
122  """Initialize the wfs component of the supervisor as a WfsCompass
123  """
124  if self.tel is not None:
125  self.wfs = WfsCompass(self.context, self.config, self.tel)
126  else:
127  raise ValueError("Configuration not loaded or Telescope not initilaized")
128 
129  def _init_rtc(self):
130  """Initialize the rtc component of the supervisor as a RtcCompass
131  """
132  if self.wfs is not None:
133  self.rtc = RtcCompass(self.context, self.config, self.tel, self.wfs, self.dms, self.atmos, cacao=self.cacao)
134  else:
135  raise ValueError("Configuration not loaded or Telescope not initilaized")
136 
137  def next(self, *, move_atmos: bool = True, nControl: int = 0,
138  tar_trace: Iterable[int] = None, wfs_trace: Iterable[int] = None,
139  do_control: bool = True, apply_control: bool = True, compute_tar_psf: bool = True) -> None:
140  """Iterates the AO loop, with optional parameters.
141 
142  Overload the GenericSupervisor next() method to handle the GEO controller
143  specific raytrace order operations
144 
145  Parameters :
146  move_atmos: (bool), optional: move the atmosphere for this iteration. Default is True
147 
148  nControl: (int, optional): Controller number to use. Default is 0 (single control configuration)
149 
150  tar_trace: (List, optional): list of targets to trace. None is equivalent to all (default)
151 
152  wfs_trace: (List, optional): list of WFS to trace. None is equivalent to all (default)
153 
154  do_control : (bool, optional) : Performs RTC operations if True (Default)
155 
156  apply_control: (bool): (optional) if True (default), apply control on DMs
157 
158  compute_tar_psf : (bool, optional) : If True (default), computes the PSF at the end of the iteration
159  """
160  if (self.config.p_controllers is not None and
161  self.config.p_controllers[nControl].type == scons.ControllerType.GEO):
162  if tar_trace is None and self.target is not None:
163  tar_trace = range(len(self.config.p_targets))
164  if wfs_trace is None and self.wfs is not None:
165  wfs_trace = range(len(self.config.p_wfss))
166 
167  if move_atmos and self.atmos is not None:
168  self.atmos.move_atmos()
169 
170  if tar_trace is not None:
171  for t in tar_trace:
172  if self.atmos.is_enable:
173  self.target.raytrace(t, tel=self.tel, atm=self.atmos, ncpa=False)
174  else:
175  self.target.raytrace(t, tel=self.tel, ncpa=False)
176 
177  if do_control and self.rtc is not None:
178  self.rtc.do_control(nControl, sources=self.target.sources)
179  self.target.raytrace(t, dms=self.dms, ncpa=True, reset=False)
180  if apply_control:
181  self.rtc.apply_control(nControl)
182  if self.cacao:
183  self.rtc.publish()
184  if compute_tar_psf:
185  for tar_index in tar_trace:
186  self.target.comp_tar_image(tar_index)
187  self.target.comp_strehl(tar_index)
188 
189  self.iter += 1
190 
191  else:
192  GenericSupervisor.next(self, move_atmos=move_atmos, nControl=nControl,
193  tar_trace=tar_trace, wfs_trace=wfs_trace,
194  do_control=do_control, apply_control=apply_control, compute_tar_psf=compute_tar_psf)
195 # ___ _ __ _ __ __ _ _ _
196 # / __|_ __ ___ __(_)/ _(_)__ | \/ |___| |_| |_ ___ __| |___
197 # \__ \ '_ \/ -_) _| | _| / _| | |\/| / -_) _| ' \/ _ \/ _` (_-<
198 # |___/ .__/\___\__|_|_| |_\__| |_| |_\___|\__|_||_\___/\__,_/__/
199 # |_|
200 
202  self, cb_count: int, projection_matrix : np.ndarray, sub_sample: int = 1, controller_index: int = 0,
203  tar_index: int = 0, see_atmos: bool = True, cube_data_type: str = None,
204  cube_data_file_path: str = "", ncpa: int = 0, ncpa_wfs: np.ndarray = None,
205  ref_slopes: np.ndarray = None, ditch_strehl: bool = True):
206  """ Used to record a synchronized circular buffer AO loop data.
207 
208  Parameters:
209  cb_count: (int) : the number of iterations to record.
210 
211  projection_matrix : (np.ndarray) : projection matrix on modal basis to compute residual coefficients
212 
213  sub_sample: (int) : sub sampling of the data (default=1, I.e no subsampling)
214 
215  controller_index: (int) :
216 
217  tar_index: (int) : target number
218 
219  see_atmos: (int) : used for the next function to enable or not the Atmos
220 
221  cube_data_type: (int) : if specified ("tarPhase" or "psfse") returns the target phase or short exposure PSF data cube in the output variable
222 
223  cube_data_file_path: (int) : if specified it will also save the target phase cube data (full path on the server)
224 
225  ncpa: (int) : !!experimental!!!: Used only in the context of PYRWFS + NCPA compensation on the fly (with optical gain)
226  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.
227 
228  ncpa_wfs: (int) : the ncpa phase as seen from the wfs array with dims = size of Mpupil
229 
230  ref_slopes: (int) : the reference slopes to use.
231 
232  ditch_strehl: (int) : resets the long exposure SR computation at the beginning of the Circular buffer (default= True)
233 
234  Return:
235  slopes: (int) : the slopes CB
236 
237  volts: (int) : the volts applied to the DM(s) CB
238 
239  ai: (int) : the modal coefficient of the residual phase projected on the currently used modal Basis
240 
241  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)
242 
243  sthrel_se_list: (int) : The SR short exposure evolution during CB recording
244 
245  sthrel_le_list: (int) : The SR long exposure evolution during CB recording
246 
247  g_ncpa_list: (int) : the gain applied to the NCPA (PYRWFS CASE) if NCPA is set to True
248 
249  cube_data: (int) : the tarPhase or psfse cube data (see cube_data_type)
250  """
251  slopes_data = None
252  volts_data = None
253  cube_data = None
254  ai_data = None
255  k = 0
256  sthrel_se_list = []
257  sthrel_le_list = []
258  g_ncpa_list = []
259 
260  # Resets the target so that the PSF LE is synchro with the data
261  # Doesn't reset it if Ditch_strehl == False (used for real time gain computation)
262  if ditch_strehl:
263  for i in range(len(self.config.p_targets)):
264  self.target.reset_strehl(i)
265 
266  # Starting CB loop...
267  for j in range(cb_count):
268  print(j, end="\r")
269  if (ncpa):
270  if (j % ncpa == 0):
271  ncpa_diff = ref_slopes[None, :]
272  ncpa_turbu = self.calibration.do_imat_phase(controller_index,
273  -ncpa_wfs[None, :, :], noise=False,
274  with_turbu=True)
275  g_ncpa = float(
276  np.sqrt(
277  np.dot(ncpa_diff, ncpa_diff.T) / np.dot(
278  ncpa_turbu, ncpa_turbu.T)))
279  if (g_ncpa > 1e18):
280  g_ncpa = 0
281  print('Warning NCPA ref slopes gain too high!')
282  g_ncpa_list.append(g_ncpa)
283  self.rtc.set_ref_slopes(-ref_slopes * g_ncpa)
284  else:
285  g_ncpa_list.append(g_ncpa)
286  print('NCPA ref slopes gain: %4.3f' % g_ncpa)
287  self.rtc.set_ref_slopes(-ref_slopes / g_ncpa)
288 
289  self.atmos.enable_atmos(see_atmos)
290  self.next()
291  for t in range(len(self.config.p_targets)):
292  self.target.comp_tar_image(t)
293 
294  srse, srle, _, _ = self.target.get_strehl(tar_index)
295  sthrel_se_list.append(srse)
296  sthrel_le_list.append(srle)
297  if (j % sub_sample == 0):
298  ai_vector = self.calibration.compute_modal_residuals(projection_matrix, selected_actus=self.basis.selected_actus)
299  if (ai_data is None):
300  ai_data = np.zeros((len(ai_vector), int(cb_count / sub_sample)))
301  ai_data[:, k] = ai_vector
302 
303  slopes_vector = self.rtc.get_slopes(controller_index)
304  if (slopes_data is None):
305  slopes_data = np.zeros((len(slopes_vector),
306  int(cb_count / sub_sample)))
307  slopes_data[:, k] = slopes_vector
308 
309  volts_vector = self.rtc.get_command(controller_index) # get_command or get_voltages ?
310  if (volts_data is None):
311  volts_data = np.zeros((len(volts_vector),
312  int(cb_count / sub_sample)))
313  volts_data[:, k] = volts_vector
314 
315  if (cube_data_type):
316  if (cube_data_type == "tarPhase"):
317  dataArray = self.target.get_tar_phase(tar_index, pupil=True)
318  elif (cube_data_type == "psfse"):
319  dataArray = self.target.get_tar_image(tar_index, expo_type="se")
320  else:
321  raise ValueError("unknown dataData" % cube_data_type)
322  if (cube_data is None):
323  cube_data = np.zeros((*dataArray.shape,
324  int(cb_count / sub_sample)))
325  cube_data[:, :, k] = dataArray
326  k += 1
327  if (cube_data_file_path != ""):
328  print("Saving tarPhase cube at: ", cube_data_file_path)
329  pfits.writeto(cube_data_file_path, cube_data, overwrite=True)
330 
331  psf_le = self.target.get_tar_image(tar_index, expo_type="le")
332  return slopes_data, volts_data, ai_data, psf_le, sthrel_se_list, sthrel_le_list, g_ncpa_list, cube_data
shesha.supervisor.compassSupervisor.CompassSupervisor.cacao
cacao
Instantiates a CompassSupervisor object.
Definition: compassSupervisor.py:120
shesha.supervisor.compassSupervisor.CompassSupervisor.tel
tel
Definition: compassSupervisor.py:135
shesha.supervisor.optimizers
User layer for optimizing AO supervisor loop.
Definition: shesha/shesha/supervisor/optimizers/__init__.py:1
shesha.supervisor.components
Definition: shesha/shesha/supervisor/components/__init__.py:1
shesha.ao.basis
Functions for modal basis (DM basis, KL, Btt, etc...)
Definition: basis.py:1
shesha.supervisor.genericSupervisor.GenericSupervisor.config
config
Definition: genericSupervisor.py:109
shesha.supervisor.compassSupervisor.CompassSupervisor.record_ao_circular_buffer
def record_ao_circular_buffer(self, int cb_count, np.ndarray projection_matrix, 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)
Used to record a synchronized circular buffer AO loop data.
Definition: compassSupervisor.py:290
shesha.constants
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
shesha.supervisor.genericSupervisor.GenericSupervisor.context
context
Definition: genericSupervisor.py:108
shesha.supervisor.genericSupervisor.GenericSupervisor
This class defines generic methods and behavior of a supervisor It is not intended to be instantiated...
Definition: genericSupervisor.py:53
shesha.supervisor.compassSupervisor.CompassSupervisor.calibration
calibration
Definition: compassSupervisor.py:123
shesha.supervisor.compassSupervisor.CompassSupervisor.dms
dms
Definition: compassSupervisor.py:149
shesha.supervisor.genericSupervisor.GenericSupervisor.iter
iter
Init the a supervisor.
Definition: genericSupervisor.py:117
shesha.supervisor.compassSupervisor.CompassSupervisor.target
target
Definition: compassSupervisor.py:157
shesha.supervisor.compassSupervisor.CompassSupervisor.atmos
atmos
Definition: compassSupervisor.py:142
shesha.supervisor.compassSupervisor.CompassSupervisor
This class implements generic supervisor to handle compass simulation.
Definition: compassSupervisor.py:57
shesha.supervisor.compassSupervisor.CompassSupervisor.__init__
def __init__(self, config, bool cacao=False)
Definition: compassSupervisor.py:119
shesha.supervisor.compassSupervisor.CompassSupervisor.rtc
rtc
Definition: compassSupervisor.py:177
shesha.supervisor.compassSupervisor.CompassSupervisor.wfs
wfs
Definition: compassSupervisor.py:167
shesha.supervisor.compassSupervisor.CompassSupervisor.basis
basis
Definition: compassSupervisor.py:122
shesha.supervisor.compassSupervisor.CompassSupervisor.next
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)
Iterates the AO loop, with optional parameters.
Definition: compassSupervisor.py:201
shesha.supervisor.genericSupervisor
Definition: genericSupervisor.py:1