COMPASS  5.0.0
End-to-end AO simulation tool using GPU acceleration
gamora.py
1 """
2 GAMORA (Gpu Accelerated Module fOr psf Reconstruction Algorithms)
3 
4 Python module for GPU accelerated PSF reconstruction using Vii functions and ROKET file
5 
6 Note: GPU devices used are hardcoded here. Change gpudevices if needed
7 """
8 import numpy as np
9 import matplotlib.pyplot as plt
10 import h5py
11 from shesha.sutra_wrap import carmaWrap_context, Gamora
12 from scipy.sparse import csr_matrix
13 from sys import stdout
14 import time
15 from guardians import drax
16 
17 plt.ion()
18 
19 gpudevices = np.array([0, 1, 2, 3], dtype=np.int32)
20 c = carmaWrap_context.get_instance_ngpu(gpudevices.size, gpudevices)
21 
22 
23 def psf_rec_Vii(filename, err=None, fitting=True, covmodes=None, cov=None):
24  """
25  PSF reconstruction using Vii functions with GPU acceleration.
26 
27  :parameters:
28  filename: (str): path to the ROKET file
29  err: (np.ndarray[ndim=2, dtype=np.float32]): (optionnal) Buffers of command error
30  fitting: (bool): (optional) Add the fitting error to the PSF or not (True by default)
31  covmodes: (np.ndarray[ndim=2, dtype=np.float32]): (optionnal) Error covariance matrix in the modal space
32  cov: (np.ndarray[ndim=2, dtype=np.float32]): (optionnal) Error covariance matrix in the DM space
33 
34  :return:
35  otftel: (np.ndarray[ndim=2, dtype=np.float32]): OTF of the perfect telescope
36  otf2: (np.ndarray[ndim=2, dtype=np.float32]): OTF due to residual phase error
37  psf: (np.ndarray[ndim=2, dtype=np.float32]): LE PSF
38  gpu: (Gamora): Gamora GPU object (for manipulation and debug)
39  """
40 
41  f = h5py.File(filename, 'r')
42  spup = drax.get_pup(filename)
43  # Sparse IF matrix
44  IF, T = drax.get_IF(filename)
45  # Covariance matrix
46  P = f["P"][:]
47  print("Projecting error buffer into modal space...")
48  if ((err is None) and (cov is None)):
49  err = drax.get_err(filename)
50  err = P.dot(err)
51  print("Computing covariance matrix...")
52  if (cov is None):
53  if (covmodes is None):
54  covmodes = err.dot(err.T) / err.shape[1]
55  else:
56  covmodes = (P.dot(covmodes)).dot(P.T)
57  else:
58  covmodes = cov
59  print("Done")
60  Btt = f["Btt"][:]
61 
62  # Scale factor
63  scale = float(2 * np.pi / f.attrs["_Param_target__Lambda"][0])
64  # Init GPU
65  gpu = Gamora(c, c.active_device, "Vii", Btt.shape[0], covmodes.shape[0],
66  f["noise"][:].shape[1], IF.data, IF.indices, IF.indptr, IF.data.size, T,
67  spup, spup.shape[0],
68  np.where(spup)[0].size, scale, Btt, covmodes)
69  # Launch computation
70  # gamora.set_eigenvals(e)
71  # gamora.set_covmodes(V)
72  tic = time.time()
73  gpu.psf_rec_Vii()
74 
75  otftel = np.array(gpu.d_otftel)
76  otf2 = np.array(gpu.d_otfVii)
77 
78  otftel /= otftel.max()
79  if (list(f.keys()).count("psfortho") and fitting):
80  print("\nAdding fitting to PSF...")
81  psfortho = f["psfortho"][:]
82  otffit = np.real(np.fft.fft2(psfortho))
83  otffit /= otffit.max()
84  psf = np.fft.fftshift(np.real(np.fft.ifft2(otffit * otf2)))
85  else:
86  psf = np.fft.fftshift(np.real(np.fft.ifft2(otftel * otf2)))
87 
88  psf *= (psf.shape[0] * psf.shape[0] / float(np.where(spup)[0].shape[0]))
89  f.close()
90  tac = time.time()
91  print(" ")
92  print("PSF renconstruction took ", tac - tic, " seconds")
93 
94  return otftel, otf2, psf, gpu
95 
96 
97 def psf_rec_vii_cpu(filename):
98  """
99  PSF reconstruction using Vii functions (CPU version)
100 
101  :parameters:
102  filename: (str): path to the ROKET file
103 
104  :return:
105  otftel: (np.ndarray[ndim=2, dtype=np.float32]): OTF of the perfect telescope
106  otf2: (np.ndarray[ndim=2, dtype=np.float32]): OTF due to residual phase error
107  psf: (np.ndarray[ndim=2, dtype=np.float32]): LE PSF
108  """
109 
110  f = h5py.File(filename, 'r')
111  IF, T = drax.get_IF(filename)
112  ratio_lambda = 2 * np.pi / f.attrs["_Param_target__Lambda"][0]
113  # Telescope OTF
114  print("Computing telescope OTF...")
115  spup = drax.get_pup(filename)
116  mradix = 2
117  fft_size = mradix**int((np.log(2 * spup.shape[0]) / np.log(mradix)) + 1)
118  pup = np.zeros((fft_size, fft_size))
119  pup[:spup.shape[0], :spup.shape[0]] = spup
120  pupfft = np.fft.fft2(pup)
121  conjpupfft = np.conjugate(pupfft)
122  otftel = np.real(np.fft.ifft2(pupfft * conjpupfft))
123  den = 1. / otftel
124  den[np.where(np.isinf(den))] = 0
125  mask = np.ones((fft_size, fft_size))
126  mask[np.where(otftel < 1e-5)] = 0
127  otftel = otftel / otftel.max()
128  print("Done")
129  # Covariance matrix
130  print("Computing covariance matrix...")
131  err = drax.get_err(filename)
132  P = f["P"][:]
133  err = P.dot(err)
134  Btt = f["Btt"][:]
135  #modes = IF.T.dot(Btt)
136  covmodes = err.dot(err.T) / err.shape[1]
137  print("Done")
138  # Vii algorithm
139  print("Diagonalizing cov matrix...")
140  e, V = np.linalg.eig(covmodes)
141  print("Done")
142  tmp = np.zeros((fft_size, fft_size))
143  newmodek = tmp.copy()
144  ind = np.where(pup)
145  for k in range(err.shape[0]):
146  #newmodek[ind] = IF.T.dot(V[:,k])
147  #newmodek[ind] = modes.dot(V[:,k])
148  tmp2 = Btt.dot(V[:, k])
149  newmodek[ind] = IF.T.dot(tmp2[:-2])
150  newmodek[ind] += T.T.dot(tmp2[-2:])
151  term1 = np.real(np.fft.fft2(newmodek**2) * conjpupfft)
152  term2 = np.abs(np.fft.fft2(newmodek))**2
153  tmp += ((term1 - term2) * e[k])
154  print(" Computing Vii : %d/%d\r" % (k, covmodes.shape[0]), end=' ')
155  print("Vii computed")
156 
157  dphi = np.real(np.fft.ifft2(2 * tmp)) * den * mask * ratio_lambda**2
158  otf2 = np.exp(-0.5 * dphi) * mask
159  otf2 = otf2 / otf2.max()
160 
161  psf = np.fft.fftshift(np.real(np.fft.ifft2(otftel * otf2)))
162  psf *= (fft_size * fft_size / float(np.where(pup)[0].shape[0]))
163 
164  f.close()
165  return otftel, otf2, psf
166 
167 
168 def test_Vii(filename):
169  """
170  Test function comparing results and performance of GPU version
171  versus CPU version of Vii PSF reconstruction
172 
173  :parameters:
174  filename: (str): path to the ROKET file
175  """
176  a = time.time()
177  otftel_cpu, otf2_cpu, psf_cpu = psf_rec_vii_cpu(filename)
178  b = time.time()
179  otftel_gpu, otf2_gpu, psf_gpu, gamora = psf_rec_Vii(filename)
180  c = time.time()
181  cputime = b - a
182  gputime = c - b
183  print("CPU exec time : ", cputime, " s")
184  print("GPU exec time : ", gputime, " s")
185  print("Speed up : x", cputime / gputime)
186  print("---------------------------------")
187  print("precision on psf : ", np.abs(psf_cpu - psf_gpu).max() / psf_cpu.max())
188 
189 
190 def add_fitting_to_psf(filename, otf, otffit):
191  """
192  Compute the PSF including the fitting OTF
193 
194  :parameters:
195  otf: (np.ndarray[ndim=2, dtype=np.float32]): OTF
196  otffit: (np.ndarray[ndim=2, dtype=np.float32]): Fitting error OTF
197  :return:
198  psf: (np.ndarray[ndim=2, dtype=np.float32]): PSF
199 
200  """
201  print("\nAdding fitting to PSF...")
202  spup = drax.get_pup(filename)
203  psf = np.fft.fftshift(np.real(np.fft.ifft2(otffit * otf)))
204  psf *= (psf.shape[0] * psf.shape[0] / float(np.where(spup)[0].shape[0]))
205 
206  return psf
207 
208 
209 def intersample(Cvvmap, pupilImage, IFImage, pixscale, dactu, lambdaIR):
210  """
211  res = intersample( Cvvmap, pupilImage, IFImage, pixscale, dactu, lambdaIR)
212 
213  Cvvmap is the 'map' of the Cvv matrix (cov matrix of tomo error
214  expressed on volts). The "volts" unit must be used together with
215  the influence function funcInflu(x,y,dm.x0) expressed in meters.
216 
217  Then, the result of intersample is in meter^2.
218 
219  <Cvvmap> : output of Cvvmap=getMap(Cvv)
220  <pupilImage> : pupil image, of size (N,N), shall be properly zero-padded,
221  ready for FFT
222  <IFImage> : image of influence function of 1 actu. Same support
223  as pupilImage, same sampling.
224  <pixscale> : size of pixel (in pupil space, meters) of pupilImage
225  and IFImage
226  <dactu> : inter-actuator pitch in pupil space (meters)
227  <lambdaIR> : in microns
228 
229  Units of IFImage and Cvvmap shall be such that the product of Cvvmap
230  numbers and IFImage^2 is microns^2
231 
232 
233  SEE ALSO: getMap()
234 
235 
236  # pour test/debug :
237  N = 1024
238  D=39.
239  npup=300
240  pixscale = D/npup
241  dactu = 4*pixscale
242  x=(np.arange(N)-N/2)*pixscale
243  x,y = np.meshgrid(x,x,indexing='ij')
244  r2=(x**2+y**2)
245  IFImage = np.exp(-1.5* r2 / dactu**2)
246  pupilImage = generateEeltPupilReflectivity(1., N, 0.53, N/2, N/2, pixscale, 0.03, -10., softGap=1)
247  Nactu = int(np.round(D/dactu))+1
248  ncov = 2*Nactu+1
249  x=np.arange(ncov)-Nactu
250  x,y = np.meshgrid(x,x,indexing='ij')
251  r=np.sqrt(x**2+y**2)
252  Cvvmap = np.exp(-r/3)
253  Cvvmap = np.zeros((ncov, ncov))
254  Cvvmap[Nactu, Nactu]=1.
255 
256 
257  """
258 
259  print("Interpolating Dphi map")
260 
261  # image size
262  N = pupilImage.shape[0]
263 
264  # size of the side of Cvvmap (always odd number)
265  ncov = Cvvmap.shape[0]
266  if (ncov % 2) == 0:
267  ncov = 3 / 0
268  print("Fucking error")
269 
270  # nber of elements on each side of the center of Cvvmap
271  nelem = (ncov - 1) // 2
272  # compute inter-actuator distance in pixels dactupix
273  # dactupix *should* be an integer : pixscale shall be chosen in such a way
274  # that dactupix is an integer. However, for safety here, we round the
275  # result.
276  dactupix = int(np.round(dactu / pixscale))
277  # Fill MAP array with values of Cvvmap. Centre of MAP is located at
278  # index [ncmap, ncmap] (i.e. Fourier-centred)
279  MAP = np.zeros((N, N))
280  ncmap = N // 2 # central element of the MAP, in a Fourier-sense
281  i = ncmap - nelem * dactupix
282  j = ncmap + nelem * dactupix + 1
283  MAP[i:j:dactupix, i:j:dactupix] = Cvvmap
284  print("done")
285 
286  # Computing the phase correlation function
287  # One should have corr(0) = phase_variance.
288  # Computing corr(0) is done using the <v_i^2> (diagonal of Cvv).
289  # We decided that <v^2> is the average value for 1 single actuator (i.e. it's the average
290  # of the diagonal of Cvv).
291  # Then the average phase variance over the pupil equals to
292  # (1/S_pupil) * $_pupil(fi^2) * Nactu * <v^2>
293  # with S_pupil the surface, and $ is an integral. Nactu needs to be here
294  # because if it wasn't, we'd have computed the phase variance over the pupil
295  # with only 1 single actu moving.
296  # So, in our formula, we have replaced the value of (S_pupil/Nactu) by (dactu^2).
297  # The (dactu^2) needs to be expressed in pixels because our integral $(fi^2) is not
298  # a real integral : it's just summing pixels instead.
299 
300  corr = np.fft.fft2(np.abs(np.fft.fft2(IFImage))**2 * np.fft.fft2(MAP)).real / (
301  IFImage.size * dactupix**2)
302  # From correlation to Dphi
303  # Dphi(r) = 2*C(0) - 2*C(r)
304  # We take advantage we need to do a multiplication to multiply by another factor
305  # in the same line. This is to translate dphi from m^2 into rd^2
306  fact = 2 * (2 * np.pi / lambdaIR)**2
307  corr = np.fft.fftshift(corr)
308  dphi = fact * corr[0, 0] - fact * corr
309 
310  # computation of the PSF
311  FTOtel = np.fft.ifft2(np.abs(np.fft.fft2(pupilImage))**2).real
312  # FTOtel is normalized with np.sum(FTOtel)=1
313  # This ensures to get a PSF with SR=np.max(psf), when the PSF is computed
314  # using just np.fft.fft2() without other normalisation
315  FTOtel /= np.sum(FTOtel)
316  # variable mask could be omitted because FTOtel should be zero after
317  # telescope cutoff. However, numeric errors lead to FTOtel small but not
318  # zero, and multiplying with exp(dphi) with dphi undefined after
319  # telescope cutoff may lead to unexpected results.
320  mask = FTOtel > (FTOtel[0, 0] / 1e9)
321  psf = np.fft.fftshift(np.fft.fft2(np.exp(-0.5 * dphi * mask) * FTOtel).real)
322 
323  return psf
324 
325 
326 # def psf_rec_roket_file(filename, err=None):
327 # """
328 # PSF reconstruction using ROKET file. SE PSF is reconstructed
329 # for each frame and stacked to obtain the LE PSF.
330 # Used for ROKET debug only.
331 
332 # :parameters:
333 # filename: (str): path to the ROKET file
334 # err: (np.ndarray[ndim=2, dtype=np.float32]): (optionnal) Buffers of command error
335 # :return:
336 # psf: (np.ndarray[ndim=2, dtype=np.float32]): LE PSF
337 # gpu: (Gamora): Gamora GPU object (for manipulation and debug)
338 # """
339 # f = h5py.File(filename, 'r')
340 # if (err is None):
341 # err = drax.get_err(filename)
342 # spup = drax.get_pup(filename)
343 # # Sparse IF matrix
344 # IF, T = drax.get_IF(filename)
345 # # Scale factor
346 # scale = float(2 * np.pi / f.attrs["_Param_target__Lambda"][0])
347 # # Init GPU
348 # gpu = gamora_init(b"roket", err.shape[0], err.shape[1],
349 # IF.data, IF.indices, IF.indptr, T,
350 # spup, scale)
351 # # Launch computation
352 # gpu.psf_rec_roket(err)
353 # # Get psf
354 # psf = gpu.get_tar_image()
355 # f.close()
356 # return psf, gpu
357 
358 # def psf_rec_roket_file_cpu(filename):
359 # """
360 # PSF reconstruction using ROKET file (CPU version). SE PSF is reconstructed
361 # for each frame and stacked to obtain the LE PSF.
362 # Used for ROKET debug only.
363 
364 # :parameters:
365 # filename: (str): path to the ROKET file
366 # :return:
367 # psf: (np.ndarray[ndim=2, dtype=np.float32]): LE PSF
368 # """
369 
370 # f = h5py.File(filename, 'r')
371 # # Get the sum of error contributors
372 # err = drax.get_err(filename)
373 
374 # # Retrieving spupil (for file where spupil was not saved)
375 # indx_pup = f["indx_pup"][:]
376 # pup = np.zeros((f["dm_dim"].value, f["dm_dim"].value))
377 # pup_F = pup.flatten()
378 # pup_F[indx_pup] = 1.
379 # pup = pup_F.reshape(pup.shape)
380 # spup = pup[np.where(pup)[0].min():np.where(pup)[0].max() + 1,
381 # np.where(pup)[1].min():np.where(pup)[1].max() + 1]
382 # phase = spup.copy()
383 # mradix = 2
384 # fft_size = mradix**int((np.log(2 * spup.shape[0]) / np.log(mradix)) + 1)
385 # amplipup = np.zeros((fft_size, fft_size), dtype=np.complex)
386 # psf = amplipup.copy()
387 # psf = psf
388 
389 # # Sparse IF matrix
390 # IF, T = drax.get_IF(filename)
391 # # Scale factor
392 # scale = float(2 * np.pi / f.attrs["_Param_target__Lambda"][0])
393 
394 # for k in range(err.shape[1]):
395 # amplipup = np.zeros((fft_size, fft_size), dtype=np.complex)
396 # phase[np.where(spup)] = IF.T.dot(err[:-2, k])
397 # phase[np.where(spup)] += T.dot(err[-2:, k])
398 # amplipup[:phase.shape[0], :phase.shape[1]] = np.exp(-1j * phase * scale)
399 # amplipup = np.fft.fft2(amplipup)
400 # psf += np.fft.fftshift(np.abs(amplipup)**2) / \
401 # IF.shape[1] / IF.shape[1] / err.shape[1]
402 # print(" Computing and stacking PSF : %d/%d\r" % (k, err.shape[1]), end=' ')
403 # print("PSF computed and stacked")
404 # f.close()
405 # return psf
guardians.gamora.psf_rec_vii_cpu
def psf_rec_vii_cpu(filename)
PSF reconstruction using Vii functions (CPU version)
Definition: gamora.py:108
shesha.sutra_wrap
Definition: sutra_wrap.py:1
guardians.gamora.add_fitting_to_psf
def add_fitting_to_psf(filename, otf, otffit)
Compute the PSF including the fitting OTF.
Definition: gamora.py:200
guardians.gamora.test_Vii
def test_Vii(filename)
Definition: gamora.py:175
guardians.gamora.psf_rec_Vii
def psf_rec_Vii(filename, err=None, fitting=True, covmodes=None, cov=None)
PSF reconstruction using Vii functions with GPU acceleration.
Definition: gamora.py:39
guardians.gamora.intersample
def intersample(Cvvmap, pupilImage, IFImage, pixscale, dactu, lambdaIR)
res = intersample( Cvvmap, pupilImage, IFImage, pixscale, dactu, lambdaIR)
Definition: gamora.py:257