COMPASS  5.4.4
End-to-end AO simulation tool using GPU acceleration
imats.py
1 
37 
38 import numpy as np # type: ignore
39 import time
40 from typing import List # Mypy checker
41 from rich.progress import track # type: ignore
42 
43 import shesha.config as conf
44 import shesha.constants as scons
45 import shesha.init.lgs_init as lgs
46 import shesha.util.hdf5_util as h5u
47 
48 from shesha.sutra_wrap import Sensors, Dms, Rtc_FFF as Rtc
49 from shesha.constants import CONST
50 
51 from astropy.io import fits
52 
53 
54 def imat_geom(wfs: Sensors, dms: Dms, p_wfss: List[conf.Param_wfs],
55  p_dms: List[conf.Param_dm], p_controller: conf.Param_controller,
56  meth: int = 0) -> np.ndarray:
57  """ Compute the interaction matrix with a geometric method
58 
59  Args:
60 
61  wfs: (Sensors) : Sensors object
62 
63  dms: (Dms) : Dms object
64 
65  p_wfss: (list of Param_wfs) : wfs settings
66 
67  p_dms: (list of Param_dm) : dms settings
68 
69  p_controller: (Param_controller) : controller settings
70 
71  meth: (int) : (optional) method type (0 or 1)
72 
73  """
74 
75  nwfs = p_controller.nwfs.size
76  ndm = p_controller.ndm.size
77  imat_size1 = 0
78  imat_size2 = 0
79 
80  for dm in dms.d_dms:
81  dm.reset_shape()
82 
83  for nw in range(nwfs):
84  nm = p_controller.nwfs[nw]
85  imat_size1 += p_wfss[nm]._nvalid * 2
86 
87  for nmc in range(ndm):
88  nm = p_controller.ndm[nmc]
89  imat_size2 += p_dms[nm]._ntotact
90 
91  imat_cpu = np.zeros((imat_size1, imat_size2), dtype=np.float32)
92  ind = 0
93  cc = 0
94  print("Doing imat geom...")
95  for nmc in range(ndm):
96  nm = p_controller.ndm[nmc]
97  dms.d_dms[nm].reset_shape()
98  for i in track(range(p_dms[nm]._ntotact), description="DM%d" % nmc):
99  dms.d_dms[nm].comp_oneactu(i, p_dms[nm].push4imat)
100  nslps = 0
101  for nw in range(nwfs):
102  n = p_controller.nwfs[nw]
103  wfs.d_wfs[n].d_gs.raytrace(dms, rst=1)
104  wfs.d_wfs[n].slopes_geom(meth)
105  imat_cpu[nslps:nslps + p_wfss[n]._nvalid * 2, ind] = np.array(
106  wfs.d_wfs[n].d_slopes)
107  nslps += p_wfss[n]._nvalid * 2
108  imat_cpu[:, ind] = imat_cpu[:, ind] / p_dms[nm].push4imat
109  ind = ind + 1
110  cc = cc + 1
111  dms.d_dms[nm].reset_shape()
112 
113  return imat_cpu
114 
115 
116 def imat_init(ncontrol: int, rtc: Rtc, dms: Dms, p_dms: list, wfs: Sensors, p_wfss: list,
117  p_tel: conf.Param_tel, p_controller: conf.Param_controller, M2V=None,
118  dataBase: dict = {}, use_DB: bool = False) -> None:
119  """ Initialize and compute the interaction matrix on the GPU
120 
121  Args:
122 
123  ncontrol: (int) : controller's index
124 
125  rtc: (Rtc) : Rtc object
126 
127  dms: (Dms) : Dms object
128 
129  p_dms: (Param_dms) : dms settings
130 
131  wfs: (Sensors) : Sensors object
132 
133  p_wfss: (list of Param_wfs) : wfs settings
134 
135  p_tel: (Param_tel) : telescope settings
136 
137  p_controller: (Param_controller) : controller settings
138 
139  M2V:(np.array) : KL_matrix
140 
141  dataBase:(dict): (optional) dict containing paths to files to load
142 
143  use_DB:(bool) : (optional) use dataBase flag
144  """
145  # first check if wfs is using lgs
146  # if so, load new lgs spot, just for imat
147  for i in range(len(p_wfss)):
148  if (p_wfss[i].gsalt > 0):
149  # TODO: check that
150  save_profile = p_wfss[i].proftype
151  p_wfss[i].proftype = scons.ProfType.GAUSS1
152  lgs.prep_lgs_prof(p_wfss[i], i, p_tel, wfs, imat=1)
153 
154  if "imat" in dataBase:
155  imat = h5u.load_imat_from_dataBase(dataBase)
156  rtc.d_control[ncontrol].set_imat(imat)
157  else:
158  t0 = time.time()
159  if M2V is not None:
160  p_controller._M2V = M2V.copy()
161  rtc.do_imat_basis(ncontrol, dms, M2V.shape[1], M2V, p_controller.klpush, p_controller.kernconv4imat)
162  else:
163  rtc.do_imat(ncontrol, dms, p_controller.kernconv4imat)
164  print("done in %f s" % (time.time() - t0))
165  imat = np.array(rtc.d_control[ncontrol].d_imat)
166  if use_DB:
167  h5u.save_imat_in_dataBase(imat)
168  p_controller.set_imat(imat)
169 
170  # Restore original profile in lgs spots
171  for i in range(len(p_wfss)):
172  if (p_wfss[i].gsalt > 0):
173  p_wfss[i].proftype = save_profile
174  lgs.prep_lgs_prof(p_wfss[i], i, p_tel, wfs)
175 
176 
177 #write imat_ts:
178 # loop over ts directions
179 # change WFS offset to direction
180 # do imat geom
181 
182 
183 def imat_geom_ts_multiple_direction(wfs: Sensors, dms: Dms, p_ts: conf.Param_wfs,
184  p_dms: List[conf.Param_dm], p_geom: conf.Param_geom,
185  ind_TS: int, ind_dmseen: List, p_tel: conf.Param_tel,
186  x, y, meth: int = 0) -> np.ndarray:
187  """ Compute the interaction matrix with a geometric method for multiple truth sensors (with different direction)
188 
189  Args:
190  wfs: (Sensors) : Sensors object
191 
192  dms: (Dms) : Dms object
193 
194  p_ts: (Param_wfs) : truth sensor settings
195 
196  ind_TS: (int) : index of the truth sensor in Sensors (wfs)
197 
198  p_dms: (list of Param_dm) : dms settings
199 
200  ind_DMs: (list of int) : indices of used DMs
201 
202  p_controller: (Param_controller) : controller settings
203 
204  Kwargs:
205  meth: (int) : (optional) method type (0 or 1)
206  """
207  imat_size2 = 0
208  print("DMS_SEEN: ", ind_dmseen)
209  for nm in ind_dmseen:
210  imat_size2 += p_dms[nm]._ntotact
211  imat_cpu = np.ndarray((0, imat_size2))
212 
213  for i in track(range(x.size), description="TS pos"):
214  xpos = x[i]
215  ypos = y[i]
216  for k in ind_dmseen:
217  dims = p_dms[k]._n2 - p_dms[k]._n1 + 1
218  dim = p_geom._mpupil.shape[0]
219  if (dim < dims):
220  dim = dims
221  xoff = xpos * CONST.ARCSEC2RAD * \
222  p_dms[k].alt / p_tel.diam * p_geom.pupdiam
223  yoff = ypos * CONST.ARCSEC2RAD * \
224  p_dms[k].alt / p_tel.diam * p_geom.pupdiam
225  xoff = xoff + (dim - p_geom._n) / 2
226  yoff = yoff + (dim - p_geom._n) / 2
227  wfs.d_wfs[ind_TS].d_gs.remove_layer(p_dms[k].type, k)
228  wfs.d_wfs[ind_TS].d_gs.add_layer(p_dms[k].type, k, xoff, yoff)
229  imat_cpu = np.concatenate(
230  (imat_cpu, imat_geom_ts(wfs, dms, p_ts, ind_TS, p_dms, ind_dmseen,
231  meth)), axis=0)
232 
233  return imat_cpu
234 
235 
236 def imat_geom_ts(wfs: Sensors, dms: Dms, p_ts: conf.Param_wfs, ind_TS: int,
237  p_dms: List[conf.Param_dm], ind_DMs: List[int],
238  meth: int = 0) -> np.ndarray:
239  """ Compute the interaction matrix with a geometric method for a single truth sensor
240 
241  Args:
242  wfs: (Sensors) : Sensors object
243 
244  dms: (Dms) : Dms object
245 
246  p_ts: (Param_wfs) : truth sensor settings
247 
248  ind_TS: (int) : index of the truth sensor in Sensors (wfs)
249 
250  p_dms: (list of Param_dm) : dms settings
251 
252  ind_DMs: (list of int) : indices of used DMs
253 
254  p_controller: (Param_controller) : controller settings
255 
256  Kwargs:
257  meth: (int) : (optional) method type (0 or 1)
258  """
259 
260  #nwfs = 1 #p_controller.nwfs.size # as parameter list of indices for wfs if several ts (only 1 ts for now)
261  ndm = len(ind_DMs) #p_controller.ndm.size # as parameter list of indices of used dms
262  imat_size1 = p_ts._nvalid * 2 # as parameter (nvalid)
263  imat_size2 = 0
264 
265  # for nw in range(nwfs):
266  # nm = p_controller.nwfs[nw]
267  # imat_size1 += p_wfss[nm]._nvalid * 2
268 
269  for dm in dms.d_dms:
270  dm.reset_shape()
271 
272  imat_size2 = 0
273  for nm in ind_DMs:
274  imat_size2 += p_dms[nm]._ntotact
275 
276  imat_cpu = np.zeros((imat_size1, imat_size2), dtype=np.float64)
277  ind = 0
278  cc = 0
279  for nm in track(ind_DMs, description="imat geom DM"):
280  dms.d_dms[nm].reset_shape()
281  for i in track(range(p_dms[nm]._ntotact), description="imat geom actu"):
282  dms.d_dms[nm].comp_oneactu(i, p_dms[nm].push4imat)
283  wfs.d_wfs[ind_TS].d_gs.raytrace(dms, rst=1)
284  wfs.d_wfs[ind_TS].slopes_geom(meth)
285  imat_cpu[:, ind] = np.array(wfs.d_wfs[ind_TS].d_slopes)
286  imat_cpu[:, ind] = imat_cpu[:, ind] / p_dms[nm].push4imat
287  ind = ind + 1
288  cc = cc + 1
289  dms.d_dms[nm].reset_shape()
290 
291  return imat_cpu
292 
293 
294 def get_metaD(sup, p_wfs, TS_xpos=None, TS_ypos=None, ind_TS=-1, n_control=0):
295  """Create an interaction matrix for the current simulation given TS position
296 
297  Args:
298  sup : (CompassSupervisor) : current COMPASS simulation
299 
300  p_ts: (Param_wfs) : truth sensor settings
301 
302  TS_xpos : np.ndarray : TS position (x axis)
303 
304  TS_ypos : np.ndarray : TS position (y axis)
305 
306  ind_TS: (int) : index of the truth sensor in Sensors (wfs)
307 
308  n_control : (int) : index of the controller
309 
310  :return:
311  metaD : np.ndarray :interaction matrix
312  """
313  if (TS_xpos is None):
314  TS_xpos = np.array([t.xpos for t in sup.config.p_wfs_ts])
315  elif (isinstance(TS_xpos, list)):
316  TS_xpos = np.array(TS_xpos)
317  elif (isinstance(TS_xpos, int) or isinstance(TS_xpos, float)):
318  TS_xpos = np.array([TS_xpos]).astype(np.float32)
319  if (TS_xpos.size < 1):
320  TS_xpos = np.zeros((1))
321 
322  if (TS_ypos is None):
323  TS_ypos = np.array([t.ypos for t in sup.config.p_wfs_ts])
324  elif (isinstance(TS_ypos, list)):
325  TS_ypos = np.array(TS_ypos)
326  elif (isinstance(TS_ypos, int) or isinstance(TS_ypos, float)):
327  TS_ypos = np.array([TS_ypos]).astype(np.float32)
328  if (TS_ypos.size < 1):
329  TS_ypos = np.zeros((1))
330 
331  return imat_geom_ts_multiple_direction(sup.wfs._wfs, sup.dms._dms, p_wfs,
332  sup.config.p_dms, sup.config.p_geom, ind_TS,
333  sup.config.p_controllers[n_control].ndm,
334  sup.config.p_tel, TS_xpos, TS_ypos)
None imat_init(int ncontrol, Rtc rtc, Dms dms, list p_dms, Sensors wfs, list p_wfss, conf.Param_tel p_tel, conf.Param_controller p_controller, M2V=None, dict dataBase={}, bool use_DB=False)
Initialize and compute the interaction matrix on the GPU.
Definition: imats.py:144
np.ndarray imat_geom_ts_multiple_direction(Sensors wfs, Dms dms, conf.Param_wfs p_ts, List[conf.Param_dm] p_dms, conf.Param_geom p_geom, int ind_TS, List ind_dmseen, conf.Param_tel p_tel, x, y, int meth=0)
Compute the interaction matrix with a geometric method for multiple truth sensors (with different dir...
Definition: imats.py:206
np.ndarray imat_geom(Sensors wfs, Dms dms, List[conf.Param_wfs] p_wfss, List[conf.Param_dm] p_dms, conf.Param_controller p_controller, int meth=0)
Compute the interaction matrix with a geometric method.
Definition: imats.py:73
def get_metaD(sup, p_wfs, TS_xpos=None, TS_ypos=None, ind_TS=-1, n_control=0)
Create an interaction matrix for the current simulation given TS position.
Definition: imats.py:312
np.ndarray imat_geom_ts(Sensors wfs, Dms dms, conf.Param_wfs p_ts, int ind_TS, List[conf.Param_dm] p_dms, List[int] ind_DMs, int meth=0)
Compute the interaction matrix with a geometric method for a single truth sensor.
Definition: imats.py:258
Parameter classes for COMPASS.
Numerical constants for shesha and config enumerations for safe-typing.
Definition: constants.py:1
Initialization of a LGS in a Wfs object.
Definition: lgs_init.py:1
Functions for handling the database system.
Definition: hdf5_util.py:1