COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
roket_gpu.py
1 """
2 Created on Tue Jul 12 09:28:23 2016
3 
4 @author: fferreira
5 """
6 
7 import cProfile
8 import pstats as ps
9 
10 import sys, os
11 import numpy as np
12 import carmaWrap as ch
13 import shesha as ao
14 import time
15 import matplotlib.pyplot as plt
16 plt.ion()
17 import hdf5_util as h5u
18 import pandas
19 from scipy.sparse import csr_matrix
20 
21 if (len(sys.argv) < 2):
22  error = 'command line should be at least:"python -i test.py parameters_filename"\n with "parameters_filename" the path to the parameters file'
23  raise Exception(error)
24 
25 #get parameters from file
26 param_file = sys.argv[1]
27 if (param_file.split('.')[-1] == b"py"):
28  filename = param_file.split('/')[-1]
29  param_path = param_file.split(filename)[0]
30  sys.path.insert(0, param_path)
31  exec("import %s as config" % filename.split(".py")[0])
32  sys.path.remove(param_path)
33 elif (param_file.split('.')[-1] == b"h5"):
34  sys.path.insert(0, os.environ["SHESHA_ROOT"] + "/data/par/par4bench/")
35  import scao_sh_16x16_8pix as config
36  sys.path.remove(os.environ["SHESHA_ROOT"] + "/data/par/par4bench/")
37  h5u.configFromH5(param_file, config)
38 else:
39  raise ValueError("Parameter file extension must be .py or .h5")
40 
41 print("param_file is", param_file)
42 
43 #if(len(sys.argv) > 2):
44 # device=int(sys.argv[2])
45 #else:
46 # device = 0
47 if (len(sys.argv) > 2):
48  savename = sys.argv[2]
49 else:
50  savename = "roket_default.h5"
51 
52 print("save file is ", savename)
53 
60 
61 if (hasattr(config, "simul_name")):
62  if (config.simul_name is None):
63  simul_name = ""
64  else:
65  simul_name = config.simul_name
66 else:
67  simul_name = ""
68 print("simul name is", simul_name)
69 
70 matricesToLoad = {}
71 if (simul_name == b""):
72  clean = 1
73 else:
74  clean = 0
75  param_dict = h5u.params_dictionary(config)
76  matricesToLoad = h5u.checkMatricesDataBase(os.environ["SHESHA_ROOT"] + "/data/",
77  config, param_dict)
78 #initialisation:
79 # context
80 #c=ch.carmaWrap_context(device)
81 c = ch.carmaWrap_context(devices=np.array([0, 1], dtype=np.int32))
82 #c.set_active_device(device)
83 
84 # wfs
85 print("->wfs")
86 wfs, tel = ao.wfs_init(config.p_wfss, config.p_atmos, config.p_tel, config.p_geom,
87  config.p_target, config.p_loop, config.p_dms)
88 
89 # atmos
90 print("->atmos")
91 atm = ao.atmos_init(c, config.p_atmos, config.p_tel, config.p_geom, config.p_loop,
92  config.p_wfss, wfs, config.p_target, rank=0, clean=clean,
93  load=matricesToLoad)
94 
95 # dm
96 print("->dm")
97 dms = ao.dm_init(config.p_dms, config.p_wfss, wfs, config.p_geom, config.p_tel)
98 
99 # target
100 print("->target")
101 tar = ao.target_init(c, tel, config.p_target, config.p_atmos, config.p_geom,
102  config.p_tel, config.p_dms)
103 
104 print("->rtc")
105 # rtc
106 rtc = ao.rtc_init(tel, wfs, config.p_wfss, dms, config.p_dms, config.p_geom,
107  config.p_rtc, config.p_atmos, atm, config.p_tel, config.p_loop,
108  clean=clean, simul_name=simul_name, load=matricesToLoad)
109 
110 if not clean:
111  h5u.validDataBase(os.environ["SHESHA_ROOT"] + "/data/", matricesToLoad)
112 
113 print("====================")
114 print("init done")
115 print("====================")
116 print("objects initialzed on GPU:")
117 print("--------------------------------------------------------")
118 print(atm)
119 print(wfs)
120 print(dms)
121 print(tar)
122 print(rtc)
123 
124 print("----------------------------------------------------")
125 print("iter# | SE SR image | LE SR image | Fitting | LE SR phase var")
126 print("----------------------------------------------------")
127 
128 error_flag = True in [w.roket for w in config.p_wfss]
129 
130 
131 
138 def loop(n):
139  """
140  Performs the main AO loop for n interations. First, initialize buffers
141  for error breakdown computations. Then, at the end of each iteration, just
142  before applying the new DM shape, calls the error_breakdown function.
143 
144  :param n: (int) : number of iterations
145 
146  :return:
147  com : (np.array((n,nactus))) : full command buffer
148 
149  noise_com : (np.array((n,nactus))) : noise contribution for error breakdown
150 
151  alias_wfs_com : (np.array((n,nactus))) : aliasing estimation in the WFS direction
152 
153  tomo_com : (np.array((n,nactus))) : tomography error estimation
154 
155  H_com : (np.array((n,nactus))) : Filtered modes contribution for error breakdown
156 
157  trunc_com : (np.array((n,nactus))) : Truncature and sampling error of WFS
158 
159  bp_com : (np.array((n,nactus))) : Bandwidth error estimation on target
160 
161  mod_com : (np.array((n,nactus))) : Commanded modes expressed on the actuators
162 
163  fit : (float) : fitting (mean variance of the residual target phase after projection)
164 
165  SR : (float) : final strehl ratio returned by the simulation
166  """
167  if (error_flag):
168  # Initialize buffers for error breakdown
169  nactu = rtc.get_command(0).size
170  nslopes = rtc.get_centroids(0).size
171  com = np.zeros((n, nactu), dtype=np.float32)
172  noise_com = np.zeros((n, nactu), dtype=np.float32)
173  alias_wfs_com = np.copy(noise_com)
174  wf_com = np.copy(noise_com)
175  tomo_com = np.copy(noise_com)
176  trunc_com = np.copy(noise_com)
177  H_com = np.copy(noise_com)
178  mod_com = np.copy(noise_com)
179  bp_com = np.copy(noise_com)
180  fit = np.zeros(n)
181  # covm = np.zeros((nslopes,nslopes))
182  # covv = np.zeros((nactu,nactu))
183 
184  t0 = time.time()
185  for i in range(-10, n):
186  atm.move_atmos()
187 
188  if (config.p_controllers[0].type == b"geo"):
189  for t in range(config.p_target.ntargets):
190  tar.atmos_trace(t, atm, tel)
191  rtc.docontrol_geo(0, dms, tar, 0)
192  rtc.applycontrol(0, dms)
193  tar.dmtrace(0, dms)
194  else:
195  for t in range(config.p_target.ntargets):
196  tar.atmos_trace(t, atm, tel)
197  tar.dmtrace(t, dms)
198  for w in range(len(config.p_wfss)):
199  wfs.sensors_trace(w, "all", tel, atm, dms)
200  wfs.sensors_compimg(w)
201  rtc.docentroids(0)
202  rtc.docontrol(0)
203  #m = np.reshape(rtc.get_centroids(0),(nslopes,1))
204  #v = np.reshape(rtc.get_command(0),(nactu,1))
205  if (error_flag and i > -1):
206  #compute the error breakdown for this iteration
207  #covm += m.dot(m.T)
208  #covv += v.dot(v.T)
209  roket.computeBreakdown()
210  rtc.applycontrol(0, dms)
211 
212  if ((i + 1) % 100 == 0 and i > -1):
213  strehltmp = tar.get_strehl(0)
214  print(i + 1, "\t", strehltmp[0], "\t", strehltmp[1], "\t",
215  np.exp(-strehltmp[2]), "\t", np.exp(-strehltmp[3]))
216  t1 = time.time()
217  print(" loop execution time:", t1 - t0, " (", n, "iterations), ", (t1 - t0) / n,
218  "(mean) ", n / (t1 - t0), "Hz")
219  if (error_flag):
220  #Returns the error breakdown
221  SR2 = np.exp(-tar.get_strehl(0, comp_strehl=False)[3])
222  SR = tar.get_strehl(0, comp_strehl=False)[1]
223  #bp_com[-1,:] = bp_com[-2,:]
224  #SR = tar.get_strehl(0,comp_strehl=False)[1]
225  return SR, SR2
226 
227 
228 def preloop(n):
229  """
230  Performs the main AO loop for n interations. First, initialize buffers
231  for error breakdown computations. Then, at the end of each iteration, just
232  before applying the new DM shape, calls the error_breakdown function.
233 
234  :param n: (int) : number of iterations
235 
236  :return:
237  com : (np.array((n,nactus))) : full command buffer
238 
239  noise_com : (np.array((n,nactus))) : noise contribution for error breakdown
240 
241  alias_wfs_com : (np.array((n,nactus))) : aliasing estimation in the WFS direction
242 
243  tomo_com : (np.array((n,nactus))) : tomography error estimation
244 
245  H_com : (np.array((n,nactus))) : Filtered modes contribution for error breakdown
246 
247  trunc_com : (np.array((n,nactus))) : Truncature and sampling error of WFS
248 
249  bp_com : (np.array((n,nactus))) : Bandwidth error estimation on target
250 
251  mod_com : (np.array((n,nactus))) : Commanded modes expressed on the actuators
252 
253  fit : (float) : fitting (mean variance of the residual target phase after projection)
254 
255  SR : (float) : final strehl ratio returned by the simulation
256  """
257  for i in range(0, n):
258  atm.move_atmos()
259 
260  if (config.p_controllers[0].type == b"geo"):
261  for t in range(config.p_target.ntargets):
262  tar.atmos_trace(t, atm, tel)
263  rtc.docontrol_geo(0, dms, tar, 0)
264  rtc.applycontrol(0, dms)
265  else:
266  for t in range(config.p_target.ntargets):
267  tar.atmos_trace(t, atm, tel)
268  for w in range(len(config.p_wfss)):
269  wfs.sensors_trace(w, "all", tel, atm, dms)
270  wfs.sensors_compimg(w)
271  rtc.docentroids(0)
272  rtc.docontrol(0)
273 
274  rtc.applycontrol(0, dms)
275 
276 
277 
283 def compute_btt():
284  IF = rtc.get_IFsparse(1).T
285  N = IF.shape[0]
286  n = IF.shape[1]
287  #T = IF[:,-2:].copy()
288  T = rtc.get_IFtt(1)
289  #IF = IF[:,:n-2]
290  n = IF.shape[1]
291 
292  delta = IF.T.dot(IF).toarray() / N
293 
294  # Tip-tilt + piston
295  Tp = np.ones((T.shape[0], T.shape[1] + 1))
296  Tp[:, :2] = T.copy() #.toarray()
297  deltaT = IF.T.dot(Tp) / N
298  # Tip tilt projection on the pzt dm
299  tau = np.linalg.inv(delta).dot(deltaT)
300 
301  # Famille generatrice sans tip tilt
302  G = np.identity(n)
303  tdt = tau.T.dot(delta).dot(tau)
304  subTT = tau.dot(np.linalg.inv(tdt)).dot(tau.T).dot(delta)
305  G -= subTT
306 
307  # Base orthonormee sans TT
308  gdg = G.T.dot(delta).dot(G)
309  U, s, V = np.linalg.svd(gdg)
310  U = U[:, :U.shape[1] - 3]
311  s = s[:s.size - 3]
312  L = np.identity(s.size) / np.sqrt(s)
313  B = G.dot(U).dot(L)
314 
315  # Rajout du TT
316  TT = T.T.dot(T) / N #.toarray()/N
317  Btt = np.zeros((n + 2, n - 1))
318  Btt[:B.shape[0], :B.shape[1]] = B
319  mini = 1. / np.sqrt(TT)
320  mini[0, 1] = 0
321  mini[1, 0] = 0
322  Btt[n:, n - 3:] = mini
323 
324  # Calcul du projecteur actus-->modes
325  delta = np.zeros((n + T.shape[1], n + T.shape[1]))
326  #IF = rtc.get_IFsparse(1).T
327  delta[:-2, :-2] = IF.T.dot(IF).toarray() / N
328  delta[-2:, -2:] = T.T.dot(T) / N
329  P = Btt.T.dot(delta)
330 
331  return Btt.astype(np.float32), P.astype(np.float32)
332 
333 
334 def compute_cmatWithBtt(Btt, nfilt):
335  D = rtc.get_imat(0)
336  #D = ao.imat_geom(wfs,config.p_wfss,config.p_controllers[0],dms,config.p_dms,meth=0)
337  # Filtering on Btt modes
338  Btt_filt = np.zeros((Btt.shape[0], Btt.shape[1] - nfilt))
339  Btt_filt[:, :Btt_filt.shape[1] - 2] = Btt[:, :Btt.shape[1] - (nfilt + 2)]
340  Btt_filt[:, Btt_filt.shape[1] - 2:] = Btt[:, Btt.shape[1] - 2:]
341 
342  # Modal interaction basis
343  Dm = D.dot(Btt_filt)
344  # Direct inversion
345  Dmp = np.linalg.inv(Dm.T.dot(Dm)).dot(Dm.T)
346  # Command matrix
347  cmat = Btt_filt.dot(Dmp)
348 
349  return Dm.astype(np.float32), cmat.astype(np.float32)
350 
351 
352 def compute_cmatWithBtt2(Btt, nfilt):
353  D = rtc.get_imat(0)
354 
355  # Modal interaction basis
356  Dm = D.dot(Btt)
357  # Filtering on modal imat
358  DmtDm = Dm.T.dot(Dm)
359  U, s, V = np.linalg.svd(DmtDm)
360  s = 1. / s
361  s[s.shape[0] - nfilt - 2:s.shape[0] - 2] = 0.
362  DmtDm1 = U.dot(np.diag(s)).dot(U.T)
363  Dmp = DmtDm1.dot(Dm.T)
364  # Command matrix
365  cmat = Btt.dot(Dmp)
366 
367  return Dm.astype(np.float32), cmat.astype(np.float32)
368 
369 
370 
377 
378 
379 def cov_cor(P, noise, trunc, alias, H, bp, tomo):
380  cov = np.zeros((6, 6))
381  bufdict = {
382  "0": noise.T,
383  "1": trunc.T,
384  "2": alias.T,
385  "3": H.T,
386  "4": bp.T,
387  "5": tomo.T
388  }
389  for i in range(cov.shape[0]):
390  for j in range(cov.shape[1]):
391  if (j >= i):
392  tmpi = P.dot(bufdict[str(i)])
393  tmpj = P.dot(bufdict[str(j)])
394  cov[i, j] = np.sum(
395  np.mean(tmpi * tmpj, axis=1) -
396  np.mean(tmpi, axis=1) * np.mean(tmpj, axis=1))
397  else:
398  cov[i, j] = cov[j, i]
399 
400  s = np.reshape(np.diag(cov), (cov.shape[0], 1))
401  sst = np.dot(s, s.T)
402  cor = cov / np.sqrt(sst)
403 
404  return cov, cor
405 
406 
407 
413 
414 
415 def save_it(filename):
416  IF = rtc.get_IFsparse(1)
417  TT = rtc.get_IFtt(1)
418  noise_com = roket.getContributor("noise")
419  trunc_com = roket.getContributor("nonlinear")
420  alias_wfs_com = roket.getContributor("aliasing")
421  H_com = roket.getContributor("filtered")
422  bp_com = roket.getContributor("bandwidth")
423  tomo_com = roket.getContributor("tomo")
424  fit = roket.getContributor("fitting")
425 
426  tmp = (config.p_geom._ipupil.shape[0] -
427  (config.p_dms[0]._n2 - config.p_dms[0]._n1 + 1)) / 2
428  tmp_e0 = config.p_geom._ipupil.shape[0] - tmp
429  tmp_e1 = config.p_geom._ipupil.shape[1] - tmp
430  pup = config.p_geom._ipupil[tmp:tmp_e0, tmp:tmp_e1]
431  indx_pup = np.where(pup.flatten() > 0)[0].astype(np.int32)
432  dm_dim = config.p_dms[0]._n2 - config.p_dms[0]._n1 + 1
433  cov, cor = cov_cor(P, noise_com, trunc_com, alias_wfs_com, H_com, bp_com, tomo_com)
434  psf = tar.get_image(0, "le", fluxNorm=False)
435  psfortho = roket.get_tar_imageortho()
436  covv = roket.get_covv()
437  covm = roket.get_covm()
438 
439  fname = "/home/fferreira/Data/" + filename
440  pdict = {
441  "noise": noise_com.T,
442  "aliasing": alias_wfs_com.T,
443  "tomography": tomo_com.T,
444  "filtered modes": H_com.T,
445  "non linearity": trunc_com.T,
446  "bandwidth": bp_com.T,
447  "P": P,
448  "Btt": Btt,
449  "IF.data": IF.data,
450  "IF.indices": IF.indices,
451  "IF.indptr": IF.indptr,
452  "TT": TT,
453  "dm_dim": dm_dim,
454  "indx_pup": indx_pup,
455  "fitting": fit,
456  "SR": SR,
457  "SR2": SR2,
458  "cov": cov,
459  "cor": cor,
460  "psfortho": psfortho,
461  "covm": covm,
462  "covv": covv
463  }
464  h5u.save_h5(fname, "psf", config, psf)
465  #h5u.writeHdf5SingleDataset(fname,com.T,datasetName="com")
466  for k in list(pdict.keys()):
467  h5u.save_hdf5(fname, k, pdict[k])
468 
469 
470 
477 nfiltered = config.p_controllers[0].maxcond
478 niters = config.p_loop.niter
479 #config.p_loop.set_niter(niters)
480 Btt, P = compute_btt()
481 rtc.load_Btt(1, Btt.dot(Btt.T))
482 Dm, cmat = compute_cmatWithBtt(Btt, nfiltered)
483 rtc.set_cmat(0, cmat)
484 R = rtc.get_cmat(0)
485 imat = rtc.get_imat(0)
486 RD = np.dot(R, imat).astype(np.float32)
487 
488 gRD = (np.identity(RD.shape[0]) - config.p_controllers[0].gain * RD).astype(np.float32)
489 roket = ao.roket_init(rtc, wfs, tar, dms, tel, atm, 0, 1, Btt.shape[0], Btt.shape[1],
490  nfiltered, niters, Btt, P, gRD, RD)
491 #diagRD = np.diag(gRD)
492 #gRD = np.diag(diagRD)
493 #gRD=np.diag(gRD)
494 
495 #imat_geom = ao.imat_geom(wfs,config.p_wfss,config.p_controllers[0],dms,config.p_dms,meth=0)
496 #RDgeom = np.dot(R,imat_geom)
497 preloop(1000)
498 SR, SR2 = loop(niters)
499 
500 save_it(savename)
roket_gpu.preloop
def preloop(n)
Definition: roket_gpu.py:256
roket_gpu.cov_cor
def cov_cor(P, noise, trunc, alias, H, bp, tomo)
Definition: roket_gpu.py:379
roket_gpu.compute_cmatWithBtt
def compute_cmatWithBtt(Btt, nfilt)
Definition: roket_gpu.py:334
roket_gpu.compute_btt
def compute_btt()
Definition: roket_gpu.py:283
roket_gpu.loop
def loop(n)
Definition: roket_gpu.py:166
roket_gpu.save_it
def save_it(filename)
Definition: roket_gpu.py:415
roket_gpu.compute_cmatWithBtt2
def compute_cmatWithBtt2(Btt, nfilt)
Definition: roket_gpu.py:352