NN build system: characterization tests + component registry (steps 1-2)

[max-dax]

First two steps of the NN build-system refactor, tracking hackathon/NN_Build_System_Design.md. Draft so the interfaces are visible early; steps 3+ (dict batches, per-architecture dim inference replacing autocomplete_model_kwargs, new config schema) will land here as they are ready.

Step 1 — characterization tests (tests/core/test_build_model.py)

Pins current behavior of the build path before it moves: build_model_from_kwargs dispatch, autocomplete_model_kwargs (plain + GNPE), end-to-end forward passes for all three model types, GNPE added_context embedding, the unconditional flow (models-as-priors path), SVD initial-weight seeding, old nsf+embedding schema mapping, save/rebuild round trip.

Documents (without fixing) a latent bug on main: get_theta_embedding_net reads frequencies from the wrong level of theta_embedding_kwargs, so any nonzero value crashes the FMPE/score build — the fmpe example only works because it uses frequencies: 0.

Step 2 — registry + NeuralDistribution rename

• dingo/core/registry.py: generic Registry with four-step name resolution — registered short name → installed entry points (group dingo.architectures) → dotted import path → /path/to/file.py:Class. Third-party architectures no longer require editing dingo source (pip plugin or plain file).
• The three model types register themselves; build_model_from_kwargs dispatches via the registry instead of the hardcoded models_dict. A file-path plugin distribution works end-to-end (tested).
• Name lookup is case-sensitive; the historical case-insensitivity for built-in type names moved into update_model_config (the back-compat boundary), so old checkpoints still load.
• BasePosteriorModel → NeuralDistribution: "posterior" was too narrow — the same class models priors/proposals (e.g. the unconditional flow). A lasting alias BasePosteriorModel keeps existing imports and other branches working.

Compatibility

Old checkpoints and settings load unchanged (pinned by the step-1 tests, which pass unmodified against step 2). tests/core: 67 passed.

Next (step 3) — for reviewers to weigh in on

The metadata contract this build system will write — inference_parameters / context_parameters / standardization — is exactly what FlowFactor.from_model on sampler-revamp consumes; the concrete schema will be proposed here and needs sign-off from the sampler group.

🤖 Generated with Claude Code

[max-dax]

Added the concrete metadata-contract proposal (the "(1)" item from the PR description): commit 4456a03 adds inference_parameters / context_parameters / standardization as properties on NeuralDistribution — the supported interface for samplers, working on old checkpoints unchanged.

Full proposal for the sampler group (incl. the step-3 dict-batch call-signature migration plan and the exact factors.py before/after): hackathon/NN_Sampler_Interface.md. Group A sign-off on that doc is the remaining coordination point.

[max-dax]

Step 3 landed: dict batches, registry-built embeddings, autocomplete_model_kwargs deleted

Three new commits (f917def2, d54a1d33, d2f9315a) complete step 3 of the build-system design (§6). All 321 tests in tests/core + tests/gw pass.

Note on the rebase: the branch was rebased onto the updated hackathon-1. #387 had added a smaller tests/core/test_build_model.py; the conflict was resolved in favor of the more comprehensive version from this branch. One #387 test asserting case-insensitive dispatch was dropped deliberately: per the step-2 decision, registry lookups are case-sensitive, and historical case-tolerance lives only in update_model_config.

1. f917def2 — dict batches (design §4.4, interface doc §2)

• The training transform chain ends in SelectKeys (dict-preserving) instead of UnpackDict (positional flattening); loaders yield {"inference_parameters": ..., "waveform": ..., "context_parameters": ...}.
• NeuralDistribution.sample / sample_and_log_prob / log_prob / loss take a single context: dict | None instead of positional *context — exactly the signature proposed to Group A in NN_Sampler_Interface.md §2.
• FlowWrapper and ContinuousFlow route context[k] for k in context_keys into the embedding network once, at the model boundary; tensors stay positional inside modules. A missing context key raises a ValueError naming it — no more silent dependence on the loader's key order.
• This kills the GNPE-proxy detection via try/except IndexError and the overloaded data_sample[2] slot (which the transformer branch reused for position tensors — the two were mutually exclusive by accident).
• Call sites updated: core/samplers.py (plain + GNPE), SampleDataset for unconditional training.

2. d54a1d33 — new schema + registries; autocomplete_model_kwargs deleted (design §4.3/4.5/4.7/4.8)

New model-settings schema, each component a registry type plus a flat kwargs dict:

model:
  distribution:    {type: normalizing_flow, kwargs: {num_flow_steps: 30, ...}}
  embedding_net:   {type: dense_svd, kwargs: {output_dim: 128, ...}}
  # context_merger: {type: concat}   # added automatically when the data has context_parameters

• Embedding contract (core/nn/enets.py module docstring): input_keys, output_dim, complete_settings(settings, sample_batch). First registered components: dense_svd (LinearProjectionRB + DenseResidualNet) and the concat context merger. Any embedding now composes with any distribution type — transformer×FMPE will need zero extra wiring.
• complete_model_settings replaces autocomplete_model_kwargs: each architecture infers its own input dims from the named sample batch; generic theta_dim/context_dim are computed once; completed settings are saved in the checkpoint so loading never needs a data sample. Dimensions in user YAML are an error; a context_merger with no context parameters in the data is an error (no silent drops).
• initialize_network loses its embedding_kwargs-presence branches; NeuralDistribution.build_embedding_net assembles embedding + merger generically. The V_rb_list deepcopy hack is gone — initial weights are extra constructor kwargs, never in saved settings.
• save_model enforces the parameter contract promised to samplers (interface doc §1): standardization must cover inference ∪ context parameters, for built-ins and plugins alike.
• Back-compat: update_model_config is the single boundary and maps all old schemas forward (nsf+embedding, posterior_model_type/posterior_kwargs/embedding_kwargs, added_context → concat merger, input_dim → theta_dim). Old checkpoints build state-dict-identical networks — pinned by test_old_schema_builds_state_dict_compatible_network and an actual old-format .pt load round trip (test_load_old_schema_checkpoint_file).
• Deleted: the never-wired embedding_net_builder parameter, create_nsf_wrapped, create_nsf_with_rb_projection_embedding_net. Example configs, pipe defaults, and the architecture docs notebook use the new schema (old YAML still works through the shim).

3. d2f9315a — chained-NPE training-side conditioning (design §4.5)

• data.conditioning_parameters feed the same context_parameters list as GNPE proxies — two sources, one declared contract.
• ContextMergerMLP ported as the registered mlp merger: context mixed in through a learned MLP so context_dim doesn't grow; its dims are inferred at completion (no hand-written input_dim), output_dim defaults to the embedding's.
• data.time_alignment: true trains on time-aligned data via ProjectOntoDetectors(apply_time_shift=False), with the validation guard (requires ra/dec/geocent_time as conditioning, forbids them as inference targets, incompatible with gnpe_time_shifts). Tests ported from the chained branch. The chained inference side is a factorized-sampler chain (Group A's domain) and is not built here.

Deliberately deferred

• FlowFactor/GNPEFlowFactor call sites (interface doc §2, impacts 1–3): factors.py is not in hackathon-1 yet, so those edits can't land in this PR. Nothing on sampler-revamp breaks in the meantime; we'll do that migration as promised once the branches meet.
• Step 4 (architecture-owned SVD init — the initial_weights threading and the trainer's svd gate still exist, marked with a TODO) and steps 5–6 (transformer merge) follow.

Review pointers

• The schema mapping: core/utils/backward_compatibility.py::update_model_config
• The autocomplete replacement: core/posterior_models/build_model.py::complete_model_settings
• The embedding contract + registered components: core/nn/enets.py
• Checkpoint compatibility tests: tests/core/test_build_model.py (old-schema section)

🤖 Generated with Claude Code