jaxht 0.1.0__py3-none-any.whl

jaxht-0.1.0.dist-info/METADATA

@@ -0,0 +1,165 @@
+Metadata-Version: 2.4
+Name: jaxht
+Version: 0.1.0
+Summary: JAX-native spherical harmonic transforms (BK-regime first): GPU-capable, differentiable, dependency-controlled.
+Project-URL: Homepage, https://github.com/jrcheshire/jht
+Project-URL: Repository, https://github.com/jrcheshire/jht
+Project-URL: Issues, https://github.com/jrcheshire/jht/issues
+Project-URL: Changelog, https://github.com/jrcheshire/jht/blob/main/CHANGELOG.md
+Author-email: James Cheshire <cheshire@caltech.edu>
+Maintainer-email: James Cheshire <cheshire@caltech.edu>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: automatic-differentiation,cosmic-microwave-background,cosmology,gpu,healpix,jax,nufft,spherical-harmonic-transform
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.4.30
+Requires-Dist: numpy>=1.24
+Provides-Extra: dev
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: validation
+Requires-Dist: ducc0; extra == 'validation'
+Requires-Dist: healpy; extra == 'validation'
+Description-Content-Type: text/markdown
+
+# jht — JAX Harmonic Transforms
+
+JAX-native spherical harmonic transforms (map ↔ aₗₘ): **GPU-capable**, **fully
+differentiable**, and **dependency-controlled** (pure JAX + numpy at runtime — no
+compiled C++ extension, no heavyweight third-party SHT library). Scoped to the
+BICEP/Keck regime — **spin-0 and spin-2** on the **HEALPix RING** pixelization,
+ℓ_max ≲ 1000, nside ≤ ~2048 — but written cleanly so it can serve as a general
+transform dependency.
+
+It exists to serve the GPU / differentiable tier of analysis that a CPU-only C++
+transform (ducc0) structurally cannot, while *owning the numerics*. See
+[`docs/motivation.md`](docs/motivation.md) for the full decision record.
+
+## Status (2026-06-10)
+
+Phases 0–4 **complete and validated** (190 tests pass + 8 GPU-gated skips,
+CPU/float64):
+
+- **On-grid transforms** — spin-0 & spin-2 synthesis (`aₗₘ→map`) and the exact
+  adjoint `Sᵀ`, validated to machine precision vs healpy **and** ducc0; spin-2
+  inverse at the HEALPix floor with **no s2fft-style structural defect**.
+- **Accuracy** — jht's own ring quadrature weights + Jacobi iteration reach
+  ~1e-13 on band-limited maps (matches `healpy.map2alm(use_weights=True)`); see
+  [`docs/accuracy.md`](docs/accuracy.md).
+- **Partial-sky** — masked pseudo-aₗₘ, a cut-sky CG deconvolution, and a masked
+  Wiener filter / constrained realization (the MUSE inner solve); see
+  [`docs/masked.md`](docs/masked.md).
+- **Off-grid (NUFFT)** — `synthesis_general` / `adjoint_synthesis_general`
+  evaluate a band-limited field at **arbitrary pointings** (spin 0–3), alm- **and**
+  pointing-differentiable. The JAX-native replacement for ducc0's
+  `sht.experimental.synthesis_general` (on-grid SHT + this NUFFT = the full ducc0
+  surface bk-jax depends on); see [`docs/offgrid.md`](docs/offgrid.md).
+- **Differentiability** — native JAX autodiff (`jacfwd ≡ jacrev`, tight adjoint
+  identity), plus a convention-clean real-DOF layer `jht.diff`; see
+  [`docs/design.md`](docs/design.md) §Differentiability.
+- **GPU** — pure JAX, so the transforms run on CUDA with no code change. Measured on
+  Cannon A100/V100 (fp64): GPU==CPU parity ~1e-13 across the BK regime **including
+  nside=2048**, forward synthesis 14–60× the CPU. See **Performance** below and
+  [`docs/gpu.md`](docs/gpu.md).
+
+## Install
+
+Standard env is [pixi](https://pixi.sh):
+
+```bash
+pixi install          # CPU env (osx-arm64 / linux-64)
+pixi run test         # the gated suite
+```
+
+GPU (CUDA, linux-64 — see [`docs/gpu.md`](docs/gpu.md)):
+
+```bash
+pixi run -e gpu python scripts/gpu_check.py   # on an NVIDIA box
+```
+
+As a dependency in another project (runtime deps are just `jax` + `numpy`):
+
+```bash
+pip install jaxht        # once released on PyPI — then `import jht`
+# or track the repo directly:
+pip install "jaxht @ git+https://github.com/jrcheshire/jht.git"
+```
+
+## Quick start
+
+float64 is **opt-in per entry point** — enable it *before creating any array*
+(library code never touches the global config):
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+import jht
+
+nside, lmax, spin = 256, 512, 0
+m = jht.synthesis(alm, nside, lmax, spin=spin)          # aₗₘ -> map
+a = jht.map2alm(m, nside, lmax, spin=spin, niter=3)     # map -> aₗₘ (weighted + iterated)
+cl = jht.bandpower(a, lmax, spin=spin)                  # angular auto-power C_ℓ
+```
+
+`spin=2` takes/returns `(E, B)` aₗₘ of shape `(2, …)` and `(Q, U)` maps of shape
+`(2, npix)`. `jht.adjoint_synthesis` is the **exact unweighted transpose** `Sᵀ`
+(the operator seam / VJP), distinct from `map2alm` (the approximate inverse). For
+gradient-based work use the real-DOF layer `jht.synthesis_real` /
+`jht.analysis_real` (plain ℝⁿ→ℝᵐ, no complex-conjugate convention subtlety).
+
+## Conventions
+
+healpy m-major triangular aₗₘ packing, orthonormal Yₗₘ with the Condon–Shortley
+phase, HEALPix-internal (COSMO) polarization — verified against healpy 1.19.0 and
+ducc0 0.41.0. Pinned in [`docs/design.md`](docs/design.md).
+
+## Accuracy tiers (the contract)
+
+jht targets the **GPU/differentiable tier** where the HEALPix ~1e-3 sampling
+floor is acceptable; weights + iteration close it to ~1e-13 on band-limited
+inputs. It is **not** a drop-in for ducc's purity-critical (~1e-4 E→B-leakage)
+production path. Tolerances are a-priori and gate-driven, never relaxed without
+sign-off. Residual mismatches are logged in
+[`DISCREPANCIES.md`](DISCREPANCIES.md).
+
+## Performance
+
+Pure JAX runs unchanged on CUDA. Measured on Cannon A100 (incl. a 20 GB MIG) / V100, fp64:
+
+- **GPU==CPU parity ~1e-13** across the BK regime, **including nside=2048** (synthesis and `map2alm`).
+- **Forward synthesis 14–60×** the 8-core CPU; fp64/fp32 ≈ 2.2×.
+- **Off-grid forward** ~0.5–0.9 s at ℓ_max=1000 (independent of the number of points; recursion-bound), with the pointing gradient ~1× a forward.
+- **nside=2048** compiles and runs on GPU — a ~20 GB slice holds synthesis + `map2alm`; the one-time compile is multi-minute (jit-cached).
+
+The recurring GPU lesson: fp64/complex scatters are catastrophic on GPU, so jht packs and assembles via **gathers**. CPU perf model + memory in [`docs/performance.md`](docs/performance.md); GPU detail in [`docs/gpu.md`](docs/gpu.md).
+
+## Using jht as a dependency
+
+jht is standalone and consumer-agnostic. The operator/grad seam a downstream
+needs (e.g. to use jht *in place of ducc0*) — and the accuracy boundary — are
+documented in [`docs/consumers.md`](docs/consumers.md). Any backend-selection
+wiring lives in the consumer, not here.
+
+## Docs
+
+- [`docs/design.md`](docs/design.md) — technical design, conventions, the crux, differentiability.
+- [`docs/accuracy.md`](docs/accuracy.md) — the accuracy contract + ring-weight algorithm.
+- [`docs/masked.md`](docs/masked.md) — partial-sky estimators.
+- [`docs/performance.md`](docs/performance.md) — CPU perf model + memory.
+- [`docs/gpu.md`](docs/gpu.md) — the GPU env, the x64 requirement, the harness.
+- [`docs/consumers.md`](docs/consumers.md) — the downstream-dependency seam.
+- [`docs/motivation.md`](docs/motivation.md) — why jht exists.
+- [`ROADMAP.md`](ROADMAP.md) — phased plan + gates.

jaxht-0.1.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+jht/__init__.py,sha256=SyhUjXRnAQWP8-H0ASXw2UGWwaizlIwh6UYq2IWFMHA,3128
+jht/_nufft.py,sha256=rHz7023H2mMMMuKtERFVUUrxT_bCf4Oou-J9hRVq5oI,7078
+jht/_recursion.py,sha256=MjPnWrafuR2Ru3nt4Gb45-hdqAaRhgfi_UuqVXfo_Aw,24951
+jht/_reference.py,sha256=ZrWTDSKQgrqzN1id0WjsvLmrjx4TcvQzqjYq_ykVf88,5459
+jht/analysis.py,sha256=BS4M-GwUNvUjyRLUdAtFxT93fAmSej3-MnqcjgSgOis,2291
+jht/diff.py,sha256=UNBw-j4lR7npGiG1grMmZSocTu_pm4ryIFgrjQABNaY,4130
+jht/healpix.py,sha256=uZrV5cXNgSpwU8r65UwmnveJC8E9mDfO8iGKJ_6FoHs,14837
+jht/masked.py,sha256=DBG-gqz_uwg1kIuRt2Dm8eF9Xsx13E-sJPXc4tav4KY,17320
+jht/offgrid.py,sha256=3t2kesjOtH4Cm2R0Hf47FtmxnKh1VJZoUQdWe9ANENM,9375
+jht/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jht/weights.py,sha256=WzY2ucCm_Ty6c3VODIwYOy27j6umMZGUlpUZ3pZw9Oc,4613
+jaxht-0.1.0.dist-info/METADATA,sha256=DTyXdwtqi6i3-llXZLorkpSxibtwXJhaielN9qcgk7Y,7855
+jaxht-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+jaxht-0.1.0.dist-info/licenses/LICENSE,sha256=N8ykVEnf_WjCA9Ymw5GPMitqZrmMx911pDxpl-xBHq4,1071
+jaxht-0.1.0.dist-info/RECORD,,

jaxht-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

jaxht-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 James Cheshire
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jht/__init__.py

@@ -0,0 +1,86 @@
+"""jht -- JAX-native spherical harmonic transforms.
+
+A clean-room, pure-JAX implementation of the forward and inverse spherical
+harmonic transform (map <-> a_lm) for **spin-0 and spin-2** fields on the
+**HEALPix RING** pixelization, in the BICEP/Keck angular regime
+(l_max <~ 1000, nside <= ~2048).  GPU-capable, fully differentiable under JAX's
+native autodiff, and dependency-controlled (runtime deps = jax + numpy only).
+
+Conventions (verified vs healpy 1.19.0 / ducc0 0.41.0; see ``docs/design.md``):
+healpy m-major triangular a_lm packing, orthonormal Y_lm with the
+Condon-Shortley phase, HEALPix-internal (COSMO) polarization.
+
+float64 is **opt-in per entry point** -- library code never touches the global
+config.  Enable it (before creating any array, on CPU or GPU) with::
+
+    import jax
+    jax.config.update("jax_enable_x64", True)
+
+Quick start::
+
+    import jax; jax.config.update("jax_enable_x64", True)
+    import jht
+    m = jht.synthesis(alm, nside, lmax, spin=0)        # a_lm -> map
+    a = jht.map2alm(m, nside, lmax, spin=0, niter=3)   # map -> a_lm (weighted)
+
+See ``README.md`` for the tour, ``docs/consumers.md`` for the downstream seam
+(e.g. using jht as a ducc0 replacement), and ``docs/gpu.md`` for the GPU story.
+"""
+
+from __future__ import annotations
+
+from .analysis import bare_analysis, map2alm
+from .diff import analysis_real, bandpower, synthesis_real
+from .healpix import (
+    adjoint_synthesis,
+    alm_metric_weight,
+    alm_size,
+    synthesis,
+)
+from .masked import (
+    alm_to_real,
+    constrained_realization,
+    deconvolve,
+    n_dof,
+    pseudo_alm,
+    real_to_alm,
+    wiener,
+)
+from .offgrid import adjoint_synthesis_general, synthesis_general
+from .weights import pixel_weights, ring_weights
+
+__version__ = "0.1.0"  # single source of truth; pyproject reads this via hatch
+__author__ = "James Cheshire"
+__email__ = "cheshire@caltech.edu"
+
+__all__ = [
+    "__version__",
+    # core on-grid transforms
+    "synthesis",  # S:  a_lm -> map
+    "adjoint_synthesis",  # S^T: map -> a_lm (exact unweighted transpose / VJP)
+    # off-grid (arbitrary points / NUFFT) -- spin 0-3, alm- AND pointing-differentiable
+    "synthesis_general",  # a_lm -> field at arbitrary (theta, phi)
+    "adjoint_synthesis_general",  # its exact transpose
+    # analysis (approximate inverse)
+    "bare_analysis",  # A0 = S^T W
+    "map2alm",  # A0 + Jacobi iteration
+    # quadrature weights
+    "ring_weights",
+    "pixel_weights",
+    # partial-sky / masked
+    "pseudo_alm",  # masked pseudo-a_lm (zero-fill)
+    "deconvolve",  # cut-sky CG deconvolution
+    "wiener",  # Wiener filter / MUSE inner solve (Cl prior + N^-1)
+    "constrained_realization",  # posterior draw (constrained realization)
+    # differentiable (real-DOF) interface
+    "synthesis_real",  # S o T^-1 : R^n -> map
+    "analysis_real",  # T o map2alm : map -> R^n
+    "bandpower",  # angular auto-power C_ell
+    # real-DOF isometry T
+    "alm_to_real",
+    "real_to_alm",
+    "n_dof",
+    # a_lm layout / inner-product helpers
+    "alm_size",
+    "alm_metric_weight",  # the (2 - delta_m0) metric G (the bk-jax 2*conj bridge)
+]

jht/_nufft.py

@@ -0,0 +1,171 @@
+"""Pure-JAX 2D NUFFT (ES kernel) -- the off-grid engine for ``jht.offgrid``.
+
+Type-2 (uniform Fourier modes -> nonuniform points) and its exact transpose,
+type-1.  Both are differentiable in the nonuniform point coordinates under JAX's
+native autodiff: the integer window index is frozen with ``stop_gradient`` (it is
+genuinely locally constant), so the gradient flows analytically through the
+smooth ES-kernel weights -- no custom rule, ``jacfwd == jacrev`` preserved.
+
+The kernel is the FINUFFT "exponential of semicircle"
+``psi(t) = exp(beta (sqrt(1 - (2t/W)^2) - 1))`` (Barnett, Magland & af Klinteberg
+2019, arXiv:1808.06736).  The accuracy-vs-(sigma, W) trade was pinned empirically
+in the M0 spike (abs err vs the exact direct sum, lmax<=512): W=14/sigma=2 reaches
+<=6e-10; W=12 ~5e-9; W=16 sits at the fp64 roundoff floor.
+
+Convention: modes are *centered*, ``k = -N//2 .. -N//2 + N - 1``; points
+``x, y`` live on the ``[0, 2pi)`` torus.  ``nufft2d2`` evaluates
+``f_j = sum_{k,m} c[k,m] exp(i (k x_j + m y_j))``.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+TWO_PI = 2.0 * np.pi
+
+# requested epsilon -> (W, sigma).  Keyed by *achievable* accuracy; the cheapest
+# (smallest W) entry meeting the request is selected.  Tuned in scripts/offgrid_spike.py.
+_KERNEL_DB: dict[float, tuple[int, float]] = {
+    1e-4: (6, 2.0),
+    1e-6: (8, 2.0),
+    1e-8: (10, 2.0),
+    1e-10: (14, 2.0),
+    1e-12: (16, 2.0),
+}
+
+
+def _eps_to_w_sigma(eps: float) -> tuple[int, float]:
+    # Keys are achievable accuracies; pick the loosest (cheapest -> largest key)
+    # that still meets the request.  If eps is tighter than the DB's best, fall
+    # back to min(key) -- the smallest key is the most-accurate (largest-W) kernel.
+    ok = [e for e in _KERNEL_DB if e <= eps * 1.0000001]
+    return _KERNEL_DB[max(ok)] if ok else _KERNEL_DB[min(_KERNEL_DB)]
+
+
+def _next_size(n: float) -> int:
+    """Next even FFT-friendly size (a 2,3,5-smooth integer) >= n."""
+    n = int(np.ceil(n))
+    if n % 2:
+        n += 1
+    while True:
+        m = n
+        for p in (2, 3, 5):
+            while m % p == 0:
+                m //= p
+        if m == 1:
+            return n
+        n += 2
+
+
+def _beta(W: int, sigma: float) -> float:
+    return float(np.pi * W * (1.0 - 0.5 / sigma))
+
+
+def _es_kernel(t, beta: float, W: int):
+    """ES window ``exp(beta(sqrt(1-(2t/W)^2)-1))``; 0 outside ``|t| <= W/2``."""
+    z2 = (2.0 * t / W) ** 2
+    val = jnp.exp(beta * (jnp.sqrt(jnp.clip(1.0 - z2, 0.0, 1.0)) - 1.0))
+    return jnp.where(z2 <= 1.0, val, 0.0)
+
+
+def _correction(N: int, n: int, W: int, beta: float) -> np.ndarray:
+    """Continuous-FT deconvolution factor ``K^(k)`` for centered modes (real, even)."""
+    ks = np.arange(-(N // 2), -(N // 2) + N)
+    t, wq = np.polynomial.legendre.leggauss(2 * W + 24)
+    t, wq = t * (W / 2.0), wq * (W / 2.0)
+    psi = np.exp(beta * (np.sqrt(np.clip(1.0 - (2.0 * t / W) ** 2, 0.0, 1.0)) - 1.0))
+    u = TWO_PI / n
+    return (psi[None, :] * np.cos(ks[:, None] * u * t[None, :]) * wq[None, :]).sum(1)
+
+
+class NufftPlan(NamedTuple):
+    """Static (host-side) tables for a 2D NUFFT at a given (Nk, Nm, eps)."""
+
+    Nk: int
+    Nm: int
+    nk: int  # oversampled grid size (theta/k axis)
+    nm: int  # oversampled grid size (phi/m axis)
+    W: int
+    beta_k: float
+    beta_m: float
+    corr_k: np.ndarray  # (Nk,)
+    corr_m: np.ndarray  # (Nm,)
+    kbins: np.ndarray  # (Nk,)  centered modes mod nk
+    mbins: np.ndarray  # (Nm,)  centered modes mod nm
+    fill_src: np.ndarray  # (nk*nm,) gather index: oversampled cell -> source mode in c.ravel()
+    fill_mask: np.ndarray  # (nk*nm,) True where a centered mode lands on this cell
+
+
+@lru_cache(maxsize=None)
+def nufft_plan(Nk: int, Nm: int, eps: float) -> NufftPlan:
+    W, sigma = _eps_to_w_sigma(eps)
+    nk, nm = _next_size(sigma * Nk), _next_size(sigma * Nm)
+    bk, bm = _beta(W, sigma), _beta(W, sigma)
+    kbins = np.arange(-(Nk // 2), -(Nk // 2) + Nk) % nk
+    mbins = np.arange(-(Nm // 2), -(Nm // 2) + Nm) % nm
+    # Inverse fill index so nufft2d2 builds the oversampled grid by a GATHER, not an
+    # fp64/complex SCATTER (`C.at[].set`) -- the scatter was 97% of the off-grid forward
+    # on GPU (~25000x the gather; see scripts/profile_offgrid_forward.py).  Centered modes
+    # at sigma>=2 are collision-free in the oversampled grid, so the gather reproduces the
+    # scatter exactly (asserted below).
+    dest = (kbins[:, None] * nm + mbins[None, :]).ravel()
+    if np.unique(dest).size != Nk * Nm:
+        raise ValueError("NUFFT grid-build indices collide; gather build invalid")
+    flat = np.full(nk * nm, -1, dtype=np.int64)
+    flat[dest] = np.arange(Nk * Nm)
+    return NufftPlan(
+        Nk, Nm, nk, nm, W, bk, bm,
+        _correction(Nk, nk, W, bk),
+        _correction(Nm, nm, W, bm),
+        kbins, mbins,
+        np.where(flat >= 0, flat, 0), flat >= 0,
+    )
+
+
+def _stencil(coord, n: int, W: int, beta: float):
+    """``(idx (Npts,W) mod n, weights (Npts,W))`` for the W-wide ES stencil.
+
+    The window base is ``stop_gradient``-frozen (locally constant); the weights
+    are smooth in ``coord`` -> native AD gives the analytic point-derivative.
+    """
+    xi = coord * n / TWO_PI
+    base = jax.lax.stop_gradient(jnp.ceil(xi - W / 2.0)).astype(jnp.int64)
+    off = jnp.arange(W)
+    ell = base[:, None] + off[None, :]  # (Npts, W) true grid coords (un-wrapped)
+    w = _es_kernel(xi[:, None] - ell, beta, W)
+    return ell % n, w
+
+
+def nufft2d2(plan: NufftPlan, coeffs, x, y):
+    """Type-2: ``f_j = sum_{k,m} coeffs[k,m] exp(i(k x_j + m y_j))``.
+
+    ``coeffs`` (Nk, Nm) complex (centered modes); ``x, y`` (Npts,) on ``[0,2pi)``.
+    """
+    c = coeffs / jnp.asarray(plan.corr_k)[:, None] / jnp.asarray(plan.corr_m)[None, :]
+    # Build the oversampled grid by a GATHER (the inverse fill index lives in the plan);
+    # the equivalent fp64/complex `.at[].set` scatter was 97% of the GPU forward.
+    C = jnp.where(
+        jnp.asarray(plan.fill_mask), c.ravel()[jnp.asarray(plan.fill_src)], 0.0 + 0.0j
+    ).reshape(plan.nk, plan.nm)
+    g = jnp.fft.ifft2(C) * (plan.nk * plan.nm)
+    lk, wk = _stencil(x, plan.nk, plan.W, plan.beta_k)
+    lm, wm = _stencil(y, plan.nm, plan.W, plan.beta_m)
+    gg = g[lk[:, :, None], lm[:, None, :]]  # (Npts, W, W)
+    return jnp.sum(wk[:, :, None] * wm[:, None, :] * gg, axis=(1, 2))
+
+
+def nufft2d1(plan: NufftPlan, vals, x, y):
+    """Type-1: ``coeffs[k,m] = sum_j vals_j exp(-i(k x_j + m y_j))`` -- transpose of type-2."""
+    lk, wk = _stencil(x, plan.nk, plan.W, plan.beta_k)
+    lm, wm = _stencil(y, plan.nm, plan.W, plan.beta_m)
+    contrib = (wk[:, :, None] * wm[:, None, :]) * jnp.asarray(vals)[:, None, None]
+    g = jnp.zeros((plan.nk, plan.nm), dtype=jnp.complex128)
+    g = g.at[lk[:, :, None], lm[:, None, :]].add(contrib)
+    C = jnp.fft.fft2(g)
+    cc = C[jnp.asarray(plan.kbins)[:, None], jnp.asarray(plan.mbins)[None, :]]
+    return cc / jnp.asarray(plan.corr_k)[:, None] / jnp.asarray(plan.corr_m)[None, :]