Repeated-axis selectors and sparse 2D parameter tables

[jc-macdonald]

Motivation

Sparse rate matrices indexed by the same axis twice are a recurring modelling pattern that op_system currently can't express compactly. The motivating case is the SMH Round 19 immune-ladder vaccination uptake (flepimop2-demo/configs/SMH_R19_op_system.yml, transitions section 5):

- from: X[age, vax=unvaccinated, loc, imm=X0]
  to:   X[age, vax=vaccinated,   loc, imm=X3]
  rate: eta_X0toX3 * nu[age]
- from: X[age, vax=unvaccinated, loc, imm=X0]
  to:   X[age, vax=vaccinated,   loc, imm=X4]
  rate: eta_X0toX4 * nu[age]
# ... 15 more

This expands to:

• 17 transition records, one per nonzero (imm_from, imm_to) pair on the imm ladder.
• 17 sibling parameter declarations eta_X{i}toX{j} whose joint structure (a sparse 11×11 imm × imm matrix) is invisible to op_system, the parameter plugin, and inference.

Other models will hit the same shape: contact matrices C[age, age], mobility/mixing kernels M[loc, loc], mutation/cross-immunity kernels K[strain, strain].

Proposed feature

Three composable additions:

1. Independent placeholders decoupled from axis names ($-prefix)

Today, X[age, vax] introduces wildcards whose names are the axis names, and rate-expression placeholders bind by axis name. This couples the placeholder identity to the axis identity and forbids using the same axis twice.

Introduce $-prefixed placeholders that are independent identifiers bound to a coord on a specified axis:

- from: X[age, vax=unvaccinated, loc, imm=$i]
  to:   X[age, vax=vaccinated,   loc, imm=$j]
  rate: eta[imm=$i, imm=$j] * nu[age]

$i and $j are two independent free variables, each ranging over imm.coords. Existing axis-name wildcards (X[age, vax]) continue to work as a shorthand equivalent to X[age=$age, vax=$vax].

2. Repeated-axis support in selectors

Allow the same axis name to appear more than once on a selector, provided each occurrence is bound to a distinct $-placeholder (or pinned to a literal coord). parse_selector currently rejects X[age, age=a0] outright; relax that to allow eta[imm=$i, imm=$j] while still rejecting unbound duplicates like X[imm, imm].

This requires auditing every consumer that builds an axis → coord mapping (gather, apply_along, initial_state expansion, operator framework) to ensure repeated axes don't conflate roles.

3. New parameter module: sparse_table

A parameter module whose values are indexed by a tuple of axis-coord pairs:

parameter:
  eta:
    module: sparse_table
    indices: [imm, imm]              # axis appears twice
    values:
      [X0, X3]: 0.0
      [X0, X4]: 0.0
      [X1, X4]: 0.0
      # ... 14 more entries

The support set (declared keys) is structural metadata available at normalization time. Transitions whose rate references eta[imm=$i, imm=$j] iterate exactly the support set rather than the full Cartesian product:

transitions:
  - from: X[age, vax=unvaccinated, loc, imm=$i]
    to:   X[age, vax=vaccinated,   loc, imm=$j]
    rate: eta[imm=$i, imm=$j] * nu[age]

The expander recognizes that eta is a sparse table over ($i, $j) and emits one transition per declared (i, j) key. Equivalently the user could write where: eta[imm=$i, imm=$j] != 0 explicitly; sparse-table presence in the rate would be a sugar for that.

A companion dense_table module (full Cartesian product) is the obvious complement for non-sparse matrices like contact matrices.

4. Plugin compatibility

The parameter plugin (when it lands) needs to populate sparse_table modules from CSV / parquet readins; the legacy R19 layout is one file per (i, j) pair, which maps naturally to one sparse-table entry each. Inference treats a single sparse_table as one structured target with priors over the support set rather than 17 independent scalars.

Scope and risk

Touches:

• src/op_system/_templates.py — selector grammar ($-placeholders, repeated-axis support).
• src/op_system/_normalize.py — _collect_transition_wildcard_axes, _render_transition_combo, _expand_single_transition, _expand_initial_state_templates, gather/apply_along consumers.
• New parameter module(s) and registry hook for sparse_table / dense_table.
• Documentation: README + spec reference updates.
• Migration of flepimop2-demo/configs/SMH_R19_op_system.yml (17 transitions + 17 eta_* scalars → 1 template + 1 sparse table).

Risks / open questions:

• Grammar lock-in. The $-prefix syntax is a one-shot decision; worth validating against 2-3 example configs (R19 ladder, contact matrix, mutation kernel) before shipping.
• Repeated-axis ripples. Every consumer of axis-coord assignments must be audited. Some may need to switch from dict[axis, coord] to list[(axis, coord)] internally.
• Sparse iteration semantics. Support set is structural (fixed at normalization), not runtime. Document this explicitly so users don't expect a rate's runtime zero values to prune transitions.
• Inference design coupling. Treating a sparse table as one target with structured priors is a real win for inference, but the precise interface should land with or after the parameter-inference work, not before.

Recommended sequencing

1. Land the design here (grammar + module shape) on paper. Don't write code yet.
2. Wait for the parameter-inference work to stabilize so we can co-design the inference interface.
3. Implement in a dedicated PR; migrate R19 in the same PR as a working example.
4. Follow up with contact-matrix / mixing-kernel example configs.

Alternatives considered

• Status quo (Option A): Leave R19 with 17 explicit records. Honest representation of sparse irregular data, but doesn't generalize and leaves eta invisible to inference as a joint object.
• for_each: template list (Option B): Compress the transitions block to one template + a row list, but the parameter block still has 17 disconnected scalars. Smaller change but doesn't address the inference-interface or matrix-parameter cases.

Option C (this issue) is the right long-term answer specifically because contact matrices and mutation kernels make repeated-axis matrix parameters near-certain in future models.

[jc-macdonald]

Design proposal: repeated-axis routing via axis:alias, lowered at the template level

This is step 1 of the sequencing above: a design on paper, no code yet. It covers the grammar, semantics, parameter shapes, and lowering, and it includes measurements, because the lowering decides whether this also fixes #206.

Summary of decisions proposed

1. Grammar: reuse the existing axis:alias binding (already used as contact_kernel[age, age:ap] inside apply_along) instead of adding $-placeholders. One binding syntax, no new sigil.
2. Semantics: a transition whose from/to bind the same axis under two aliases is a routing transition. It moves mass along that axis with a matrix-valued rate.
3. Lowering: routing transitions and aliased apply_along reductions are lowered once per template to a tensor contraction along the routed axis. They are never expanded per cell or per matrix entry.
4. Parameters: matrix parameters are ordinary shaped parameters over [axis, axis] (optionally with time first). A sparse table is dense storage plus a declared support. The support is metadata for validation and inference; it is not used to expand transitions.

Evidence: expansion is the bottleneck, not vectorization

The routing inflow[j] = sum_i K[i, j] X[i] is already expressible today in an expr spec, and it is exact (zero error against einsum, mass conserved to 1e-12):

equations:
  X[age, vax, loc, imm]: >-
    apply_along(K[imm:i, imm] * X[age, vax, loc, imm:i], imm=i, kernel=sum)
    - X[age, vax, loc, imm] * apply_along(K[imm, imm:j], imm=j, kernel=sum)

But normalization expands it per cell with the reduction unrolled. Timings for age 15 × vax 2 × loc 3 × imm n, NumPy backend:

n (imm)	cells	normalize	compile	chars per cell equation
11	990	0.38 s	0.03 s	554
21	1,890	2.17 s	0.07 s	1,054
31	2,790	4.40 s	0.12 s	1,554
41	3,690	6.33 s	0.20 s	2,054
61	5,490	14.65 s	0.41 s	3,054

Vectorization itself takes milliseconds; it re-discovers the structure that per-cell expansion flattened. The same cost appears with coordinate-pinned transitions (#206: an Erlang-3 immune ladder with 31 coordinates took 261 s against 15 s for the production spec). So the win comes from not expanding per cell, which the "one transition per support entry" expansion in the original proposal would not deliver.

1. Grammar

A second role of an axis is written axis:alias, in selectors and in parameter subscripts:

transitions:
  # Waning along the immune axis with a generator G (off-diagonal rates).
  - from: X[age, risk, vax, loc, imm:i]
    to:   X[age, risk, vax, loc, imm:j]
    rate: waning_rate * G[imm:i, imm:j]

  # R19 vaccination uptake: status flip plus routing along imm.
  - from: X[age, risk, vax=unvaccinated, loc, imm:i]
    to:   X[age, risk, vax=vaccinated,   loc, imm:j]
    rate: eta_scale * nu[time, age, risk, loc] * eta[time, imm:i, imm:j]

Rules:

• An alias is local to one transition (or one apply_along call).
• The same axis may appear once in a selector. from binds one alias and to binds another; that pair defines the routed axis. At most one routed axis per transition.
• Unaliased axes (age, loc, ...) keep today's meaning: the same coordinate on both sides, or a pinned coordinate. This is exactly today's shorthand, unchanged.
• A parameter subscript may use each alias of the routed axis at most once: G[imm:i, imm:j], and eta[time, imm:i, imm:j] for a time-varying matrix.
• Reject unbound duplicates (X[imm, imm]), aliases that appear only on one side, and aliases in from/to without a matching rate subscript. Each error names the transition and the alias.

$i stays unused. If a future case needs placeholders that are independent of any axis, they can be added then without conflicting with this.

2. Semantics

For a routing transition with per-capita rate r(...) * K[i, j] from source slice S to target slice T along axis a:

dX_S[..., i] -= r * X_S[..., i] * sum_j K[i, j]
dX_T[..., j] += r * sum_i K[i, j] * X_S[..., i]

Mass is conserved by construction. Two proposed rules:

• Self-routing is dropped. When S == T (no other axis changes), the diagonal K[i, i] would move a cell into itself. It contributes nothing and is excluded, so a generator's diagonal is ignored rather than double-counted. When S != T (for example the vax flip), the diagonal is a real move: rung 9 unvaccinated to rung 9 vaccinated.
• Rates are per-capita and nonnegative, as for ordinary transitions. K holds rates or routing weights, not a generator with a negative diagonal.

The routed axis must not be a factorize_axes block axis. It can be any other categorical or ordinal axis.

3. Parameters

• Dense matrices are shaped parameters over [imm, imm] or [time, imm, imm]. Parameter modules can already return repeated-axis shapes (COVID19_USA_noise_experiments' parametric plugin evaluates expressions over [age, age] with positional names age_0, age_1), and time-varying shaped parameters already go through op_system's time interpolation.
• sparse_table (provider module) stores a dense array plus support: [[X0, X3], [X0, X4], ...]. Off-support entries are structurally zero. op_system validates at normalization that referenced matrices declare [axis, axis] shapes. The support is exposed as metadata (for example on requested_parameters, or via Public spec validation API and CLI with vectorizer diagnostics and cost summary #205's validate_spec) so inference can place priors on the support entries only. The R19 layout (one file per (i, j)) maps onto one table with one column per support pair and a time index.
• This keeps the inference-interface decision separate, as Repeated-axis selectors and sparse 2D parameter tables #88 recommends: op_system only needs the shape and the support.

4. Lowering

• Normalization: keep a routing transition as a template-level routing term; don't expand it per cell. Cells untouched by routing keep today's per-cell expansion, so existing specs are unaffected.
• IR: add a contraction node, route(K, X_S, axis=a, source_alias=i, target_alias=j), lowered for the vectorized path to xp.tensordot/matmul along a, and to the equivalent per-cell sum only on the scalar fallback. Aliased apply_along(K[a:i, a] * X[..., a:i], a=i, kernel=sum) lowers to the same node, so existing contact-kernel aliases get the speedup too.
• Time-varying matrices use the existing time-interpolation wrapper; the node receives K(t) with shape [a, a].
• Block axes: with factorize_axes, the contraction runs on the block-stripped template unchanged, because the routed axis is required to differ from the block axis.

Consumers to audit, per the original risk list: _templates.parse_selector; transition expansion in _normalize (_collect_transition_wildcard_axes, _render_transition_combo, _expand_single_transition); _ir_expand/_ir_lower for the new node; _vectorize plan construction; initial-state expansion (unchanged, since aliases aren't allowed there); and the operator framework (unchanged).

5. Validation against the three example cases

Case	Today	With this design
R19 ladder vaccination	17 pinned transitions + 17 eta_* scalars; one pinned transition per pair on finer ladders	1 routing transition + 1 sparse_table eta[time, imm, imm]
Contact matrix	apply_along(C[age, age:ap] * ...), expanded per cell	same syntax, lowered to one contraction
Mutation / cross-immunity K[strain, strain]	not expressible compactly	from: X[strain:s] to: X[strain:t] rate: mu * K[strain:s, strain:t]
Immune waning generator (#206)	advection operator, or an engine operator	from: X[imm:i] to: X[imm:j] rate: v * G[imm:i, imm:j], compiled by op_system

With this, the axis_kernel operator from #208 remains useful for engines that apply routing outside the compiled right-hand side, but transitions-kind specs no longer need it.

6. Open questions for maintainers

1. Is dropping self-routing when S == T the right default, or should a nonzero diagonal there be an error?
2. Should a routing transition allow more than one routed axis (for example imm and strain at once)? Proposal: no in v1; compose two transitions instead.
3. Where should the sparse support live: a provider-module attribute read by op_system at normalization, or a spec-level declaration (parameters: {eta: {support: ...}})? The first keeps op_system parameter-agnostic.
4. Should the scalar (non-vectorized) path support routing at all, or should routing require the vectorized path, as axis-indexed specs already do (feat: error when axis-indexed spec fails vectorization instead of silent scalar fallback #160)?

7. Proposed PR sequence

1. Lowering first, no new grammar: template-level contraction for aliased apply_along. This speeds up existing contact kernels immediately and is independently testable against the current per-cell results.
2. Routing transitions: the axis:alias grammar in from/to, desugared onto the node from step 1, with validation errors and validate_spec reporting.
3. Provider side: sparse_table/dense_table modules in flepimop2-op_system and the R19 migration as the worked example.
4. Inference interface for tables, after the parameter-inference work settles.

Step 1 carries no grammar lock-in, so it could start before questions 1–4 are settled.

[jc-macdonald]

Decisions on the open questions

1. Diagonal when nothing else changes: ignored, not an error. A move from a compartment into itself can't change the dynamics, and op_system sees matrix values only at runtime. Ignoring it also lets a generator with a negative diagonal work as-is, because only off-diagonal rates are used. Exception: if a sparse table explicitly declares (i, i) as a nonzero entry for such a transition, that is reported as an error, since it's known up front and almost certainly a mistake. When another axis also changes (for example the vax flip), the diagonal is a real move and is kept.
2. One routed axis per transition in v1. Multiple routed axes are rejected with a clear error. Models that need joint routing can combine the axes into one; this can be added later without breaking specs.
3. Sparse support lives with the parameter (the flepimop2-op_system table module), not in the op_system spec. op_system stays parameter-agnostic and exposes the compiled routing structure (parameter, axis, whether the diagonal is ignored). The provider or flepimop2 check validates declared entries against it, and inference reads them from the parameter.
4. Routing requires the vectorized path, consistent with feat: error when axis-indexed spec fails vectorization instead of silent scalar fallback #160. The Add validated axis_kernel operators with shared reference semantics #208 reference functions serve as the test oracle.

Proceeding with step 1 (template-level contraction for aliased apply_along), which adds no grammar.