feat: metamatrix graph backend, transport/decentering, default to metamath - #138
Open
vhaasteren wants to merge 93 commits into
Open
feat: metamatrix graph backend, transport/decentering, default to metamath#138vhaasteren wants to merge 93 commits into
vhaasteren wants to merge 93 commits into
Conversation
…atrix.ipynb, but I have not checked the math against the standard implementation
…_example-metamatrix.ipynb
…oefficient likelihood
…, visualize graphs with mm.visualize_graph
… folding; seems to be OK for individual pulsars
…etamatrix.ipynb work with the subgraph patch. Memory use and compilation time are much better.
…signals.py with a patch to set `.inv` so that the HD prior uses the faster inversion method. Now the HD likelihood is the same speed as the old code on my laptop (55ms for 67 pulsar HD).
Reduce CPU memory for models that have precomputations by computing design matrices on-the-fly for each pulsar and then discarding them. Keep GPU memory in check by adding a `jax.device_put` context in `fold_constants`.
…W that can work with ExtSignal. The new ExtSignal class that allows for a Fourier determinstic signal to carry its own basis.
… that compare old matrix.py methodology with new metamath/metamatrix methodology.
…nt likelihoods, etc. a BIG update. Include relevant testing suites that haven't been committed yet. Including the "_patch.py" that allows us to switch between the old matrix.py and the new metamath.py-based classes. etc.
… called separately from matrix.py and metamath.py. The plan is to move these towards independent paths that users can switch between for v1.0
Adds ds.config(kernels='matrix'|'metamath') switch that rebinds the top-level discovery.PulsarLikelihood / GlobalLikelihood / ArrayLikelihood between likelihood.py (matrix.py-backed, default) and the new likelihood_metamath.py (a verbatim copy for now, to be migrated class-by-class to metamath equivalents). Parity tests now exercise three routes per row via a shared build_routes helper: the matrix.py reference, metamath via the existing _patch.py monkeypatch, and metamath-native via the new ds.config kernel switch. Eager forcing of cached_properties inside each route keeps closure capture correct across patch / config boundaries. 43/43 parity tests pass; mh_native is trivially equal to matrix today and is where regressions will surface as likelihood_metamath.py migrates.
Two pieces: 1. ds.config(kernels='metamath') now installs a persistent matrix→metamath swap on `discovery.matrix` (signals.py's matrix.NoiseMatrix*, etc. resolve to metamath equivalents) via the new src/discovery/_kernel_switch.py. The mappings live there now; tests/metamatrix/_patch.py re-exports the context-manager form for the mh_patched test route. 2. likelihood_metamath.py is refactored to compose with metamath kernels directly: PulsarLikelihood / ArrayLikelihood / GlobalLikelihood call metamath.CompoundGP / WoodburyKernel / CompoundDelay / VectorWoodburyKernel rather than the matrix.* factories. Dead matrix-side fallbacks under PulsarLikelihood.conditional are dropped. The ArrayLikelihood.clogL decentering bridge (which used .solve_2d under matrix.py) simplifies to the graph-only path. To preserve the configurability that `matrix.config(backend=..., factor= ..., regularize_FtNmF=...)` provides (numpy vs jax, float32 vs float64, cholesky vs LU), kernel_helpers.py forwards the configured aliases — jnp, jsp, jnparray, jnpsplit, jnpnormal, matrix_factor, matrix_solve, matrix_norm, regularize_FtNmF, SM_algorithm, single_precision — dynamically from matrix.py via PEP 562 module __getattr__. likelihood_ metamath uses `import kernel_helpers as kh` and reads `kh.jnp`, etc., so `matrix.config(...)` continues to drive the underlying primitives in both backends. Remaining matrix.* references in likelihood_metamath are the GP marker types (matrix.ConstantGP/VariableGP, used for isinstance dispatch on signals.py outputs — to be removed by the signals.py migration step), matrix.CompoundGlobalGP (untested edge case), and matrix.make_logdet_ estimator / matrix.cgsolve inside cglogL (the open Tier-3 design question, not part of the parity suite). All 43 parity rows pass across matrix / mh_patched / mh_native routes.
ConstantKernel, VariableKernel, ConstantMatrix, VariableMatrix, NoiseMatrix (base), GP, ConstantGP, VariableGP, GlobalVariableGP now live in kernel_helpers.py. matrix.py re-imports them so matrix.X keeps working for existing call sites (signals.py, optimal.py, etc.). likelihood_metamath.py drops the matrix.* isinstance dispatch on the marker types and uses kh.Kernel / kh.ConstantGP / kh.VariableGP directly. matrix is still imported there for the Tier-3 cglogL helpers (matrix.cgsolve, matrix.make_logdet_estimator) and matrix.CompoundGlobalGP (an untested edge case). 43/43 parity tests pass.
Move the numerical backend config (config(), dynamic globals, rngkey, and the optional matfree cgsolve/make_logdet_estimator block) and the kernel-primitive markers + indexed Sherman-Morrison helpers out of matrix.py / kernel_helpers.py into a new path-neutral utils.py. - matrix.py re-exports these from utils so legacy matrix.<name> access keeps working for signals.py / likelihood.py. - metamath.py and likelihood_metamath.py read markers and config aliases from utils instead of kernel_helpers. - kernel_helpers.py deleted; its __getattr__ matrix-proxy is no longer needed since config now lives in the neutral module. This breaks the kernel_helpers -> matrix back-reference, giving both kernel paths a common substrate -- a prerequisite for eventually deleting matrix.py. Full suite green (43 parity + 29 matrix-path tests). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…easurement noise
Production metamath path no longer monkeypatches matrix.*; signals.py builds
kernels through an explicit factory selected by ds.config(kernels=...).
- New _kernels.py: factory resolving kernel constructors to matrix or metamath
per set_mode(). metamath mode uses a canonical _METAMATH map (which
_kernel_switch now imports as _PATCHES, so the test mh_patched route and the
factory can't drift); matrix mode falls through to matrix.* at call time so
the test monkeypatch still flows through.
- __init__.config() sets the factory mode instead of apply_patches/restore.
- signals.py / deterministic.py: all matrix.* references repointed to utils.*
(backend + GP markers) and kernels.* (constructors); unused matrix imports
dropped.
- Closed the VectorNoiseMatrix12D_var gap: going explicit forced the vector
dispatcher to resolve to metamath's NoiseMatrix12D (previously the leaf-class
monkeypatch covered it implicitly).
Collapse the measurement-noise constructors (strangler):
- New measurement_noise.py: makenoise_measurement{,_simple} in collapsed form,
the _novar/_var class dispatch replaced by variant-agnostic factory entry
points _kernels.NoiseMatrix1D / NoiseMatrixSM (choose the class from whether
the noise arg is an array or a callable).
- signals.py re-exports them (keeps signals.makenoise_measurement and
ds.makenoise_measurement resolving) at module end to avoid a circular import.
- Equivalence was verified by a temporary parity test (matrix-isolation:
collapsed-vs-original both matrix mode; plus collapsed-metamath vs matrix
oracle) across white fixed/var, tnequad, ecorr-SM fixed/var, and simple
configs; the migration test is removed now that signals adopts the collapsed
path.
Full suite green (43 parity + 29 matrix-path tests).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
likelihood_metamath.py no longer imports `matrix`: - cgsolve / make_logdet_estimator / jnp / jsp now come from `utils` (kh.*), where they have lived since Phase 0. - CompoundGlobalGP is reached through the `_kernels` factory (kernels.*). It is not in the metamath map, so it falls through to matrix.CompoundGlobalGP unchanged -- behavior-preserving for this untested edge case. Relocating / porting CompoundGlobalGP out of matrix.py is deferred to Phase 4. - GP marker isinstance dispatch already used utils types; stale comments referencing matrix.ConstantGP / matrix.ExtSignal corrected. Full suite green (43 parity + 29 matrix-path tests). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Swept production + notebook (examples/, docs/tutorials/) kernel-constructor usage against parity routes. Results in docs/components/phase3_coverage.md. Closed gaps by adding single-pulsar parity rows: - fftcov_2d: exercises the 2D variable covariance noise path (NoiseMatrix2D_var via the NoiseMatrix12D_var dispatcher), used by makegp_fftcov in os_example / numpyro_example but previously reached by no 1D-PSD row. - delay: exercises CompoundDelay (makedelay), used by the CW examples. Mapped NoiseMatrix2D_novar -> mh.NoiseMatrix2D in the factory. Two carry-overs to Phase 4, both documented: - fourier_variance_fixed: strict xfail. The NoiseMatrix2D_novar *kernel* maps, but an all-constant 2D GP prior is unsupported in the metamath likelihood (CompoundGP._build_mixed_logprior requires gp.index). strict=True flips it to a failure once Phase 4 supports the path. Not used by any notebook. - CompoundGlobalGP: factory fallthrough to matrix.py, no metamath port; to be relocated/ported in Phase 4 (the matrix.py deletion forces it). Integration check: cw_extsignal_example.ipynb model build + clogL/logL at its test point (5 psrs, CW + HD global + decenter) agrees matrix vs metamath to ~1e-9. Full suite green (45 passed, 1 xfailed parity + 29 matrix-path tests). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Phase 4 now makes the metamath path feature-complete without removing anything -- port the all-constant 2D GP prior path (flip the fourier_variance_fixed xfail) and relocate/port CompoundGlobalGP, keeping matrix.py/likelihood.py as the oracle. This is the checkpoint for external testing of the metamath path. Phase 5 is the deletion (matrix.py, likelihood.py, _kernel_switch, factory collapse, promote likelihood_metamath), done only after Phase 4 is signed off and externally exercised. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
metamath.CompoundGP built the per-GP coefficient log-prior (and returned Phi=None) for any mixed 1D/2D compound, which required gp.index -- absent on marginalized constant GPs. So an all-constant mixed compound (e.g. a constant timing GP + a constant fourier-variance GP, NoiseMatrix2D_novar) errored in the metamath path. Now the coefficient-log-prior / Phi=None branch is gated to the *vector* decentered path (commongp/globalgp with sampled coefficients); a mixed but marginalized compound builds a real combined dense Phi via _mixed_dense_Phi, promoting 1D diagonal blocks to dense and block-diagonalising (mirrors matrix.CompoundGP's mixed-const/var handling). Handles static and callable Phi blocks. Closes Phase-3 carry-over C: test_pulsar test_logL[fourier_variance_fixed] flips from strict xfail to passing; xfail marker removed. Parity suite green (46 passed, 0 xfailed). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The globalgp-as-list path combined global GPs via matrix.CompoundGlobalGP,
which reads matrix-only GP attributes (gp.Phi.params/.make_inv) and so couldn't
consume metamath GP priors -- metamath-native errored ('NoiseMatrix2D' has no
attribute 'params').
Add signals.CompoundGlobalGP: a backend-agnostic reimplementation that builds
the combined block-structured prior through the _kernels factory and
utils.GlobalVariableGP, reading only mode-neutral attributes (gp.Phi.getN,
gp.Phi.getN.params, gp.Phi_inv) -- the surface makeglobalgp_fourier populates in
either mode. Both likelihood.py and likelihood_metamath.py now call it; the
metamath path no longer touches matrix.CompoundGlobalGP (now dead code, removed
with matrix.py in Phase 5). Drop the now-unused _kernels import from
likelihood_metamath.
Parity route added: test_global.py global_compound (HD + monopole as a list),
covering logL and conditional. matrix vs metamath agree exactly.
Closes Phase-3 carry-over D. With 4a (carry-over C), Phase 4 is complete: no
xfails remain. Full suite green (48 parity + 29 matrix-path).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ture/metamatrix The Phase exit plan and coverage notes were committed under the Sphinx docs/components/ tree, which is the wrong home for development/architecture notes. Move the two tracked ones (exit_plan.md, phase3_coverage.md) to dev_architecture/metamatrix/ and update the metamatrix_architecture.md path reference in metamath.py's module docstring. (The remaining planning markdowns in that directory stay untracked for now.) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…g docs Add a README.md indexing the metamatrix development notes -- summarizing each markdown file and what it was used for -- and commit the three previously untracked planning docs (metamatrix_architecture, mixed_phi_compoundgp_plan, branch_diff_vs_upstream_metamatrix) so we have a record. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ts (route b)
Replace the dense two-perturbation outer-logdet form (S0_out / mid matmuls +
slogdet LU, both ~npsr^3) with the algebraically-equal increment built from the
Cholesky logdets already on hand:
d_ld_out = (lP + lS_cf) - (lPr + lSr)
where the full outer logdet is lP + lS = logdet(I + Phi_gw G~) (cf.
globalwoodbury_fused). lS_cf is the current outer capacitance logdet, previously
discarded as `_`; lP is live from Pinv; lPr/lSr fold to f64 constants. The
big-minus-big (each side ~1e4) is done in f64 (combine_f64), with lS_cf staying
f32 upstream (born from the f32 Cholesky).
This deletes the dominant dense work: refdelta now reuses the single current
outer factorization `cf` it already pays for `nu`. Also drops the now-dead
Phi/Phir/dPhi nodes; Pcov is unused and prunes away (left in the signature so
routing/tests don't ripple -- removable follow-up).
Verification (CPU): f64 parity refd==normal 2.3e-10/4.7e-10; f32 accuracy
7.1e-4/8.0e-4 (unchanged from ~5e-4 floor -- the budgeted ~1e-3 regression did
not materialize); FLOPs vs normal f64 1.38x (N=6) / 1.51x (N=12), down from ~74x;
14/14 refdelta tests pass.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…fdelta Route (b) (prev commit) removed the only consumer of the current-Phi_gw covariance leaf, so Pcov was dead and merely pruned away. Remove it from the kernel signature, the make_kernelproduct routing (was passing self.P.getN), and the two test call sites (were passing Pgw.getN), plus the docstring/comment. No behavioral change: harness f64 parity 2.3e-10/4.7e-10 and f32 accuracy 7.1e-4/8.0e-4 are identical to before; 14/14 refdelta tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
New executed notebook docs/advanced/single_precision.ipynb (in the Advanced Topics toctree) showing how to build the HD model three ways -- float64, blanket float32, and reference-delta float32 -- with the standard timing model (project=False), then comparing their precision against the float64 truth. Includes a short (3-equation) statement of how the reference splits the likelihood (Phi = Phi_ref + dPhi; logL = logL_ref + dlogL; the O(1) increment). Outputs are baked in (nb_execution_mode='off'): on 12 psr the errors vs truth are float64 1.4e-8, blanket f32 7.0e-2, reference-delta f32 1.1e-3. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The single-precision tutorial (and any markdown math) used $...$ / $$...$$, but MyST leaves dollar-math off by default, so it rendered as literal text. Enable myst_enable_extensions = ['dollarmath', 'amsmath']. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add a "Single precision (float32) on GPU" section to docs/metamatrix.md covering Half A (f64 final combine), Half B (reference+delta, incl. the fused HD twins and the outer-logdet increment speedup), and timing-model projection -- with the precision table (blanket f32 ~7e-2 vs reference+delta ~1e-3) and the ~1.5x-of-f64 speed result (down from ~74x). Links the new tutorial and the ADRs. Also corrects the now-stale "Not yet cross-checked" bullets: GPU has been exercised (GH200) for the single-precision HD benchmark, and that benchmark is the one matrix-vs-metamath / f64-vs-f32 timing comparison that exists. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Built models can now describe themselves without re-deriving structure from logL.params (which hides fixed white noise, basis-only GPs, and deterministic signals). - model.summary() / summary_frame(): per-signal table (kind, basis shape, free/fixed params with prior ranges) with independent show_free/show_fixed toggles and an optional signals[i] access column - model.tree(): composition tree (one node per signal, with live handles); tree(literal=True) mirrors the nested Woodbury kernel (model.N.N...), showing concat=True fusion and column slices - __repr__ on GP/signal objects and on the likelihoods (tree-based, with a many-pulsar guard) - factories tag components for introspection (gpname/index/orfnames, white-noise measurement metadata); likelihoods retain signals + concat Pure introspection, no kernel-math change; identical under matrix/metamath. Tested across all recipes in both backends (tests/metamatrix/test_summary.py). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
# Conflicts: # src/discovery/likelihood.py # src/discovery/params.py
…anch PR1 of feature_metamatrix_cleanup (§6, D14/D15/D16). - makesampler_nuts: the old body iterated `kwargs.items()` and tested each (key, value) tuple for membership in an argname list, so the test was always false and EVERY kwarg was silently discarded. Route NUTS/MCMC kwargs by name and raise TypeError on unknown ones; potential_fn is rejected because a positional model is always supplied. - run_nuts_with_checkpoints: always create outdir (the old guard only ran when outdir was not already a Path), attach to_df via the new _ensure_sampler_to_df, and correct the docstring's JSON/pickle claim. - utils.make_uind: return an empty index table for a zero-column exposure matrix instead of taking the max of an empty array (an ECORR selection that matches no TOAs). matrix.py imports this definition, so both routes are covered by the one fix. - matrix.WoodburyKernel_varP: add make_kernelsolve_simple; an all-variable single-pulsar `conditional` crashed with AttributeError. cho_factor/cho_solve are called directly rather than via the configurable matrix_factor alias because the returned factor is handed to callers under a lower-Cholesky contract that an LU config -- or cho_factor's lower=False default -- would silently violate. Tests cover both kernel routes; the varP conditional is checked against the certified metamath route, not merely asserted to run.
…aph_params PR2 of feature_metamatrix_cleanup (§3, §7.3, §4.2; D2/D19/D5). - PulsarLikelihood(concat=False) with >=2 variable GPs now raises unless marginalize_all_but_last=True. The chained construction overwrites `.index` per iteration, so every variable GP but the last was silently marginalized; the error names both the surviving block and the shadowed ones. Added to the legacy likelihood.py too, which has the same overwrite behavior and stays user-reachable until Phase 5. - PulsarLikelihood.clogL's fallback is ffunc-wrapped, so the property always returns a `(params) -> value` callable carrying `.params` regardless of branch. ffunc is a no-op for callables, so the primary branch is unchanged. - metamatrix.graph_params: pure-introspection union of `.params` over a graph, recursing into GraphLeafs and callable-attached `.graph` values. This is what PR4's clogl_form='auto' routing will consult; it must not fold or evaluate. The coefficient-frontend assertion in the shadow-guard tests is scoped to metamath: `clogL` over CHAINED variable GPs is unsupported on the legacy matrix route (WoodburyKernel_varNP reaches for `make_solve_1d` on its inner WoodburyKernel_varP, which does not define it). That gap predates this guard.
… column
PR3 of feature_metamatrix_cleanup (§7, §10.1; D17/D18/D20).
§7.2 assembly hygiene (D18): `conditional` / `clogL` / `logL` / `cglogL` each
rebuilt the kernel assembly and wrote it to `self.vsm` / `self.ys` as a side
effect, with variations -- so whichever property you touched first decided what
the others saw (`conditional` even short-circuited on an already-set
`self.vsm`). Two cached helpers replace that:
_marginal_assembly -- logL / cglogL / conditional; commongp only, P_ref
attached when reference= is set
_coefficient_assembly -- clogL; globalgp folded in, reference= NOT consulted
Assigning `reference` post-construction now invalidates the dependent caches
(the single-precision opt-in workflow does exactly that); without this the
refdelta leaves would be decided by whichever property was touched first --
the very staleness the assemblies exist to prevent.
§7.1 summary column (D17): a `coefficients` column reporting how each block
appears to the COEFFICIENT frontend, derived from the assembled kernel's
`.index` rather than from GP type. This is what makes concat shadowing legible:
a shadowed variable GP reports `marginalized`, not `sampled`. Values: sampled
(k) / marginalized / projected for GP rows, `kernel` for white noise,
`deterministic` for ExtSignals, `—` for residual and delay rows.
§7.4 (D20): the cglogL parity test is committed but SKIPS. Investigating its
missing coverage showed cglogL is broken on both routes, predating this work:
metamath lacks VectorWoodburyKernel.make_kernelterms (globalgp) and returns an
unwrapped graph (commongp-only); the matrix globalgp branch dies in the CG
stack on a JAX API change; and the CG extras (jaxopt + matfree) are undeclared
and absent. Repair is out of scope; docs/metamatrix.md now records exactly
what is broken instead of calling it merely "not cross-checked".
test_reference_wiring reads the assembly instead of the removed `self.vsm`.
…ting PR4 of feature_metamatrix_cleanup (§4; D4/D5/D6). - metamath.vectorresidualcomponent: the FtNmF-free twin of vectorgpcomponent. Same leaf contract, same 'logp'/'staged' outputs; different algebra (r_i = y_i - F_i c[i] - sum ExtSignals, then a single N-solve of r). With varying white noise nothing of shape (n_toa, k) is pushed through the solve. - VectorWoodburyKernel._coefficient_leaves extracts the shared prior/fold/means plumbing so cross and residual forms build identical coefficient leaves; make_residualproduct is the residual twin of make_kernelproduct_gpcomponent. - ArrayLikelihood(clogl_form="auto"|"cross"|"residual"), default "auto" = residual iff any per-pulsar noise solve has free parameters (pure introspection via metamatrix.graph_params, no folding). Resolved form exposed as clogl_form_resolved. reference= is deliberately not consulted by clogL. - _kernels.require_metamath guards the residual branch (D1). Tests (test_residual_clogl.py): cross==residual to 1e-16 (constant N, ExtSignal, means); residual vs a dense NumPy oracle to 1e-9 under varying EFAC; a deterministic topology gate proving no noise solve consumes a GP design matrix (plus a discrimination test showing the cross form DOES); the exact Gaussian marginalization identity; an informational k-scaling benchmark (residual ~3.5x faster at k=60, flat in k); reference= coexistence; the metamath-only guard.
…lete closure
PR5a of feature_metamatrix_cleanup (§5; D7-D13, D24). Also carries PR4's clogL
routing wiring in likelihood_metamath.py (the residual/cross dispatch), on which
PR5a composes transport reparams.
New module src/discovery/transport.py -- a boundary module (D8) that bakes
constants at construction and participates in graphs only through the
FuncLeaf/reparam contract; kernel methods stay graph-pure:
- TransportBlock + adapters gp_block / globalgp_curn_block. Every block
carries a MANDATORY conditioner precision -- the exact live prior precision
for a diagonal GP, or the explicitly-named inverse-marginal-variance CURN
view of a dense global prior (D11). No floors, no ridges, no defaults (D9).
- reference_noise(psr) / reference_noise_frozen(kernel, params0): one-method
frozen solves. Freezing requires COMPLETE params0 and a metamath kernel --
the direct fix for the closure's silent `params={}` anti-pattern.
- Transport (xi -> (q, ldJ)) and ArrayTransport (batched, equal-width, D24),
with true .params (D13), an eager validate() PD gate, and diagnostics().
Failure semantics are honest (§5.7): construction and validate raise;
runtime apply is NaN-propagating JAX. Cholesky is called directly so an LU
config cannot silently change the transform's meaning.
ArrayLikelihood(transport=...) accepts a prebuilt transport (eagerly
compatibility-checked); decenter=True becomes sugar that builds one from the
commongp (+ globalgp CURN) blocks with per-pulsar frozen-noise references. The
in-likelihood decenter closure is DELETED.
decenter=True on varying white noise now raises a diagnosed error (the frozen
reference is incomplete) instead of silently assuming constant N; callers build
the transport explicitly with reference_noise(psr) and pass transport=.
Tests (test_transport.py): closure parity is the deletion gate -- a standalone
replica of the deleted closure (§10.5's legacy_decenter_transform) certifies the
transport's clogL to rtol=1e-12 (~1e-15 in practice) and ldJ to 1e-12 across 20
draws on both decenter recipes under the default Cholesky config; the raw
coefficients match to ~1e-10 (numpy-baked-vs-jax-baked BLAS divergence on an
ill-conditioned solve, which washes out of the likelihood). Plus Jacobian,
reparam-.params, anti-ridge scale matching, adapter correctness, the full
validation-error matrix incl. a collinear rank-deficiency PD failure, inverse-map
geometry, parameter-dependent-basis rejection, and §10.2-§10.4 (free-EFAC,
legal-extreme PD, non-picklability, a checkpointed sampler run).
PR6 of feature_metamatrix_cleanup (§8, D21).
__init__.py sets _KERNELS = "metamath" and calls config(kernels="metamath") once
at import, so the factory mode and the likelihood-class bindings agree in
lock-step (before the flip they agreed only by the coincidence of both defaults
being "matrix"). Every user now gets the certified graph path by default;
config(kernels='matrix') remains selectable until Phase 5. Rollback is the
one-line default plus the import-time config call.
Test-suite mode-independence:
- tests/conftest.py gains an autouse fixture that snapshots and restores the
kernel mode around every test, so the suite is mode-independent regardless of
any per-test reset convention (teardowns that hardcoded config('matrix') were
correct only while matrix was the default).
- tests/metamatrix/_routes.py sets the matrix reference route's mode EXPLICITLY
and restores the module default; otherwise, after the flip, the "matrix"
reference would silently be built under metamath and every parity comparison
would become a self-comparison.
- test_matrix.py, test_measurement_noise.py, and the nanograv outlier model test
are pinned to matrix mode: they exercise matrix.py internals directly
(concrete kernel classes, solve_1d/solve_2d, likelihood.PulsarLikelihood +
matrix-only kernel APIs) that the factory only produces under matrix mode.
These stay valid until the Phase 5 deletion.
Full discovery suite green under the metamath default (391 passed, 1 skipped
before these three pins; all green after). docs/metamatrix.md records the new
default and the rollback. The remaining gate item -- one real-analysis sign-off
by the branch maintainer (§8 gate 4) -- is the maintainer's to give.
The clogL cross form (vectorgpcomponent) sums each ExtSignal independently and omits the cross-terms between distinct ExtSignals; the residual form (vectorresidualcomponent) includes them via a single N^-1 solve of the full residual and is the correct joint likelihood. Only affects >=2 non-orthogonal ExtSignals (0/1 is exact) and is usually swamped by an improper-prior dynamic range in practice, but clogl_form="auto" can route the same model to different values depending on whether white noise is fixed vs free. Pre-existing defect in the cross form; documented for a maintainer decision (fixing it touches the certified parity path).
Two things in the transport, committed together:
1. Transport.fingerprint() / ArrayTransport.fingerprint() — a stable structural
digest (sha256 over diagnostics() structure, no params) so a saved run's
transport can be reconciled without serializing an opaque closure. Consumed
by nltiming's dynamic run manifest.
2. Replace the vacuous anti-ridge test. The old test_scale_matching used the
same noise for the frozen reference N0 and the live noise solve, so the
diagnostics target H structurally cancels the precision against A (H/A is
forced to 1 for ANY precision), and the real-data curvature (~2.6e16) swamped
p/alpha under rtol=1e-8 — a floored/ignored conditioner passed. Root cause is
a spec inconsistency: §5.8 diagnostics builds H from the block's own
precision, so it cannot reveal a ridge, yet §5.11 test 4 assumed it could.
Replaced with two honest tests:
- test_anti_ridge_whitening_at_exact_precision_only: computes the transformed
posterior curvature M/A directly with O(1)-scale numbers and an INDEPENDENT
true prior precision, so exact conditioner -> 1.0 and 100x-wrong ->
(lam+p)/(lam+alpha) are far apart. Verified it FAILS under a ridge-floor
mutation that the old test passed.
- test_diagnostics_metric_reports_reference_noise_mismatch: honest coverage of
diagnostics(noise_solve=...) as a reference-noise (N0 vs N) mismatch check.
…PR5b) Enable external timing blocks and template-subtracted/soft-clamped centering for joint full-basis sampling, and expose PulsarLikelihood.sampled_gps.
…ient GP) A proper ConstantGP with identity coefficient covariance (ones(k)), for the z-prior W_m marginal block of the nonlinear-timing geometry stack: unlike makegp_improper (constant=1e40), it is a genuinely proper prior (c ~ N(0, I)) with log-determinant retained, no projection, and no column normalization (the caller passes the unnormalized W_m basis). Includes a 2-D shape guard and tests.
Retain the frozen reference-noise operator on Transport as _reference_noise and add reference_noise_quadratic(v) -> v^T N0^-1 v and reference_noise_standard_deviation() -> sqrt(diag(N0)), used by the nltiming geometry certifier to standardize a residual remainder without reconstructing white/ECORR noise from a notebook dict. _FrozenSolve now carries the exact diag(N0); reference_noise(psr) supplies toaerrs^2 and reference_noise_frozen computes it for NoiseMatrix / Sherman-Morrison ECORR NoiseMatrixSM references (best-effort: other kernels leave the diagonal unavailable rather than failing construction). Tests cover the diagonal, ECORR exposure term, and the transport helpers.
Add PowerLawParameterization plus make_powerlaw_pivot: sample log10_A_pivot at a pivot frequency instead of 1/yr, decorrelating amplitude and slope via the affine unit-Jacobian map log10_A_ref = log10_A_pivot + 0.5*gamma*log10(f_pivot/f_ref). sensitivity_weighted_pivot_frequency computes log(f_pivot) as the sensitivity- weighted mean of log(f_j) with weights tr(F_j^T N0^-1 F_j), and fourier_sensitivity_weights derives those weights from a Fourier basis and a frozen reference-noise operator. Public amplitude name is log10_A_pivot; reference_log10_amplitude decodes log10_A at 1/yr. No ambiguous name reuse. §11.1.
…oise Address review 8.2: exercise fourier_sensitivity_weights and the sensitivity- weighted pivot with discovery reference_noise_frozen (real _FrozenSolve.solve over a measurement-noise kernel) and a real Fourier basis, not just a toy diagonal solver.
Fold in transport fingerprinting, PR5b (array_block/ExtSignal centering/softclip), nltiming-facing helpers (standard-normal GP, pivot power-law, reference-noise quadratic/diagonal), and the ExtSignal cross-term bug note.
vhaasteren
marked this pull request as draft
July 20, 2026 07:41
Implements the discovery side of the marginalized-dynamic-decentering mode (feature_marginalized_dynamic_decentering.md §4): the eta-dependent transport of a single external (timing) block whitened against the LIVE marginalized covariance C(eta), for the case where the RN/DM/timing GP coefficients are analytically integrated out and only the small sampled timing block is sampled. - MarginalTransport + marginal_transport factory (D1-D9, D21, D22): apply/split/ as_reparam/validate/diagnostics/fingerprint mirror the joint Transport; per-eval products come from the single kernel.make_kernelsolve(y_t, W_s) graph (fold_constants caches the frozen-N0 cross-products; no duplicated TNT route). Single external block only; no clamps/softclip (D6); metamath-only (D5); conditioner precision 1.0 (unit-normal chart prior, D4). Schema "discovery-marginal-transport-v1". - D8b live-kernel geometry hooks: live_kernel_quadratic via one make_solve graph; live_kernel_standard_deviation via _live_kernel_diagonal, a recursive Woodbury- stack diagonal walk (NoiseMatrix / NoiseMatrixSM / WoodburyKernel; nested recurse on .N; WoodburyProjKernel raises per D22; else TypeError). The 1e40 improper term is float64-only (D22, commented). - T-D1..T-D6: dense-Woodbury oracle (stable with the 1e40 improper block), Jacobian == ldJ, params propagation, GLS centering sign, failure semantics (callable y / no make_kernelsolve / multi-key block / matrix mode / negative precision), and live-kernel hooks vs dense v^T C^-1 v and sqrt(diag C). Transport suite green (57 passed).
…diag dispatch (PR-M1 polish)
Bring in main's chromatic/solar models, unique-.params fix, Fourier-basis renames, and sampler/test hardening while keeping the metamatrix relocation of make_uind and measurement-noise constructors. Port add_equad/tnequad onto measurement_noise.makenoise_measurement_simple and adapt incoming matrix-route assertions for the metamath default.
Importing discovery always loaded metamatrix, which pulled in sympy even though Matrix was never used. Remove the dead dependency so the default test env can collect again.
vhaasteren
marked this pull request as ready for review
July 24, 2026 10:34
WoodburyKernel_novar stores a SciPy factorization and later feeds it to jax.scipy.linalg.cho_solve when residuals are parameter-dependent (e.g. deterministic delays). JAX treats `lower` as a static argument and rejects numpy bool arrays such as array(False), which some SciPy builds return.
Replace the branch-status metamatrix page with a user guide and a detailed developer guide, drop the root ExtSignal bug note (tracked as nanograv#137), and remove decision-ID / phase / ADR process markers from comments in the metamatrix PR surface. Keep dev_architecture notes intact for local design history.
Move lasting metamatrix/single-precision design docs under docs/design/, delete session handoffs and lab harnesses, and retarget code doclinks.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR lands the long-running metamatrix work: a graph-based kernel/likelihood path (
metamatrix/metamath/likelihood_metamath) certified against the legacymatrix.pyroute, then completed with residual-form coefficient likelihoods, a free-standing transport/decentering object, and a default flip to the graph backend. This work was mostly designed by @vallis and @meyers-academic .Intentionally left out of this PR for review: deletion of the legacy
matrix.py/likelihood.pypath. Both backends remain selectable viadiscovery.config(kernels='matrix'|'metamath')so review can add tests/goldens before the cutover.Core (metamatrix completion)
tests/metamatrix/)ArrayLikelihood.clogLwithclogl_form="auto"|"cross"|"residual"Transport/ArrayTransportreplace the in-likelihood decenter closure;decenter=Trueis sugar over transportconcat=False+ multiple variable GPs)coefficientscolumn; assembly hygiene / call-order invariancemakesampler_nutskwargs, empty-basismake_uind,WoodburyKernel_varP.make_kernelsolve_simple_KERNELS = "metamath"Also in this PR (post–default-flip)
array_block, ExtSignal-subtracted centering,softclip;PulsarLikelihood.sampled_gpsTransport.fingerprint()/ArrayTransport.fingerprint()for nltiming run manifestsreference_noise_quadratic,reference_noise_standard_deviationmakegp_standard_normal(proper (c \sim N(0,I)))make_powerlaw_pivot, sensitivity-weighted (f_{\rm pivot}`)BUG_extsignal_cross_terms.md(known open issue; see below)summary(), single-precision / refdelta opt-in work, developer docsExplicitly deferred
matrix.py/ legacylikelihood.py, collapseconfig(kernels=...), switch parity to goldens — after review + any extra tests we agree onKnown open issue
clogLcross form omits inter-ExtSignal cross-terms when ≥2 non-orthogonal ExtSignals are present; residual form is correct.clogl_form="auto"can therefore disagree with itself when white noise is fixed vs free. Workaround until fixed:clogl_form="residual"for multi-ExtSignal models. SeeBUG_extsignal_cross_terms.md.Test plan
pytest tests/metamatrix/ -q(parity, residual clogL, transport, summary, call-order, cglogL)pytest tests/test_samplers_numpyro.py tests/test_likelihood.py tests/test_matrix.py -q(bug-port coverage)pytest tests/test_powerlaw_pivot.py tests/test_standard_normal_gp.py -qdiscovery.config()default is"metamath"andkernels='matrix'still builds a modelNotes for reviewers
main)