2 GAMORA (Gpu Accelerated Module fOr psf Reconstruction Algorithms)
4 Python module for GPU accelerated PSF reconstruction using Vii functions and ROKET file
6 Note: GPU devices used are hardcoded here. Change gpudevices if needed
9 import matplotlib.pyplot
as plt
12 from scipy.sparse
import csr_matrix
13 from sys
import stdout
15 from guardians
import drax
19 gpudevices = np.array([0, 1, 2, 3], dtype=np.int32)
20 c = carmaWrap_context.get_instance_ngpu(gpudevices.size, gpudevices)
23 def psf_rec_Vii(filename, err=None, fitting=True, covmodes=None, cov=None):
25 PSF reconstruction using Vii functions with GPU acceleration.
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
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)
41 f = h5py.File(filename,
'r')
42 spup = drax.get_pup(filename)
44 IF, T = drax.get_IF(filename)
47 print(
"Projecting error buffer into modal space...")
48 if ((err
is None)
and (cov
is None)):
49 err = drax.get_err(filename)
51 print(
"Computing covariance matrix...")
53 if (covmodes
is None):
54 covmodes = err.dot(err.T) / err.shape[1]
56 covmodes = (P.dot(covmodes)).dot(P.T)
63 scale = float(2 * np.pi / f.attrs[
"_Param_target__Lambda"][0])
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,
68 np.where(spup)[0].size, scale, Btt, covmodes)
75 otftel = np.array(gpu.d_otftel)
76 otf2 = np.array(gpu.d_otfVii)
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)))
86 psf = np.fft.fftshift(np.real(np.fft.ifft2(otftel * otf2)))
88 psf *= (psf.shape[0] * psf.shape[0] / float(np.where(spup)[0].shape[0]))
92 print(
"PSF renconstruction took ", tac - tic,
" seconds")
94 return otftel, otf2, psf, gpu
99 PSF reconstruction using Vii functions (CPU version)
102 filename: (str): path to the ROKET file
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
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]
114 print(
"Computing telescope OTF...")
115 spup = drax.get_pup(filename)
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))
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()
130 print(
"Computing covariance matrix...")
131 err = drax.get_err(filename)
136 covmodes = err.dot(err.T) / err.shape[1]
139 print(
"Diagonalizing cov matrix...")
140 e, V = np.linalg.eig(covmodes)
142 tmp = np.zeros((fft_size, fft_size))
143 newmodek = tmp.copy()
145 for k
in range(err.shape[0]):
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")
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()
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]))
165 return otftel, otf2, psf
170 Test function comparing results and performance of GPU version
171 versus CPU version of Vii PSF reconstruction
174 filename: (str): path to the ROKET file
179 otftel_gpu, otf2_gpu, psf_gpu, gamora =
psf_rec_Vii(filename)
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())
192 Compute the PSF including the fitting OTF
195 otf: (np.ndarray[ndim=2, dtype=np.float32]): OTF
196 otffit: (np.ndarray[ndim=2, dtype=np.float32]): Fitting error OTF
198 psf: (np.ndarray[ndim=2, dtype=np.float32]): PSF
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]))
209 def intersample(Cvvmap, pupilImage, IFImage, pixscale, dactu, lambdaIR):
211 res = intersample( Cvvmap, pupilImage, IFImage, pixscale, dactu, lambdaIR)
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.
217 Then, the result of intersample is in meter^2.
219 <Cvvmap> : output of Cvvmap=getMap(Cvv)
220 <pupilImage> : pupil image, of size (N,N), shall be properly zero-padded,
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
226 <dactu> : inter-actuator pitch in pupil space (meters)
227 <lambdaIR> : in microns
229 Units of IFImage and Cvvmap shall be such that the product of Cvvmap
230 numbers and IFImage^2 is microns^2
242 x=(np.arange(N)-N/2)*pixscale
243 x,y = np.meshgrid(x,x,indexing='ij')
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
249 x=np.arange(ncov)-Nactu
250 x,y = np.meshgrid(x,x,indexing='ij')
252 Cvvmap = np.exp(-r/3)
253 Cvvmap = np.zeros((ncov, ncov))
254 Cvvmap[Nactu, Nactu]=1.
259 print(
"Interpolating Dphi map")
262 N = pupilImage.shape[0]
265 ncov = Cvvmap.shape[0]
268 print(
"Fucking error")
271 nelem = (ncov - 1) // 2
276 dactupix = int(np.round(dactu / pixscale))
279 MAP = np.zeros((N, N))
281 i = ncmap - nelem * dactupix
282 j = ncmap + nelem * dactupix + 1
283 MAP[i:j:dactupix, i:j:dactupix] = Cvvmap
300 corr = np.fft.fft2(np.abs(np.fft.fft2(IFImage))**2 * np.fft.fft2(MAP)).real / (
301 IFImage.size * dactupix**2)
306 fact = 2 * (2 * np.pi / lambdaIR)**2
307 corr = np.fft.fftshift(corr)
308 dphi = fact * corr[0, 0] - fact * corr
311 FTOtel = np.fft.ifft2(np.abs(np.fft.fft2(pupilImage))**2).real
315 FTOtel /= np.sum(FTOtel)
320 mask = FTOtel > (FTOtel[0, 0] / 1e9)
321 psf = np.fft.fftshift(np.fft.fft2(np.exp(-0.5 * dphi * mask) * FTOtel).real)