From a04b2ded610ec459d66ad43d3428b68c899f6245 Mon Sep 17 00:00:00 2001 From: AdamField118 Date: Mon, 24 Aug 2026 05:30:31 -0400 Subject: [PATCH] Cap the k-space interpolation stencil at one period --- jax_galsim/interpolant.py | 22 +++++++++++ jax_galsim/interpolatedimage.py | 23 +++++++---- tests/jax/test_interpolant_jax.py | 47 +++++++++++++++++++++++ tests/jax/test_interpolatedimage_utils.py | 6 +-- 4 files changed, 88 insertions(+), 10 deletions(-) diff --git a/jax_galsim/interpolant.py b/jax_galsim/interpolant.py index 8cfe75bb..ff4389da 100644 --- a/jax_galsim/interpolant.py +++ b/jax_galsim/interpolant.py @@ -156,6 +156,16 @@ def xval(self, x): def _xval_noraise(self, x): return self.__class__._xval(x) + def _xval_wrapped_noraise(self, x, n): + """The sum of ``xval`` over all aliases ``x + j*n``, as in + ``Interpolant::xvalWrapped`` in GalSim.""" + xdown = x - n * jnp.floor(x / n + 0.5) + if 2 * self.xrange <= n: + return self._xval_noraise(xdown) + nalias = int(math.ceil(self.xrange / n + 0.5)) + js = jnp.arange(-nalias, nalias + 1) * n + return jnp.sum(self._xval_noraise(jnp.asarray(xdown)[..., None] + js), axis=-1) + @implements(_galsim.interpolant.Interpolant.kval) def kval(self, k): if jnp.ndim(k) > 1: @@ -420,6 +430,18 @@ def _uval(u): jnp.where(absu < 0.5, 1.0, 0.5), ) + def _xval_wrapped_noraise(self, x, n): + x = x * np.pi + msk = jnp.abs(x) < 1e-4 + xs = jnp.where(msk, 1.0, x) + if n % 2 == 0: + taylor = 1.0 - x * x * (1.0 / 6.0 + 1.0 / 2.0 - 1.0 / (6.0 * n * n)) + val = jnp.sin(xs) * jnp.cos(xs / n) / (n * jnp.sin(xs / n)) + else: + taylor = 1.0 - (1.0 / 6.0) * x * x * (1.0 - 1.0 / (n * n)) + val = jnp.sin(xs) / (n * jnp.sin(xs / n)) + return jnp.where(msk, taylor, val) + def urange(self): """The maximum extent of the interpolant in Fourier space (in 2pi/pixels).""" return 0.5 diff --git a/jax_galsim/interpolatedimage.py b/jax_galsim/interpolatedimage.py index 589c27a9..a2a6bd9e 100644 --- a/jax_galsim/interpolatedimage.py +++ b/jax_galsim/interpolatedimage.py @@ -1095,12 +1095,24 @@ def _kValue_arr( return jnp.where(msk, val * xint_val * pfac, 0.0) +def _kval_offsets(ixrange, n): + """The stencil offsets for one axis of a k-space image of period ``n``. + + Interpolants wider than the period are summed over their aliases by + ``_xval_wrapped_noraise``, so the stencil need never exceed one period. + """ + if ixrange < n: + irange = ixrange // 2 + return jnp.arange(-irange, irange + 1) + return jnp.arange(n) - n // 2 + + @partial(jax.vmap, in_axes=(0, None, None, None, None, None)) -@partial(jax.jit, static_argnames=("interp",)) +@partial(jax.jit, static_argnames=("nkx", "interp")) def _interp_weight_1d_kval(ioff, kxi, kxp, kx, nkx, interp): kxind = (kxi + ioff) % nkx _kx = kx - (kxp + ioff) - wkx = interp._xval_noraise(_kx) + wkx = interp._xval_wrapped_noraise(_kx, nkx) return wkx, kxind.astype(jnp.int32) @@ -1131,11 +1143,8 @@ def _draw_with_interpolant_kval(kx, ky, kxmin, kymin, zp, interp): kyp = kyi + kymin nky = zp.shape[0] - irange = interp.ixrange // 2 - iinds = jnp.arange(-irange, irange + 1) - wkx, kxind = _interp_weight_1d_kval( - iinds, + _kval_offsets(interp.ixrange, nkx), kxi, kxp, kx, @@ -1144,7 +1153,7 @@ def _draw_with_interpolant_kval(kx, ky, kxmin, kymin, zp, interp): ) wky, kyind = _interp_weight_1d_kval( - iinds, + _kval_offsets(interp.ixrange, nky), kyi, kyp, ky, diff --git a/tests/jax/test_interpolant_jax.py b/tests/jax/test_interpolant_jax.py index 45d328fb..a126ddca 100644 --- a/tests/jax/test_interpolant_jax.py +++ b/tests/jax/test_interpolant_jax.py @@ -524,3 +524,50 @@ def test_interpolant_jax_unit_integrals(): np.testing.assert_array_equal(short, full[:5]) np.testing.assert_array_equal(med, full[:8]) np.testing.assert_array_equal(long, full[:10]) + + +@pytest.mark.parametrize( + "interp", + [ + galsim.Delta(), + galsim.Nearest(), + galsim.Linear(), + galsim.Cubic(), + galsim.Quintic(), + galsim.Lanczos(3), + galsim.Lanczos(5), + ], + ids=lambda x: str(x).replace("galsim.", "").replace("(", "").replace(")", ""), +) +@pytest.mark.parametrize("n", [2, 3, 4, 5, 8, 9, 32]) +def test_interpolant_jax_xval_wrapped(interp, n): + """The wrapped kernel is the sum of xval over all aliases x + j*n.""" + gs = getattr(ref_galsim, interp.__class__.__name__) + gs = gs(interp.n) if isinstance(interp, galsim.Lanczos) else gs() + + x = np.linspace(-12.0, 12.0, 97) + js = np.arange(-200, 201) * n + expected = gs.xval(np.ravel(x[:, None] + js[None, :])).reshape(x.size, js.size) + expected = expected.sum(axis=1) + + np.testing.assert_allclose( + interp._xval_wrapped_noraise(x, n), expected, rtol=0, atol=1e-10 + ) + + +@pytest.mark.parametrize("n", [8, 9, 32, 128]) +def test_interpolant_jax_sinc_xval_wrapped(n): + """The sinc kernel wraps to the Dirichlet kernel, which the truncated + alias sum converges to as the window grows.""" + x = np.linspace(-12.0, 12.0, 97) + got = galsim.SincInterpolant()._xval_wrapped_noraise(x, n) + + prev = None + for m in (10**3, 10**4, 10**5): + js = np.arange(-m, m + 1) * n + expected = np.sinc(x[:, None] + js[None, :]).sum(axis=1) + err = np.max(np.abs(np.asarray(got) - expected)) + if prev is not None: + assert err < prev + prev = err + assert err < 1e-5 diff --git a/tests/jax/test_interpolatedimage_utils.py b/tests/jax/test_interpolatedimage_utils.py index b45ffbd1..fdef029d 100644 --- a/tests/jax/test_interpolatedimage_utils.py +++ b/tests/jax/test_interpolatedimage_utils.py @@ -6,12 +6,13 @@ import pytest import jax_galsim -from jax_galsim.interpolant import ( # SincInterpolant, +from jax_galsim.interpolant import ( Cubic, Lanczos, Linear, Nearest, Quintic, + SincInterpolant, ) from jax_galsim.interpolatedimage import ( _draw_with_interpolant_kval, @@ -66,8 +67,7 @@ def test_interpolatedimage_utils_draw_with_interpolant_xval(interp): [ Nearest(), Linear(), - # this is really slow right now and I am not sure why will fix later - # SincInterpolant(), + SincInterpolant(), Linear(), Cubic(), Quintic(),