Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Zernike enhancements #1327

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Zernike enhancements #1327

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1553081
Add xycoef method to DoubleZernike
jmeyers314 Jan 23, 2025
9dff874
Enable higher order annular Zernikes
jmeyers314 Jan 23, 2025
65f072f
Remove py3.8 from CI
jmeyers314 Jan 23, 2025
b3f8aba
Mike's comments on test_large_j
jmeyers314 Feb 6, 2025
24f9bcc
Rewrite xy_contribution as cached recursive
jmeyers314 Feb 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" ]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
16 changes: 15 additions & 1 deletion galsim/zernike.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -387,7 +387,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
Expand Down Expand Up @@ -1215,6 +1215,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__:
Expand Down
66 changes: 66 additions & 0 deletions tests/test_zernike.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down Expand Up @@ -1553,5 +1568,56 @@ 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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we expect this to converge particularly well when r>1, right? So should probably limit this to have x**2 + y**2 < 1.
When I tried it, I was able to use significantly lower tolerances:

(10, 1.e-12)
(20, 1.e-12)
(40, 1.e-10)
(60, 1.e-8)
(80, 1.e-6)
(100, 2.e-3)

So it is still becoming less accurate as j get very large. But not as catastrophically bad as this test implies.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The ones that are least accurate are always ones with r close to 1. Like 0.95 or so. This is probably a clue about where the recursion is going unstable. I might play around with the recursion formula to see if I can harden it up at all.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, limiting to 0.5 < r < 1.0 might even be appropriate. Though, OTOH, these are just 2d polynomials. They're orthogonal over an annulus but defined everywhere.

The particular polynomials here are proportional to r^n, so that might explain why the large r are trickier. Other values of j will presumably produce polynomials with larger amplitudes near r=0.

In case it's helpful while you're looking at recursions, here's the closed-form table I got out of Mathematica up to Z28:
Screenshot 2025-01-24 at 6 26 12 AM

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the problem is the conversion of the coef array from rrsq format to xy format. The xy version has lots of very large coefficients with alternating sign. The native rho/rhosq version is much more stable numerically.

If I change evalCartesian to:

        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

then all the tests pass at 1.e-12 tolerance, except Z5151, which needed 1.e-11. (It's also massively faster than the current implementation -- I guess the rrsq_to_xy function must be a slow step?)

I don't know if this means we should actually make the above change though. There might be other downsides (e.g. not getting to use the C++ layer horner2d implementation, at least without a little more work). But I suppose we could at least provide another method that would evaluate this way for users who want more accuracy at very high Zernike order.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting. I sped up rrsq_to_xy significantly by rewriting it as a cached recursive function. Now the tradeoff seems to be roughly:

The xy array is about 2x the effort to compute as the rrsq array, which makes sense since the xy is derived from the rrsq. However, most of this effort is cached so I don't think there's a significant difference after the initial Zernike constructions.

The xy approach (so c++ horner) is about 10x faster when there are ~100 evaluation points. However, on my machine (M3 Pro) the complex-python approach actually runs faster when the number of evaluation points is ~10_000 or more (!). Maybe this is because the rrsq array is ~2x smaller than the xy array?

Given the improved accuracy and the above, I tentatively vote make the rrsq approach the default. I want to benchmark some donut fitting first though.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmmm.... Maybe the horner step depends on jmax too. I just ran some donut image forward models (jmax ~65 and npoints ~10000) and they were ~4x faster using the xy approach. So maybe need to leave that as an expert option.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Either way is fine with me. As long as there is a public API way to get the calculation that is accurate for high order.

I think most use cases won't care about the accuracy as much as the speed. But it seems important to have a way to get the accurate one.


R_outer = 1.0
R_inner = 0.5
eps = R_inner/R_outer

test_vals = [
(10, 1e-12), # Z66
(20, 1e-12), # Z231
(40, 1e-9), # Z861
]
if run_slow:
test_vals += [
(60, 1e-6), # Z1891
(80, 1e-3), # Z3321
# (100, 10), # Z5151 # This one is catastrophic failure!
]

print()
for n, tol in test_vals:
j = np.sum(np.arange(n+2))
n, m = galsim.zernike.noll_to_zern(j)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Somewhat confusing reuse of variable n. Probably should rename the other one. (k maybe?)

print(f"Z{j} => (n, m) = ({n}, {n})")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
print(f"Z{j} => (n, m) = ({n}, {n})")
print(f"Z{j} => (n, m) = ({n}, {m})")

Probably followed by assert n == abs(m) since AFAIU, that is the intent of these choices for j.

coefs = np.zeros(j+1)
coefs[j] = 1.0
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

np.testing.assert_allclose(
zk(x, y),
analytic_zk(x, y),
atol=tol, rtol=tol
)


if __name__ == "__main__":
runtests(__file__)
Loading