Cap the k-space interpolation stencil at one period

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

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,

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

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(),