prtools 1.3.2__tar.gz → 1.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {prtools-1.3.2/prtools.egg-info → prtools-1.4.0}/PKG-INFO +1 -1
- {prtools-1.3.2 → prtools-1.4.0}/prtools/__init__.py +11 -2
- prtools-1.4.0/prtools/convolve.py +348 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/fourier.py +1 -1
- {prtools-1.3.2 → prtools-1.4.0}/prtools/jax.py +25 -6
- {prtools-1.3.2 → prtools-1.4.0}/prtools/zernike.py +91 -30
- {prtools-1.3.2 → prtools-1.4.0/prtools.egg-info}/PKG-INFO +1 -1
- {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/SOURCES.txt +1 -0
- {prtools-1.3.2 → prtools-1.4.0}/LICENSE +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/NOTICE +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/README.rst +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/array.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/__init__.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/_base.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/_jax.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/_numpy.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/cost.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/misc.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/segmented.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/shape.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/sparse.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools/stats.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/dependency_links.txt +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/requires.txt +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/top_level.txt +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/pyproject.toml +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/setup.cfg +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/tests/test_array.py +0 -0
- {prtools-1.3.2 → prtools-1.4.0}/tests/test_fourier.py +0 -0
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
__version__ = '1.
|
|
1
|
+
__version__ = '1.4.0'
|
|
2
2
|
|
|
3
3
|
from prtools.backend import __backend__, use
|
|
4
4
|
|
|
@@ -15,6 +15,16 @@ from prtools.array import (
|
|
|
15
15
|
medfix
|
|
16
16
|
)
|
|
17
17
|
|
|
18
|
+
from prtools.convolve import (
|
|
19
|
+
fftconv,
|
|
20
|
+
gauss,
|
|
21
|
+
gauss_blur,
|
|
22
|
+
gauss_kernel,
|
|
23
|
+
pixel_kernel,
|
|
24
|
+
pixelate,
|
|
25
|
+
sinc,
|
|
26
|
+
)
|
|
27
|
+
|
|
18
28
|
from prtools.cost import sserror
|
|
19
29
|
|
|
20
30
|
from prtools.fourier import dft2, idft2
|
|
@@ -38,7 +48,6 @@ from prtools.shape import (
|
|
|
38
48
|
hexagon,
|
|
39
49
|
rectangle,
|
|
40
50
|
spider,
|
|
41
|
-
gauss,
|
|
42
51
|
sin,
|
|
43
52
|
waffle
|
|
44
53
|
)
|
|
@@ -0,0 +1,348 @@
|
|
|
1
|
+
from prtools.backend import numpy as np
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def fftconv(array, kernel, normalize_kernel=True, fft_array=True,
|
|
5
|
+
fft_kernel=True):
|
|
6
|
+
r"""Convolve an array with a kernel using the FFT.
|
|
7
|
+
|
|
8
|
+
The colvolution is computed as
|
|
9
|
+
|
|
10
|
+
.. math::
|
|
11
|
+
|
|
12
|
+
u*v = \mathcal{F}^{-1}\left\{\mathcal{F}\{u\}\cdot\mathcal{F}\{v\}\right\}
|
|
13
|
+
|
|
14
|
+
Parameters
|
|
15
|
+
----------
|
|
16
|
+
array : array_like
|
|
17
|
+
Array to be convolved with `kernel`.
|
|
18
|
+
kernel : array_like
|
|
19
|
+
Convolution kernel. Should have the same shape as `array`.
|
|
20
|
+
normalize_kernel : bool, optional
|
|
21
|
+
If True (default), kernel is normalized so that `kernel =
|
|
22
|
+
kernel / np.sum(kernel)`.
|
|
23
|
+
fft_array : bool, optional
|
|
24
|
+
If True (default), the array is assumed to be provided in the spatial
|
|
25
|
+
domain. If False, the array is assumed to be provided in the frequency
|
|
26
|
+
domain.
|
|
27
|
+
fft_kernel : bool, optional
|
|
28
|
+
If True (default), the kernel is assumed to be provided in the spatial
|
|
29
|
+
domain. If False, the kernel is assumed to be provided in the frequency
|
|
30
|
+
domain.
|
|
31
|
+
|
|
32
|
+
Returns
|
|
33
|
+
-------
|
|
34
|
+
ndarray
|
|
35
|
+
|
|
36
|
+
See also
|
|
37
|
+
--------
|
|
38
|
+
:func:`~prtools.gauss_blur`
|
|
39
|
+
:func:`~prtools.pixelate`
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
a = np.fft.fft2(array) if fft_array else array
|
|
43
|
+
k = np.fft.fft2(array) if fft_kernel else kernel
|
|
44
|
+
|
|
45
|
+
if normalize_kernel:
|
|
46
|
+
k /= np.sum(k)
|
|
47
|
+
|
|
48
|
+
return np.fft.ifft2(a*k).real
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def gauss_blur(img, sigma, oversample=1):
|
|
52
|
+
"""Blur an image using a Gaussian filter using the FFT.
|
|
53
|
+
|
|
54
|
+
Parameters
|
|
55
|
+
----------
|
|
56
|
+
img : array_like
|
|
57
|
+
Array to be blurred
|
|
58
|
+
sigma : float or (2,) array_like
|
|
59
|
+
Standard deviation of Gaussian. Providing two values allows for
|
|
60
|
+
non-symmetric Gaussian interpreted as `(sigma_row, sigma_col)`
|
|
61
|
+
oversample : float, optional
|
|
62
|
+
Oversampling factor of `img`. Default is 1.
|
|
63
|
+
|
|
64
|
+
Returns
|
|
65
|
+
-------
|
|
66
|
+
ndarray
|
|
67
|
+
|
|
68
|
+
Examples
|
|
69
|
+
--------
|
|
70
|
+
.. plot::
|
|
71
|
+
:include-source:
|
|
72
|
+
:context: reset
|
|
73
|
+
:scale: 50
|
|
74
|
+
|
|
75
|
+
>>> img = prtools.circle((256,256), 100)
|
|
76
|
+
>>> img_blurred = prtools.gauss_blur(img, sigma=3)
|
|
77
|
+
>>> fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(5,2))
|
|
78
|
+
>>> ax[0].imshow(img, cmap='gray')
|
|
79
|
+
>>> ax[0].set_title('Input image')
|
|
80
|
+
>>> ax[1].imshow(img_blurred, cmap='gray')
|
|
81
|
+
>>> ax[1].set_title('Blurred image')
|
|
82
|
+
"""
|
|
83
|
+
img = np.asarray(img)
|
|
84
|
+
kernel = gauss_kernel(img.shape, sigma, oversample, fftshift=False)
|
|
85
|
+
return fftconv(img, kernel, normalize_kernel=True, fft_array=True,
|
|
86
|
+
fft_kernel=False)
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def pixelate(img, oversample=1):
|
|
90
|
+
"""Apply the aperture effects of an idealized square pixel on a discretely
|
|
91
|
+
sampled image using the FFT.
|
|
92
|
+
|
|
93
|
+
Parameters
|
|
94
|
+
----------
|
|
95
|
+
img : array_like
|
|
96
|
+
Input image
|
|
97
|
+
oversample : float, optional
|
|
98
|
+
Oversampling factor of `img`. Default is 1.
|
|
99
|
+
|
|
100
|
+
Returns
|
|
101
|
+
-------
|
|
102
|
+
out : ndarray
|
|
103
|
+
Image with pixel sampling effects applied.
|
|
104
|
+
|
|
105
|
+
Notes
|
|
106
|
+
-----
|
|
107
|
+
To avoid the introduction of numerical artifacts, this function should be
|
|
108
|
+
performed on data that is at least 2x oversampled.
|
|
109
|
+
|
|
110
|
+
Examples
|
|
111
|
+
--------
|
|
112
|
+
.. plot::
|
|
113
|
+
:include-source:
|
|
114
|
+
:context: reset
|
|
115
|
+
:scale: 50
|
|
116
|
+
|
|
117
|
+
>>> img = prtools.circle((17,17), 5)
|
|
118
|
+
>>> img_pixelate = prtools.pixelate(img)
|
|
119
|
+
>>> fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(5,2))
|
|
120
|
+
>>> ax[0].imshow(img, cmap='gray')
|
|
121
|
+
>>> ax[0].set_title('Input image')
|
|
122
|
+
>>> ax[1].imshow(img_pixelate, cmap='gray')
|
|
123
|
+
>>> ax[1].set_title('Pixelated image')
|
|
124
|
+
|
|
125
|
+
Note that the pixelation process preserves image power:
|
|
126
|
+
|
|
127
|
+
.. code:: pycon
|
|
128
|
+
|
|
129
|
+
>>> print(np.sum(img), np.sum(img_pixelate))
|
|
130
|
+
64.0, 64.0
|
|
131
|
+
"""
|
|
132
|
+
img = np.asarray(img)
|
|
133
|
+
kernel = pixel_kernel(img.shape, oversample=oversample, fftshift=False)
|
|
134
|
+
return fftconv(img, kernel, normalize_kernel=False, fft_array=True,
|
|
135
|
+
fft_kernel=False)
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def gauss(x1, x2, sigma, oversample=1, indexing='ij', normalize=False):
|
|
139
|
+
"""2D Gaussian function
|
|
140
|
+
|
|
141
|
+
Parameters
|
|
142
|
+
----------
|
|
143
|
+
x1, x2 : array_like
|
|
144
|
+
1-D arrays representing the grid coordinates
|
|
145
|
+
sigma : float or (2,) array_like
|
|
146
|
+
Standard deviation of Gaussian. Providing two values allows for
|
|
147
|
+
non-symmetric Gaussian interpreted as `(sigma_row, sigma_col)`
|
|
148
|
+
oversample : float, optional
|
|
149
|
+
Oversampling factor. Defailt is 1.
|
|
150
|
+
indexing : {'ij', 'xy'}, optional
|
|
151
|
+
Matrix ('ij', default) or cartesian ('xy') indexing of output.
|
|
152
|
+
normalize : bool, optional
|
|
153
|
+
If True, the output is normalized such that its sum is equal to 1.
|
|
154
|
+
If False (default), the output has max equal to one.
|
|
155
|
+
|
|
156
|
+
Returns
|
|
157
|
+
-------
|
|
158
|
+
ndarray
|
|
159
|
+
|
|
160
|
+
Examples
|
|
161
|
+
--------
|
|
162
|
+
.. plot::
|
|
163
|
+
:include-source:
|
|
164
|
+
:context: reset
|
|
165
|
+
:scale: 50
|
|
166
|
+
|
|
167
|
+
>>> r = np.arange(-128,128)
|
|
168
|
+
>>> c = np.arange(-128,128)
|
|
169
|
+
>>> g = prtools.gauss(r, c, sigma=20)
|
|
170
|
+
>>> plt.imshow(g, cmap='gray')
|
|
171
|
+
"""
|
|
172
|
+
xx1, xx2 = np.meshgrid(x1, x2, indexing=indexing)
|
|
173
|
+
sigma = np.broadcast_to(sigma, (2,)) / oversample
|
|
174
|
+
g = np.exp(-((xx1**2/(2*sigma[0]**2)) + (xx2**2/(2*sigma[1]**2))))
|
|
175
|
+
if normalize:
|
|
176
|
+
g = g / (2*np.pi * np.prod(sigma))
|
|
177
|
+
return g
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
def sinc(x1, x2, indexing='ij'):
|
|
181
|
+
r"""Normalized 2D sinc function
|
|
182
|
+
|
|
183
|
+
Parameters
|
|
184
|
+
----------
|
|
185
|
+
x1, x2 : array_like
|
|
186
|
+
1-D arrays representing the grid coordinates
|
|
187
|
+
indexing : {'ij', 'xy'}, optional
|
|
188
|
+
Matrix ('ij', default) or cartesian ('xy') indexing of output.
|
|
189
|
+
normalize : bool, optional
|
|
190
|
+
If True, the output is normalized such that its sum is equal to 1.
|
|
191
|
+
If False (default), the output has max equal to one.
|
|
192
|
+
|
|
193
|
+
Returns
|
|
194
|
+
-------
|
|
195
|
+
ndarray
|
|
196
|
+
|
|
197
|
+
See also
|
|
198
|
+
--------
|
|
199
|
+
:func:`~prtools.pixel_kernel`
|
|
200
|
+
|
|
201
|
+
Notes
|
|
202
|
+
-----
|
|
203
|
+
This function uses ``numpy.sinc``, which computes the sinc function as
|
|
204
|
+
|
|
205
|
+
.. math::
|
|
206
|
+
\mbox{sinc}(x) = \frac{\sin \pi x}{\pi x}
|
|
207
|
+
|
|
208
|
+
Examples
|
|
209
|
+
--------
|
|
210
|
+
.. plot::
|
|
211
|
+
:include-source:
|
|
212
|
+
:context: reset
|
|
213
|
+
:scale: 50
|
|
214
|
+
|
|
215
|
+
>>> r = np.linspace(-3,3,512)
|
|
216
|
+
>>> c = np.linspace(-3,3,512)
|
|
217
|
+
>>> s = prtools.sinc(r, c)
|
|
218
|
+
>>> plt.imshow(s, cmap='gray', extent=[r.min(), r.max(), c.min(), c.max()])
|
|
219
|
+
"""
|
|
220
|
+
xx1, xx2 = np.meshgrid(x1, x2, indexing=indexing)
|
|
221
|
+
return np.sinc(xx1) * np.sinc(xx2)
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def gauss_kernel(shape, sigma, oversample=1, fftshift=True):
|
|
225
|
+
"""2D Gaussian filter kernel
|
|
226
|
+
|
|
227
|
+
This function returns the transfer function of a normalized 2D Gaussian
|
|
228
|
+
function.
|
|
229
|
+
|
|
230
|
+
Parameters
|
|
231
|
+
----------
|
|
232
|
+
shape : int or tuple of int
|
|
233
|
+
Shape of the kernel
|
|
234
|
+
sigma : float
|
|
235
|
+
Standard deviation of Gaussian. Providing two values allows for
|
|
236
|
+
non-symmetric Gaussian interpreted as `(sigma_row, sigma_col)`
|
|
237
|
+
oversample : float, optional
|
|
238
|
+
Oversampling factor to represent in the kernel. Default is 1.
|
|
239
|
+
fftshift : bool, optional
|
|
240
|
+
If True (default), the kernel is FFT-shifted before it is returned.
|
|
241
|
+
|
|
242
|
+
Returns
|
|
243
|
+
-------
|
|
244
|
+
ndarray
|
|
245
|
+
|
|
246
|
+
See also
|
|
247
|
+
--------
|
|
248
|
+
:func:`~prtools.fftconv`
|
|
249
|
+
:func:`~prtools.gauss_blur`
|
|
250
|
+
:func:`~prtools.gauss`
|
|
251
|
+
|
|
252
|
+
Examples
|
|
253
|
+
--------
|
|
254
|
+
.. plot::
|
|
255
|
+
:include-source:
|
|
256
|
+
:context: reset
|
|
257
|
+
:scale: 50
|
|
258
|
+
|
|
259
|
+
>>> blur_kernel = prtools.gauss_kernel((256,256), 1)
|
|
260
|
+
>>> plt.imshow(blur_kernel, cmap='gray')
|
|
261
|
+
|
|
262
|
+
Note this is functionally equivalent to
|
|
263
|
+
|
|
264
|
+
.. plot::
|
|
265
|
+
:include-source:
|
|
266
|
+
:context: reset
|
|
267
|
+
:scale: 50
|
|
268
|
+
|
|
269
|
+
>>> r = np.fft.fftfreq(256)
|
|
270
|
+
>>> c = np.fft.fftfreq(256)
|
|
271
|
+
>>> sigma = 1
|
|
272
|
+
>>> blur_kernel = np.fft.fftshift(prtools.gauss(r, c, 1/(2*np.pi*sigma)))
|
|
273
|
+
>>> plt.imshow(blur_kernel, cmap='gray')
|
|
274
|
+
"""
|
|
275
|
+
|
|
276
|
+
shape = np.broadcast_to(shape, (2,))
|
|
277
|
+
sigma = np.broadcast_to(sigma, (2,))
|
|
278
|
+
|
|
279
|
+
x1 = np.fft.fftfreq(shape[0])
|
|
280
|
+
x2 = np.fft.fftfreq(shape[1])
|
|
281
|
+
|
|
282
|
+
k = gauss(x1, x2, 1/(2*np.pi*sigma), oversample, indexing='ij',
|
|
283
|
+
normalize=False)
|
|
284
|
+
|
|
285
|
+
if fftshift:
|
|
286
|
+
return np.fft.fftshift(k)
|
|
287
|
+
else:
|
|
288
|
+
return k
|
|
289
|
+
|
|
290
|
+
|
|
291
|
+
def pixel_kernel(shape, oversample=1, fftshift=True):
|
|
292
|
+
r"""2D pixel MTF filter kernel
|
|
293
|
+
|
|
294
|
+
This function returns a normalized 2D sinc function sized to represent
|
|
295
|
+
the transfer function of an idealized square pixel.
|
|
296
|
+
|
|
297
|
+
Parameters
|
|
298
|
+
----------
|
|
299
|
+
shape : int or tuple of int
|
|
300
|
+
Shape of the kernel
|
|
301
|
+
oversample : float, optional
|
|
302
|
+
Oversampling factor to represent in the kernel. Default is 1.
|
|
303
|
+
fftshift : bool, optional
|
|
304
|
+
If True (default), the kernel is FFT-shifted before it is returned.
|
|
305
|
+
|
|
306
|
+
Returns
|
|
307
|
+
-------
|
|
308
|
+
ndarray
|
|
309
|
+
|
|
310
|
+
See also
|
|
311
|
+
--------
|
|
312
|
+
:func:`~prtools.fftconv`
|
|
313
|
+
:func:`~prtools.pixelate`
|
|
314
|
+
:func:`~prtools.sinc`
|
|
315
|
+
|
|
316
|
+
Examples
|
|
317
|
+
--------
|
|
318
|
+
.. plot::
|
|
319
|
+
:include-source:
|
|
320
|
+
:context: reset
|
|
321
|
+
:scale: 50
|
|
322
|
+
|
|
323
|
+
>>> pixel_mtf = prtools.pixel_kernel((256,256), 1)
|
|
324
|
+
>>> plt.imshow(pixel_mtf, cmap='gray', vmin=0)
|
|
325
|
+
|
|
326
|
+
Note this is functionally equivalent to
|
|
327
|
+
|
|
328
|
+
.. plot::
|
|
329
|
+
:include-source:
|
|
330
|
+
:context: reset
|
|
331
|
+
:scale: 50
|
|
332
|
+
|
|
333
|
+
>>> r = np.fft.fftfreq(256)
|
|
334
|
+
>>> c = np.fft.fftfreq(256)
|
|
335
|
+
>>> pixel_mtf = np.fft.fftshift(np.abs(prtools.sinc(r, c)))
|
|
336
|
+
>>> plt.imshow(pixel_mtf, cmap='gray', vmin=0)
|
|
337
|
+
"""
|
|
338
|
+
shape = np.broadcast_to(shape, (2,))
|
|
339
|
+
|
|
340
|
+
x1 = np.fft.fftfreq(shape[0]) / oversample
|
|
341
|
+
x2 = np.fft.fftfreq(shape[1]) / oversample
|
|
342
|
+
|
|
343
|
+
k = np.abs(sinc(x1, x2, indexing='ij'))
|
|
344
|
+
|
|
345
|
+
if fftshift:
|
|
346
|
+
return np.fft.fftshift(k)
|
|
347
|
+
else:
|
|
348
|
+
return k
|
|
@@ -74,7 +74,7 @@ def dft2(f, alpha, shape=None, shift=(0, 0), offset=(0, 0), unitary=True,
|
|
|
74
74
|
References
|
|
75
75
|
----------
|
|
76
76
|
[1] Soummer, et. al. Fast computation of Lyot-style coronagraph
|
|
77
|
-
|
|
77
|
+
propagation (2007)
|
|
78
78
|
"""
|
|
79
79
|
return _dftcore(f, alpha, shape, shift, offset, unitary, out, forward=True)
|
|
80
80
|
|
|
@@ -31,14 +31,22 @@ class JaxOptimizeResult:
|
|
|
31
31
|
state: Any #: Optimizer state
|
|
32
32
|
|
|
33
33
|
|
|
34
|
-
def lbfgs(
|
|
34
|
+
def lbfgs(fn, x0, gtol=None, maxiter=None, callback=None, fn_args=None,
|
|
35
|
+
fn_kwargs=None):
|
|
35
36
|
"""Minimize a scalar function of one or more variables using the L-BFGS
|
|
36
37
|
algorithm
|
|
37
38
|
|
|
38
39
|
Parameters
|
|
39
40
|
----------
|
|
40
|
-
|
|
41
|
-
The objective function to be
|
|
41
|
+
fn : callable
|
|
42
|
+
The objective function to be minimized:
|
|
43
|
+
|
|
44
|
+
.. code:: python
|
|
45
|
+
|
|
46
|
+
fn(x, *fn_args, **fn_kwargs)
|
|
47
|
+
|
|
48
|
+
where ``x`` is a 1-D array with shape (n,) and ``fn_args`` and
|
|
49
|
+
``fn_kwargs`` are optional positional and keyword arguments.
|
|
42
50
|
x0 : jax.Array
|
|
43
51
|
Initial guess
|
|
44
52
|
gtol : float
|
|
@@ -53,6 +61,10 @@ def lbfgs(fun, x0, gtol=None, maxiter=None, callback=None):
|
|
|
53
61
|
callback(intermediate_result: JaxOptimizeResult)
|
|
54
62
|
|
|
55
63
|
where ``intermediate_result`` is a :class:`JaxOptimizeResult`.
|
|
64
|
+
fn_args : iterable or None
|
|
65
|
+
Extra positional arguments passed to ``fn()``
|
|
66
|
+
fn_kwargs : dict or None
|
|
67
|
+
Extra keyword arguments passed to the ``fn()``
|
|
56
68
|
|
|
57
69
|
Returns
|
|
58
70
|
-------
|
|
@@ -71,14 +83,21 @@ def lbfgs(fun, x0, gtol=None, maxiter=None, callback=None):
|
|
|
71
83
|
if not any((gtol, maxiter)):
|
|
72
84
|
raise ValueError('At least one termination tolerance must be specified.')
|
|
73
85
|
|
|
86
|
+
fn_args = [] if fn_args is None else fn_args
|
|
87
|
+
fn_kwargs = {} if fn_kwargs is None else fn_kwargs
|
|
88
|
+
|
|
74
89
|
opt = optax.lbfgs()
|
|
75
|
-
|
|
90
|
+
value_and_grad_fn = optax.value_and_grad_from_state(fn)
|
|
76
91
|
|
|
77
92
|
def step(carry):
|
|
78
93
|
params, state = carry
|
|
79
|
-
|
|
94
|
+
# NOTE: passing *args and **kwargs to value_and_grad_fun is very
|
|
95
|
+
# poorly documented in optax (as of v0.2.6 - 10/2025) but this
|
|
96
|
+
# seems to work for now
|
|
97
|
+
value, grad = value_and_grad_fn(params, *fn_args, state=state,
|
|
98
|
+
**fn_kwargs)
|
|
80
99
|
updates, state = opt.update(
|
|
81
|
-
grad, state, params, value=value, grad=grad, value_fn=
|
|
100
|
+
grad, state, params, value=value, grad=grad, value_fn=fn)
|
|
82
101
|
if callback:
|
|
83
102
|
res = JaxOptimizeResult(
|
|
84
103
|
n=otu.tree_get(state, 'count')-1,
|
|
@@ -4,7 +4,7 @@ import prtools
|
|
|
4
4
|
from prtools.backend import numpy as np
|
|
5
5
|
|
|
6
6
|
|
|
7
|
-
def zernike(mask, index, normalize=True, rho=None, theta=None):
|
|
7
|
+
def zernike(mask, index, normalize=True, order='noll', rho=None, theta=None):
|
|
8
8
|
"""Compute the circular Zernike polynomial for a given mask.
|
|
9
9
|
|
|
10
10
|
Parameters
|
|
@@ -17,6 +17,8 @@ def zernike(mask, index, normalize=True, rho=None, theta=None):
|
|
|
17
17
|
normalize : bool, optional
|
|
18
18
|
If True (default), the output is normalized according to [1]. If False,
|
|
19
19
|
the output value ranges [-1, 1] over the mask.
|
|
20
|
+
order : {'noll', 'fringe', 'ansi'}, optional
|
|
21
|
+
Zernike ordering scheme. Default is 'noll'.
|
|
20
22
|
rho : array_like, optional
|
|
21
23
|
Radial coordinates of the mask array. :attr:`rho` should be 0 at the
|
|
22
24
|
origin and 1 at the edge of the circle.
|
|
@@ -44,9 +46,15 @@ def zernike(mask, index, normalize=True, rho=None, theta=None):
|
|
|
44
46
|
|
|
45
47
|
"""
|
|
46
48
|
mask = np.asarray(mask)
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
49
|
+
|
|
50
|
+
if order == 'noll':
|
|
51
|
+
n, m = noll_to_nm(index)
|
|
52
|
+
elif order == 'fringe':
|
|
53
|
+
n, m = fringe_to_nm(index)
|
|
54
|
+
elif order =='ansi':
|
|
55
|
+
n, m = ansi_to_nm(index)
|
|
56
|
+
else:
|
|
57
|
+
raise ValueError('Unknown order', order)
|
|
50
58
|
|
|
51
59
|
if rho is None:
|
|
52
60
|
rho, theta = zernike_coordinates(mask)
|
|
@@ -54,8 +62,6 @@ def zernike(mask, index, normalize=True, rho=None, theta=None):
|
|
|
54
62
|
if theta is None:
|
|
55
63
|
raise ValueError("Both rho and theta must be specified")
|
|
56
64
|
|
|
57
|
-
m, n = zernike_index(index)
|
|
58
|
-
|
|
59
65
|
if m == 0:
|
|
60
66
|
if n == 0:
|
|
61
67
|
Z = mask
|
|
@@ -97,26 +103,24 @@ def R(m, n, rho):
|
|
|
97
103
|
return R
|
|
98
104
|
|
|
99
105
|
|
|
100
|
-
def zernike_compose(mask, coeffs, normalize=True, rho=None, theta=None):
|
|
106
|
+
def zernike_compose(mask, coeffs, normalize=True, order='noll', rho=None, theta=None):
|
|
101
107
|
"""Create an OPD based on the supplied Zernike coefficients.
|
|
102
108
|
|
|
103
109
|
Parameters
|
|
104
110
|
----------
|
|
105
111
|
mask : array_like
|
|
106
112
|
Binary mask defining the extent to compute the Zernike polynomial over.
|
|
107
|
-
|
|
108
113
|
coeffs : array_like
|
|
109
114
|
List of coefficients corresponding to Zernike indices (Noll ordering)
|
|
110
115
|
used to create the OPD.
|
|
111
|
-
|
|
112
116
|
normalize : bool, optional
|
|
113
117
|
If True (default), the output is normalized according to [1]. If False,
|
|
114
118
|
the output value ranges [-1, 1] over the mask.
|
|
115
|
-
|
|
119
|
+
order : {'noll', 'fringe', 'ansi'}, optional
|
|
120
|
+
Zernike ordering scheme. Default is 'noll'.
|
|
116
121
|
rho : array_like, optional
|
|
117
122
|
Radial coordinates of the mask array. :attr:`rho` should be 0 at the
|
|
118
123
|
origin and 1 at the edge of the circle.
|
|
119
|
-
|
|
120
124
|
theta : array_like, optional
|
|
121
125
|
Angular coordinates of the mask array in radians.
|
|
122
126
|
|
|
@@ -159,11 +163,12 @@ def zernike_compose(mask, coeffs, normalize=True, rho=None, theta=None):
|
|
|
159
163
|
opd = np.zeros(mask.shape)
|
|
160
164
|
|
|
161
165
|
for index, coeff in enumerate(coeffs):
|
|
162
|
-
opd += coeff * zernike(mask, index+1, normalize, rho, theta)
|
|
166
|
+
opd += coeff * zernike(mask, index+1, normalize, order, rho, theta)
|
|
163
167
|
|
|
164
168
|
return opd
|
|
165
169
|
|
|
166
|
-
def zernike_basis(mask, modes, vectorize=False, normalize=True,
|
|
170
|
+
def zernike_basis(mask, modes, vectorize=False, normalize=True, order='noll',
|
|
171
|
+
rho=None, theta=None):
|
|
167
172
|
"""Compute a Zernike basis set for a given mask.
|
|
168
173
|
|
|
169
174
|
Parameters
|
|
@@ -180,6 +185,8 @@ def zernike_basis(mask, modes, vectorize=False, normalize=True, rho=None, theta=
|
|
|
180
185
|
normalize : bool, optional
|
|
181
186
|
If True (default), the output is normalized according to [1]. If False,
|
|
182
187
|
the output value ranges [-1, 1] over the mask.
|
|
188
|
+
order : {'noll', 'fringe', 'ansi'}, optional
|
|
189
|
+
Zernike ordering scheme. Default is 'noll'.
|
|
183
190
|
rho : array_like, optional
|
|
184
191
|
Radial coordinates of the mask array. :attr:`rho` should be 0 at the
|
|
185
192
|
origin and 1 at the edge of the circle.
|
|
@@ -203,7 +210,7 @@ def zernike_basis(mask, modes, vectorize=False, normalize=True, rho=None, theta=
|
|
|
203
210
|
|
|
204
211
|
basis = []
|
|
205
212
|
for index, mode in enumerate(modes):
|
|
206
|
-
basis.append(zernike(mask, mode, normalize, rho, theta))
|
|
213
|
+
basis.append(zernike(mask, mode, normalize, order, rho, theta))
|
|
207
214
|
basis = np.asarray(basis)
|
|
208
215
|
|
|
209
216
|
if vectorize:
|
|
@@ -211,9 +218,10 @@ def zernike_basis(mask, modes, vectorize=False, normalize=True, rho=None, theta=
|
|
|
211
218
|
return basis.reshape(basis.shape[0], -1)
|
|
212
219
|
else:
|
|
213
220
|
return basis
|
|
214
|
-
|
|
215
221
|
|
|
216
|
-
|
|
222
|
+
|
|
223
|
+
def zernike_fit(opd, mask, modes, normalize=True, order='noll', rho=None,
|
|
224
|
+
theta=None):
|
|
217
225
|
"""Fit a Zernike basis set to an OPD.
|
|
218
226
|
|
|
219
227
|
Parameters
|
|
@@ -227,6 +235,8 @@ def zernike_fit(opd, mask, modes, normalize=True, rho=None, theta=None):
|
|
|
227
235
|
normalize : bool, optional
|
|
228
236
|
If True (default), the output is normalized according to [1]. If False,
|
|
229
237
|
the output value ranges [-1, 1] over the mask.
|
|
238
|
+
order : {'noll', 'fringe', 'ansi'}, optional
|
|
239
|
+
Zernike ordering scheme. Default is 'noll'.
|
|
230
240
|
rho : array_like, optional
|
|
231
241
|
Radial coordinates of the mask array. :attr:`rho` should be 0 at the
|
|
232
242
|
origin and 1 at the edge of the circle.
|
|
@@ -263,14 +273,14 @@ def zernike_fit(opd, mask, modes, normalize=True, rho=None, theta=None):
|
|
|
263
273
|
opd = np.asarray(opd)
|
|
264
274
|
mask = np.asarray(mask)
|
|
265
275
|
|
|
266
|
-
basis = zernike_basis(mask, modes, True, normalize, rho, theta)
|
|
276
|
+
basis = zernike_basis(mask, modes, True, normalize, order, rho, theta)
|
|
267
277
|
|
|
268
278
|
basis = np.linalg.pinv(basis)
|
|
269
279
|
|
|
270
280
|
return np.einsum('ij,i->j', basis, opd.ravel())
|
|
271
281
|
|
|
272
282
|
|
|
273
|
-
def zernike_remove(opd, mask, modes, rho=None, theta=None):
|
|
283
|
+
def zernike_remove(opd, mask, modes, order='noll', rho=None, theta=None):
|
|
274
284
|
"""Fit and remove a Zernike basis set from an OPD.
|
|
275
285
|
|
|
276
286
|
Parameters
|
|
@@ -281,6 +291,8 @@ def zernike_remove(opd, mask, modes, rho=None, theta=None):
|
|
|
281
291
|
Binary mask defining the extent to compute the Zernike basis over.
|
|
282
292
|
modes : array_like
|
|
283
293
|
List of modes (Noll ordering) to remove.
|
|
294
|
+
order : {'noll', 'fringe', 'ansi'}, optional
|
|
295
|
+
Zernike ordering scheme. Default is 'noll'.
|
|
284
296
|
rho : array_like, optional
|
|
285
297
|
Radial coordinates of the mask array. :attr:`rho` should be 0 at the
|
|
286
298
|
origin and 1 at the edge of the circle.
|
|
@@ -303,35 +315,34 @@ def zernike_remove(opd, mask, modes, rho=None, theta=None):
|
|
|
303
315
|
opd = np.asarray(opd)
|
|
304
316
|
mask = np.asarray(mask)
|
|
305
317
|
|
|
306
|
-
coeffs = zernike_fit(opd, mask, modes, rho, theta)
|
|
307
|
-
fit_opd = zernike_compose(mask, coeffs, rho, theta)
|
|
318
|
+
coeffs = zernike_fit(opd, mask, modes, True, order, rho, theta)
|
|
319
|
+
fit_opd = zernike_compose(mask, coeffs, True, order, rho, theta)
|
|
308
320
|
|
|
309
321
|
residual = opd - fit_opd
|
|
310
322
|
|
|
311
323
|
return residual
|
|
312
324
|
|
|
313
325
|
|
|
314
|
-
def
|
|
315
|
-
"""Convert
|
|
326
|
+
def noll_to_nm(j):
|
|
327
|
+
"""Convert Noll index to (n, m) index.
|
|
316
328
|
|
|
317
329
|
Parameters
|
|
318
330
|
----------
|
|
319
331
|
j : int
|
|
320
|
-
Noll Zernike index
|
|
332
|
+
Noll Zernike index
|
|
321
333
|
|
|
322
334
|
Returns
|
|
323
335
|
-------
|
|
324
336
|
n : int
|
|
325
|
-
Radial
|
|
337
|
+
Radial degree.
|
|
326
338
|
m : int
|
|
327
|
-
Azimuthal
|
|
328
|
-
|
|
339
|
+
Azimuthal degree.
|
|
329
340
|
"""
|
|
330
341
|
if j < 1:
|
|
331
|
-
raise ValueError('
|
|
342
|
+
raise ValueError('Noll index must be greater than zero')
|
|
332
343
|
|
|
333
344
|
# find the row j is in
|
|
334
|
-
n = int(np.ceil((-1 + np.sqrt(1 + 8*j)) / 2) - 1)
|
|
345
|
+
n = int(np.ceil((-1 + np.sqrt(1 + 8 * j)) / 2) - 1)
|
|
335
346
|
if n == 0:
|
|
336
347
|
m = 0
|
|
337
348
|
else:
|
|
@@ -349,13 +360,63 @@ def zernike_index(j):
|
|
|
349
360
|
else:
|
|
350
361
|
row_m = [0]
|
|
351
362
|
|
|
352
|
-
for
|
|
363
|
+
for _ in range(n // 2):
|
|
353
364
|
row_m.append(row_m[-1] + 2)
|
|
354
365
|
row_m.append(row_m[-1])
|
|
355
366
|
|
|
356
367
|
m = row_m[r] * sign
|
|
357
368
|
|
|
358
|
-
return
|
|
369
|
+
return int(n), int(m)
|
|
370
|
+
|
|
371
|
+
|
|
372
|
+
def ansi_to_nm(j):
|
|
373
|
+
"""Convert ANSI index to (n, m) index.
|
|
374
|
+
|
|
375
|
+
Parameters
|
|
376
|
+
----------
|
|
377
|
+
j : int
|
|
378
|
+
ANSI index
|
|
379
|
+
|
|
380
|
+
Returns
|
|
381
|
+
-------
|
|
382
|
+
n : int
|
|
383
|
+
Radial degree
|
|
384
|
+
m : int
|
|
385
|
+
Azimuthal degree
|
|
386
|
+
"""
|
|
387
|
+
if j < 0:
|
|
388
|
+
raise ValueError('ANSI index must be zero or larger')
|
|
389
|
+
|
|
390
|
+
# https://en.wikipedia.org/wiki/Zernike_polynomials
|
|
391
|
+
n = np.floor((np.sqrt(8 * j + 1) - 1) / 2)
|
|
392
|
+
m = 2*j - n*(n + 2)
|
|
393
|
+
return int(n), int(m)
|
|
394
|
+
|
|
395
|
+
|
|
396
|
+
def fringe_to_nm(j):
|
|
397
|
+
"""Convert Fringe index to (n, m) index.
|
|
398
|
+
|
|
399
|
+
Parameters
|
|
400
|
+
----------
|
|
401
|
+
j : int
|
|
402
|
+
Fringe index
|
|
403
|
+
|
|
404
|
+
Returns
|
|
405
|
+
-------
|
|
406
|
+
n : int
|
|
407
|
+
Radial degree
|
|
408
|
+
m : int
|
|
409
|
+
Azimuthal degree
|
|
410
|
+
"""
|
|
411
|
+
if j < 1:
|
|
412
|
+
raise ValueError('Fringe index must be greater than zero')
|
|
413
|
+
|
|
414
|
+
# https://github.com/arielmission-space/PAOS/blob/main/paos/classes/zernike.py
|
|
415
|
+
m_n = 2 * (np.ceil(np.sqrt(j)) - 1)
|
|
416
|
+
g_s = (m_n / 2) ** 2 + 1
|
|
417
|
+
n = m_n / 2 + np.floor((j - g_s) / 2)
|
|
418
|
+
m = (m_n - n) * (1 - np.mod(j - g_s, 2) * 2)
|
|
419
|
+
return int(n), int(m)
|
|
359
420
|
|
|
360
421
|
|
|
361
422
|
def zernike_coordinates(mask, shift=None, angle=0, indexing='ij'):
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|