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