Zernike enhancements

.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
             matrix:
                 # First all python versions in basic linux
                 os: [ ubuntu-latest ]
-                py: [ 3.7, 3.8, 3.9, "3.10", 3.11, 3.12 ]
+                py: [ 3.8, 3.9, "3.10", 3.11, 3.12 ]
                 CC: [ gcc ]
                 CXX: [ g++ ]
                 FFTW_DIR: [ "/usr/local/lib" ]
@@ -107,7 +107,7 @@ jobs:
 
                 # Easier if eigen is installed with apt-get, but on at least one system, check that
                 # it gets downloaded and installed properly if it isn't installed.
-                if ${{ matrix.py != 3.7 }}; then sudo -H apt install -y libeigen3-dev; fi
+                if ${{ matrix.py != 3.8 }}; then sudo -H apt install -y libeigen3-dev; fi
 
                 # Need this for the mpeg tests
                 sudo -H apt install -y ffmpeg
@@ -189,7 +189,7 @@ jobs:
                 FFTW_DIR=$FFTW_DIR pip install -vvv .
 
             - name: Check download_cosmos only if it changed. (And only on 1 runner)
-              if: matrix.py == 3.7 && github.base_ref != ''
+              if: matrix.py == 3.8 && github.base_ref != ''
               run: |
                 git status
                 git --no-pager log --graph --pretty=oneline --abbrev-commit | head -50

galsim/zernike.py

@@ -140,7 +140,7 @@ def __noll_coef_array(jmax, obscuration):
 _noll_coef_array = LRU_Cache(__noll_coef_array)
 
 
-def _xy_contribution(rho2_power, rho_power, shape):
+def __xy_contribution(rho2_power, rho_power):
     """Convert (rho, |rho|^2) bivariate polynomial coefficients to (x, y) bivariate polynomial
     coefficients.
     """
@@ -162,25 +162,29 @@ def _xy_contribution(rho2_power, rho_power, shape):
     #
     # and so on.  We can apply these operations repeatedly to effect arbitrary powers of rho or
     # |rho|^2.
+    if rho2_power == 0 and rho_power == 0:
+        return np.array([[1]], dtype=np.complex128)
+
+    if rho2_power > 0:
+        prev = _xy_contribution(rho2_power-1, rho_power)
+        shape = (prev.shape[0] + 2, prev.shape[1] +2)
+        out = np.zeros(shape, dtype=np.complex128)
+        for (i, j) in zip(*np.nonzero(prev)):
+            val = prev[i, j]
+            out[i+2, j] += val
+            out[i, j+2] += val
+        return out
+
+    # else rho_power > 0
+    prev = _xy_contribution(rho2_power, rho_power-1)
+    shape = (prev.shape[0] + 1, prev.shape[1] + 1)
     out = np.zeros(shape, dtype=np.complex128)
-    out[0,0] = 1
-    while rho2_power >= 1:
-        new = np.zeros_like(out)
-        for (i, j) in zip(*np.nonzero(out)):
-            val = out[i, j]
-            new[i+2, j] += val
-            new[i, j+2] += val
-        rho2_power -= 1
-        out = new
-    while rho_power >= 1:
-        new = np.zeros_like(out)
-        for (i, j) in zip(*np.nonzero(out)):
-            val = out[i, j]
-            new[i+1, j] += val
-            new[i, j+1] += 1j*val
-        rho_power -= 1
-        out = new
+    for (i, j) in zip(*np.nonzero(prev)):
+        val = prev[i, j]
+        out[i+1, j] += val
+        out[i, j+1] += 1j*val
     return out
+_xy_contribution = LRU_Cache(__xy_contribution)
 
 
 def _rrsq_to_xy(coefs, shape):
@@ -190,7 +194,8 @@ def _rrsq_to_xy(coefs, shape):
 
     # Now we loop through the elements of coefs and compute their contribution to new_coefs
     for (i, j) in zip(*np.nonzero(coefs)):
-        new_coefs += (coefs[i, j]*_xy_contribution(i, j, shape)).real
+        contribution = (coefs[i, j]*_xy_contribution(i, j)).real
+        new_coefs[:contribution.shape[0], :contribution.shape[1]] += contribution
     return new_coefs
 
 
@@ -387,7 +392,7 @@ def __annular_zern_rho_coefs(n, m, eps):
             if i % 2 == 1: continue
             j = i // 2
             more_coefs = (norm**j) * binomial(-eps**2, 1, j)
-            out[0:i+1:2] += coef*more_coefs
+            out[0:i+1:2] += float(coef)*more_coefs
     elif m == n:  # Equation (25)
         norm = 1./np.sqrt(np.sum((eps**2)**np.arange(n+1)))
         out[n] = norm
@@ -716,17 +721,40 @@ def gradY(self):
         newCoef /= self.R_outer
         return Zernike(newCoef, R_outer=self.R_outer, R_inner=self.R_inner)
 
-    def __call__(self, x, y):
+    def __call__(self, x, y, robust=False):
         """Evaluate this Zernike polynomial series at Cartesian coordinates x and y.
         Synonym for `evalCartesian`.
 
+        Parameters:
+            x:       x-coordinate of evaluation points.  Can be list-like.
+            y:       y-coordinate of evaluation points.  Can be list-like.
+            robust:  If True, use a more robust method for evaluating the polynomial.
+                     This can sometimes be slower, but is usually more accurate,
+                     especially for large Noll indices.  [default: False]
+        Returns:
+            Series evaluations as numpy array.
+        """
+        if robust:
+            return self.evalCartesianRobust(x, y)
+        return self.evalCartesian(x, y)
+
+    def evalCartesianRobust(self, x, y):
+        """Evaluate this Zernike polynomial series at Cartesian coordinates x and y using a more
+        robust method than the default `evalCartesian`.
+
         Parameters:
             x:    x-coordinate of evaluation points.  Can be list-like.
             y:    y-coordinate of evaluation points.  Can be list-like.
+
         Returns:
             Series evaluations as numpy array.
         """
-        return self.evalCartesian(x, y)
+        x = np.asarray(x)
+        y = np.asarray(y)
+        rho = (x + 1j*y) / self.R_outer
+        rhosq = np.abs(rho)**2
+        ar = _noll_coef_array(len(self.coef)-1, self.R_inner/self.R_outer).dot(self.coef[1:])
+        return horner2d(rhosq, rho, ar).real
 
     def evalCartesian(self, x, y):
         """Evaluate this Zernike polynomial series at Cartesian coordinates x and y.
@@ -738,7 +766,8 @@ def evalCartesian(self, x, y):
         Returns:
             Series evaluations as numpy array.
         """
-        return horner2d(x, y, self._coef_array_xy, dtype=float)
+        ar = self._coef_array_xy
+        return horner2d(x, y, ar, dtype=float)
 
     def evalPolar(self, rho, theta):
         """Evaluate this Zernike polynomial series at polar coordinates rho and theta.
@@ -1215,6 +1244,20 @@ def __call__(self, u, v, x=None, y=None):
                     assert np.shape(x) == np.shape(u)
                     return horner4d(u, v, x, y, self._coef_array_uvxy)
 
+    def xycoef(self, u, v):
+        """Return the xy Zernike coefficients for a given uv point or points.
+
+        Parameters:
+            u, v: float or array.  UV point(s).
+
+        Returns:
+            [npoint, jmax] array.  Zernike coefficients for the given UV point(s).
+        """
+        uu, vv = np.broadcast_arrays(u, v)
+        out = np.array([z(uu, vv) for z in self._xy_series]).T
+        out[..., 0] = 0.0  # Zero out piston term
+        return out
+
     def __neg__(self):
         """Negate a DoubleZernike."""
         if 'coef' in self.__dict__:

tests/test_zernike.py

@@ -876,6 +876,21 @@ def test_dz_val():
                 atol=2e-13, rtol=0
             )
 
+        zk_coefs = dz.xycoef(*uv_vector)
+        for zk, c in zip(zk_list, zk_coefs):
+            # Zk may have trailing zeros...
+            ncoef = len(c)
+            np.testing.assert_allclose(
+                zk.coef[:ncoef],
+                c,
+                atol=2e-13, rtol=0
+            )
+        for i, (u, v) in enumerate(zip(*uv_vector)):
+            np.testing.assert_equal(
+                zk_coefs[i],
+                dz.xycoef(u, v)
+            )
+
         # Check asserts
         with assert_raises(AssertionError):
             dz(0.0, [1.0])
@@ -1553,5 +1568,70 @@ def test_dz_mean():
         )
 
 
+def test_large_j(run_slow):
+    # The analytic form for an annular Zernike of the form (n, m) = (n, n) or (n, -n)
+    # is:
+    #   r^n sincos(n theta) sqrt((2n+2) / sum_i=0^n-1 eps^(2i))
+    # where sincos is sin if n is even and cos if n is odd, and eps = R_inner/R_outer.
+
+    rng = galsim.BaseDeviate(1234).as_numpy_generator()
+    x = rng.uniform(-1.0, 1.0, size=100)
+    y = rng.uniform(-1.0, 1.0, size=100)
+
+    R_outer = 1.0
+    R_inner = 0.5
+    eps = R_inner/R_outer
+
+    test_vals = [
+        (10, 1e-12, 1e-12),  # Z66
+        (20, 1e-12, 1e-12),  # Z231
+        (40, 1e-9, 1e-12),   # Z861
+    ]
+    if run_slow:
+        test_vals += [
+            (60, 1e-6, 1e-12),  # Z1891
+            (80, 1e-3, 1e-12),  # Z3321
+            (100, None, 1e-11),   # Z5151
+            (200, None, 1e-11),   # Z20301
+        ]
+
+    print()
+    for n, tol, tol2 in test_vals:
+        j = (n+1)*(n+2)//2
+        _, m = galsim.zernike.noll_to_zern(j)
+        print(f"Z{j} => (n, m) = ({n}, {m})")
+        assert n == abs(m)
+        coefs = [0]*j+[1]
+        zk = Zernike(coefs, R_outer=R_outer, R_inner=R_inner)
+
+        def analytic_zk(x, y):
+            r = np.hypot(x, y)
+            theta = np.arctan2(y, x)
+            factor = np.sqrt((2*n+2) / np.sum([eps**(2*i) for i in range(n+1)]))
+            if m > 0:
+                return r**n * np.cos(n*theta) * factor
+            else:
+                return r**n * np.sin(n*theta) * factor
+
+        analytic_vals = analytic_zk(x, y)
+        if n < 100:
+            np.testing.assert_allclose(
+                zk(x, y),
+                analytic_vals,
+                atol=tol, rtol=tol
+            )
+
+        robust_vals = zk.evalCartesianRobust(x, y)
+        np.testing.assert_allclose(
+            robust_vals,
+            analytic_vals,
+            atol=tol2, rtol=tol2
+        )
+        np.testing.assert_equal(
+            robust_vals,
+            zk(x, y, robust=True)
+        )
+
+
 if __name__ == "__main__":
     runtests(__file__)