Add axis-aware field specifications for parameters and initial conditions

[jc-macdonald]

Summary

Add a unified fields block to op_system specs for defining quantities over declared axes and materializing them on grid coordinates. Targets both model parameters and initial conditions.

Problem

Current configs are awkward when a quantity is function-valued over an axis grid. Users must either enumerate per-bin scalar parameters or rely on external preprocessing. This inflates config size, inference dimension, and maintenance burden.

Proposal

Introduce an optional top-level fields block. Each entry contains:

Field	Description
name	Unique identifier
target	parameter or initial_state
binds	Templated symbol to populate, e.g. beta[age,loc], S[age,vax,loc]
axes	Axes this field is defined over
expr	Expression string evaluated on axis coordinates
params	Named coefficients mapped to model symbols (templating allowed)
normalize	Optional normalization mode and target

Draft YAML

axes:
  age:
    type: continuous
    coords: [0, 1, 2, ..., 100]   # or linspace(0, 100, 101)
  vax:
    type: categorical
    values: [u, v]
  loc:
    type: categorical
    values: [A, B]

fields:
  # Age-varying transmission rate: linear in age, no normalization
  - name: beta_age
    target: parameter
    binds: beta[age]
    axes: [age]
    expr: "m * age + b"
    params:
      m: beta_m
      b: beta_b
    normalize:
      mode: none

  # Initial susceptible profile: Gaussian over age, normalized to a per-(vax,loc) total
  - name: s0_age_profile
    target: initial_state
    binds: S[age,vax,loc]
    axes: [age]
    expr: "scale * gaussian(age, mu, sigma)"
    params:
      mu: s0_mu[vax,loc]
      sigma: s0_sigma[vax,loc]
      scale: s0_scale[vax,loc]
    normalize:
      mode: integral         # trapezoidal, uses continuous-axis deltas
      target: s0_total[vax,loc]
      nonnegative: true

  # Lognormal exposure kernel over age
  - name: exposure_age
    target: parameter
    binds: chi[age]
    axes: [age]
    expr: "lognormal(age, mu_log, sigma_log)"
    params:
      mu_log: chi_mu_log
      sigma_log: chi_sigma_log
    normalize:
      mode: sum

Expression design

No kind field. All shapes are written as expressions. The expression evaluator whitelist is extended with named convenience functions so users do not have to spell out common shapes manually:

• gaussian(x, mu, sigma)
• lognormal(x, mu_log, sigma_log)
• gamma_pdf(x, k, theta)
• beta_pdf(x, alpha, beta_)
• weibull(x, k, lam)
• All existing entries: exp, log, sqrt, maximum, minimum, clip, where, trig, etc.

These are resolved to sandboxed, JAX-compatible implementations during compilation, not arbitrary dynamic eval.

Normalization modes

• none
• sum
• integral — continuous axes only, using trapezoidal axis deltas

Constraints

1. Backward compatible with existing initial_state mapping behavior and all existing template expansion.
2. integral normalization rejects categorical-only normalization axes.
3. Duplicate concrete bind targets after expansion are validation errors.
4. Expression evaluation is sandboxed via AST whitelist, not arbitrary dynamic eval.
5. Convenience functions must be JAX-compatible to preserve differentiability.
6. Likelihood-family configuration remains out of scope for op_system.

Out of scope

• Observation likelihood specification.
• Mixture-model convenience syntax.

Estimated size

• Files touched: 7–9
• Net LOC: ~350–550

Implementation checklist

• Add fields schema and validator in spec normalization
• Extend expression evaluator whitelist with convenience functions
• Add bind expansion and concrete target assignment
• Add normalization logic for none / sum / integral
• Wire materialized values into adapter for parameter and initial_state targets
• Tests: validation, expansion, materialization math, normalization, backward compat
• Docs and usage examples

Suggested follow-up issues

• Additional convenience functions on request.
• Mixture-model convenience syntax.
• Richer normalization constraints and positivity transforms.