-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathkl_modes.py
180 lines (133 loc) · 5.62 KB
/
kl_modes.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
import numpy as np
import zern
import matplotlib.pyplot as pl
from tqdm import tqdm
import scipy.special as sp
def _even(x):
return x%2 == 0
def _zernike_parity( j, jp):
return _even(j-jp)
class KL(object):
def __init__(self):
pass
# tmp = np.load('kl/kl_data.npy')
# self.noll_KL = tmp[:,0].astype('int')
# self.noll_Z = tmp[:,1].astype('int')
# self.cz = tmp[:,2]
def precalculate(self, npix_image, first_noll=1):
"""
Precalculate KL modes. We skip the first mode, which is just the
aperture. The second and third mode (first and second one in the return)
are tip-tilt using standard Zernike modes. The rest are KL modes
obtained by the diagonalization of the Kolmogorov Noll's matrix
Parameters
----------
npix_image : int
Number of pixels on the pupil plane
Returns
-------
float array
KL modes
"""
self.npix_image = npix_image
print("Computing KL modes...")
Z_machine = zern.ZernikeNaive(mask=[])
x = np.linspace(-1, 1, npix_image)
xx, yy = np.meshgrid(x, x)
rho = np.sqrt(xx ** 2 + yy ** 2)
theta = np.arctan2(yy, xx)
aperture_mask = rho <= 1.0
self.KL = np.zeros((self.n_modes_max, self.npix_image, self.npix_image))
for mode in (range(self.n_modes_max)):
j = mode + first_noll
# if (i <= 2):
# n, m = zern.zernIndex(i + first_noll)
# Z = Z_machine.Z_nm(n, m, rho, theta, True, 'Jacobi')
# self.KL[i,:,:] = Z * aperture_mask
# else:
indx = np.where(self.noll_KL == j)[0]
tmp = vh[mode,:]
print(mode, self.cz[indx], tmp[np.abs(tmp) > 0])
for i in range(len(indx)):
jz = self.noll_Z[indx[i]]
cz = self.cz[indx[i]]
n, m = zern.zernIndex(jz)
Z = Z_machine.Z_nm(n, m, rho, theta, True, 'Jacobi')
self.KL[mode,:,:] += cz * Z * aperture_mask
return self.KL
def precalculate_covariance(self, npix_image, n_modes_max, first_noll=1, overfill=1.0):
"""
Precalculate KL modes. We skip the first mode, which is just the
aperture. The second and third mode (first and second one in the return)
are tip-tilt using standard Zernike modes. The rest are KL modes
obtained by the diagonalization of the Kolmogorov Noll's matrix
Parameters
----------
npix_image : int
Number of pixels on the pupil plane
n_modes_max : int
Maximum number of nodes to consider
first_noll : int
First Noll index to consider. j=1 is the aperture. j=2/3 are the tip-tilts
Returns
-------
float array
KL modes
"""
self.npix_image = npix_image
self.first_noll = first_noll - 1
self.n_modes_max = n_modes_max + first_noll
print("Computing Kolmogorov covariance...")
covariance = np.zeros((self.n_modes_max, self.n_modes_max))
for j in range(self.n_modes_max):
n, m = zern.zernIndex(j+1)
for jpr in range(self.n_modes_max):
npr, mpr = zern.zernIndex(jpr+1)
deltaz = (m == mpr) and (_zernike_parity(j, jpr) or m == 0)
if (deltaz):
phase = (-1.0)**(0.5*(n+npr-2*m))
t1 = np.sqrt((n+1)*(npr+1))
t2 = sp.gamma(14./3.0) * sp.gamma(11./6.0)**2 * (24.0/5.0*sp.gamma(6.0/5.0))**(5.0/6.0) / (2.0*np.pi**2)
Kzz = t2 * t1 * phase
t1 = sp.gamma(0.5*(n+npr-5.0/3.0))
t2 = sp.gamma(0.5*(n-npr+17.0/3.0)) * sp.gamma(0.5*(npr-n+17.0/3.0)) * sp.gamma(0.5*(n+npr+23.0/3.0))
covariance[j,jpr] = Kzz * t1 / t2
covariance[0,:] = 0.0
covariance[:,0] = 0.0
covariance[0,0] = 1.0
print("Diagonalizing Kolmogorov covariance...")
u, s, vh = np.linalg.svd(covariance)
vh[np.abs(vh) < 1e-10] = 0.0
print("Computing KL modes...")
Z_machine = zern.ZernikeNaive(mask=[])
x = np.linspace(-1, 1, npix_image)
xx, yy = np.meshgrid(x, x)
rho = overfill * np.sqrt(xx ** 2 + yy ** 2)
theta = np.arctan2(yy, xx)
aperture_mask = rho <= 1.0
self.KL = np.zeros((self.n_modes_max, self.npix_image, self.npix_image))
noll_Z = 1 + np.arange(self.n_modes_max)
for mode in tqdm(range(self.n_modes_max)):
cz = vh[mode,:]
ind = np.where(cz != 0)[0]
for i in range(len(ind)):
jz = noll_Z[ind[i]]
coeff = cz[ind[i]]
n, m = zern.zernIndex(jz)
Z = Z_machine.Z_nm(n, m, rho, theta, True, 'Jacobi')
self.KL[mode,:,:] += coeff * Z * aperture_mask
self.KL = self.KL[self.first_noll+1:,:,:]
return self.KL
if (__name__ == '__main__'):
tmp = KL()
tmp.precalculate_covariance(npix_image=128, n_modes_max=25, first_noll=3)
f, ax = pl.subplots(nrows=4, ncols=4)
for i in range(16):
ax.flat[i].imshow(tmp.KL[i, :, :])
pl.show()
mat = np.zeros((20,20))
for i in range(20):
for j in range(20):
mat[i,j] = np.sum(tmp.KL[i,:,:]*tmp.KL[j,:,:])
#pl.imshow(np.log(np.abs(mat)))
#pl.show()