COMPASS  5.4.4
End-to-end AO simulation tool using GPU acceleration
wfsCompass.py
1 
37 from shesha.init.wfs_init import wfs_init
38 import shesha.util.utilities as util
39 import shesha.ao.wfs as wfs_util
40 from shesha.supervisor.components.sourceCompass import SourceCompass
41 import numpy as np
42 from typing import List
43 
45  """ WFS handler for compass simulation
46 
47  Attributes:
48  sources : (List) : List of SutraSource instances used for raytracing
49 
50  _wfs : (sutraWrap.Wfs) : SutraSensors instance
51 
52  _context : (carmaContext) : CarmaContext instance
53 
54  _config : (config module) : Parameters configuration structure module
55  """
56  def __init__(self, context, config, tel):
57  """ Initialize a wfsCompass component for wfs related supervision
58 
59  Args:
60  context : (carmaContext) : CarmaContext instance
61 
62  config : (config module) : Parameters configuration structure module
63 
64  tel : (TelescopeCompass) : A TelescopeCompass instance
65  """
66  self._context_context = context
67  self._config_config = config # Parameters configuration coming from supervisor init
68  print("->wfs init")
69  self._wfs_wfs = wfs_init(self._context_context, tel._tel, self._config_config.p_wfss,
70  self._config_config.p_tel, self._config_config.p_geom, self._config_config.p_dms,
71  self._config_config.p_atmos)
72  self.sourcessourcessources = [wfs.d_gs for wfs in self._wfs_wfs.d_wfs]
73 
74  def get_wfs_image(self, wfs_index : int) -> np.ndarray:
75  """ Get an image from the WFS (wfs[0] by default), or from the centroider handling the WFS
76  to get the calibrated image
77 
78  Args:
79  wfs_index : (int) : index of the WFS (or the centroider) to request an image
80 
81  Returns:
82  image : (np.ndarray) : WFS image
83  """
84  if self._config_config.p_wfss[wfs_index].fakecam:
85  return np.array(self._wfs_wfs.d_wfs[wfs_index].d_camimg)
86  else:
87  return np.array(self._wfs_wfs.d_wfs[wfs_index].d_binimg)
88 
89  def set_wfs_image(self, wfs_index : int, img: np.ndarray):
90  """ Set an image in the WFS (wfs[0] by default)
91  Args:
92  wfs_index : (int) : index of the WFS (or the centroider) to request an image
93  img: (np.ndarray) : Image to set
94  """
95  self._wfs_wfs.d_wfs[wfs_index].set_binimg(img, img.size)
96 
97  def set_pyr_modulation_points(self, wfs_index : int, cx: np.ndarray, cy: np.ndarray,
98  *, weights: np.ndarray = None) -> None:
99  """ Set pyramid modulation positions
100 
101  Args:
102  wfs_index : (int) : WFS index
103 
104  cx : (np.ndarray) : X positions of the modulation points [arcsec]
105 
106  cy : (np.ndarray) : Y positions of the modulation points [arcsec]
107 
108  Kwargs:
109  weights : (np.ndarray) : Weights to apply on each modulation point contribution
110  """
111  pyr_npts = len(cx)
112  pwfs = self._config_config.p_wfss[wfs_index]
113  pwfs.set_pyr_npts(pyr_npts)
114  pwfs.set_pyr_cx(cx)
115  pwfs.set_pyr_cy(cy)
116  if weights is None:
117  self._wfs_wfs.d_wfs[wfs_index].set_pyr_modulation_points(cx, cy, pyr_npts)
118  else:
119  self._wfs_wfs.d_wfs[wfs_index].set_pyr_modulation_points(
120  cx, cy, weights, pyr_npts)
121 
122  def set_pyr_modulation_ampli(self, wfs_index: int, pyr_mod: float) -> float:
123  """ Set pyramid circular modulation amplitude value - in lambda/D units.
124 
125  Compute new modulation points corresponding to the new amplitude value
126  and upload them.
127  WARNING : if you are using slopes-based centroider with the PWFS,
128  also update the centroider scale (rtc.set_scale) with the returned
129  value
130 
131  Args:
132  wfs_index : (int) : WFS index
133 
134  pyr_mod : (float) : new pyramid modulation amplitude value
135 
136  Returns:
137  scale : (float) : scale factor
138  """
139  p_wfs = self._config_config.p_wfss[wfs_index]
140 
141  cx, cy, scale, pyr_npts = wfs_util.comp_new_pyr_ampl(wfs_index, pyr_mod,
142  self._config_config.p_wfss,
143  self._config_config.p_tel)
144  p_wfs.set_pyr_ampl(pyr_mod)
145  self.set_pyr_modulation_pointsset_pyr_modulation_points(wfs_index, cx, cy)
146 
147  if (len(p_wfs._halfxy.shape) == 2):
148  print("PYR modulation set to: %f L/D using %d points" % (pyr_mod, pyr_npts))
149  elif (len(p_wfs._halfxy.shape) == 3):
150  newhalfxy = np.tile(p_wfs._halfxy[0, :, :], (pyr_npts, 1, 1))
151  print("Loading new modulation arrays")
152  self._wfs_wfs.d_wfs[wfs_index].set_phalfxy(
153  np.exp(1j * newhalfxy).astype(np.complex64).T)
154  print("Done. PYR modulation set to: %f L/D using %d points" % (pyr_mod,
155  pyr_npts))
156  else:
157  raise ValueError("Error unknown p_wfs._halfxy shape")
158 
159  return scale
160 
161  def set_pyr_multiple_stars_source(self, wfs_index: int, coords: List,
162  *, weights: List = None, pyr_mod: float = 3.,
163  niters: int = None) -> None:
164  """ Sets the Pyramid modulation points with a multiple star system
165 
166  Args:
167  wfs_index : (int) : WFS index
168 
169  coords : (list) : list of couples of length n, coordinates of the n stars in lambda/D
170 
171  Kwargs:
172  weights : (list) : list of weights to apply on each modulation points. Default is None
173 
174  pyr_mod : (float): modulation amplitude of the pyramid in lambda/D. Default is 3
175 
176  niters : (int) : number of iteration. Default is None
177  """
178  if niters is None:
179  perim = pyr_mod * 2 * np.pi
180  niters = int((perim // 4 + 1) * 4)
181  print(niters)
182  scale_circ = self._config_config.p_wfss[wfs_index]._pyr_scale_pos * pyr_mod
183  temp_cx = []
184  temp_cy = []
185  for k in coords:
186  temp_cx.append(scale_circ * \
187  np.sin((np.arange(niters)) * 2. * np.pi / niters) + \
188  k[0] * self._config_config.p_wfss[wfs_index]._pyr_scale_pos)
189  temp_cy.append(scale_circ * \
190  np.cos((np.arange(niters)) * 2. * np.pi / niters) + \
191  k[1] * self._config_config.p_wfss[wfs_index]._pyr_scale_pos)
192  cx = np.concatenate(np.array(temp_cx))
193  cy = np.concatenate(np.array(temp_cy))
194  #Gives the arguments to the simulation
195  if weights is not None:
196  w = []
197  for k in weights:
198  w += niters * [k]
199  weights = np.array(w)
200  self.set_pyr_modulation_pointsset_pyr_modulation_points(wfs_index, cx, cy, weights=weights)
201 
202  def set_pyr_disk_source_hexa(self, wfs_index: int, radius: float) -> None:
203  """ Create disk object by packing PSF in a given radius, using hexagonal packing
204  and set it as modulation pattern
205 
206  There is no modulation
207 
208  Args:
209  wfs_index : (int) : WFS index
210 
211  radius : (float) : radius of the disk object in lambda/D
212  """
213  #Vectors used to generate the hexagonal paving
214  gen_xp, gen_yp = np.array([1,
215  0.]), np.array([np.cos(np.pi / 3),
216  np.sin(np.pi / 3)])
217  n = 1 + int(1.2 * radius)
218  mat_circ = []
219  for k in range(-n, n):
220  for l in range(-n, n):
221  coord = k * gen_xp + l * gen_yp
222  if np.sqrt(coord[0]**2 + coord[1]**2) <= radius:
223  mat_circ.append(coord)
224  mat_circ = np.array(mat_circ)
225  cx, cy = mat_circ[:, 0], mat_circ[:, 1]
226  self.set_pyr_modulation_pointsset_pyr_modulation_points(wfs_index, cx, cy)
227 
228 
229  def set_pyr_disk_source(self, wfs_index: int, radius: float, *, density: float = 1.) -> None:
230  """ Create disk object by packing PSF in a given radius, using square packing
231  and set it as modulation pattern
232 
233  There is no modulation
234 
235  Args:
236  wfs_index : (int) : WFS index
237 
238  radius : (float) : radius of the disk object in lambda/D
239 
240  Kwargs:
241  density : (float) : Spacing between the packed PSF in the disk object, in lambda/D.
242  Default is 1
243  """
244  cx, cy = util.generate_circle(radius, density)
245  cx = cx.flatten() * self._config_config.p_wfss[wfs_index]._pyr_scale_pos
246  cy = cy.flatten() * self._config_config.p_wfss[wfs_index]._pyr_scale_pos
247  self.set_pyr_modulation_pointsset_pyr_modulation_points(wfs_index, cx, cy)
248 
249  def set_pyr_square_source(self, wfs_index: int, radius: float, *, density: float = 1.) -> None:
250  """ Create a square object by packing PSF in a given radius, using square packing
251  and set it as modulation pattern
252 
253  There is no modulation
254 
255  Args:
256  wfs_index : (int) : WFS index
257 
258  radius : (float) : radius of the disk object in lambda/D
259 
260  Kwargs:
261  density : (float) : Spacing between the packed PSF in the disk object, in lambda/D.
262  Default is 1
263  """
264  cx, cy = util.generate_square(radius, density)
265  cx = cx.flatten() * self._config_config.p_wfss[wfs_index]._pyr_scale_pos
266  cy = cy.flatten() * self._config_config.p_wfss[wfs_index]._pyr_scale_pos
267  self.set_pyr_modulation_pointsset_pyr_modulation_points(wfs_index, cx, cy)
268 
269 
270  def set_pyr_pseudo_source(self, wfs_index: int, radius: float, *,
271  additional_psf: int = 0, density: float = 1.) -> None:
272  """ TODO : DESCRIPTION
273 
274  Args:
275  wfs_index : (int) : WFS index
276 
277  radius : (float) : TODO : DESCRIPTION
278 
279  Kwargs:
280  additional_psf : (int) : TODO : DESCRIPTION
281 
282  density : (float) :TODO : DESCRIPTION
283  """
284  cx, cy, weights, _, _ = util.generate_pseudo_source(radius, additional_psf,
285  density)
286  cx = cx.flatten() * self._config_config.p_wfss[wfs_index]._pyr_scale_pos
287  cy = cy.flatten() * self._config_config.p_wfss[wfs_index]._pyr_scale_pos
288  self.set_pyr_modulation_pointsset_pyr_modulation_points(wfs_index, cx, cy, weights=weights)
289 
290  def set_fourier_mask(self, wfs_index : int, new_mask: np.ndarray) -> None:
291  """ Set a mask in the Fourier Plane of the given WFS
292 
293  Args:
294  wfs_index : (int) : WFS index
295 
296  new_mask : (ndarray) : mask to set
297  """
298  if new_mask.shape != self._config_config.p_wfss[wfs_index].get_halfxy().shape:
299  print('Error : mask shape should be {}'.format(
300  self._config_config.p_wfss[wfs_index].get_halfxy().shape))
301  else:
302  self._wfs_wfs.d_wfs[wfs_index].set_phalfxy(
303  np.exp(1j * np.fft.fftshift(new_mask)).astype(np.complex64).T)
304 
305  def set_noise(self, wfs_index : int, noise: float, *, seed: int = 1234) -> None:
306  """ Set noise value of WFS wfs_index
307 
308  Args:
309  wfs_index : (int) : WFS index
310 
311  noise : (float) : readout noise value in e-
312 
313  Kwargs:
314  seed : (int) : RNG seed. The seed used will be computed as seed + wfs_index
315  Default is 1234
316  """
317  self._wfs_wfs.d_wfs[wfs_index].set_noise(noise, int(seed + wfs_index))
318  print("Noise set to: %f on WFS %d" % (noise, wfs_index))
319 
320  def set_gs_mag(self, wfs_index : int, mag : float) -> None:
321  """ Change the guide star magnitude for the given WFS
322 
323  Args:
324  wfs_index : (int) : WFS index
325 
326  mag : (float) : New magnitude of the guide star
327  """
328  wfs = self._wfs_wfs.d_wfs[wfs_index]
329  if (self._config_config.p_wfss[0].type == "pyrhr"):
330  r = wfs.comp_nphot(self._config_config.p_loop.ittime,
331  self._config_config.p_wfss[wfs_index].optthroughput,
332  self._config_config.p_tel.diam, self._config_config.p_tel.cobs,
333  self._config_config.p_wfss[wfs_index].zerop, mag)
334  else:
335  r = wfs.comp_nphot(self._config_config.p_loop.ittime,
336  self._config_config.p_wfss[wfs_index].optthroughput,
337  self._config_config.p_tel.diam, self._config_config.p_wfss[wfs_index].nxsub,
338  self._config_config.p_wfss[wfs_index].zerop, mag)
339  if (r == 0):
340  print("GS magnitude is now %f on WFS %d" % (mag, wfs_index))
341  self._config_config.p_wfss[wfs_index].set_gsmag(mag)
342  self._config_config.p_wfss[wfs_index]._nphotons = wfs.nphot
343 
344  def compute_wfs_image(self, wfs_index : int, *, noise: bool = True) -> None:
345  """ Computes the image produced by the WFS from its phase screen
346 
347  Args:
348  wfs_index : (int): WFS index
349 
350  Kwargs:
351  noise : (bool) : Flag to enable noise for image computation. Default is True
352  """
353  self._wfs_wfs.d_wfs[wfs_index].comp_image(noise)
354 
355  def reset_noise(self) -> None:
356  """ Reset all the WFS RNG to their original state
357  """
358  for wfs_index, p_wfs in enumerate(self._config_config.p_wfss):
359  self._wfs_wfs.d_wfs[wfs_index].set_noise(p_wfs.noise, 1234 + wfs_index)
360 
361  def reset_image(self, *, wfs_index : int = None):
362  """ Reset the WFS image
363 
364  Kwargs:
365  wfs_index : (int): WFS index. If not provided, will reset all WFS
366  """
367  if wfs_index is not None:
368  self._wfs_wfs.d_wfs[wfs_index].d_binimg.reset()
369  return
370 
371  for _,swfs in enumerate(self._wfs_wfs.d_wfs):
372  swfs.d_binimg.reset()
373 
374  def get_ncpa_wfs(self, wfs_index : int) -> np.ndarray:
375  """ Return the current NCPA phase screen of the WFS path
376 
377  Args:
378  wfs_index : (int) : Index of the WFS
379 
380  Returns:
381  ncpa : (np.ndarray) : NCPA phase screen
382  """
383  return np.array(self._wfs_wfs.d_wfs[wfs_index].d_gs.d_ncpa_phase)
384 
385  def get_wfs_phase(self, wfs_index : int) -> np.ndarray:
386  """ Return the WFS phase screen of WFS number wfs_index
387 
388  Args:
389  wfs_index : (int) : Index of the WFS
390 
391  Returns:
392  phase : (np.ndarray) : WFS phase screen
393  """
394  return np.array(self._wfs_wfs.d_wfs[wfs_index].d_gs.d_phase)
395 
396  def get_pyrhr_image(self, wfs_index : int) -> np.ndarray:
397  """ Get an high resolution image from the PWFS
398 
399  Args:
400  wfs_index : (int) : Index of the WFS
401 
402  Returns:
403  image : (np.ndarray) : PWFS high resolution image
404 
405  """
406  return np.array(self._wfs_wfs.d_wfs[wfs_index].d_hrimg)
407 
408  def set_ncpa_wfs(self, wfs_index : int, ncpa: np.ndarray) -> None:
409  """ Set the additional fixed NCPA phase in the WFS path.
410  ncpa must be of the same size of the mpupil support
411 
412  Args:
413  wfs_index : (int) : WFS index
414 
415  ncpa : (ndarray) : NCPA phase screen to set [µm]
416  """
417  self._wfs_wfs.d_wfs[wfs_index].d_gs.set_ncpa(ncpa)
418 
419  def set_wfs_phase(self, wfs_index : int, phase : np.ndarray) -> None:
420  """ Set the phase screen seen by the WFS
421 
422  Args:
423  wfs_index : (int) : WFS index
424 
425  phase : (np.ndarray) : phase screen to set
426  """
427  self._wfs_wfs.d_wfs[wfs_index].d_gs.set_phase(phase)
428 
429  def set_wfs_pupil(self, wfs_index : int, pupil : np.ndarray) -> None:
430  """ Set the pupil seen by the WFS
431  Other pupils remain unchanged, i.e. DM and target can see an other
432  pupil than the WFS after this call.
433  <pupil> must have the same shape than p_geom._mpupil support
434 
435  Args:
436  wfs_index : (int) : WFS index
437 
438  pupil : (np.ndarray) : new pupil to set
439  """
440  old_mpup = self._config_config.p_geom._mpupil
441  dimx = old_mpup.shape[0]
442  dimy = old_mpup.shape[1]
443  if ((pupil.shape[0] != dimx) or (pupil.shape[1] != dimy)):
444  print("Error pupil shape on wfs %d must be: (%d,%d)" % (wfs_index, dimx,
445  dimy))
446  else:
447  self._wfs_wfs.d_wfs[wfs_index].set_pupil(pupil.copy())
448 
449  def get_pyr_focal_plane(self, wfs_index : int) -> np.ndarray:
450  """ Returns the psf on the top of the pyramid.
451  pyrhr WFS only
452 
453  Args:
454  wfs_index : (int) : WFS index
455 
456  Returns:
457  focal_plane : (np.ndarray) : psf on the top of the pyramid
458  """
459  return np.fft.fftshift(np.array(self._wfs_wfs.d_wfs[wfs_index].d_pyrfocalplane))
Source handler for compass simulation.
sources
(List) : List of SutraSource instances
WFS handler for compass simulation.
Definition: wfsCompass.py:48
def __init__(self, context, config, tel)
Definition: wfsCompass.py:76
np.ndarray get_wfs_phase(self, int wfs_index)
Return the WFS phase screen of WFS number wfs_index.
Definition: wfsCompass.py:410
None reset_noise(self)
Reset all the WFS RNG to their original state.
Definition: wfsCompass.py:374
np.ndarray get_pyr_focal_plane(self, int wfs_index)
Returns the psf on the top of the pyramid.
Definition: wfsCompass.py:475
None set_pyr_modulation_points(self, int wfs_index, np.ndarray cx, np.ndarray cy, *np.ndarray weights=None)
Set pyramid modulation positions.
Definition: wfsCompass.py:121
float set_pyr_modulation_ampli(self, int wfs_index, float pyr_mod)
Set pyramid circular modulation amplitude value - in lambda/D units.
Definition: wfsCompass.py:149
def set_wfs_image(self, int wfs_index, np.ndarray img)
Set an image in the WFS (wfs[0] by default)
Definition: wfsCompass.py:105
sources
(List) : List of SutraSource instances used for raytracing
Definition: wfsCompass.py:83
None set_pyr_square_source(self, int wfs_index, float radius, *float density=1.)
Create a square object by packing PSF in a given radius, using square packing and set it as modulatio...
Definition: wfsCompass.py:278
None set_pyr_multiple_stars_source(self, int wfs_index, List coords, *List weights=None, float pyr_mod=3., int niters=None)
Sets the Pyramid modulation points with a multiple star system.
Definition: wfsCompass.py:188
np.ndarray get_wfs_image(self, int wfs_index)
Get an image from the WFS (wfs[0] by default), or from the centroider handling the WFS to get the cal...
Definition: wfsCompass.py:94
np.ndarray get_pyrhr_image(self, int wfs_index)
Get an high resolution image from the PWFS.
Definition: wfsCompass.py:422
None set_pyr_disk_source_hexa(self, int wfs_index, float radius)
Create disk object by packing PSF in a given radius, using hexagonal packing and set it as modulation...
Definition: wfsCompass.py:223
None set_wfs_phase(self, int wfs_index, np.ndarray phase)
Set the phase screen seen by the WFS.
Definition: wfsCompass.py:443
None set_wfs_pupil(self, int wfs_index, np.ndarray pupil)
Set the pupil seen by the WFS Other pupils remain unchanged, i.e.
Definition: wfsCompass.py:456
None set_gs_mag(self, int wfs_index, float mag)
Change the guide star magnitude for the given WFS.
Definition: wfsCompass.py:344
None compute_wfs_image(self, int wfs_index, *bool noise=True)
Computes the image produced by the WFS from its phase screen.
Definition: wfsCompass.py:369
None set_pyr_disk_source(self, int wfs_index, float radius, *float density=1.)
Create disk object by packing PSF in a given radius, using square packing and set it as modulation pa...
Definition: wfsCompass.py:256
None set_noise(self, int wfs_index, float noise, *int seed=1234)
Set noise value of WFS wfs_index.
Definition: wfsCompass.py:333
None set_fourier_mask(self, int wfs_index, np.ndarray new_mask)
Set a mask in the Fourier Plane of the given WFS.
Definition: wfsCompass.py:312
None set_pyr_pseudo_source(self, int wfs_index, float radius, *int additional_psf=0, float density=1.)
TODO DESCRIPTION.
Definition: wfsCompass.py:298
def reset_image(self, *int wfs_index=None)
Reset the WFS image.
Definition: wfsCompass.py:383
None set_ncpa_wfs(self, int wfs_index, np.ndarray ncpa)
Set the additional fixed NCPA phase in the WFS path.
Definition: wfsCompass.py:433
np.ndarray get_ncpa_wfs(self, int wfs_index)
Return the current NCPA phase screen of the WFS path.
Definition: wfsCompass.py:399
On the fly modification of the WFS.
Definition: ao/wfs.py:1
Initialization of a Sensors object.
Definition: wfs_init.py:1
Basic utilities function.
Definition: utilities.py:1