Make the SFNO local filter configurable and add its two-branch dhconv form

[mcgibbon]

The local (DISCO) filter used by local_blocks hardcoded a 3x3 Morlet basis and derived its support radius from a mode-count heuristic, so its receptive field could not be set at all. This makes those choices configurable and adds the two-branch form that turns the DISCO filter into the analog of the dhconv spectral filter, so the receptive field becomes a continuous knob where the spectral filter has none.

Under the real-SHT storage convention (only m >= 0 stored, m < 0 implied by conjugate symmetry), SpectralConvS2's complex per-l weight w_l = a_l + i b_l acts over the full m range as a_l + i b_l sign(m). The real part is genuinely isotropic; the imaginary part is an axis-tied zonal phase shift that no isotropic kernel can express, since any real span of radius-only kernels is mirror-symmetric. By Funk-Hecke, convolving with a real radius-only kernel scales f_lm by a real number identically in m, so the filter decomposes exactly as

dhconv  ==  K_a * f + H[K_b * f]

with two learned real radius-only kernels and one fixed, parameter-free H (multiplication by -i sign(m); the sign of the profile K_b learns is absorbed into it). H touches only longitude structure, so it is a per-latitude-row FFT phase flip and needs no spherical harmonic transform. Because H stays zonally global while theta_cutoff shrinks the K branches, locality and chirality become independent knobs.

Defaults reproduce the previous hardcoded filter bit-for-bit, down to the order in which it draws its weights, so an existing local_blocks config builds the same numbers apart from the activation fix below. Its weight moves from blocks.<i>.filter.filter.conv.weight to blocks.<i>.filter.filter.branches.0.weight, so an existing local_blocks checkpoint no longer loads; no released checkpoint uses that path.

Changes:

• fme.core.models.conditional_sfno.local_filter: new module holding LocalFilterConfig (kernel_shape, basis_type, theta_cutoff, basis_norm_mode, two_branch, match_spectral_init), the LocalFilter module, and zonal_quarter_cycle_shift (the H operator). kernel_shape: lmax and theta_cutoff: global are resolution-independent spellings of "one degree of freedom per total wavenumber" and "pi"; kernel_shape: 2lmax and friends ask for a multiple of that count, which the second bullet below shows is what a global filter actually needs.
• match_spectral_init rescales the branch weights to SpectralConvS2's output magnitude. SpectralConvS2's 1/sqrt(channels) init is unit-gain (measured output RMS 1.07 for unit-RMS input); DISCO's 1/sqrt(channels * kernel_size) is ~50x weaker at kernel_size = lmax, which would otherwise confound a comparison of training dynamics with an initialization difference. The scale is computed analytically from the precomputed convolution tensor, so it is deterministic.
• fme.core.disco: expose isotropic morlet in BasisType (it existed but was not config-selectable, and is now reachable from LocalNetConfig too); declare isotropy per family as FilterBasis.is_isotropic, since for piecewise-linear bases it depends on kernel_shape rather than the family; add kernel_shape_for_basis_count; move _compute_cutoff_radius here as compute_cutoff_radius, next to the basis vocabulary it is keyed to, and BasisType with it.
• fme.core.models.conditional_sfno.sfnonet.FourierNeuralOperatorBlock: local blocks now get the same post-filter activation as linear blocks. They previously got none, so a local block differed from a linear block by a missing nonlinearity as well as by its filter. This changes the computation of any model using local_blocks; no training config does, and the path has no released checkpoints.
• SFNONetConfig.local_filter and NoiseConditionedSFNOBuilder.local_filter plumb the config through.
• Test fixtures that build local_blocks: the three spatial-parallel step regression baselines (sm_with_atmos_corr, sm_noise_conditioned, multi_call) are regenerated for the renamed weight and the new activation, and NoiseConditionedSFNO_module_config.yaml records the new local_filter key.

two_branch requires branches that are purely radial at the given kernel_shape and rejects spatial model parallelism, since H is an FFT over the full longitude circle. match_spectral_init likewise needs the fork's precomputed convolution tensor, which the model-parallel torch-harmonics path does not expose.

Two properties of the radial basis worth flagging for review. Both are measured by least-squares fitting the two-branch filter to a random dhconv operator over the whole (l, m) plane, at 4 degrees (45x90, lmax = 45) and reproduced at the 16x32 test grid:

• Counting degrees of freedom is not enough, and conditioning is why. isotropic morlet at kernel_shape: lmax has an l-transfer condition number of 2e11 with only 25 of 45 usable directions, so it cannot realize arbitrary per-l profiles in practice despite nominally spanning them. Piecewise-linear radial bumps at (2 * lmax - 1, 1) give the same 45 degrees of freedom at condition number 81 with exact profile representability.
• Counting degrees of freedom is not enough for a second, independent reason: lmax radial modes span every per-l profile but do not behave isotropically. Each basis function's l-transfer drifts with m (the quadrature is only approximately rotation-equivariant on a lat-lon grid), and cancelling that drift costs degrees of freedom beyond the ones spent spanning the profile. Best achievable relative error: 18% at lmax, 4.0% at 2lmax, 1.7% at 3lmax, 0.9% at 4lmax, with the residual concentrated at high l and the fitted coefficient norm falling as the basis grows -- so oversampling is a better-conditioned fit, not a finer cancellation. Cost is linear in the count, since the DISCO contraction saves a (batch, channels, kernel_size, nlat, nlon) activation.

The fits are done over the full (l, m) plane rather than against the l-transfer measured at m = 0. Fitting at m = 0 looks natural, since the target is m-independent, but it spends every degree of freedom on one m slice and lands about 5x worse; an earlier revision of this branch did that and concluded the ~30% ceiling was a property of the discretization.

• Tests added
• If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

[mcgibbon]

Claude (pre-review agent): independent pre-review, read cold against main at 11a7b6c5f. I have fix authority on this PR, so findings below are already applied and pushed (66714a0d..4213e7a3).

Audit passes

• CI / test-suite ❌→fixed. The branch was red on all three test jobs, four failures, all caused by this change and all invisible from the diff alone:
  • test_latest_module_backwards_compatibility[NoiseConditionedSFNO] — local_filter is a new key in the builder config, which that test rejects until the cached copy is regenerated.
  • test_step_regression[{sm_with_atmos_corr,sm_noise_conditioned,multi_call}] — all three cached step configurations in fme/core/distributed/parallel_tests/ set local_blocks: [0], so they break twice over: the weight renames from filter.filter.conv.weight to filter.filter.branches.0.weight, and the new activation changes their outputs. Regenerated single-rank per AGENTS.md, then verified under torchrun at model/H=2,W=1, model/H=1,W=2 and torch/1,1. Full -m parallel suite green in all three.
• Correctness of the decomposition ✅, verified independently rather than from the derivation. With SpectralConvS2's per-l weight set purely imaginary: an m = 0 field maps to exactly zero, a Nyquist-only field is identically zero on the grid (so its imaginary part is a null direction as claimed), and on a general field SpectralConvS2 equals -H[isht(sht(f) * b_l)] to 1.5e-7. The sign lands on the second branch's learned profile; I noted that in the description.
• _set_spectral_matched_init algebra ✅. Var(out) = Var(w) * groupsize * sum_k g_k^2 holds even though the per-basis outputs are correlated across k, because the weights are independent and zero-mean, so the cross terms drop. sum(psi_k^2) / nlat_out is the right normalization given that psi is stored per output latitude with longitude handled by the FFT contraction. Measured: a single matched branch has output std 1.016 on unit-variance input, and each of two branches 0.708, against the predicted 1 and 0.707.
• Silent-failure pass ✅ with one note. Every new branch fails loudly: unknown kernel_shape spelling, theta_cutoff out of (0, pi], non-radial basis under two_branch, zero total basis gain, and the hasattr probe for the fork's convolution tensor. Note: basis_type: isotropic morlet and zernike silently use kernel_shape[0] and ignore kernel_shape[1], which becomes reachable from config for the first time here via LocalNetConfig; documented rather than made an error, since LocalNetConfig.kernel_shape is typed as a 2-tuple.
• Behavior-change pass ✅. Defaults are bit-identical to main: I built SpectralFilterLayer(filter_type="local") on both checkouts under one seed and got the same weights and the same output to the last printed digit, so the RNG draw order is unchanged too. BasisType widens by one member, which additively widens LocalNetConfig and the LocalNet registry builder. The local_filter=None fallback in SpectralFilterLayer and FourierNeuralOperatorBlock duplicates the dataclass default but only reaches direct constructions (e.g. the block benchmark), where it gives the historical filter.
• Differential-test pass ✅ after fixes. Three tolerances were not doing the work their docstrings claimed:
  • test_isotropic_branch_commutes_with_longitude_reflection bounded the error at 1e-2 and called the symmetry approximate. It is exact to float32 round-off — measured 2e-7 — so the bound would have passed with the branch substantially anisotropic. Tightened to 1e-5 and corrected the comment.
  • test_match_spectral_init... asserted unscaled < 0.5 * matched, measured 0.419 vs 0.462: 10% from red, on a scale (the ~50x of the description needs kernel_size 45 at 512 channels, not 32 at 16) the test does not measure. Loosened to 0.7 and both numbers now print on failure.
  • Added a parallel test for the two_branch spatial-parallelism guard, which is what stops a tiled run from computing the phase shift per longitude tile and silently returning a wrong filter.
  • The remaining margins are honest: reproduction error 0.044 against a 0.07 bound; errors["2lmax"] < errors["lmax"] / 4 is 0.0438 vs 0.0539, the tightest of them, and it is guarding exactly the 4x claim. test_oversampling... does guard its claim (all four assertions fail if the basis is shrunk back to lmax); I re-marked it medium_duration since it runs in 2s and deserves to run in the --fast loop.
• kernel_shape parsing ✅. ^(\d+)?lmax$ accepts lmax/2lmax, rejects lmax2, 2 lmax, all, and 0lmax separately; each case is covered. resolved_kernel_shape for an explicit list is unvalidated, but every bad list raises inside the basis constructor, so nothing is silent.
• AI-artifact pass ✅ after fixes. Dropped a type: ignore[arg-type] in resolved_theta_cutoff by ordering the branches so mypy narrows the union (AGENTS.md wants a justification for each one). The precision comment in zonal_quarter_cycle_shift claimed the filter runs with autocast disabled; only the spectral path does that, so it now states the cast on its own terms. Distributed.get_instance() was called twice in one constructor.
• Split heuristic ⚠️ see the question below.
• Writing style ✅ on the description and the new docstrings.

Description

Updated: the state-dict key rename is now stated as a checkpoint break on the local_blocks path, "reproduce exactly" is qualified as bit-for-bit including draw order, the regenerated fixtures are listed, and the K_b sign convention and the LocalNetConfig reach of the BasisType widening are noted.

Open question

dd9cce0c4 (the activation change) is the one commit here that alters existing behaviour, and it is why three regression baselines shared with unrelated step configurations had to be regenerated. It is legitimately part of the point of the PR — a local block that differs from a linear block by a missing nonlinearity is not the comparison you want — but if you would rather review the baseline churn on its own, it splits cleanly into its own PR ahead of this one. I left it here.

Does it do what it says, and is it reviewable?

Yes. The claim I would have most expected to be overstated — that the two-branch form is the exact continuum decomposition of dhconv — holds under direct measurement, including the two edge wavenumbers, and the init algebra is right for a reason (weight independence) that is easy to get wrong. The DOF-counting-is-not-enough findings are the substance and they are carried by tests rather than only by prose. Review load is concentrated in one new module plus its tests; the disco changes are a move plus one new property per family.

Ready for review, after the fixes above. Still marked draft.

[mcgibbon]

Claude (pre-review agent): one more fix after the comment above, plus the verification record.

07a2cb66 rejects a boolean theta_cutoff. bool is an int subclass, so theta_cutoff: true passed the (0, pi] range check and silently configured a 1 radian support radius. true is a plausible thing to write for an option whose name does not read as a number, so this failed the wrong way round. Covered in test_config_rejects_invalid_values.

Also loaded the new config through the registry from YAML end to end, which was otherwise only exercised by direct dataclass construction: kernel_shape: 2lmax with theta_cutoff: global round-trips, an int cutoff is accepted, and a bad kernel_shape spelling, a bad cutoff string, an unknown key and a bad basis_type each raise.

Local verification of the final branch: pytest -n 4 . green apart from three failures that predate the branch and are absent from CI (missing graphcast deps, torchrun subprocess paths in a worktree); -m parallel green at every configuration in cpu_test_all_parallel I could run here — model at H2W1, H1W2, H2W2 and torch/model with no spatial ranks at 2 and 3 ranks; pre-commit clean on every changed file.