Neural fields across the stack: what exists, what is missing, and the plan

1Motivation¶

Report 18 reviews neural fields in the geosciences for 2025 and 2026. Its conclusion for GeoStack is that neural fields are worth building for 1) gridless reconstruction (fitting tracks, swaths and pixels directly), 2) uncertainty-aware gap-filling, and 3) differentiable state parameterizations in variational assimilation. However, they are not worth building as an archival compression format, where SZ3 still wins.

We surveyed geonnax (v0.0.9), pyrox (pyrox-nn, pyrox-gp, pyrox-lgm) and vardax (v0.1.9) on 23 October 2026 and found that the stack already has most of the parts, and what it lacks is the plumbing between them. So this page records what exists, what is missing, where each missing piece belongs under the dependency rules, and the order in which we build it. The work is Wave 6.

2Decision¶

1. Deterministic neural-field parts go in geonnax: encoders, architectures and coordinate-derivative operators. geonnax keeps its rules (no numpyro, no kernels, no losses, one example per call).

2. Probabilistic parts go in pyrox-nn, as wrappers which swap named geonnax fields for NumPyro sites, plus an estimator whose likelihood takes per-observation noise and an observation operator.

3. Assimilation parts go in vardax: a point-observation batch, a nonlinear StateParameterization, an optax minimiser, a PDE residual cost and a space-time neural-field 4D-Var model.

4. Acceptance tests are existing case studies: case 9 (SSH mapping, with DUACS/OI as the oracle), case 11 (SST gap-filling, with the GP as the oracle) and overview stage 3 (methane L2 → L3, with kriging as the oracle).

5. Compression is a research bet, which we score against SZ3/ZFP at matched maximum error before anything is built for it.

3Where a neural field flows¶

  observations  y ∈ ℝ^{N_obs}  (tracks, swaths, pixels)
  geostack: GeoPoints / GeoTracks, to_jax(pad_to=N) + valid mask   (CS-F3)
        │
        ▼
  geonnax · encode         γ(s, t), s = (λ, φ)
    SphericalHarmonicEncoder, SlepianEncoder, Cartesian3DEncoder,
    CyclicEncoder, GeoContextEncoder
    ✗ Gaussian Fourier features, HEALPix grid / feature grid
        │
        ▼
  geonnax · represent      x(s, t; θ) = g_net(γ(s, t); θ)
    SIREN, MultiScaleSIREN, FourierNet, GaborNet,
    ConditionedINR + FiLM / HyperLinear, HyperSIREN
    ✗ separable, latent table + shift modulation, LIIF, BNF core
        │
        ├── ✗ geonnax · differentiate   ∇x, ∇²x; on S²: ∇_S² x, ∇²_S² x
        │
        ▼
  pyrox-nn · make Bayesian  θ ~ p_prior(θ)
    BayesianSIREN, BayesianFourierNet, BayesianNeuralField,
    EnsembleMAP, EnsembleVI, BNFEstimator
        │
        ▼
  vardax · observe          ŷ = ℋ[x(·; θ)]
    AveragingKernel, MultiInstrumentFusion
    ✗ point evaluation at observation coordinates
        │
        ▼
  likelihood / loss         p_lik(y | ŷ; Σ_ε)  or  ℒ(θ)
    ✗ pyrox-nn: per-observation σ, ℋ inside the likelihood
    ✗ vardax: θ as the 4D-Var control, PDE residual term
        │
        ▼
  geostack · score          case 9, case 11, methane stage 3

4What exists, and the gap¶

5The plan by library¶

geonnax (theme 1, theme 2)¶

The derivative operators are operators and not losses, so they keep geonnax’s “no losses” rule, and the residual costs live in vardax. The BayesNF core moves down from pyrox-nn for two reasons: 1) so that pyrox-nn’s BNF becomes a wrapped core like every other Bayesian network there, and 2) so that it can take spherical-harmonic features.

pyrox-nn (theme)¶

pyrox (GeoML tier 2) may not import vardax (tier 3), so the likelihood takes the observation operator as a plain callable, which vardax’s AveragingKernel satisfies.

vardax (theme)¶

The point-observation batch also fixes the classical models on along-track data, so it pays off without any neural field. somax stays a forbidden import, so the PDE residuals enter as callables, as the forward models do today.

The space-time neural-field 4D-Var loss is written over the network weights $\gvect{\theta} \in \mathbb{R}^{D_\theta}$, with the field $x(\vect{s}, t; \gvect{\theta}) \in \mathbb{R}$ at the coordinates $\vect{s} = (\lambda, \phi)$ and time $t$. Each observation $y_n \in \mathbb{R}$, $n = 1, \ldots, N_\text{obs}$, sits at $(\vect{s}_n, t_n)$ with error standard deviation $\sigma_n \in \mathbb{R}_{>0}$. The PDE residual operator $\mathcal{R}[x]$ vanishes on the true flow and is evaluated at the collocation points $(\vect{s}_j, t_j)$, $j = 1, \ldots, N_c$, and $\beta_\text{PDE}, \beta_\theta \in \mathbb{R}_{\ge 0}$ are the weights of the two penalties:

There is no background term (the $\gmat{\Sigma}_x$ term of classical 4D-Var), as the network’s spectral bias plays its role. We write it as a loss $\mathcal{L}$ rather than a cost $\mathcal{J}$ because the control is the parameters $\gvect{\theta}$ and not a state.

geostack (theme)¶

Carrier wrappers stay out of geonnax (L1 only). The GeoTracks → field operator for case 9 goes behind vardax[geo], and those for case 11 and stage 3 go behind pyrox.adapters.pipekit.

6Stages and gates¶

  A · plumbing
    vardax#83 point obs · pyrox#316 per-obs σ · geonnax#132 diffops · geostack#84 stat dim
    ── GATE: SH-SIREN on case 9 tracks matches OI RMSE;
             BNF + σₙ on case 11 ≤ GP CRPS
        │
        ▼
  B · uncertainty + methane
    pyrox#317, #318, #319, #320, #321 · geonnax#133, #134, #135
    ── GATE: stage-3 XCH₄ gap-fill beats kriging;
             90% intervals cover 85–95%
        │
        ▼
  C · physics
    vardax#84, #85, #86, #87, #88 · geonnax#136 · pyrox#322, #323
    ── GATE: reproduce Oh (2025) on 2-D Kolmogorov flow;
             case 9 with a QG residual beats the stage-A field
        │
        ▼
  D · research bets
    geonnax#137–#141 · geostack#86
    ── GATE: each beats its stage A–C baseline on one anchor,
             or it stays a notebook

The gates are scored in geostack#85. Stage A is low risk to start now, as every item in it also helps models without neural fields.

7Open questions¶

8Relationships¶

• Wave epic: geostack#82

• Literature: report 18

• Case studies: case 9, case 11; stage 3 in the overview

• Layers and tiers: levels across the stack, dependency graph