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