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.
Files changed (29) hide show
  1. {prtools-1.3.2/prtools.egg-info → prtools-1.4.0}/PKG-INFO +1 -1
  2. {prtools-1.3.2 → prtools-1.4.0}/prtools/__init__.py +11 -2
  3. prtools-1.4.0/prtools/convolve.py +348 -0
  4. {prtools-1.3.2 → prtools-1.4.0}/prtools/fourier.py +1 -1
  5. {prtools-1.3.2 → prtools-1.4.0}/prtools/jax.py +25 -6
  6. {prtools-1.3.2 → prtools-1.4.0}/prtools/zernike.py +91 -30
  7. {prtools-1.3.2 → prtools-1.4.0/prtools.egg-info}/PKG-INFO +1 -1
  8. {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/SOURCES.txt +1 -0
  9. {prtools-1.3.2 → prtools-1.4.0}/LICENSE +0 -0
  10. {prtools-1.3.2 → prtools-1.4.0}/NOTICE +0 -0
  11. {prtools-1.3.2 → prtools-1.4.0}/README.rst +0 -0
  12. {prtools-1.3.2 → prtools-1.4.0}/prtools/array.py +0 -0
  13. {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/__init__.py +0 -0
  14. {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/_base.py +0 -0
  15. {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/_jax.py +0 -0
  16. {prtools-1.3.2 → prtools-1.4.0}/prtools/backend/_numpy.py +0 -0
  17. {prtools-1.3.2 → prtools-1.4.0}/prtools/cost.py +0 -0
  18. {prtools-1.3.2 → prtools-1.4.0}/prtools/misc.py +0 -0
  19. {prtools-1.3.2 → prtools-1.4.0}/prtools/segmented.py +0 -0
  20. {prtools-1.3.2 → prtools-1.4.0}/prtools/shape.py +0 -0
  21. {prtools-1.3.2 → prtools-1.4.0}/prtools/sparse.py +0 -0
  22. {prtools-1.3.2 → prtools-1.4.0}/prtools/stats.py +0 -0
  23. {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/dependency_links.txt +0 -0
  24. {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/requires.txt +0 -0
  25. {prtools-1.3.2 → prtools-1.4.0}/prtools.egg-info/top_level.txt +0 -0
  26. {prtools-1.3.2 → prtools-1.4.0}/pyproject.toml +0 -0
  27. {prtools-1.3.2 → prtools-1.4.0}/setup.cfg +0 -0
  28. {prtools-1.3.2 → prtools-1.4.0}/tests/test_array.py +0 -0
  29. {prtools-1.3.2 → prtools-1.4.0}/tests/test_fourier.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: prtools
3
- Version: 1.3.2
3
+ Version: 1.4.0
4
4
  Summary: Utility functions for image-based phase retrieval
5
5
  Author: Andy Kee
6
6
  License: MIT License
@@ -1,4 +1,4 @@
1
- __version__ = '1.3.2'
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
- propagation (2007)
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(fun, x0, gtol=None, maxiter=None, callback=None):
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
- fun : callable
41
- The objective function to be minimied
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
- value_and_grad_fun = optax.value_and_grad_from_state(fun)
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
- value, grad = value_and_grad_fun(params, state=state)
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=fun)
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
- #out = np.zeros_like(mask)
48
- #mask_slice = util.boundary_slice(mask, pad=0)
49
- #mask = mask[mask_slice]
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, rho=None, theta=None):
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
- def zernike_fit(opd, mask, modes, normalize=True, rho=None, theta=None):
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 zernike_index(j):
315
- """Convert a 1-D Noll index to a 2-D radiul and azimuthal index.
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 Zernike index.
337
+ Radial degree.
326
338
  m : int
327
- Azimuthal Zernike index.
328
-
339
+ Azimuthal degree.
329
340
  """
330
341
  if j < 1:
331
- raise ValueError('Zernike index j must be a positive integer')
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 i in range(int(np.floor(n / 2))):
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 m, n
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'):
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: prtools
3
- Version: 1.3.2
3
+ Version: 1.4.0
4
4
  Summary: Utility functions for image-based phase retrieval
5
5
  Author: Andy Kee
6
6
  License: MIT License
@@ -4,6 +4,7 @@ README.rst
4
4
  pyproject.toml
5
5
  prtools/__init__.py
6
6
  prtools/array.py
7
+ prtools/convolve.py
7
8
  prtools/cost.py
8
9
  prtools/fourier.py
9
10
  prtools/jax.py
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