From 6be9f07a5e85a03aececba5aabfe1423c6ca3f91 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Tue, 8 Feb 2022 17:45:37 -0600 Subject: [PATCH 01/13] DEV: Prototype some zernike polynomials --- src/tike/zernike.py | 45 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 src/tike/zernike.py diff --git a/src/tike/zernike.py b/src/tike/zernike.py new file mode 100644 index 00000000..7bd475f0 --- /dev/null +++ b/src/tike/zernike.py @@ -0,0 +1,45 @@ +import numpy as np + + +def Z(m, n, ρ, θ): + assert np.all(0 <= ρ <= 1), "Radii must be in range [0, 1]." + assert n >= 0, "Radial degree must be non-negative." + _m_ = np.abs(m) + assert _m_ <= n, "Angular frequency must be less than radial degree." + return N(m, n) * R(_m_, n, ρ) * np.exp(1j * _m_ * θ) + + +def N(m, n): + """Zernike normalization factor.""" + # @StevenHenke1 this must be floating point division? + return np.sqrt(2 * (n + 1) / (1 + (m == 0))) + + +def R(m, n, ρ): + """Zernike radial polynomial.""" + if (n - m) % 2: + return 0 + else: + # Initialize with k=0 case because this term will always be included + sign = 1 + b0 = 1 + b1 = 1 + result = ρ**n + # @StevenHenke2 these must be integer division? + # @StevenHenke3 Does this sum include k = (n - m) // 2 ? + for k in range(1, (n - m) // 2 + 1): + sign = -sign + b0 *= bino(n - k, k) + b1 *= bino(n - 2 * k, (n - m) // 2 - k) + result += sign * b0 * b1 * ρ**(n - 2 * k) + return result + + +def bino(n, i): + """One product term of the binomial coefficient.""" + # @StevenHenke4 these must be integer division? + assert i >= 0 + if i == 0: + return 1 + else: + return (n - i + 1) // i From d69838663b3babe4eac011515c4a1b3bc819ac9e Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Tue, 8 Feb 2022 17:57:43 -0600 Subject: [PATCH 02/13] DOC: Add more documentation to zernike module --- src/tike/zernike.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 7bd475f0..430d6c11 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -1,7 +1,18 @@ +"""Provide functions to generate complex coefficients of Zernike polynomials +on a discrete grid.""" + import numpy as np -def Z(m, n, ρ, θ): +def Z(m: int, n: int, ρ: np.array, θ: np.array) -> np.array: + """Return the coefficients of the Zernike[m,n] polynomial. + + Parameters + ---------- + ρ: radius + θ: angle + """ + assert np.all(0 <= ρ <= 1), "Radii must be in range [0, 1]." assert n >= 0, "Radial degree must be non-negative." _m_ = np.abs(m) @@ -9,13 +20,13 @@ def Z(m, n, ρ, θ): return N(m, n) * R(_m_, n, ρ) * np.exp(1j * _m_ * θ) -def N(m, n): +def N(m: int, n: int) -> int: """Zernike normalization factor.""" # @StevenHenke1 this must be floating point division? return np.sqrt(2 * (n + 1) / (1 + (m == 0))) -def R(m, n, ρ): +def R(m: int, n: int, ρ: np.array) -> np.array: """Zernike radial polynomial.""" if (n - m) % 2: return 0 @@ -35,7 +46,7 @@ def R(m, n, ρ): return result -def bino(n, i): +def bino(n: int, i: int) -> int: """One product term of the binomial coefficient.""" # @StevenHenke4 these must be integer division? assert i >= 0 From 217e66c8c0468850208e25a5b50cf8dfd87fa380 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Fri, 18 Feb 2022 19:12:35 -0600 Subject: [PATCH 03/13] Try oseudo-zernikes --- src/tike/zernike.py | 78 ++++++++++++++++++++++++------------------- tests/test_zernike.py | 21 ++++++++++++ 2 files changed, 65 insertions(+), 34 deletions(-) create mode 100644 tests/test_zernike.py diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 430d6c11..273fe78e 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -1,56 +1,66 @@ """Provide functions to generate complex coefficients of Zernike polynomials on a discrete grid.""" +from math import factorial import numpy as np -def Z(m: int, n: int, ρ: np.array, θ: np.array) -> np.array: +def Z(m: int, n: int, radius: np.array, angle: np.array) -> np.array: """Return the coefficients of the Zernike[m,n] polynomial. Parameters ---------- - ρ: radius - θ: angle + m : int + angular frequency + n : int + radial degree + radius: float [0, 1] + radius + angle: float radians + angle """ - - assert np.all(0 <= ρ <= 1), "Radii must be in range [0, 1]." assert n >= 0, "Radial degree must be non-negative." _m_ = np.abs(m) assert _m_ <= n, "Angular frequency must be less than radial degree." - return N(m, n) * R(_m_, n, ρ) * np.exp(1j * _m_ * θ) + polynomial = R(_m_, n, radius) * np.exp(1j * m * angle) + polynomial[np.logical_or(radius < 0, radius > 1)] = np.nan + return polynomial def N(m: int, n: int) -> int: """Zernike normalization factor.""" - # @StevenHenke1 this must be floating point division? return np.sqrt(2 * (n + 1) / (1 + (m == 0))) -def R(m: int, n: int, ρ: np.array) -> np.array: +def R(m: int, n: int, radius: np.array) -> np.array: """Zernike radial polynomial.""" - if (n - m) % 2: - return 0 - else: - # Initialize with k=0 case because this term will always be included - sign = 1 - b0 = 1 - b1 = 1 - result = ρ**n - # @StevenHenke2 these must be integer division? - # @StevenHenke3 Does this sum include k = (n - m) // 2 ? - for k in range(1, (n - m) // 2 + 1): - sign = -sign - b0 *= bino(n - k, k) - b1 *= bino(n - 2 * k, (n - m) // 2 - k) - result += sign * b0 * b1 * ρ**(n - 2 * k) - return result - - -def bino(n: int, i: int) -> int: - """One product term of the binomial coefficient.""" - # @StevenHenke4 these must be integer division? - assert i >= 0 - if i == 0: - return 1 - else: - return (n - i + 1) // i + # Initialize with k=0 case because this term will always be included + sign = 1 + result = bino(n, m, 0) * radius**n + for k in range(1, n - m + 1): + sign = -sign + result += sign * bino(n, m, k) * radius**(n - k) + return result + + +def bino(n: int, m: int, k: int) -> int: + """Return the approximate binomial coeffient (a b).""" + return int( + factorial(2 * n + 1 - k) / factorial(k) / factorial(n - m - k) / + factorial(n + m - k + 1)) + + +def mode(size: int, n: int) -> np.array: + endpoint = 1.0 - 1 / (2 * size) + x = np.linspace(-endpoint, endpoint, size, endpoint=True) + coords = np.meshgrid(x, x, indexing='ij') + radius = np.linalg.norm(coords, axis=0) + theta = np.arctan(coords[0] / coords[1]) + + basis = [] + for _n in range(0, n): + for m in range(-_n, _n + 1): + basis.append(Z(m, n, radius, theta)) + + basis = np.stack(basis, axis=0) + return basis diff --git a/tests/test_zernike.py b/tests/test_zernike.py new file mode 100644 index 00000000..b8a01b1f --- /dev/null +++ b/tests/test_zernike.py @@ -0,0 +1,21 @@ +import unittest +import os + +import tike.zernike +import tike.view +import matplotlib.pyplot as plt + +testdir = os.path.dirname(__file__) + + +class TestZernike(unittest.TestCase): + + def test_zernike(self): + + fname = os.path.join(testdir, 'result', 'zernike') + os.makedirs(fname, exist_ok=True) + for i, Z in enumerate(tike.zernike.mode(256, 3)): + plt.figure() + tike.view.plot_complex(Z) + plt.savefig(os.path.join(fname, f"zernike-{i}.png")) + plt.close() From 0f80f4f01bc0ac6529200b98d82169bb1c710652 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Mon, 21 Feb 2022 18:27:32 -0600 Subject: [PATCH 04/13] Pretty confident that zernikes are correct now --- src/tike/zernike.py | 40 +++++++++++++++++++++++++--------------- tests/test_zernike.py | 36 +++++++++++++++++++++++++++++++++--- 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 273fe78e..85cd07c0 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -22,12 +22,15 @@ def Z(m: int, n: int, radius: np.array, angle: np.array) -> np.array: assert n >= 0, "Radial degree must be non-negative." _m_ = np.abs(m) assert _m_ <= n, "Angular frequency must be less than radial degree." - polynomial = R(_m_, n, radius) * np.exp(1j * m * angle) - polynomial[np.logical_or(radius < 0, radius > 1)] = np.nan + if m < 0: + polynomial = R(_m_, n, radius) * np.sin(m * angle) + else: + polynomial = R(_m_, n, radius) * np.cos(m * angle) + polynomial[np.logical_or(radius < 0, radius > 1)] = 0 return polynomial -def N(m: int, n: int) -> int: +def N(m: int, n: int) -> float: """Zernike normalization factor.""" return np.sqrt(2 * (n + 1) / (1 + (m == 0))) @@ -35,32 +38,39 @@ def N(m: int, n: int) -> int: def R(m: int, n: int, radius: np.array) -> np.array: """Zernike radial polynomial.""" # Initialize with k=0 case because this term will always be included - sign = 1 - result = bino(n, m, 0) * radius**n - for k in range(1, n - m + 1): + sign = -1 + result = 0 * radius + for k in range(0, (n - m) // 2 + 1): sign = -sign - result += sign * bino(n, m, k) * radius**(n - k) + result += sign * bino(n, m, k) * radius**(n - 2 * k) return result def bino(n: int, m: int, k: int) -> int: """Return the approximate binomial coeffient (a b).""" return int( - factorial(2 * n + 1 - k) / factorial(k) / factorial(n - m - k) / - factorial(n + m - k + 1)) + factorial(n - k) / factorial(k) / factorial((n + m) // 2 - k) / + factorial((n - m) // 2 - k)) -def mode(size: int, n: int) -> np.array: - endpoint = 1.0 - 1 / (2 * size) +def zernike_basis(size: int, degree: int) -> np.array: + """Return all circular Zernike basis for radial degree up to n.""" + endpoint = 1.0 - 1.0 / (2 * size) x = np.linspace(-endpoint, endpoint, size, endpoint=True) coords = np.meshgrid(x, x, indexing='ij') radius = np.linalg.norm(coords, axis=0) - theta = np.arctan(coords[0] / coords[1]) + theta = np.arctan2(coords[0], coords[1]) basis = [] - for _n in range(0, n): - for m in range(-_n, _n + 1): - basis.append(Z(m, n, radius, theta)) + for m, n in valid_zernike_indices(degree): + basis.append(Z(m, n, radius, theta)) basis = np.stack(basis, axis=0) return basis + + +def valid_zernike_indices(degree): + for n in range(0, degree): + for m in range(-n, n + 1): + if (n - abs(m)) % 2 == 0: + yield m, n diff --git a/tests/test_zernike.py b/tests/test_zernike.py index b8a01b1f..b983f259 100644 --- a/tests/test_zernike.py +++ b/tests/test_zernike.py @@ -4,6 +4,7 @@ import tike.zernike import tike.view import matplotlib.pyplot as plt +import numpy as np testdir = os.path.dirname(__file__) @@ -14,8 +15,37 @@ def test_zernike(self): fname = os.path.join(testdir, 'result', 'zernike') os.makedirs(fname, exist_ok=True) - for i, Z in enumerate(tike.zernike.mode(256, 3)): + for i, Z in enumerate(tike.zernike.zernike_basis(256, degree=10)): plt.figure() - tike.view.plot_complex(Z) - plt.savefig(os.path.join(fname, f"zernike-{i}.png")) + tike.view.plot_complex(Z, rmin=-1, rmax=1) + plt.savefig(os.path.join(fname, f"zernike-{i:02d}.png")) plt.close() + + def _radial_template(self, m=0): + fname = os.path.join(testdir, 'result', 'zernike') + os.makedirs(fname, exist_ok=True) + + radius = np.linspace(0, 1, 200) + + plt.figure() + labels = [] + for n in range(0, 9): + if (n + m) % 2 == 0: + v = tike.zernike.R(m, n, radius) + plt.plot( + radius, + v, + ) + labels.append(n) + plt.legend(labels) + plt.savefig(os.path.join(fname, f"radial-function-{m}.png")) + plt.close() + + def test_radial(self): + self._radial_template(0) + + def test_radial_1(self): + self._radial_template(1) + + def test_radial_2(self): + self._radial_template(2) From c148165cd950b4346118e22da7e78200191ee687 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Mon, 21 Feb 2022 18:54:02 -0600 Subject: [PATCH 05/13] Use efficient bonomials --- src/tike/zernike.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 85cd07c0..e859a948 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -42,15 +42,18 @@ def R(m: int, n: int, radius: np.array) -> np.array: result = 0 * radius for k in range(0, (n - m) // 2 + 1): sign = -sign - result += sign * bino(n, m, k) * radius**(n - 2 * k) + b0 = bino(n - k, k) + b1 = bino(n - 2 * k, (n - m) // 2 - k) + result += sign * b0 * b1 * radius**(n - 2 * k) return result -def bino(n: int, m: int, k: int) -> int: +def bino(a: int, b: int) -> int: """Return the approximate binomial coeffient (a b).""" - return int( - factorial(n - k) / factorial(k) / factorial((n + m) // 2 - k) / - factorial((n - m) // 2 - k)) + result = 1 + for i in range(1, b + 1): + result *= (a - i + 1) / i + return result def zernike_basis(size: int, degree: int) -> np.array: From 3f82d2988ee8f8a6b5044e71af2bb2c0cb848264 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Tue, 22 Feb 2022 13:01:24 -0600 Subject: [PATCH 06/13] DOC; Update documentation --- src/tike/zernike.py | 74 ++++++++++++++++++++++++++++++++----------- tests/test_zernike.py | 3 +- 2 files changed, 58 insertions(+), 19 deletions(-) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index e859a948..61a88d71 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -1,27 +1,30 @@ -"""Provide functions to generate complex coefficients of Zernike polynomials -on a discrete grid.""" +"""Provide functions to evaluate Zernike polynomials on a discrete grid.""" -from math import factorial import numpy as np def Z(m: int, n: int, radius: np.array, angle: np.array) -> np.array: - """Return the coefficients of the Zernike[m,n] polynomial. + """Return values of Zernike[m,n] polynomial at given radii, angles. + + Values outside valid radius will be zero. Parameters ---------- m : int - angular frequency + Angular frequency of the polynomial. n : int - radial degree + Radial degree of the polynomial. radius: float [0, 1] - radius + The radial coordinates of the evaluated polynomial. angle: float radians - angle + The angular coordinates of the evaluated polynomial. + """ - assert n >= 0, "Radial degree must be non-negative." + if n < 0: + raise ValueError("Radial degree must be non-negative.") _m_ = np.abs(m) - assert _m_ <= n, "Angular frequency must be less than radial degree." + if _m_ > n: + raise ValueError("Angular frequency must be less than radial degree.") if m < 0: polynomial = R(_m_, n, radius) * np.sin(m * angle) else: @@ -36,19 +39,38 @@ def N(m: int, n: int) -> float: def R(m: int, n: int, radius: np.array) -> np.array: - """Zernike radial polynomial.""" + """Return the values of the Zernike radial polynomial at the given radii. + + This polynomial matches Figure 3 in Lakshminarayanan & Fleck (2011). + + Parameters + ---------- + m : int + Angular frequency of the polynomial. + n : int + Radial degree of the polynomial. + radius: float [0, 1] + The radial coordinates of the evaluated polynomial. + + References + ---------- + Vasudevan Lakshminarayanan & Andre Fleck (2011): Zernike polynomials: a + guide, Journal of ModernOptics, 58:7, 545-561 + http://dx.doi.org/10.1080/09500340.2011.554896 + + """ # Initialize with k=0 case because this term will always be included sign = -1 result = 0 * radius for k in range(0, (n - m) // 2 + 1): sign = -sign - b0 = bino(n - k, k) - b1 = bino(n - 2 * k, (n - m) // 2 - k) + b0 = _bino(n - k, k) + b1 = _bino(n - 2 * k, (n - m) // 2 - k) result += sign * b0 * b1 * radius**(n - 2 * k) return result -def bino(a: int, b: int) -> int: +def _bino(a: int, b: int) -> int: """Return the approximate binomial coeffient (a b).""" result = 1 for i in range(1, b + 1): @@ -56,8 +78,23 @@ def bino(a: int, b: int) -> int: return result -def zernike_basis(size: int, degree: int) -> np.array: - """Return all circular Zernike basis for radial degree up to n.""" +def basis(size: int, degree: int) -> np.array: + """Return all circular Zernike basis up to given radial degree. + + Parameters + ---------- + size : int + The width of the discrete basis in pixel. + degree : int + The maximum radial degree of the polynomial (not inclusive). The number + of degrees included in the set of bases. + + Returns + ------- + basis : (degree, size, size) + The Zernike bases. + + """ endpoint = 1.0 - 1.0 / (2 * size) x = np.linspace(-endpoint, endpoint, size, endpoint=True) coords = np.meshgrid(x, x, indexing='ij') @@ -65,14 +102,15 @@ def zernike_basis(size: int, degree: int) -> np.array: theta = np.arctan2(coords[0], coords[1]) basis = [] - for m, n in valid_zernike_indices(degree): + for m, n in valid_indices(degree): basis.append(Z(m, n, radius, theta)) basis = np.stack(basis, axis=0) return basis -def valid_zernike_indices(degree): +def valid_indices(degree: int) -> tuple: + """Enumerate all valid zernike indices (m,n) up to the given degree.""" for n in range(0, degree): for m in range(-n, n + 1): if (n - abs(m)) % 2 == 0: diff --git a/tests/test_zernike.py b/tests/test_zernike.py index b983f259..7de167b6 100644 --- a/tests/test_zernike.py +++ b/tests/test_zernike.py @@ -15,7 +15,7 @@ def test_zernike(self): fname = os.path.join(testdir, 'result', 'zernike') os.makedirs(fname, exist_ok=True) - for i, Z in enumerate(tike.zernike.zernike_basis(256, degree=10)): + for i, Z in enumerate(tike.zernike.basis(256, degree=10)): plt.figure() tike.view.plot_complex(Z, rmin=-1, rmax=1) plt.savefig(os.path.join(fname, f"zernike-{i:02d}.png")) @@ -38,6 +38,7 @@ def _radial_template(self, m=0): ) labels.append(n) plt.legend(labels) + plt.ylim([-1, 1]) plt.savefig(os.path.join(fname, f"radial-function-{m}.png")) plt.close() From 9d0041b785d72066f715d21eac1a8de6a7d4e6fa Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Tue, 22 Feb 2022 13:08:48 -0600 Subject: [PATCH 07/13] Compatability for cupy --- src/tike/zernike.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 61a88d71..4632d460 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -97,7 +97,7 @@ def basis(size: int, degree: int) -> np.array: """ endpoint = 1.0 - 1.0 / (2 * size) x = np.linspace(-endpoint, endpoint, size, endpoint=True) - coords = np.meshgrid(x, x, indexing='ij') + coords = np.stack(np.meshgrid(x, x, indexing='ij'), axis=0) radius = np.linalg.norm(coords, axis=0) theta = np.arctan2(coords[0], coords[1]) From d18873dd225b6708ec8f4ebe8fdb616ec2d15493 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Thu, 24 Feb 2022 13:45:43 -0600 Subject: [PATCH 08/13] Zernike transform prototype --- src/tike/zernike.py | 8 ++++---- tests/data/probe.png | Bin 0 -> 3245 bytes tests/test_zernike.py | 45 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 49 insertions(+), 4 deletions(-) create mode 100644 tests/data/probe.png diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 4632d460..d31347e6 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -78,7 +78,7 @@ def _bino(a: int, b: int) -> int: return result -def basis(size: int, degree: int) -> np.array: +def basis(size: int, degree_min: int, degree_max: int) -> np.array: """Return all circular Zernike basis up to given radial degree. Parameters @@ -102,16 +102,16 @@ def basis(size: int, degree: int) -> np.array: theta = np.arctan2(coords[0], coords[1]) basis = [] - for m, n in valid_indices(degree): + for m, n in valid_indices(degree_min, degree_max): basis.append(Z(m, n, radius, theta)) basis = np.stack(basis, axis=0) return basis -def valid_indices(degree: int) -> tuple: +def valid_indices(degree_min: int, degree_max: int) -> tuple: """Enumerate all valid zernike indices (m,n) up to the given degree.""" - for n in range(0, degree): + for n in range(degree_min, degree_max): for m in range(-n, n + 1): if (n - abs(m)) % 2 == 0: yield m, n diff --git a/tests/data/probe.png b/tests/data/probe.png new file mode 100644 index 0000000000000000000000000000000000000000..cce3eb970cf02af386ae85b6f48189f3c676540e GIT binary patch literal 3245 zcmV;e3{vxnP)u*i~x-9Xp{X|TjlgccMLStJmwq5`riNCyyMRKN{r z#kQhYJ6-HH-G1tHnQ8x$&a~%En`G{id%y2}mh+r*o`7)57^5MPLUBUzMhi|9<+)f# zQd?~zrG(}RpHNAuppq~miNI^Z6bUF3oa2uQ#wllnN_?Lafw$Ty8(B$1C7Gs-8lyN@ zjw^!$@P&|S!2;Zn%24AB93N^*fjq4GWzYI4Dd@_g}K`D_$TnJo((72(DI~&`~;!50*OXNhfMY;&s zgbZ=xbQZ8skLArnY5eChR1QfA8%K^3YK=}ZW3@vjL^_IZMF^EpW~d0qY7maq#CHH? zAylPK5=5lYQ}j+Kjk<)!8pD}X#!@X&E<*4*S{TwmhDZjbLvx~t7@kF1QH_=hjf|PA zlxnFI2wEkC8x#ACl|cSzM~LS{_JLG6l#P;2+@@=qzdgGzI9OpAe0@1Rm(xRy$fM6I%0W zk~Nm2hMHw@s?byPDUQcL%OxKx9C`sH0t1>6KY$TcS>$5VIB$OEq>kpw%$~MtQ)g#q zTQ#L=oKsuR1ebdxAiw73Z@9k=D@0{2%xviL3 zk}}ItTYS;_`J@Vk3xEdzE5MC|*fCTfj;gKWS|&C3ZFqEO=XX~3OzoJoeDSQ#nLSO- zah=Nx<>%j%kM&Ft24t*<%qfn(3nWoj*U{VA-9CHS?$FZvw{lx3{K7aqUGq+BxT)$=IqLsrN zmrsfc<5Ew700aQ#n;`?G@eD%5Ofu;ji?w3=%Jr+(&)x9tZuhP853cY{rrSuoDvWt!EQ`AtAvS555;!o+;yu4hrjjQiPt`O`_|o$?p!qyZhEDW%NhntKN}?K$<%jq~pv z8QOO0_-7C9zJBnHcTQft`tp{J)+yuTxX7XcQUF_lZcutm0#Pno&f9~VkzP11A^$Q2r&g-7qhVf=;1l>>yXY&3qdCs4vtnx zkuO(OWB0rbdk!Bv_Qs{d2cMd{ddJAN6>FBhaPW;ACok?nCGU{Q(!EUHQ z$eEE+Yh&Y_vrTgs4xc!E^W^DsyALdQWMRjOC6f*h?z?{S<>yBR*LC(}wTY3=$B_bx zsApCd7Q}&Dl}J< zET>DcOSGmDAQ=$MAe0jnXNyUkHqKZ#|L~dHH*W3N`{<@`49)-I_glYtaK+Az$4?#E zv1wp_eSK0_AuKVgcr(Kfb(q!>n=TV&W;$g1(?}AtHcCOW`1EIaf;4j^6nr$6h&g{n)2(jx6XMdZvEL z@Vf40FFtwo-o0OcZ|l;9tvQ}pN}~uafEzHRSLav&-brF}4Y{moc7OkypZxNlpZ#j{ zvqRITKelZ9&c|2pxH|g&-MbeLJu=|7(W)iNoR8VsVsSJY7AhVEG_cU`>v9PLh(~9MCF48ejATiBw6<*#s3K~KXUEs2cQ4p zH}Bni>(Kg#a}(w*ZcHZU#+5=5CDu8eqypS31CA7k_hZmo{L8qyhQ{*jdA;)%ZrOL~ z+J|3#{namiduj99)k`-$F?DK3#b~A+1~+sMH!_u$`uN-P~f6@ zO>h5>-3!~-JpaOh!_S_%uz%OG83Qf#7UtV3lX*f15dJ6c1~E|&N;tZxQZ2bS$7M9J zZQ8tn<=x9S9@yJA^U0lk3wm3+#>aUp6(Y!BPrN%9a7zIUEdF3}g+(%P$--u%9m1-7 z{j6DCGuql3Dsvxhm|UvnOrdC%voHua4L_h}16tS66%Z@2!q7N&B82>xswuL`6>AOE zV%eFAebc8q?dr!nC81<+_ZlH66&e9P@S_BSExccaXo0LkFm9<5H`E-=QEl?{_NMk) zLu2NOEf#Ll`VoM`=3$5~c$)`_g~i6M!L+<00C|@>NmSBCt`Ve^|A-wAXRD?rD;E7zE zhZif<*VpT)2-l5dKqDMyAzJ7&)&p36Ul%%N+Cyl*aZp{f+6rcz&P6D%#uS!>HBlu+ z!%-52Ilv$x6XKmJECyOIHj`ee1@U=Db_U$HRIkT{clOd4Y+3oU>b{|L?o!Fu?5LJ zE(stmpxnftg|aJr3B2K8w%SUp0FBNQ7*#GC=?eoV2@!qzhHY>kCO^>>FU zgh9i;C>1nOctQc4P`^09zk1i;8=_Kgc>Kx>8zu7@=F1u!fb|xMQu=?d4|9k&0rAtU z#F*PANLCP^4GF#d(D$exy-v!zsGr9_Er(#r`{j<*#^Df|0o>xK~Nq}5E1sS zD2mWdUspuvmpjeIBETyE7!LwS483uufEB)kslpT17o-!S-4-beX;K4NVsA!aA!E#Y fG`|u1ZH@gGZ7 0).astype('float32') + + # x = tike.linalg.lstsq(basis, f0, weights=w) + x, _, _, _ = np.linalg.lstsq(basis, f0, rcond=1e-4) + + f1 = basis @ x + f1 = f1.reshape(size, size) + plt.imsave(os.path.join(fname, f'basis-{d:02d}.png'), f1) + + plt.figure() + plt.title(f"basis weights for {d} degree polynomials") + plt.bar(list(range(x.size)), x.flatten()) + plt.savefig(os.path.join(fname, f'basis-w-{d:02d}.png')) + plt.close() From d9c743aa99d32390cfc485a3d32abaa0ca747660 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Thu, 18 Jan 2024 14:46:21 -0600 Subject: [PATCH 09/13] NEW: Wrap zenike functions in Operator class --- src/tike/operators/cupy/__init__.py | 3 +- src/tike/operators/cupy/zernike.py | 63 +++++++++++++++++++ src/tike/zernike.py | 30 ++++++--- tests/operators/test_zernike.py | 51 ++++++++++++++++ tests/test_zernike.py | 95 +++++++++++++++++++++++------ 5 files changed, 217 insertions(+), 25 deletions(-) create mode 100644 src/tike/operators/cupy/zernike.py create mode 100644 tests/operators/test_zernike.py diff --git a/src/tike/operators/cupy/__init__.py b/src/tike/operators/cupy/__init__.py index cb2a9d6f..a67a72f5 100644 --- a/src/tike/operators/cupy/__init__.py +++ b/src/tike/operators/cupy/__init__.py @@ -11,11 +11,12 @@ from .convolution import * from .flow import * from .lamino import * -from .operator import * from .objective import * +from .operator import * from .pad import * from .patch import * from .propagation import * from .ptycho import * from .rotate import * from .shift import * +from .zernike import * diff --git a/src/tike/operators/cupy/zernike.py b/src/tike/operators/cupy/zernike.py new file mode 100644 index 00000000..310f6f6c --- /dev/null +++ b/src/tike/operators/cupy/zernike.py @@ -0,0 +1,63 @@ +"""Defines an inverse-Zernike transform operator.""" + +__author__ = "Daniel Ching" +__copyright__ = "Copyright (c) 2024, UChicago Argonne, LLC." + +import numpy.typing as npt +import numpy as np +import tike.zernike +import tike.linalg + +from .operator import Operator + + +class Zernike(Operator): + """Reconstruct an image from coefficients and zernike basis using CuPy. + + Take an (..., W) array of zernike coefficients and reconstruct an image + from them. + + Parameters + ---------- + size : int + The pixel width and height of the reconstruction. + weights: (..., W) complex64 + The zernike coefficients + + .. versionadded:: 0.25.5 + + """ + + def fwd( + self, + weights: npt.NDArray[np.csingle], + size: int, + degree_max: int, + **kwargs, + ) -> npt.NDArray[np.csingle]: + basis = tike.zernike.basis( + size=size, + degree_min=0, + degree_max=degree_max, + xp=self.xp, + ) + basis /= tike.linalg.norm(basis, axis=(-2, -1), keepdims=True) + # (..., W, 1, 1) * (W, size, size) + return np.sum(weights[..., None, None] * basis, axis=-3) + + def adj( + self, + images: npt.NDArray[np.csingle], + size: int, + degree_max: int, + **kwargs, + ) -> npt.NDArray[np.csingle]: + basis = tike.zernike.basis( + size=size, + degree_min=0, + degree_max=degree_max, + xp=self.xp, + ) + basis /= tike.linalg.norm(basis, axis=(-2, -1), keepdims=True) + # (..., 1, size, size) * (W, size, size) + return np.sum(images[..., None, :, :] * basis, axis=(-2, -1)) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index d31347e6..07d12900 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -66,7 +66,7 @@ def R(m: int, n: int, radius: np.array) -> np.array: sign = -sign b0 = _bino(n - k, k) b1 = _bino(n - 2 * k, (n - m) // 2 - k) - result += sign * b0 * b1 * radius**(n - 2 * k) + result += sign * b0 * b1 * radius ** (n - 2 * k) return result @@ -78,7 +78,13 @@ def _bino(a: int, b: int) -> int: return result -def basis(size: int, degree_min: int, degree_max: int) -> np.array: +def _bino1(a: int, b: int, xp=np) -> int: + """Return the approximate binomial coeffient (a b).""" + result = np.arange(a, a - b, -1) / np.arange(1, b + 1) + return np.prod(result) + + +def basis(size: int, degree_min: int, degree_max: int, xp=np) -> np.array: """Return all circular Zernike basis up to given radial degree. Parameters @@ -96,16 +102,16 @@ def basis(size: int, degree_min: int, degree_max: int) -> np.array: """ endpoint = 1.0 - 1.0 / (2 * size) - x = np.linspace(-endpoint, endpoint, size, endpoint=True) - coords = np.stack(np.meshgrid(x, x, indexing='ij'), axis=0) - radius = np.linalg.norm(coords, axis=0) - theta = np.arctan2(coords[0], coords[1]) + x = xp.linspace(-endpoint, endpoint, size, endpoint=True) + coords = xp.stack(xp.meshgrid(x, x, indexing="ij"), axis=0) + radius = xp.linalg.norm(coords, axis=0) + theta = xp.arctan2(coords[0], coords[1]) basis = [] for m, n in valid_indices(degree_min, degree_max): basis.append(Z(m, n, radius, theta)) - basis = np.stack(basis, axis=0) + basis = xp.stack(basis, axis=0) return basis @@ -115,3 +121,13 @@ def valid_indices(degree_min: int, degree_max: int) -> tuple: for m in range(-n, n + 1): if (n - abs(m)) % 2 == 0: yield m, n + + +def degree_from_num_coeffients(num_coefficients: int) -> int: + coefficient_count = 0 + for n in range(0, 999_999): + for m in range(-n, n + 1): + if (n - abs(m)) % 2 == 0: + coefficient_count += 1 + if coefficient_count >= num_coefficients: + return n + 1, coefficient_count diff --git a/tests/operators/test_zernike.py b/tests/operators/test_zernike.py new file mode 100644 index 00000000..2b5ccc7f --- /dev/null +++ b/tests/operators/test_zernike.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +import unittest + +import numpy as np +from tike.operators import Zernike +import tike.precision +import tike.linalg + +from .util import random_complex, OperatorTests + +__author__ = "Daniel Ching" +__copyright__ = "Copyright (c) 2020, UChicago Argonne, LLC." +__docformat__ = "restructuredtext en" + + +class TestZernike(unittest.TestCase, OperatorTests): + """Test the Zernike operator.""" + + def setUp(self): + self.nscan = 21 + self.nprobe = 11 + self.nbasis = 128 + self.size = 16 + self.degree_max, self.nbasis = tike.zernike.degree_from_num_coeffients( + self.nbasis + ) + + basis = tike.zernike.basis(size=3, degree_min=0, degree_max=self.degree_max) + assert basis.shape == (self.nbasis, 3, 3), basis.size + + self.operator = Zernike() + self.operator.__enter__() + self.xp = self.operator.xp + + np.random.seed(0) + images = random_complex(self.nscan, self.nprobe, self.size, self.size) + weights = random_complex(self.nscan, self.nprobe, self.nbasis) + + self.m = self.xp.asarray(weights) + self.m_name = "weights" + self.kwargs = { + "size": self.size, + "degree_max": self.degree_max, + } + + self.d = self.xp.asarray(images) + self.d_name = "images" + + print(self.operator) diff --git a/tests/test_zernike.py b/tests/test_zernike.py index 82fe12a1..28ff9d62 100644 --- a/tests/test_zernike.py +++ b/tests/test_zernike.py @@ -11,19 +11,23 @@ class TestZernike(unittest.TestCase): - def test_zernike(self): - - fname = os.path.join(testdir, 'result', 'zernike') + fname = os.path.join(testdir, "result", "zernike") os.makedirs(fname, exist_ok=True) - for i, Z in enumerate(tike.zernike.basis(256, degree=10)): + for i, Z in enumerate( + tike.zernike.basis( + 256, + degree_min=0, + degree_max=6, + ) + ): plt.figure() tike.view.plot_complex(Z, rmin=-1, rmax=1) plt.savefig(os.path.join(fname, f"zernike-{i:02d}.png")) plt.close() def _radial_template(self, m=0): - fname = os.path.join(testdir, 'result', 'zernike') + fname = os.path.join(testdir, "result", "zernike") os.makedirs(fname, exist_ok=True) radius = np.linspace(0, 1, 200) @@ -53,45 +57,102 @@ def test_radial_2(self): self._radial_template(2) def test_transform(self): - - fname = os.path.join(testdir, 'result', 'zernike') + fname = os.path.join(testdir, "result", "zernike") os.makedirs(fname, exist_ok=True) - # import libimage + import libimage - # f0 = libimage.load('cryptomeria', 256) - f0 = plt.imread(os.path.join(testdir, 'data', 'probe.png')) + f0 = libimage.load("cryptomeria", 256) + # f0 = plt.imread(os.path.join(testdir, "data", "probe.png")) size = f0.shape[-1] - plt.imsave(os.path.join(fname, 'basis-0.png'), f0) + plt.imsave(os.path.join(fname, "basis-0.png"), f0, vmin=0, vmax=1.0) f0 = f0.reshape(size * size, 1) _basis = [] - for d in range(0, 100): + for d in range(0, 64): _basis.append( tike.zernike.basis( size, degree_min=d, degree_max=d + 1, - )) + ) + ) basis = np.concatenate(_basis, axis=0) - print(f'degree {d} - {len(basis)}') + print(f"degree {d} - {len(basis)}") basis = np.moveaxis(basis, 0, -1) basis = basis.reshape(size * size, -1) # weight only pixels inside basis - w = (basis[..., 0] > 0).astype('float32') + w = (basis[..., 0] > 0).astype("float32") # x = tike.linalg.lstsq(basis, f0, weights=w) x, _, _, _ = np.linalg.lstsq(basis, f0, rcond=1e-4) f1 = basis @ x f1 = f1.reshape(size, size) - plt.imsave(os.path.join(fname, f'basis-{d:02d}.png'), f1) + plt.imsave(os.path.join(fname, f"basis-{d:02d}.png"), f1, vmin=0, vmax=1.0) plt.figure() plt.title(f"basis weights for {d} degree polynomials") plt.bar(list(range(x.size)), x.flatten()) - plt.savefig(os.path.join(fname, f'basis-w-{d:02d}.png')) + plt.savefig(os.path.join(fname, f"basis-w-{d:02d}.png")) plt.close() + + def test_transform1(self): + fname = os.path.join(testdir, "result", "zernike") + os.makedirs(fname, exist_ok=True) + + import libimage + + f0 = libimage.load("cryptomeria", 256) + print(f0.max()) + # f0 = plt.imread(os.path.join(testdir, "data", "probe.png")) + size = f0.shape[-1] + plt.imsave(os.path.join(fname, "basis1-0.png"), f0, vmin=0, vmax=1.0) + + f0 = f0.reshape(1, size * size) + + _basis = [] + + print(f"This image has {size * size} pixels.") + + for d in range(0, 64): + more_basis = tike.zernike.basis( + size, + degree_min=d, + degree_max=d + 1, + ) + _basis.append( + # Normalize the basis (size-dependent) + more_basis + / tike.linalg.norm( + more_basis, + axis=(-2, -1), + keepdims=True, + ) + ) + + basis = np.concatenate(_basis, axis=0) + + print(f"Adding degree {d} - {len(basis)} total basis functions") + basis = basis.reshape(-1, size * size) + + # print(basis.shape, f0.shape) + # y = tike.linalg.inner(basis[0], basis[-1], axis=-1, keepdims=True) + # print(f"orthogonality {y}") + x = tike.linalg.inner(f0, basis, axis=-1, keepdims=True) + # print(x.shape) + f1 = x.T @ basis + f1 = f1.reshape(size, size) + print(f1.max()) + plt.imsave(os.path.join(fname, f"basis1-{d:02d}.png"), f1, vmin=0, vmax=1.0) + + plt.figure() + plt.title(f"basis weights for {d} degree polynomials") + plt.bar(list(range(x.size)), x.flatten()) + plt.savefig(os.path.join(fname, f"basis1-w-{d:02d}.png")) + plt.close() + + # print(f"{x.flatten()[:16]}") From 8300a434f5759533cf699b3fdbe5b3f72f9f48f7 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Thu, 18 Jan 2024 16:35:05 -0600 Subject: [PATCH 10/13] REF: Zernike operator is not scaled --- src/tike/operators/cupy/zernike.py | 2 -- src/tike/zernike.py | 39 ++++++++++++++++++++++++------ tests/operators/test_zernike.py | 6 ++++- 3 files changed, 37 insertions(+), 10 deletions(-) diff --git a/src/tike/operators/cupy/zernike.py b/src/tike/operators/cupy/zernike.py index 310f6f6c..98c630df 100644 --- a/src/tike/operators/cupy/zernike.py +++ b/src/tike/operators/cupy/zernike.py @@ -41,7 +41,6 @@ def fwd( degree_max=degree_max, xp=self.xp, ) - basis /= tike.linalg.norm(basis, axis=(-2, -1), keepdims=True) # (..., W, 1, 1) * (W, size, size) return np.sum(weights[..., None, None] * basis, axis=-3) @@ -58,6 +57,5 @@ def adj( degree_max=degree_max, xp=self.xp, ) - basis /= tike.linalg.norm(basis, axis=(-2, -1), keepdims=True) # (..., 1, size, size) * (W, size, size) return np.sum(images[..., None, :, :] * basis, axis=(-2, -1)) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 07d12900..33f3acd5 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -1,4 +1,23 @@ -"""Provide functions to evaluate Zernike polynomials on a discrete grid.""" +"""Provide functions to evaluate Zernike polynomials on a discrete grid. + + +References +---------- +@article{Niu_2022, +doi = {10.1088/2040-8986/ac9e08}, +url = {https://dx.doi.org/10.1088/2040-8986/ac9e08}, +year = {2022}, +month = {nov}, +publisher = {IOP Publishing}, +volume = {24}, +number = {12}, +pages = {123001}, +author = {Kuo Niu and Chao Tian}, +title = {Zernike polynomials and their applications}, +journal = {Journal of Optics}, +abstract = {The Zernike polynomials are a complete set of continuous functions orthogonal over a unit circle. Since first developed by Zernike in 1934, they have been in widespread use in many fields ranging from optics, vision sciences, to image processing. However, due to the lack of a unified definition, many confusing indices have been used in the past decades and mathematical properties are scattered in the literature. This review provides a comprehensive account of Zernike circle polynomials and their noncircular derivatives, including history, definitions, mathematical properties, roles in wavefront fitting, relationships with optical aberrations, and connections with other polynomials. We also survey state-of-the-art applications of Zernike polynomials in a range of fields, including the diffraction theory of aberrations, optical design, optical testing, ophthalmic optics, adaptive optics, and image analysis. Owing to their elegant and rigorous mathematical properties, the range of scientific and industrial applications of Zernike polynomials is likely to expand. This review is expected to clear up the confusion of different indices, provide a self-contained reference guide for beginners as well as specialists, and facilitate further developments and applications of the Zernike polynomials.} +} +""" import numpy as np @@ -26,16 +45,18 @@ def Z(m: int, n: int, radius: np.array, angle: np.array) -> np.array: if _m_ > n: raise ValueError("Angular frequency must be less than radial degree.") if m < 0: - polynomial = R(_m_, n, radius) * np.sin(m * angle) - else: - polynomial = R(_m_, n, radius) * np.cos(m * angle) - polynomial[np.logical_or(radius < 0, radius > 1)] = 0 - return polynomial + return np.sqrt(2 * (n + 1)) * R(_m_, n, radius) * np.sin(m * angle) + if m == 0: + return np.sqrt(n + 1) * R(_m_, n, radius) + if m > 0: + return np.sqrt(2 * (n + 1)) * R(_m_, n, radius) * np.cos(m * angle) def N(m: int, n: int) -> float: """Zernike normalization factor.""" - return np.sqrt(2 * (n + 1) / (1 + (m == 0))) + if m == 0: + return np.sqrt(n + 1) + return np.sqrt(2 * (n + 1)) def R(m: int, n: int, radius: np.array) -> np.array: @@ -67,6 +88,7 @@ def R(m: int, n: int, radius: np.array) -> np.array: b0 = _bino(n - k, k) b1 = _bino(n - 2 * k, (n - m) // 2 - k) result += sign * b0 * b1 * radius ** (n - 2 * k) + result[radius > 1] = 0.0 return result @@ -124,6 +146,9 @@ def valid_indices(degree_min: int, degree_max: int) -> tuple: def degree_from_num_coeffients(num_coefficients: int) -> int: + """ + There are a total of (n + 1)(n + 2)/2 linearly independent polynomials for a degree ⩽ n. + """ coefficient_count = 0 for n in range(0, 999_999): for m in range(-n, n + 1): diff --git a/tests/operators/test_zernike.py b/tests/operators/test_zernike.py index 2b5ccc7f..24b5cf3c 100644 --- a/tests/operators/test_zernike.py +++ b/tests/operators/test_zernike.py @@ -11,7 +11,7 @@ from .util import random_complex, OperatorTests __author__ = "Daniel Ching" -__copyright__ = "Copyright (c) 2020, UChicago Argonne, LLC." +__copyright__ = "Copyright (c) 2024, UChicago Argonne, LLC." __docformat__ = "restructuredtext en" @@ -49,3 +49,7 @@ def setUp(self): self.d_name = "images" print(self.operator) + + @unittest.skip('FIXME: This operator is not scaled.') + def test_scaled(self): + pass From 144b1eb268a4d745a4b197117f113f9bbf4cfd0f Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Thu, 18 Jan 2024 17:00:45 -0600 Subject: [PATCH 11/13] REF: Use einsum instead of sequential function calls --- src/tike/operators/cupy/zernike.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/tike/operators/cupy/zernike.py b/src/tike/operators/cupy/zernike.py index 98c630df..0e6ffc3d 100644 --- a/src/tike/operators/cupy/zernike.py +++ b/src/tike/operators/cupy/zernike.py @@ -41,8 +41,8 @@ def fwd( degree_max=degree_max, xp=self.xp, ) - # (..., W, 1, 1) * (W, size, size) - return np.sum(weights[..., None, None] * basis, axis=-3) + # (..., W) @ (W, size, size) + return np.einsum("...c,cwh->...wh", weights, basis) def adj( self, @@ -57,5 +57,5 @@ def adj( degree_max=degree_max, xp=self.xp, ) - # (..., 1, size, size) * (W, size, size) - return np.sum(images[..., None, :, :] * basis, axis=(-2, -1)) + # (..., size, size) @ (W, size, size) + return np.einsum("...wh,cwh->...c", images, basis) From c49039bae3ecdd05e587c860b694f7b40367058c Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Thu, 18 Jan 2024 17:48:00 -0600 Subject: [PATCH 12/13] REF: Helper functions for estimating basis size --- src/tike/zernike.py | 26 ++++++++++++++------------ tests/operators/test_zernike.py | 7 +++---- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index 33f3acd5..c85cec0c 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -18,6 +18,7 @@ abstract = {The Zernike polynomials are a complete set of continuous functions orthogonal over a unit circle. Since first developed by Zernike in 1934, they have been in widespread use in many fields ranging from optics, vision sciences, to image processing. However, due to the lack of a unified definition, many confusing indices have been used in the past decades and mathematical properties are scattered in the literature. This review provides a comprehensive account of Zernike circle polynomials and their noncircular derivatives, including history, definitions, mathematical properties, roles in wavefront fitting, relationships with optical aberrations, and connections with other polynomials. We also survey state-of-the-art applications of Zernike polynomials in a range of fields, including the diffraction theory of aberrations, optical design, optical testing, ophthalmic optics, adaptive optics, and image analysis. Owing to their elegant and rigorous mathematical properties, the range of scientific and industrial applications of Zernike polynomials is likely to expand. This review is expected to clear up the confusion of different indices, provide a self-contained reference guide for beginners as well as specialists, and facilitate further developments and applications of the Zernike polynomials.} } """ +import typing import numpy as np @@ -137,7 +138,10 @@ def basis(size: int, degree_min: int, degree_max: int, xp=np) -> np.array: return basis -def valid_indices(degree_min: int, degree_max: int) -> tuple: +def valid_indices( + degree_min: int, + degree_max: int, +) -> typing.Generator[typing.Tuple[int, int], None, None]: """Enumerate all valid zernike indices (m,n) up to the given degree.""" for n in range(degree_min, degree_max): for m in range(-n, n + 1): @@ -145,14 +149,12 @@ def valid_indices(degree_min: int, degree_max: int) -> tuple: yield m, n -def degree_from_num_coeffients(num_coefficients: int) -> int: - """ - There are a total of (n + 1)(n + 2)/2 linearly independent polynomials for a degree ⩽ n. - """ - coefficient_count = 0 - for n in range(0, 999_999): - for m in range(-n, n + 1): - if (n - abs(m)) % 2 == 0: - coefficient_count += 1 - if coefficient_count >= num_coefficients: - return n + 1, coefficient_count +def num_basis_less_than_degree(degree_max: int) -> int: + """Return number of zernike basis in degrees < degree_max (strictly).""" + # And odd times even number is even and always cleanly divisible by 2. + return (degree_max) * (degree_max + 1) // 2 + + +def degree_max_from_num_basis(num_basis: int) -> int: + """Return the max degree (non-inclusive) required to have at least some number of basis.""" + return int(np.ceil(0.5 * (-1 + np.sqrt(8 * num_basis)))) diff --git a/tests/operators/test_zernike.py b/tests/operators/test_zernike.py index 24b5cf3c..05a4639b 100644 --- a/tests/operators/test_zernike.py +++ b/tests/operators/test_zernike.py @@ -23,9 +23,8 @@ def setUp(self): self.nprobe = 11 self.nbasis = 128 self.size = 16 - self.degree_max, self.nbasis = tike.zernike.degree_from_num_coeffients( - self.nbasis - ) + self.degree_max = tike.zernike.degree_max_from_num_basis(self.nbasis) + self.nbasis = tike.zernike.num_basis_less_than_degree(self.degree_max) basis = tike.zernike.basis(size=3, degree_min=0, degree_max=self.degree_max) assert basis.shape == (self.nbasis, 3, 3), basis.size @@ -50,6 +49,6 @@ def setUp(self): print(self.operator) - @unittest.skip('FIXME: This operator is not scaled.') + @unittest.skip("FIXME: This operator is not scaled.") def test_scaled(self): pass From 2230723a38bafca751096b3a0466f7927fb45dd1 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Fri, 19 Jan 2024 12:23:22 -0600 Subject: [PATCH 13/13] REF: Smooth the edge of the zernike basis --- src/tike/zernike.py | 5 ++++- tests/test_zernike.py | 24 +++++++++++------------- 2 files changed, 15 insertions(+), 14 deletions(-) diff --git a/src/tike/zernike.py b/src/tike/zernike.py index c85cec0c..90a8bc51 100644 --- a/src/tike/zernike.py +++ b/src/tike/zernike.py @@ -89,7 +89,10 @@ def R(m: int, n: int, radius: np.array) -> np.array: b0 = _bino(n - k, k) b1 = _bino(n - 2 * k, (n - m) // 2 - k) result += sign * b0 * b1 * radius ** (n - 2 * k) - result[radius > 1] = 0.0 + # Smooth the sharp edges of the polynomial with a supergaussian window + # Higher smoothing degree makes the window edge sharper + smoothing_degree = 32 + result *= np.exp(-(radius ** (2 * smoothing_degree))) return result diff --git a/tests/test_zernike.py b/tests/test_zernike.py index 28ff9d62..a5d8b025 100644 --- a/tests/test_zernike.py +++ b/tests/test_zernike.py @@ -11,7 +11,7 @@ class TestZernike(unittest.TestCase): - def test_zernike(self): + def test_zernike_preview(self): fname = os.path.join(testdir, "result", "zernike") os.makedirs(fname, exist_ok=True) for i, Z in enumerate( @@ -60,10 +60,9 @@ def test_transform(self): fname = os.path.join(testdir, "result", "zernike") os.makedirs(fname, exist_ok=True) - import libimage - - f0 = libimage.load("cryptomeria", 256) - # f0 = plt.imread(os.path.join(testdir, "data", "probe.png")) + # import libimage + # f0 = libimage.load("cryptomeria", 256) + f0 = plt.imread(os.path.join(testdir, "data", "probe.png")) size = f0.shape[-1] plt.imsave(os.path.join(fname, "basis-0.png"), f0, vmin=0, vmax=1.0) @@ -85,10 +84,10 @@ def test_transform(self): basis = np.moveaxis(basis, 0, -1) basis = basis.reshape(size * size, -1) # weight only pixels inside basis - w = (basis[..., 0] > 0).astype("float32") + # w = (basis[..., 0] > 0).astype("float32") - # x = tike.linalg.lstsq(basis, f0, weights=w) - x, _, _, _ = np.linalg.lstsq(basis, f0, rcond=1e-4) + # x = tike.linalg.lstsq(basis, f0, )#weights=w) + x, _, _, _ = np.linalg.lstsq(basis, f0, rcond=1e-9) f1 = basis @ x f1 = f1.reshape(size, size) @@ -104,11 +103,10 @@ def test_transform1(self): fname = os.path.join(testdir, "result", "zernike") os.makedirs(fname, exist_ok=True) - import libimage - - f0 = libimage.load("cryptomeria", 256) - print(f0.max()) - # f0 = plt.imread(os.path.join(testdir, "data", "probe.png")) + # import libimage + # f0 = libimage.load("cryptomeria", 256) + # print(f0.max()) + f0 = plt.imread(os.path.join(testdir, "data", "probe.png")) size = f0.shape[-1] plt.imsave(os.path.join(fname, "basis1-0.png"), f0, vmin=0, vmax=1.0)