COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
utilities.py
1 
37 
38 import importlib
39 import sys, os
40 
41 import numpy as np
42 
43 
44 def rebin(a, shape):
45  """
46  TODO: docstring
47 
48  """
49 
50  sh = shape[0], a.shape[0] // shape[0], shape[1], a.shape[1] // shape[1]
51  return np.mean(a.reshape(sh), axis=(1, 3))
52 
53 
54 def fft_goodsize(s):
55  """find best size for a fft from size s
56 
57  :parameters:
58 
59  s: (int) size
60  """
61  return 2**(int(np.log2(s)) + 1)
62 
63 
64 def bin2d(data_in, binfact):
65  """
66  Returns the input 2D array "array", binned with the binning factor "binfact".
67  The input array X and/or Y dimensions needs not to be a multiple of
68  "binfact"; The final/edge pixels are in effect replicated if needed.
69  This routine prepares the parameters and calls the C routine _bin2d.
70  The input array can be of type int, float or double.
71  Last modified: Dec 15, 2003.
72  Author: F.Rigaut
73  SEE ALSO: _bin2d
74 
75  :parmeters:
76 
77  data_in: (np.ndarray) : data to binned
78 
79  binfact: (int) : binning factor
80  """
81  if (binfact < 1):
82  raise ValueError("binfact has to be >= 1")
83 
84  nx = data_in.shape[0]
85  ny = data_in.shape[1]
86  fx = int(np.ceil(nx / float(binfact)))
87  fy = int(np.ceil(ny / float(binfact)))
88 
89  data_out = np.zeros((fx, fy), dtype=data_in.dtype)
90 
91  for i1 in range(fx):
92  for j1 in range(fy):
93  for i2 in range(binfact):
94  for j2 in range(binfact):
95  i = i1 * binfact + i2
96  j = j1 * binfact + j2
97  if (i >= nx):
98  i = nx - 1
99  if (j >= ny):
100  j = ny - 1
101  data_out[i1, j1] += data_in[i, j]
102 
103  return data_out
104 
105 
106 def pad_array(A, N):
107  """
108  TODO: docstring
109 
110  """
111 
112  S = A.shape
113  D1 = (N - S[0]) // 2
114  D2 = (N - S[1]) // 2
115  padded = np.zeros((N, N))
116  padded[D1:D1 + S[0], D2:D2 + S[1]] = A
117  return padded
118 
119 
120 def dist(dim, xc=-1, yc=-1):
121  """
122  TODO: docstring
123 
124  """
125 
126  if (xc < 0):
127  xc = int(dim / 2.)
128  if (yc < 0):
129  yc = int(dim / 2.)
130 
131  dx = np.tile(np.arange(dim) - xc, (dim, 1))
132  dy = np.tile(np.arange(dim) - yc, (dim, 1)).T
133 
134  d = np.sqrt(dx**2 + dy**2)
135  return d
136 
137 
138 def makegaussian(size, fwhm, xc=-1, yc=-1, norm=0):
139  """
140  Returns a centered gaussian of specified size and fwhm.
141  norm returns normalized 2d gaussian
142 
143  :param size: (int) :
144 
145  :param fwhm: (float) :
146 
147  :param xc: (float) : (optional) center position on x axis
148 
149  :param yc: (float) : (optional) center position on y axis
150 
151  :param norm: (int) : (optional) normalization
152  """
153  tmp = np.exp(-(dist(size, xc, yc) / (fwhm / 1.66))**2.)
154  if (norm > 0):
155  tmp = tmp / (fwhm**2. * 1.140075)
156  return tmp
157 
158 
159 def load_config_from_file(filename_path: str):
160  """
161  Load the parameters from the parameters file
162 
163  Parameters:
164  filename_path: (str): path to the parameters file
165 
166  Return:
167  config : (config) : a config module
168  """
169  path = os.path.dirname(os.path.abspath(filename_path))
170  filename = os.path.basename(filename_path)
171  name, ext = os.path.splitext(filename)
172 
173  if (ext == ".py"):
174  if (path not in sys.path):
175  sys.path.insert(0, path)
176 
177  return load_config_from_module(name)
178 
179  # exec("import %s as wao_config" % filename)
180  sys.path.remove(path)
181  elif importlib.util.find_spec(filename_path) is not None:
182  return load_config_from_module(filename_path)
183  else:
184  raise ValueError("Config file must be .py or a module")
185 
186 
187 def load_config_from_module(filepath: str):
188  """
189  Load the parameters from the parameters module
190 
191  Parameters:
192  filename_path: (str): path to the parameters file
193 
194  Return:
195  config : (config) : a config module
196  """
197  filename = filepath.split('.')[-1]
198  print("loading: %s" % filename)
199 
200  config = importlib.import_module(filepath)
201  del sys.modules[config.__name__] # Forced reload
202  config = importlib.import_module(filepath)
203 
204  if hasattr(config, 'par'):
205  config = getattr("config.par.par4bench", filename)
206 
207  # Set missing config attributes to None
208  if not hasattr(config, 'p_loop'):
209  config.p_loop = None
210  if not hasattr(config, 'p_geom'):
211  config.p_geom = None
212  if not hasattr(config, 'p_tel'):
213  config.p_tel = None
214  if not hasattr(config, 'p_atmos'):
215  config.p_atmos = None
216  if not hasattr(config, 'p_dms'):
217  config.p_dms = None
218  if not hasattr(config, 'p_targets'):
219  config.p_targets = None
220  if not hasattr(config, 'p_wfss'):
221  config.p_wfss = None
222  if not hasattr(config, 'p_centroiders'):
223  config.p_centroiders = None
224  if not hasattr(config, 'p_controllers'):
225  config.p_controllers = None
226 
227  if not hasattr(config, 'simul_name'):
228  config.simul_name = None
229 
230  return config
231 
232 def generate_square(radius: float, density: float = 1.):
233  """ Generate modulation points positions following a square pattern
234 
235  Parameters:
236  radius : (float) : half the length of a side in lambda/D
237 
238  density : (float), optional) : number of psf per lambda/D. Default is 1
239 
240  Return:
241  cx : (np.ndarray) : X-positions of the modulation points
242 
243  cy : (np.ndarray) : Y-positions of the modulation points
244  """
245  x = np.linspace(-radius, radius, 1 + 2 * int(radius * density))
246  cx, cy = np.meshgrid(x, x, indexing='ij')
247  cx = cx.flatten()
248  cy = cy.flatten()
249  return (cx, cy)
250 
251 def generate_circle(radius: float, density: float = 1.):
252  """ Generate modulation points positions following a circular pattern
253 s
254  Parameters:
255  radius : (float) : half the length of a side in lambda/D
256 
257  density : (float), optional) : number of psf per lambda/D. Default is 1
258 
259  Return:
260  cx : (np.ndarray) : X-positions of the modulation points
261 
262  cy : (np.ndarray) : Y-positions of the modulation points
263  """
264  cx, cy = generate_square(radius, density)
265  r = cx * cx + cy * cy <= radius**2
266  return (cx[r], cy[r])
267 
268  def generate_pseudo_source(radius: float, additional_psf=0, density=1.):
269  """ Used to generate a pseudo source for PYRWFS
270 
271  Parameters:
272  radius : (float) : TODO description
273 
274  additional_psf : (int) :TODO description
275 
276  density : (float, optional) :TODO description
277 
278  Return:
279  ox : TODO description & explicit naming
280 
281  oy : TODO description & explicit naming
282 
283  w : TODO description & explicit naming
284 
285  xc : TODO description & explicit naming
286 
287  yc : TODO description & explicit naming
288  """
289  struct_size = (1 + 2 * additional_psf)**2
290  center_x, center_y = generate_square(additional_psf, density)
291  center_weight = (1 + 2 * int(additional_psf * density))**2 * [1]
292  center_size = 1 + 2 * int(additional_psf * density)
293 
294  weight_edge = [(1 + 2 * int(radius * density) - center_size) // 2]
295  xc, yc = generate_circle(radius, density)
296  for k in range(additional_psf):
297  line_length = np.sum(yc == (k + 1))
298  print(line_length)
299  weight_edge.append((line_length - center_size) // 2)
300 
301  edge_dist = (radius + additional_psf) // 2
302  V_edge_x = []
303  V_edge_y = []
304  V_edge_weight = []
305  for m in [-1, 1]:
306  V_edge_x.append(0)
307  V_edge_y.append(m * edge_dist)
308  V_edge_weight.append(weight_edge[0])
309  for k, val in enumerate(weight_edge[1:]):
310  for l in [-1, 1]:
311  for m in [-1, 1]:
312  V_edge_x.append(l * (k + 1) * density)
313  V_edge_y.append(m * edge_dist)
314  V_edge_weight.append(val)
315  H_edge_x = []
316  H_edge_y = []
317  H_edge_weight = []
318  for m in [-1, 1]:
319  H_edge_x.append(m * edge_dist)
320  H_edge_y.append(0)
321  H_edge_weight.append(weight_edge[0])
322  for k, val in enumerate(weight_edge[1:]):
323  for l in [-1, 1]:
324  for m in [-1, 1]:
325  H_edge_x.append(m * edge_dist)
326  H_edge_y.append(l * (k + 1) * density)
327  H_edge_weight.append(val)
328  pup_cent_x = []
329  pup_cent_y = []
330  pup_cent_weight = 4 * [(len(xc) - 2 * np.sum(H_edge_weight) - struct_size) / 4]
331  pup_cent_dist = int(edge_dist // np.sqrt(2))
332  for l in [-1, 1]:
333  for m in [-1, 1]:
334  pup_cent_x.append(l * pup_cent_dist)
335  pup_cent_y.append(m * pup_cent_dist)
336  ox = np.concatenate((center_x, V_edge_x, H_edge_x, pup_cent_x))
337  oy = np.concatenate((center_y, V_edge_y, H_edge_y, pup_cent_y))
338  w = np.concatenate((center_weight, V_edge_weight, H_edge_weight,
339  pup_cent_weight))
340  return (ox, oy, w, xc, yc)
341 
342 def first_non_zero(array: np.ndarray, axis: int,
343  invalid_val: int = -1) -> np.ndarray:
344  """ Find the first non zero element of an array
345 
346  Parameters:
347  array : (np.ndarray) : input array
348 
349  axis : (int) : axis index
350 
351  invalid_val : (int, optional) : Default is -1
352 
353  Return:
354  non_zeros_pos : (np.ndarray) : Index of the first non-zero element
355  for each line or column following the axis
356  """
357  mask = array != 0
358  return np.where(mask.any(axis=axis), mask.argmax(axis=axis), invalid_val)
359 
360 # def rotate3d(im, ang, cx=-1, cy=-1, zoom=1.0):
361 # """Rotates an image of an angle "ang" (in DEGREES).
362 
363 # The center of rotation is cx,cy.
364 # A zoom factor can be applied.
365 
366 # (cx,cy) can be omitted :one will assume one rotates around the
367 # center of the image.
368 # If zoom is not specified, the default value of 1.0 is taken.
369 
370 # modif dg : allow to rotate a cube of images with one angle per image
371 
372 # :parameters:
373 
374 # im: (np.ndarray[ndim=3,dtype=np.float32_t]) : array to rotate
375 
376 # ang: (np.ndarray[ndim=1,dtype=np.float32_t]) : rotation angle (in degrees)
377 
378 # cx: (float) : (optional) rotation center on x axis (default: image center)
379 
380 # cy: (float) : (optional) rotation center on x axis (default: image center)
381 
382 # zoom: (float) : (opional) zoom factor (default =1.0)
383 # """
384 
385 # # TODO test it
386 # if (zoom == 0):
387 # zoom = 1.0
388 # if (ang.size == 1):
389 # if (zoom == 1.0 and ang[0] == 0.):
390 # return im
391 
392 # ang *= np.pi / 180
393 
394 # nx = im.shape[1]
395 # ny = im.shape[2]
396 
397 # if (cx < 0):
398 # cx = nx / 2 + 1
399 # if (cy < 0):
400 # cy = ny / 2 + 1
401 
402 # x = np.tile(np.arange(nx) - cx + 1, (ny, 1)).T / zoom
403 # y = np.tile(np.arange(ny) - cy + 1, (nx, 1)) / zoom
404 
405 # rx = np.zeros((nx, ny, ang.size)).astype(np.int64)
406 # ry = np.zeros((nx, ny, ang.size)).astype(np.int64)
407 # wx = np.zeros((nx, ny, ang.size)).astype(np.float32)
408 # wy = np.zeros((nx, ny, ang.size)).astype(np.float32)
409 
410 # ind = np.zeros((nx, ny, ang.size)).astype(np.int64)
411 
412 # imr = np.zeros((im.shape[0], im.shape[1], im.shape[2])).\
413 # astype(np.float32)
414 
415 # for i in range(ang.size):
416 # matrot = np.array([[np.cos(ang[i]), -np.sin(ang[i])],
417 # [np.sin(ang[i]), np.cos(ang[i])]], dtype=np.float32)
418 # wx[:, :, i] = x * matrot[0, 0] + y * matrot[1, 0] + cx
419 # wy[:, :, i] = x * matrot[0, 1] + y * matrot[1, 1] + cy
420 
421 # nn = np.where(wx < 1)
422 # wx[nn] = 1.
423 # nn = np.where(wy < 1)
424 # wy[nn] = 1.
425 
426 # nn = np.where(wx > (nx - 1))
427 # wx[nn] = nx - 1
428 # nn = np.where(wy > (ny - 1))
429 # wy[nn] = ny - 1
430 
431 # rx = wx.astype(np.int64) # partie entiere
432 # ry = wy.astype(np.int64)
433 # wx -= rx # partie fractionnaire
434 # wy -= ry
435 
436 # ind = rx + (ry - 1) * nx
437 # if (ang.size > 1):
438 # for i in range(ang.size):
439 # ind[:, :, i] += i * nx * ny
440 
441 # imr.flat = \
442 # (im.flatten()[ind.flatten()] *
443 # (1 - wx.flatten()) +
444 # im.flatten()[ind.flatten() + 1] * wx.flatten())\
445 # * (1 - wy.flatten()) + \
446 # (im.flatten()[ind.flatten() + nx] * (1 - wx.flatten()) +
447 # im.flatten()[ind.flatten() + nx + 1] * wx.flatten()) * wy.flatten()
448 
449 # return imr
450 
451 # def rotate(im, ang, cx=-1, cy=-1, zoom=1.0):
452 # """Rotates an image of an angle "ang" (in DEGREES).
453 
454 # The center of rotation is cx,cy.
455 # A zoom factor can be applied.
456 
457 # (cx,cy) can be omitted :one will assume one rotates around the
458 # center of the image.
459 # If zoom is not specified, the default value of 1.0 is taken.
460 
461 # :parameters:
462 
463 # im: (np.ndarray[ndim=3,dtype=np.float32_t]) : array to rotate
464 
465 # ang: (float) : rotation angle (in degrees)
466 
467 # cx: (float) : (optional) rotation center on x axis (default: image center)
468 
469 # cy: (float) : (optional) rotation center on x axis (default: image center)
470 
471 # zoom: (float) : (opional) zoom factor (default =1.0)
472 # """
473 # # TODO merge it with rotate3d or see if there is any np.rotate or other...
474 # if (zoom == 0):
475 # zoom = 1.0
476 # if (zoom == 1.0 and ang == 0.):
477 # return im
478 
479 # ang *= np.pi / 180
480 # nx = im.shape[1]
481 # ny = im.shape[2]
482 
483 # if (cx < 0):
484 # cx = nx / 2 + 1
485 # if (cy < 0):
486 # cy = ny / 2 + 1
487 
488 # x = np.tile(np.arange(nx) - cx + 1, (ny, 1)).T / zoom
489 # y = np.tile(np.arange(ny) - cy + 1, (nx, 1)) / zoom
490 
491 # ind = np.zeros((nx, ny, ang.size))
492 
493 # imr = np.zeros((im.shape[0], im.shape[1], im.shape[2])).\
494 # astype(np.float32)
495 
496 # matrot = np.array([[np.cos(ang), -np.sin(ang)], [np.sin(ang), np.cos(ang)]])
497 
498 # wx = x * matrot[0, 0] + y * matrot[1, 0] + cx
499 # wy = x * matrot[0, 1] + y * matrot[1, 1] + cy
500 
501 # wx = np.clip(wx, 1., nx - 1)
502 # wy = np.clip(wy, 1., ny - 1)
503 # wx, rx = np.modf(wx)
504 # wy, ry = np.modf(wy)
505 
506 # ind = rx + (ry - 1) * nx
507 
508 # imr = (im[ind] * (1 - wx) + im[ind + 1] * wx) * (1 - wy) + \
509 # (im[ind + nx] * (1 - wx) + im[ind + nx + 1] * wx) * wy
510 
511 # return imr
shesha.util.utilities.first_non_zero
np.ndarray first_non_zero(np.ndarray array, int axis, int invalid_val=-1)
Find the first non zero element of an array.
Definition: utilities.py:355
shesha.util.utilities.fft_goodsize
def fft_goodsize(s)
find best size for a fft from size s
Definition: utilities.py:60
shesha.util.utilities.makegaussian
def makegaussian(size, fwhm, xc=-1, yc=-1, norm=0)
Definition: utilities.py:152
shesha.util.utilities.generate_square
def generate_square(float radius, float density=1.)
Generate modulation points positions following a square pattern.
Definition: utilities.py:244
shesha.util.utilities.dist
def dist(dim, xc=-1, yc=-1)
TODO docstring.
Definition: utilities.py:124
shesha.util.utilities.pad_array
def pad_array(A, N)
TODO docstring.
Definition: utilities.py:110
shesha.util.utilities.load_config_from_file
def load_config_from_file(str filename_path)
Load the parameters from the parameters file.
Definition: utilities.py:168
shesha.util.utilities.generate_circle
def generate_circle(float radius, float density=1.)
Generate modulation points positions following a circular pattern s.
Definition: utilities.py:263
shesha.util.utilities.rebin
def rebin(a, shape)
TODO docstring.
Definition: utilities.py:48
shesha.util.utilities.bin2d
def bin2d(data_in, binfact)
Definition: utilities.py:80
shesha.util.utilities.load_config_from_module
def load_config_from_module(str filepath)
Load the parameters from the parameters module.
Definition: utilities.py:196