ENH add Interpolant classes

dev/notebooks/lanczos_interp_devel.ipynb

@@ -0,0 +1,151 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "dad774e5-17de-4dd5-9383-69cd716601e2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import galsim\n",
+    "import numpy as np\n",
+    "import proplot as pplt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aa6372a5-0c8b-468b-8d18-30ec2127b1a5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axs = pplt.subplots(figsize=(7, 7))\n",
+    "\n",
+    "def _s(x, c, m):\n",
+    "    return 1.0 / (1.0 + np.exp(-(x-m)/c))\n",
+    "\n",
+    "def _approx(order, tol):\n",
+    "    plaw = np.exp(4.8 + 9.2e-4 * np.log(tol/1e-8)**2 - 0.33 * np.log(tol/1e-8) - 0.66 * np.log(order))\n",
+    "    sval = _s(np.log(order), np.log(15), 0.1) * (1.0 - _s(np.log(tol), np.log(1e-2), 0.1))\n",
+    "    return np.maximum(sval * 1.5 + (1.0 - sval) * plaw, 0.55)\n",
+    "\n",
+    "m = []\n",
+    "orders = np.arange(1, 31, dtype=int)\n",
+    "tols = np.logspace(-8, 0, 10)\n",
+    "\n",
+    "vals = np.zeros((tols.shape[0], orders.shape[0]))\n",
+    "for j, tol in enumerate(tols):\n",
+    "    yv = []\n",
+    "    for i in orders:\n",
+    "        lz = galsim.interpolant.Lanczos(i, gsparams=galsim.GSParams(kvalue_accuracy=tol))\n",
+    "        yv.append(lz._i.urange())\n",
+    "    vals[j, :] = np.array(yv)\n",
+    "    m.append(axs.loglog(orders, yv, label=\"log10(tol) = %d\" % np.log10(tol)))\n",
+    "    axs.loglog(\n",
+    "        orders, \n",
+    "        _approx(orders, tol), \n",
+    "        \":\", \n",
+    "        color=\"k\",\n",
+    "    )\n",
+    "axs.legend(m, loc=\"ur\", ncols=4, frameon=False)\n",
+    "axs.format(xlabel=\"Lanczos order\", ylabel=\"urange [$2\\pi$/pixels]\")\n",
+    "        "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "382e847c-1858-45b4-a169-d8f8a2ad9694",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "print(\"_posflux_conserve_dc = {\")\n",
+    "for i in orders:\n",
+    "    lz = galsim.interpolant.Lanczos(i, conserve_dc=True, gsparams=galsim.GSParams(kvalue_accuracy=tol))\n",
+    "    print(\"    %d: %r,\" % (i, lz.positive_flux))\n",
+    "print(\"}\")\n",
+    "\n",
+    "print(\"_posflux_no_conserve_dc = {\")\n",
+    "for i in orders:\n",
+    "    lz = galsim.interpolant.Lanczos(i, conserve_dc=False, gsparams=galsim.GSParams(kvalue_accuracy=tol))\n",
+    "    print(\"    %d: %r,\" % (i, lz.positive_flux))\n",
+    "print(\"}\")\n",
+    "\n",
+    "print(\"_negflux_conserve_dc = {\")\n",
+    "for i in orders:\n",
+    "    lz = galsim.interpolant.Lanczos(i, conserve_dc=True, gsparams=galsim.GSParams(kvalue_accuracy=tol))\n",
+    "    print(\"    %d: %r,\" % (i, lz.negative_flux))\n",
+    "print(\"}\")\n",
+    "\n",
+    "print(\"_negflux_no_conserve_dc = {\")\n",
+    "for i in orders:\n",
+    "    lz = galsim.interpolant.Lanczos(i, conserve_dc=False, gsparams=galsim.GSParams(kvalue_accuracy=tol))\n",
+    "    print(\"    %d: %r,\" % (i, lz.negative_flux))\n",
+    "print(\"}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2bae35ce-c659-4360-9778-2a3770ec8870",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "with np.printoptions(floatmode=\"unique\"):\n",
+    "\n",
+    "    print(\"_unit_integrals_no_conserve_dc = {\")\n",
+    "    for i in orders:\n",
+    "        lz = galsim.interpolant.Lanczos(i, conserve_dc=False, gsparams=galsim.GSParams(kvalue_accuracy=tol))\n",
+    "        print(\"    %d: jnp.%s, dtype=float),\" % (i, repr(lz.unit_integrals())[:-1]))\n",
+    "    print(\"}\")\n",
+    "\n",
+    "    print(\"_unit_integrals_conserve_dc = {\")\n",
+    "    for i in orders:\n",
+    "        lz = galsim.interpolant.Lanczos(i, conserve_dc=True, gsparams=galsim.GSParams(kvalue_accuracy=tol))\n",
+    "        print(\"    %d: jnp.%s, dtype=float),\" % (i, repr(lz.unit_integrals())[:-1]))\n",
+    "    print(\"}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6059101-1f9c-43fc-ba2c-d45e9d9b5c3f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8af5c04d-9643-443c-9053-e066a1a328e6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python [conda env:work]",
+   "language": "python",
+   "name": "conda-env-work-py"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

jax_galsim/__init__.py

@@ -65,4 +65,19 @@
 
 # Shear
 from .shear import Shear, _Shear
+
+# Interpolations
+from .interpolant import (
+    Interpolant,
+    Delta,
+    Nearest,
+    SincInterpolant,
+    Linear,
+    Cubic,
+    Quintic,
+    Lanczos,
+)
+
+# packages kept separate
+from . import bessel
 from . import fits

jax_galsim/bessel.py

@@ -0,0 +1,103 @@
+import galsim as _galsim
+import jax
+import jax.numpy as jnp
+from jax._src.numpy.util import _wraps
+
+
+# the code here for Si, f, g and _si_small_pade is taken from galsim/src/math/Sinc.cpp
+@jax.jit
+def _f_pade(x, x2):
+    # fmt: off
+    y = 1. / x2
+    f = (
+            (1. +  # noqa: W504, E126, E226
+             y*(7.44437068161936700618e2 +  # noqa: W504, E126, E226
+                y*(1.96396372895146869801e5 +  # noqa: W504, E126, E226
+                   y*(2.37750310125431834034e7 +  # noqa: W504, E126, E226
+                      y*(1.43073403821274636888e9 +  # noqa: W504, E126, E226
+                         y*(4.33736238870432522765e10 +  # noqa: W504, E126, E226
+                            y*(6.40533830574022022911e11 +  # noqa: W504, E126, E226
+                               y*(4.20968180571076940208e12 +  # noqa: W504, E126, E226
+                                  y*(1.00795182980368574617e13 +  # noqa: W504, E126, E226
+                                     y*(4.94816688199951963482e12 +  # noqa: W504, E126, E226
+                                        y*(-4.94701168645415959931e11)))))))))))  # noqa: W504, E126, E226
+            / (x*(1. +  # noqa: W504, E126, E226
+                  y*(7.46437068161927678031e2 +  # noqa: W504, E126, E226
+                     y*(1.97865247031583951450e5 +  # noqa: W504, E126, E226
+                        y*(2.41535670165126845144e7 +  # noqa: W504, E126, E226
+                           y*(1.47478952192985464958e9 +  # noqa: W504, E126, E226
+                              y*(4.58595115847765779830e10 +  # noqa: W504, E126, E226
+                                 y*(7.08501308149515401563e11 +  # noqa: W504, E126, E226
+                                    y*(5.06084464593475076774e12 +  # noqa: W504, E126, E226
+                                       y*(1.43468549171581016479e13 +  # noqa: W504, E126, E226
+                                          y*(1.11535493509914254097e13)))))))))))  # noqa: W504, E126, E226
+    )
+    # fmt: on
+    return f
+
+
+@jax.jit
+def _g_pade(x, x2):
+    # fmt: off
+    y = 1. / x2
+    g = (
+            y*(1. +  # noqa: W504, E126, E226
+               y*(8.1359520115168615e2 +  # noqa: W504, E126, E226
+                  y*(2.35239181626478200e5 +  # noqa: W504, E126, E226
+                     y*(3.12557570795778731e7 +  # noqa: W504, E126, E226
+                        y*(2.06297595146763354e9 +  # noqa: W504, E126, E226
+                           y*(6.83052205423625007e10 +  # noqa: W504, E126, E226
+                              y*(1.09049528450362786e12 +  # noqa: W504, E126, E226
+                                 y*(7.57664583257834349e12 +  # noqa: W504, E126, E226
+                                    y*(1.81004487464664575e13 +  # noqa: W504, E126, E226
+                                       y*(6.43291613143049485e12 +  # noqa: W504, E126, E226
+                                          y*(-1.36517137670871689e12)))))))))))  # noqa: W504, E126, E226
+            / (1. +  # noqa: W504, E126, E226
+               y*(8.19595201151451564e2 +  # noqa: W504, E126, E226
+                  y*(2.40036752835578777e5 +  # noqa: W504, E126, E226
+                     y*(3.26026661647090822e7 +  # noqa: W504, E126, E226
+                        y*(2.23355543278099360e9 +  # noqa: W504, E126, E226
+                           y*(7.87465017341829930e10 +  # noqa: W504, E126, E226
+                              y*(1.39866710696414565e12 +  # noqa: W504, E126, E226
+                                 y*(1.17164723371736605e13 +  # noqa: W504, E126, E226
+                                    y*(4.01839087307656620e13 +  # noqa: W504, E126, E226
+                                       y*(3.99653257887490811e13))))))))))  # noqa: W504, E126, E226
+    )
+    # fmt: on
+    return g
+
+
+@jax.jit
+def _si_small_pade(x, x2):
+    # fmt: off
+    return (
+                x*(1. +  # noqa: W504, E126, E226
+                   x2*(-4.54393409816329991e-2 +  # noqa: W504, E126, E226
+                       x2*(1.15457225751016682e-3 +  # noqa: W504, E126, E226
+                           x2*(-1.41018536821330254e-5 +  # noqa: W504, E126, E226
+                               x2*(9.43280809438713025e-8 +  # noqa: W504, E126, E226
+                                   x2*(-3.53201978997168357e-10 +  # noqa: W504, E126, E226
+                                       x2*(7.08240282274875911e-13 +  # noqa: W504, E126, E226
+                                           x2*(-6.05338212010422477e-16))))))))  # noqa: W504, E126, E226
+                / (1. +  # noqa: W504, E126, E226
+                   x2*(1.01162145739225565e-2 +  # noqa: W504, E126, E226
+                       x2*(4.99175116169755106e-5 +  # noqa: W504, E126, E226
+                           x2*(1.55654986308745614e-7 +  # noqa: W504, E126, E226
+                               x2*(3.28067571055789734e-10 +  # noqa: W504, E126, E226
+                                   x2*(4.5049097575386581e-13 +  # noqa: W504, E126, E226
+                                       x2*(3.21107051193712168e-16)))))))  # noqa: W504, E126, E226
+    )
+    # fmt: on
+
+
+@_wraps(_galsim.bessel.si)
+@jax.jit
+def si(x):
+    x2 = x * x
+    return jnp.where(
+        x2 > 16.0,
+        jnp.sign(x) * (jnp.pi / 2)
+        - _f_pade(x, x2) * jnp.cos(x)
+        - _g_pade(x, x2) * jnp.sin(x),
+        _si_small_pade(x, x2),
+    )