COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
roket.py
1 """
2 ROKET (erROr breaKdown Estimation Tool)
3 
4 Computes the error breakdown during a COMPASS simulation
5 and saves it in a HDF5 file
6 Error contributors are bandwidth, tomography, noise, aliasing,
7 WFS non linearity, filtered modes and fitting
8 Saved file contained temporal buffers of those contributors
9 """
10 
11 import cProfile
12 import pstats as ps
13 
14 import sys, os
15 import numpy as np
16 from shesha.supervisor.compassSupervisor import CompassSupervisor
17 from shesha.util.rtc_util import centroid_gain
18 from shesha.ao.tomo import create_nact_geom
19 from shesha.ao.basis import compute_btt, compute_cmat_with_Btt
20 import shesha.constants as scons
21 import time
22 import matplotlib.pyplot as pl
23 pl.ion()
24 import shesha.util.hdf5_util as h5u
25 import pandas
26 from scipy.sparse import csr_matrix
27 
28 
30  """
31  ROKET class
32  Inherits from CompassSupervisor
33  """
34 
35  def __init__(self, str=None, N_preloop=1000, gamma=1.):
36  """
37  Initializes an instance of Roket class
38 
39  :parameters:
40  str: (str): (optional) path to a parameter file
41  N_preloop: (int): (optional) number of iterations before starting error breakdown estimation
42  gamma: (float): (optional) centroid gain
43  """
44  super().__init__(str)
45  self.N_preloop = N_preloop
46  self.gamma = gamma
47 
48  def init_config(self):
49  """
50  Initializes the COMPASS simulation and the ROKET buffers
51  """
52  super().init_config()
53  self.iter_number = 0
54  self.n = self.config.p_loop.niter + self.N_preloop
55  self.nfiltered = int(self.config.p_controllers[0].maxcond)
56  self.nactus = self.get_command(0).size
57  self.nslopes = self.get_slopes(0).size
58  self.com = np.zeros((self.n, self.nactus), dtype=np.float32)
59  self.noise_com = np.zeros((self.n, self.nactus), dtype=np.float32)
60  self.alias_wfs_com = np.copy(self.noise_com)
61  self.alias_meas = np.zeros((self.n, self.nslopes), dtype=np.float32)
62  self.wf_com = np.copy(self.noise_com)
63  self.tomo_com = np.copy(self.noise_com)
64  self.trunc_com = np.copy(self.noise_com)
65  self.trunc_meas = np.copy(self.alias_meas)
66  self.H_com = np.copy(self.noise_com)
67  self.mod_com = np.copy(self.noise_com)
68  self.bp_com = np.copy(self.noise_com)
69  self.fit = np.zeros(self.n)
70  self.psf_ortho = self.get_tar_image(0) * 0
71  self.centroid_gain = 0
72  self.centroid_gain2 = 0
73  self.slopes = np.zeros((self.n, self.nslopes), dtype=np.float32)
74  #gamma = 1.0
75  self.config.p_loop.set_niter(self.n)
76  self.IFpzt = self.get_influ_basis_sparse(1)
77  self.TT = self.get_tt_influ_basis(1)
78 
79  self.Btt, self.P = compute_btt(self.IFpzt.T, self.TT)
80  tmp = self.Btt.dot(self.Btt.T)
81  self._sim.rtc.d_control[1].load_Btt(tmp[:-2, :-2], tmp[-2:, -2:])
82  compute_cmat_with_Btt(self._sim.rtc, self.Btt, self.nfiltered)
83  self.cmat = self.get_command_matrix(0)
84  self.D = self.get_interaction_matrix(0)
85  self.RD = np.dot(self.cmat, self.D)
86  self.gRD = np.identity(
87  self.RD.
88  shape[0]) - self.config.p_controllers[0].gain * self.gamma * self.RD
89 
90  self.Nact = create_nact_geom(self.config.p_dms[0])
91 
92  def next(self, **kwargs):
93  """
94  function next
95  Iterates the AO loop, with optional parameters
96 
97  :param move_atmos (bool): move the atmosphere for this iteration, default: True
98  :param nControl (int): Controller number to use, default 0 (single control configurations)
99  :param tar_trace (None or list[int]): list of targets to trace. None equivalent to all.
100  :param wfs_trace (None or list[int]): list of WFS to trace. None equivalent to all.
101  :param apply_control (bool): (optional) if True (default), apply control on DMs
102  """
103  self._sim.next(apply_control=False)
104  self.error_breakdown()
105  self._sim.applyControl(0)
106  self.iter_number += 1
107 
108  def loop(self, monitoring_freq=100, **kwargs):
109  """
110  Performs the AO loop for n iterations
111 
112  :parameters:
113  monitoring_freq: (int): (optional) Loop monitoring frequency [frames] in the terminal
114  """
115  print("-----------------------------------------------------------------")
116  print("iter# | SE SR | LE SR | FIT SR | PH SR | ETR (s) | Framerate (Hz)")
117  print("-----------------------------------------------------------------")
118  t0 = time.time()
119  for i in range(self.n):
120  self.next(**kwargs)
121  if ((i + 1) % monitoring_freq == 0):
122  framerate = (i + 1) / (time.time() - t0)
123  self._sim.comp_tar_image(0)
124  self._sim.compStrehl(0)
125  strehltmp = self.get_strehl(0)
126  etr = (self.n - i) / framerate
127  print("%d \t %.2f \t %.2f\t %.2f \t %.2f \t %.1f \t %.1f" %
128  (i + 1, strehltmp[0], strehltmp[1], np.exp(-strehltmp[2]),
129  np.exp(-strehltmp[3]), etr, framerate))
130  t1 = time.time()
131 
132  print(" loop execution time:", t1 - t0, " (", self.n, "iterations), ",
133  (t1 - t0) / self.n, "(mean) ", self.n / (t1 - t0), "Hz")
134 
135  #self.tar.comp_image(0)
136  SRs = self.get_strehl(0)
137  self.SR2 = np.exp(SRs[3])
138  self.SR = SRs[1]
139 
140  def error_breakdown(self):
141  """
142  Compute the error breakdown of the AO simulation. Returns the error commands of
143  each contributors. Suppose no delay (for now) and only 2 controllers : the main one, controller #0, (specified on the parameter file)
144  and the geometric one, controller #1 (automatically added if roket is asked in the parameter file)
145  Commands are computed by applying the loop filter on various kind of commands : (see schema_simulation_budget_erreur_v2)
146 
147  - Ageom : Aliasing contribution on WFS direction
148  Obtained by computing commands from DM orthogonal phase (projection + slopes_geom)
149 
150  - B : Projection on the target direction
151  Obtained as the commmands output of the geometric controller
152 
153  - C : Wavefront
154  Obtained by computing commands from DM parallel phase (RD*B)
155 
156  - E : Wavefront + aliasing + ech/trunc + tomo
157  Obtained by performing the AO loop iteration without noise on the WFS
158 
159  - F : Wavefront + aliasing + tomo
160  Obtained by performing the AO loop iteration without noise on the WFS and using phase deriving slopes
161 
162  - G : tomo
163 
164  Note : rtc.get_err returns to -CMAT.slopes
165  """
166  g = self.config.p_controllers[0].gain
167  Dcom = self.get_command(0)
168  Derr = self.get_err(0)
169  self.com[self.iter_number, :] = Dcom
170  tarphase = self.get_tar_phase(0)
171  self.slopes[self.iter_number, :] = self.get_slopes(0)
172 
173 
176  if (self.config.p_wfss[0].type == scons.WFSType.SH):
177  ideal_img = np.array(self._sim.wfs.d_wfs[0].d_binimg_notnoisy)
178  binimg = np.array(self._sim.wfs.d_wfs[0].d_binimg)
179  if (self.config.p_centroiders[0].type == scons.CentroiderType.TCOG
180  ): # Select the same pixels with or without noise
181  invalidpix = np.where(binimg <= self.config.p_centroiders[0].thresh)
182  ideal_img[invalidpix] = 0
183  self.set_centroider_threshold(0, -1e16)
184  self._sim.wfs.d_wfs[0].set_binimg(ideal_img, ideal_img.size)
185  elif (self.config.p_wfss[0].type == scons.centroiderType.PYRHR):
186  ideal_pyrimg = np.array(self._sim.wfs.d_wfs[0].d_binimg_notnoisy)
187  self._sim.wfs.d_wfs[0].set_binimg(ideal_pyrimg, ideal_pyrimg.size)
188 
189  self._sim.doCentroids(0)
190  if (self.config.p_centroiders[0].type == scons.CentroiderType.TCOG):
191  self.set_centroider_threshold(0, config.p_centroiders[0].thresh)
192 
193  self._sim.do_control(0)
194  E = self.get_err(0)
195  E_meas = self.get_slopes(0)
196  # Apply loop filter to get contribution of noise on commands
197  if (self.iter_number + 1 < self.config.p_loop.niter):
198  self.noise_com[self.iter_number + 1, :] = self.gRD.dot(
199  self.noise_com[self.iter_number, :]) + g * (Derr - E)
200 
203  self._sim.doCentroidsGeom(0)
204  self._sim.do_control(0)
205  F = self.get_err(0)
206  F_meas = self.get_slopes(0)
207  self.trunc_meas[self.iter_number, :] = E_meas - F_meas
208  # Apply loop filter to get contribution of sampling/truncature on commands
209  if (self.iter_number + 1 < self.config.p_loop.niter):
210  self.trunc_com[self.iter_number + 1, :] = self.gRD.dot(
211  self.trunc_com[self.iter_number, :]) + g * (E - self.gamma * F)
212  self.centroid_gain += centroid_gain(E, F)
213  self.centroid_gain2 += centroid_gain(Derr, F)
214 
217  self._sim.do_control(1, 0, wfs_direction=True)
218  self._sim.applyControl(1)
219  for w in range(len(self.config.p_wfss)):
220  self._sim.raytrace_wfs(w, "dm", rst=False)
221  """
222  wfs.sensors_compimg(0)
223  if(config.p_wfss[0].type == scons.WFSType.SH):
224  ideal_img = wfs.get_binimgNotNoisy(0)
225  binimg = wfs.get_binimg(0)
226  if(config.p_centroiders[0].type == scons.CentroiderType.TCOG): # Select the same pixels with or without noise
227  invalidpix = np.where(binimg <= config.p_centroiders[0].thresh)
228  ideal_img[self.iter_numbernvalidpix] = 0
229  rtc.setthresh(0,-1e16)
230  wfs.set_binimg(0,ideal_img)
231  elif(config.p_wfss[0].type == scons.centroiderType.PYRHR):
232  ideal_pyrimg = wfs.get_binimg_notnoisy(0)
233  wfs.set_pyrimg(0,ideal_pyrimg)
234  """
235  self._sim.doCentroidsGeom(0)
236  self._sim.do_control(0)
237  Ageom = self.get_err(0)
238  self.alias_meas[self.iter_number, :] = self.get_slopes(0)
239  if (self.iter_number + 1 < self.config.p_loop.niter):
240  self.alias_wfs_com[self.iter_number + 1, :] = self.gRD.dot(
241  self.alias_wfs_com[self.iter_number, :]) + self.gamma * g * (
242  Ageom) # - (E-F))
243 
244 
247  self._sim.raytraceTar(0, "atmos")
248  self._sim.do_control(1, 0, wfs_direction=False)
249  B = self.get_command(1)
250 
251 
254  self._sim.applyControl(1)
255  self._sim.raytraceTar(0, "dm", rst=False)
256 
257  self._sim.comp_tar_image(0, compLE=False)
258  self._sim.compStrehl(0)
259  self.fit[self.iter_number] = self.get_strehl(0)[2]
260  if (self.iter_number >= self.N_preloop):
261  self.psf_ortho += self.get_tar_image(0, 'se')
262 
263 
266  modes = self.P.dot(B)
267  modes_filt = modes.copy() * 0.
268  modes_filt[-self.nfiltered - 2:-2] = modes[-self.nfiltered - 2:-2]
269  self.H_com[self.iter_number, :] = self.Btt.dot(modes_filt)
270  modes[-self.nfiltered - 2:-2] = 0
271  self.mod_com[self.iter_number, :] = self.Btt.dot(modes)
272 
273 
276  C = self.mod_com[self.iter_number, :] - self.mod_com[self.iter_number - 1, :]
277 
278  self.bp_com[self.iter_number, :] = self.gRD.dot(
279  self.bp_com[self.iter_number - 1, :]) - C
280 
281 
285  for w in range(len(self.config.p_wfss)):
286  self._sim.raytrace_wfs(w, "atmos")
287 
288  self._sim.do_control(1, 0, wfs_direction=True)
289  G = self.get_command(1)
290  modes = self.P.dot(G)
291  modes[-self.nfiltered - 2:-2] = 0
292  self.wf_com[self.iter_number, :] = self.Btt.dot(modes)
293 
294  G = self.mod_com[self.iter_number, :] - self.wf_com[self.iter_number, :]
295  if (self.iter_number + 1 < self.config.p_loop.niter):
296  self.tomo_com[self.iter_number + 1, :] = self.gRD.dot(
297  self.tomo_com[self.iter_number, :]) - g * self.gamma * self.RD.dot(G)
298 
299  # Without anyone noticing...
300  self._sim.tar.d_targets[0].set_phase(tarphase)
301  self._sim.rtc.d_control[0].set_com(Dcom, Dcom.size)
302 
303  def save_in_hdf5(self, savename):
304  """
305  Saves all the ROKET buffers + simuation parameters in a HDF5 file
306 
307  :parameters:
308  savename: (str): name of the output file
309  """
310  tmp = (self.config.p_geom._ipupil.shape[0] -
311  (self.config.p_dms[0]._n2 - self.config.p_dms[0]._n1 + 1)) // 2
312  tmp_e0 = self.config.p_geom._ipupil.shape[0] - tmp
313  tmp_e1 = self.config.p_geom._ipupil.shape[1] - tmp
314  pup = self.config.p_geom._ipupil[tmp:tmp_e0, tmp:tmp_e1]
315  indx_pup = np.where(pup.flatten() > 0)[0].astype(np.int32)
316  dm_dim = self.config.p_dms[0]._n2 - self.config.p_dms[0]._n1 + 1
317  self.cov_cor()
318  psf = self.get_tar_image(0, "le")
319 
320  fname = os.getenv("DATA_GUARDIAN") + savename
321  pdict = {
322  "noise":
323  self.noise_com[self.N_preloop:, :].T,
324  "aliasing":
325  self.alias_wfs_com[self.N_preloop:, :].T,
326  "tomography":
327  self.tomo_com[self.N_preloop:, :].T,
328  "filtered modes":
329  self.H_com[self.N_preloop:, :].T,
330  "non linearity":
331  self.trunc_com[self.N_preloop:, :].T,
332  "bandwidth":
333  self.bp_com[self.N_preloop:, :].T,
334  "wf_com":
335  self.wf_com[self.N_preloop:, :].T,
336  "P":
337  self.P,
338  "Btt":
339  self.Btt,
340  "IF.data":
341  self.IFpzt.data,
342  "IF.indices":
343  self.IFpzt.indices,
344  "IF.indptr":
345  self.IFpzt.indptr,
346  "TT":
347  self.TT,
348  "dm_dim":
349  dm_dim,
350  "indx_pup":
351  indx_pup,
352  "fitting":
353  np.mean(self.fit[self.N_preloop:]),
354  "SR":
355  self.SR,
356  "SR2":
357  self.SR2,
358  "cov":
359  self.cov,
360  "cor":
361  self.cor,
362  "psfortho":
363  np.fft.fftshift(self.psf_ortho) /
364  (self.config.p_loop.niter - self.N_preloop),
365  "centroid_gain":
366  self.centroid_gain / (self.config.p_loop.niter - self.N_preloop),
367  "centroid_gain2":
368  self.centroid_gain2 /
369  (self.config.p_loop.niter - self.N_preloop),
370  "dm.xpos":
371  self.config.p_dms[0]._xpos,
372  "dm.ypos":
373  self.config.p_dms[0]._ypos,
374  "R":
375  self.get_command_matrix(0),
376  "D":
377  self.get_interaction_matrix(0),
378  "Nact":
379  self.Nact,
380  "com":
381  self.com[self.N_preloop:, :].T,
382  "slopes":
383  self.slopes[self.N_preloop:, :].T,
384  "alias_meas":
385  self.alias_meas[self.N_preloop:, :].T,
386  "trunc_meas":
387  self.trunc_meas[self.N_preloop:, :].T
388  }
389  h5u.save_h5(fname, "psf", self.config, psf)
390  for k in list(pdict.keys()):
391  h5u.save_hdf5(fname, k, pdict[k])
392 
393  def cov_cor(self):
394  """
395  Computes covariance matrix and correlation matrix between all the contributors
396  """
397  self.cov = np.zeros((6, 6))
398  self.cor = np.zeros((6, 6))
399  bufdict = {
400  "0": self.noise_com.T,
401  "1": self.trunc_com.T,
402  "2": self.alias_wfs_com.T,
403  "3": self.H_com.T,
404  "4": self.bp_com.T,
405  "5": self.tomo_com.T
406  }
407  for i in range(self.cov.shape[0]):
408  for j in range(self.cov.shape[1]):
409  if (j >= i):
410  tmpi = self.P.dot(bufdict[str(i)])
411  tmpj = self.P.dot(bufdict[str(j)])
412  self.cov[i, j] = np.sum(
413  np.mean(tmpi * tmpj, axis=1) -
414  np.mean(tmpi, axis=1) * np.mean(tmpj, axis=1))
415  else:
416  self.cov[i, j] = self.cov[j, i]
417 
418  s = np.reshape(np.diag(self.cov), (self.cov.shape[0], 1))
419  sst = np.dot(s, s.T)
420  ok = np.where(sst)
421  self.cor[ok] = self.cov[ok] / np.sqrt(sst[ok])
422 
423 
424 
429 if __name__ == "__main__":
430  if (len(sys.argv) < 2):
431  error = 'command line should be at least:"python -i test.py parameters_filename"\n with "parameters_filename" the path to the parameters file'
432  raise Exception(error)
433 
434  #get parameters from file
435  param_file = sys.argv[1]
436 
437  if (len(sys.argv) > 2):
438  savename = sys.argv[2]
439  else:
440  savename = "roket_default.h5"
441 
442  roket = Roket(param_file)
443  roket.init_config()
444  roket.loop()
445  roket.save_in_hdf5(savename)
shesha.supervisor.compassSupervisor
Initialization and execution of a COMPASS supervisor.
Definition: compassSupervisor.py:1
guardians.roket.Roket.next
def next(self, **kwargs)
Definition: roket.py:102
guardians.roket.Roket.centroid_gain
centroid_gain
Definition: roket.py:71
guardians.roket.Roket.cmat
cmat
Definition: roket.py:83
guardians.roket.Roket.cov
cov
Definition: roket.py:397
guardians.roket.Roket.cor
cor
Definition: roket.py:398
guardians.roket.Roket.gRD
gRD
Definition: roket.py:86
guardians.roket.Roket.IFpzt
IFpzt
Definition: roket.py:76
guardians.roket.Roket.noise_com
noise_com
Definition: roket.py:59
shesha.ao.basis
Functions for modal basis (DM basis, KL, Btt, etc...)
Definition: basis.py:1
guardians.roket.Roket.cov_cor
def cov_cor(self)
Computes covariance matrix and correlation matrix between all the contributors.
Definition: roket.py:396
guardians.roket.Roket.init_config
def init_config(self)
Initializes the COMPASS simulation and the ROKET buffers.
Definition: roket.py:51
guardians.roket.Roket.alias_meas
alias_meas
Definition: roket.py:61
shesha.util.hdf5_util
Functions for handling the database system.
Definition: hdf5_util.py:1
shesha.supervisor.genericSupervisor.GenericSupervisor.config
config
Definition: genericSupervisor.py:109
shesha.constants
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
guardians.roket.Roket.tomo_com
tomo_com
Definition: roket.py:63
guardians.roket.Roket
Definition: roket.py:33
guardians.roket.Roket.com
com
Definition: roket.py:58
guardians.roket.Roket.alias_wfs_com
alias_wfs_com
Definition: roket.py:60
guardians.roket.Roket.iter_number
iter_number
Definition: roket.py:53
guardians.roket.Roket.error_breakdown
def error_breakdown(self)
Definition: roket.py:165
guardians.roket.Roket.fit
fit
Definition: roket.py:69
guardians.roket.Roket.bp_com
bp_com
Definition: roket.py:68
guardians.roket.Roket.H_com
H_com
Definition: roket.py:66
guardians.roket.Roket.D
D
Definition: roket.py:84
guardians.roket.Roket.trunc_com
trunc_com
Definition: roket.py:64
guardians.roket.Roket.gamma
gamma
Definition: roket.py:46
guardians.roket.Roket.slopes
slopes
Definition: roket.py:73
guardians.roket.Roket.nfiltered
nfiltered
Definition: roket.py:55
guardians.roket.Roket.SR2
SR2
Definition: roket.py:137
guardians.roket.Roket.centroid_gain2
centroid_gain2
Definition: roket.py:72
guardians.roket.Roket.nslopes
nslopes
Definition: roket.py:57
guardians.roket.Roket.n
n
Definition: roket.py:54
shesha.util.rtc_util
Some utilities functions for RTC.
Definition: rtc_util.py:1
guardians.roket.Roket.trunc_meas
trunc_meas
Definition: roket.py:65
guardians.roket.Roket.save_in_hdf5
def save_in_hdf5(self, savename)
Saves all the ROKET buffers + simuation parameters in a HDF5 file.
Definition: roket.py:309
guardians.roket.Roket.mod_com
mod_com
Definition: roket.py:67
guardians.roket.Roket.SR
SR
Definition: roket.py:138
shesha.supervisor.compassSupervisor.CompassSupervisor
This class implements generic supervisor to handle compass simulation.
Definition: compassSupervisor.py:57
guardians.roket.Roket.loop
def loop(self, monitoring_freq=100, **kwargs)
Performs the AO loop for n iterations.
Definition: roket.py:114
guardians.roket.Roket.RD
RD
Definition: roket.py:85
guardians.roket.Roket.__init__
def __init__(self, str=None, N_preloop=1000, gamma=1.)
Initializes an instance of Roket class.
Definition: roket.py:43
guardians.roket.Roket.N_preloop
N_preloop
Definition: roket.py:45
guardians.roket.Roket.psf_ortho
psf_ortho
Definition: roket.py:70
guardians.roket.Roket.nactus
nactus
Definition: roket.py:56
guardians.roket.Roket.TT
TT
Definition: roket.py:77
guardians.roket.Roket.wf_com
wf_com
Definition: roket.py:62
shesha.ao.tomo
Computation of tomographic reconstructor.
Definition: tomo.py:1
guardians.roket.Roket.P
P
Definition: roket.py:79
guardians.roket.Roket.Nact
Nact
Definition: roket.py:90