COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
basis.py
1 
37 
38 import numpy as np
39 
40 from shesha.sutra_wrap import Dms, Rtc_FFF as Rtc
41 
42 import shesha.config as conf
43 import shesha.constants as scons
44 
45 from scipy.sparse import csr_matrix
46 
47 from typing import List
48 from tqdm import trange
49 
50 
51 def compute_KL2V(p_controller: conf.Param_controller, dms: Dms, p_dms: list,
52  p_geom: conf.Param_geom, p_atmos: conf.Param_atmos,
53  p_tel: conf.Param_tel):
54  """ Compute the Karhunen-Loeve to Volt matrix
55  (transfer matrix between the KL space and volt space for a pzt dm)
56 
57  :parameters:
58 
59  p_controller: (Param_controller) : p_controller settings
60 
61  dms : (shesha_dms) : Dms object
62 
63  p_dms: (list of Param_dm) : dms settings
64 
65  p_geom : (Param_geom) : geometry parameters
66 
67  p_atmos : (Param_atmos) : atmos parameters
68 
69  p_tel : (Param_tel) : telescope parameters
70 
71  :return:
72 
73  KL2V : (np.array(np.float32,dim=2)) : KL to Volt matrix
74  """
75  ntotact = np.array([p_dms[i]._ntotact for i in range(len(p_dms))], dtype=np.int64)
76  KL2V = np.zeros((np.sum(ntotact), np.sum(ntotact)), dtype=np.float32)
77 
78  indx_act = 0
79  nTT = 0
80 
81  for i in range(p_controller.ndm.size):
82  ndm = p_controller.ndm[i]
83  if (p_dms[ndm].type == scons.DmType.PZT):
84  tmp = (p_geom._ipupil.shape[0] - (p_dms[ndm]._n2 - p_dms[ndm]._n1 + 1)) // 2
85  tmp_e0 = p_geom._ipupil.shape[0] - tmp
86  tmp_e1 = p_geom._ipupil.shape[1] - tmp
87  pup = p_geom._ipupil[tmp:tmp_e0, tmp:tmp_e1]
88  indx_valid = np.where(pup.flatten("F") > 0)[0].astype(np.int32)
89  p2m = p_tel.diam / p_geom.pupdiam
90  norm = -(p2m * p_tel.diam / (2 * p_atmos.r0))**(5. / 3)
91 
92  dms.d_dms[ndm].compute_KLbasis(p_dms[ndm]._xpos, p_dms[ndm]._ypos,
93  indx_valid, indx_valid.size, norm, 1.0)
94 
95  KL2V[indx_act:indx_act + ntotact[ndm], indx_act:indx_act + ntotact[ndm]] = \
96  np.fliplr(dms.d_dms[ndm].d_KLbasis)
97  indx_act += ntotact[ndm]
98  elif (p_dms[ndm].type == scons.DmType.TT):
99  nTT += 1
100  if (p_controller.nmodes is not None and
101  p_controller.nmodes < KL2V.shape[1] - 2 * nTT):
102  KL2V = KL2V[:, :p_controller.nmodes]
103  else:
104  KL2V = KL2V[:, :KL2V.shape[1] - 2 * nTT]
105  if (nTT > 1):
106  raise ValueError("More than 1 TipTilt found! Stupid")
107  if (nTT != 0):
108  KL2V[:, :KL2V.shape[1] - 2] = KL2V[:, 2:]
109  KL2V[:, KL2V.shape[1] - 2:] = np.zeros((np.sum(ntotact), 2), dtype=np.float32)
110  KL2V[np.sum(ntotact) - 2:, KL2V.shape[1] - 2:] = np.identity(2, dtype=np.float32)
111 
112  return KL2V
113 
114 
115 def compute_dm_basis(g_dm, p_dm: conf.Param_dm, p_geom: conf.Param_geom):
116  """ Compute a the DM basis as a sparse matrix :
117  - push on each actuator
118  - get the corresponding dm shape
119  - apply pupil mask and store in a column
120 
121  :parameters:
122  g_dm: (Dm) : Dm object
123 
124  p_dm: (Param_dm) : dm settings
125 
126  p_geom: (Param_geom) : geom settings
127 
128  :return:
129 
130  IFbasis = (csr_matrix) : DM IF basis
131  """
132  tmp = (p_geom._ipupil.shape[0] - (p_dm._n2 - p_dm._n1 + 1)) // 2
133  tmp_e0 = p_geom._ipupil.shape[0] - tmp
134  tmp_e1 = p_geom._ipupil.shape[1] - tmp
135  pup = p_geom._ipupil[tmp:tmp_e0, tmp:tmp_e1]
136  indx_valid = np.where(pup.flatten("F") > 0)[0].astype(np.int32)
137 
138  #IFbasis = np.ndarray((indx_valid.size, p_dm._ntotact), dtype=np.float32)
139  for i in trange(p_dm._ntotact):
140  g_dm.reset_shape()
141  g_dm.comp_oneactu(i, 1.0)
142  shape = np.array(g_dm.d_shape)
143  IFvec = csr_matrix(shape.flatten("F")[indx_valid])
144  if (i == 0):
145  val = IFvec.data
146  col = IFvec.indices
147  row = np.append(0, IFvec.getnnz())
148  else:
149  val = np.append(val, IFvec.data)
150  col = np.append(col, IFvec.indices)
151  row = np.append(row, row[-1] + IFvec.getnnz())
152  g_dm.reset_shape()
153  IFbasis = csr_matrix((val, col, row))
154  return IFbasis
155 
156 
157 def compute_IFsparse(g_dm: Dms, p_dms: list, p_geom: conf.Param_geom):
158  """ Compute the influence functions of all DMs as a sparse matrix :
159  - push on each actuator
160  - get the corresponding dm shape
161  - apply pupil mask and store in a column
162 
163  :parameters:
164 
165  g_dm: (Dms) : Dms object
166 
167  p_dms: (Param_dms) : dms settings
168 
169  p_geom: (Param_geom) : geom settings
170 
171  :return:
172 
173  IFbasis = (csr_matrix) : DM IF basis
174  """
175  ndm = len(p_dms)
176  for i in range(ndm):
177  IFi = compute_dm_basis(g_dm.d_dms[i], p_dms[i], p_geom)
178  if (i == 0):
179  val = IFi.data
180  col = IFi.indices
181  row = IFi.indptr
182  else:
183  val = np.append(val, IFi.data)
184  col = np.append(col, IFi.indices)
185  row = np.append(row, row[-1] + IFi.indptr[1:])
186  IFsparse = csr_matrix((val, col, row))
187  return IFsparse
188 
189 
190 def command_on_Btt(rtc: Rtc, dms: Dms, p_dms: list, p_geom: conf.Param_geom, nfilt: int):
191  """ Compute a command matrix in Btt modal basis (see error breakdown) and set
192  it on the sutra_rtc. It computes by itself the volts to Btt matrix.
193 
194  :parameters:
195 
196  rtc: (Rtc) : rtc object
197 
198  dms: (Dms): dms object
199 
200  p_dms: (list of Param_dm): dms settings
201 
202  p_geom: (Param_geom): geometry settings
203 
204  nfilt: (int): number of modes to filter
205  """
206 
207  IFs = compute_IFsparse(dms, p_dms, p_geom).T
208  n = IFs.shape[1]
209  IFtt = IFs[:, -2:].copy().toarray()
210  IFpzt = IFs[:, :n - 2]
211 
212  Btt, P = compute_btt(IFpzt, IFtt)
213  compute_cmat_with_Btt(rtc, Btt, nfilt)
214 
215 
216 def compute_cmat_with_Btt(rtc: Rtc, Btt: np.ndarray, nfilt: int):
217  """ Compute a command matrix on the Btt basis and load it in the GPU
218 
219  :parameters:
220 
221  rtc: (Rtc): rtc object
222 
223  Btt: (np.ndarray[ndim=2, dtype=np.float32]) : volts to Btt matrix
224 
225  nfilt: (int): number of modes to filter
226  """
227  D = np.array(rtc.d_control[0].d_imat)
228  #D = ao.imat_geom(wfs,config.p_wfss,config.p_controllers[0],dms,config.p_dms,meth=0)
229  # Filtering on Btt modes
230  Btt_filt = np.zeros((Btt.shape[0], Btt.shape[1] - nfilt))
231  Btt_filt[:, :Btt_filt.shape[1] - 2] = Btt[:, :Btt.shape[1] - (nfilt + 2)]
232  Btt_filt[:, Btt_filt.shape[1] - 2:] = Btt[:, Btt.shape[1] - 2:]
233 
234  # Modal interaction basis
235  Dm = D.dot(Btt_filt)
236  # Direct inversion
237  Dmp = np.linalg.inv(Dm.T.dot(Dm)).dot(Dm.T)
238  # Command matrix
239  cmat = Btt_filt.dot(Dmp)
240  rtc.d_control[0].set_cmat(cmat.astype(np.float32))
241 
242  return cmat.astype(np.float32)
243 
244 
245 def command_on_KL(rtc: Rtc, dms: Dms, p_controller: conf.Param_controller,
246  p_dms: List[conf.Param_dm], p_geom: conf.Param_geom,
247  p_atmos: conf.Param_atmos, p_tel: conf.Param_tel, nfilt: int):
248  """ Compute a command matrix in KL modal basis and set
249  it on the sutra_rtc. It computes by itself the volts to KL matrix.
250 
251  :parameters:
252 
253  rtc: (Rtc) : rtc object
254 
255  dms: (Dms): dms object
256 
257  p_dms: (list of Param_dm): dms settings
258 
259  p_geom: (Param_geom): geometry settings
260 
261  p_atmos : (Param_atmos) : atmos parameters
262 
263  p_tel : (Param_tel) : telescope parameters
264 
265  nfilt: (int): number of modes to filter
266  """
267  KL2V = compute_KL2V(p_controller, dms, p_dms, p_geom, p_atmos, p_tel)
268  return compute_cmat_with_KL(rtc, KL2V, nfilt)
269 
270 
271 def compute_cmat_with_KL(rtc: Rtc, KL2V: np.ndarray, nfilt: int):
272  """ Compute a command matrix on the KL basis and load it in the GPU
273 
274  :parameters:
275 
276  rtc: (Rtc): rtc object
277 
278  KL2V: (np.ndarray[ndim=2, dtype=np.float32]) : volts to KL matrix
279 
280  nfilt: (int): number of modes to filter
281  """
282  D = np.array(rtc.d_control[0].d_imat)
283  KL2V_filt = np.zeros((KL2V.shape[0], KL2V.shape[1] - nfilt))
284  KL2V_filt[:, :KL2V_filt.shape[1] - 2] = KL2V[:, :KL2V.shape[1] - (nfilt + 2)]
285  KL2V_filt[:, KL2V_filt.shape[1] - 2:] = KL2V[:, KL2V.shape[1] - 2:]
286 
287  # Modal interaction basis
288  Dm = D.dot(KL2V_filt)
289  # Direct inversion
290  Dmp = np.linalg.inv(Dm.T.dot(Dm)).dot(Dm.T)
291  # Command matrix
292  cmat = KL2V_filt.dot(Dmp)
293  rtc.d_control[0].set_cmat(cmat.astype(np.float32))
294 
295  return cmat.astype(np.float32)
296 
297 
298 def compute_fourier(nActu: int, pitch: float, actu_x_pos: np.ndarray,
299  actu_y_pos: np.ndarray, periodic='n'):
300  '''
301  Values you are looking for are:
302  config.p_dm0.nact
303  config.p_dm0._pitch
304  config.p_dm0._i1
305  config.p_dm0._j1
306  '''
307  # Offset xpos and ypos to get integer indices.
308  # Compute nact x nact x nact x nact Fourier basis # Periodic condition n / n-1 as option
309  # Extract values for actuators - flatten
310 
311  # Periodic may be 'n' or 'n-1'
312  # Will work only for squared pitch mirrors so far
313  if periodic == 'n':
314  n = nActu
315  elif periodic == 'n-1':
316  n = nActu - 1
317  else:
318  raise ValueError('periodic can only be "n" or "n-1" to set boundary condition.')
319  xnorm = (np.round((actu_x_pos - np.min(actu_x_pos)) / pitch).astype(np.int32)) % n
320  ynorm = (np.round((actu_y_pos - np.min(actu_y_pos)) / pitch).astype(np.int32)) % n
321  totActu = len(xnorm)
322 
323  data = np.zeros((n, n, n, n), dtype=np.float32)
324  for i in range(n):
325  for j in range(n):
326  data[i, j, i, j] = 1.
327 
328  data = np.fft.fftn(data, axes=(2, 3))
329 
330  takeSine = np.zeros((n, n), dtype=bool) # Where to take sine instead of cosine
331  takeSine[0, n // 2 + 1:] = 1 # Half of first line
332  if n % 2 == 0:
333  takeSine[n // 2, n // 2 + 1:] = 1 # Half of waffle line
334 
335  takeSine[n // 2 + 1:, :] = 1 # Bottom half
336 
337  data[takeSine] *= 1j
338  data = data.real
339 
340  # Extract actuators
341  actuPush = data[:, :, xnorm, ynorm]
342 
343  # Add a renorm ?
344 
345  return actuPush
346 
347 
348 def compute_btt(IFpzt, IFtt, influ_petal=None, return_delta=False):
349  """ Returns Btt to Volts and Volts to Btt matrices
350 
351  :parameters:
352 
353  IFpzt : (csr_matrix) : influence function matrix of pzt DM, sparse and arrange as (Npts in pup x nactus)
354 
355  IFtt : (np.ndarray(ndim=2,dtype=np.float32)) : Influence function matrix of the TT mirror arrange as (Npts in pup x 2)
356 
357  influ_petal : (np.ndarray) : Influence function matrix of M4 petals.
358  Default is None, if set, the Btt produced is also orthogonal
359  to petal modes, then only driven by petal DM
360 
361  return_delta : (bool, optional) : If True, returns delta instead of P. Default is False
362 
363  :returns:
364 
365  Btt : (np.ndarray(ndim=2,dtype=np.float32)) : Btt to Volts matrix
366 
367  P : (np.ndarray(ndim=2,dtype=np.float32)) : Volts to Btt matrix
368  """
369  N = IFpzt.shape[0]
370  n = IFpzt.shape[1]
371  if (n > N):
372  raise ValueError("Influence functions must be arrange as (Npts_pup x nactus)")
373 
374  delta = IFpzt.T.dot(IFpzt).toarray() / N
375 
376  # Tip-tilt + piston
377  Tp = np.ones((IFtt.shape[0], IFtt.shape[1] + 1))
378  Tp[:, :2] = IFtt.copy(
379  ) # THIS IS NOT A SPARSE OBJECT !!!!! STOP PUTTING .toarray() HERE PLEASE !!!
380  deltaT = IFpzt.T.dot(Tp) / N
381  # Tip tilt projection on the pzt dm
382  tau = np.linalg.inv(delta).dot(deltaT)
383  nfilt = 3 # Piston + tip + tilt
384 
385  if influ_petal is not None:
386  # Petal basis generation (orthogonal to global piston)
387  nseg = influ_petal.toarray().shape[0]
388  petal_modes = -1 / (nseg - 1) * np.ones((nseg, (nseg - 1)))
389  petal_modes += nseg / (nseg - 1) * np.eye(nseg)[:, 0:(
390  nseg - 1)] # petal modes within the petal dm space
391  tau_petal = IFpzt.T.dot(influ_petal).toarray().dot(petal_modes.T)
392  tau = np.concatenate((tau_petal, np.linalg.inv(delta).dot(deltaT)), axis=1)
393  nfilt = 8
394  # Famille generatrice sans tip tilt
395  G = np.identity(n)
396  tdt = tau.T.dot(delta).dot(tau)
397  subTT = tau.dot(np.linalg.inv(tdt)).dot(tau.T).dot(delta)
398  G -= subTT
399 
400  # Base orthonormee sans TT
401  gdg = G.T.dot(delta).dot(G)
402  U, s, V = np.linalg.svd(gdg)
403  U = U[:, :U.shape[1] - nfilt]
404  s = s[:s.size - nfilt]
405  L = np.identity(s.size) / np.sqrt(s)
406  B = G.dot(U).dot(L)
407 
408  # Rajout du TT
409  TT = IFtt.T.dot(IFtt) / N
410  Btt = np.zeros((n + 2, n - 1))
411  Btt[:B.shape[0], :B.shape[1]] = B
412  mini = 1. / np.sqrt(np.abs(TT))
413  mini[0, 1] = 0
414  mini[1, 0] = 0
415  Btt[n:, -2:] = mini
416  if influ_petal is not None:
417  Btt[:n, -7:-2] = tau_petal
418 
419  # Calcul du projecteur actus-->modes
420  Delta = np.zeros((n + IFtt.shape[1], n + IFtt.shape[1]))
421  #IFpzt = rtc.get_IFpztsparse(1).T
422  Delta[:-2, :-2] = delta
423  Delta[-2:, -2:] = TT
424  if return_delta:
425  P = Delta
426  else:
427  P = Btt.T.dot(Delta)
428 
429  return Btt.astype(np.float32), P.astype(np.float32)
shesha.ao.basis.compute_dm_basis
def compute_dm_basis(g_dm, conf.Param_dm p_dm, conf.Param_geom p_geom)
Compute a the DM basis as a sparse matrix :
Definition: basis.py:131
shesha.ao.basis.compute_KL2V
def compute_KL2V(conf.Param_controller p_controller, Dms dms, list p_dms, conf.Param_geom p_geom, conf.Param_atmos p_atmos, conf.Param_tel p_tel)
Compute the Karhunen-Loeve to Volt matrix (transfer matrix between the KL space and volt space for a ...
Definition: basis.py:72
shesha.sutra_wrap
Definition: sutra_wrap.py:1
shesha.constants
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
shesha.config
Parameter classes for COMPASS.
Definition: shesha/shesha/config/__init__.py:1
shesha.ao.basis.command_on_KL
def command_on_KL(Rtc rtc, Dms dms, conf.Param_controller p_controller, List[conf.Param_dm] p_dms, conf.Param_geom p_geom, conf.Param_atmos p_atmos, conf.Param_tel p_tel, int nfilt)
Compute a command matrix in KL modal basis and set it on the sutra_rtc.
Definition: basis.py:264
shesha.ao.basis.compute_IFsparse
def compute_IFsparse(Dms g_dm, list p_dms, conf.Param_geom p_geom)
Compute the influence functions of all DMs as a sparse matrix :
Definition: basis.py:174
shesha.ao.basis.command_on_Btt
def command_on_Btt(Rtc rtc, Dms dms, list p_dms, conf.Param_geom p_geom, int nfilt)
Compute a command matrix in Btt modal basis (see error breakdown) and set it on the sutra_rtc.
Definition: basis.py:205
shesha.ao.basis.compute_cmat_with_Btt
def compute_cmat_with_Btt(Rtc rtc, np.ndarray Btt, int nfilt)
Compute a command matrix on the Btt basis and load it in the GPU.
Definition: basis.py:226
shesha.ao.basis.compute_fourier
def compute_fourier(int nActu, float pitch, np.ndarray actu_x_pos, np.ndarray actu_y_pos, periodic='n')
Definition: basis.py:305
shesha.ao.basis.compute_cmat_with_KL
def compute_cmat_with_KL(Rtc rtc, np.ndarray KL2V, int nfilt)
Compute a command matrix on the KL basis and load it in the GPU.
Definition: basis.py:281
shesha.ao.basis.compute_btt
def compute_btt(IFpzt, IFtt, influ_petal=None, return_delta=False)
Returns Btt to Volts and Volts to Btt matrices.
Definition: basis.py:368