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