plan: EP analytic updates — four work packages (exact PriorFactors, projection contract, linear-Gaussian, conjugate table)

[Jammy2211]

Overview

Implementation plan for the Phase 6 scoping verdict (#1337) of the EP framework review. Plan only — nothing here is implemented. Written to be executable by a fresh session with no prior context (read #1337's verdict and autofit/graphical/README.md from PR #1334 first). Umbrella: PyAutoMind/research/graphical_ep/ep_framework_review.md.

Background in one paragraph: EP's exact-update machinery exists (ExactFactorFit; has_exact_projection / calc_exact_update / calc_exact_projection hooks on Factor, populated from the wrapped callable by Factor._set_factor; auto-selection in EPOptimiser.from_meanfield) and regression tests prove user-supplied analytic projections work (test_autofit/graphical/regression/test_exact.py::ProbitModel, test_linear_regression.py::LinearModel). But the declarative path (FactorGraphModel.optimise) never uses any of it: PriorFactor wraps the bound method prior.factor which strips the hooks, and AbstractDeclarativeFactor._make_ep_optimiser constructs EPOptimiser(...) directly instead of via from_meanfield. Every declarative EP fit therefore runs a numerical optimiser once per free parameter per cycle for closed-form conjugate updates (in-code TODO at autofit/graphical/declarative/factor/prior.py:22).

---

WP1 — Exact PriorFactor updates + the analytic-projection contract (≈2–3 days)

Phase 6 candidates 1+2 combined: they share tests and documentation.

Step 1: forward the exact hooks through PriorFactor

autofit/graphical/declarative/factor/prior.py. Two options; recommended: (a).

• (a) Override the three methods on PriorFactor (it subclasses FactorKW, whose has_exact_projection/calc_exact_projection/calc_exact_update in factor_graphs/abstract.py:386-415 consult self._has_exact_projection etc.): implement has_exact_projection(mean_field) → delegate to self.prior.message.has_exact_projection(mean_field[self.prior]) and calc_exact_update(mean_field) → wrap self.prior.message.calc_exact_update(...) results back into a MeanField keyed by self.prior (mirror the nested_filter packaging in abstract.py:404-415). This avoids touching the callable/call semantics.
• (b) Pass the message object as the factor callable (the AbstractMessage.as_factor pattern) — riskier: changes PriorFactor.__call__ semantics (message.factor = per-value logpdf vs message.__call__ = summed logpdf) and the class doubles as an Analysis for classic searches; only pursue if (a) hits packaging friction.

Caution — transformed messages: for UniformPrior/LogUniformPrior-backed priors the message is a TransformedMessage. has_exact_projection must verify the cavity shares the same transform stack before claiming exactness (same validity condition documented on TransformedMessage.kl in PR #1334); base-message natural-parameter multiplication is then exact. If stack comparison is awkward, restrict WP1 exactness to untransformed messages first (NormalMessage and subclasses) and record the extension as follow-up — that already covers GaussianPrior, the dominant prior-passing family.

Step 2: auto-select exact optimisers declaratively

autofit/graphical/declarative/abstract.py:_make_ep_optimiser (line ~149): switch to EPOptimiser.from_meanfield(model_approx, default_optimiser=optimiser, factor_optimisers={explicit ones}, ep_history=..., paths=...) — from_meanfield already implements the has-exact→ExactFactorFit selection (expectation_propagation/optimiser.py:228-259). optimise() (line ~167) already builds self.mean_field_approximation(); construct it once and pass to both. Add an escape hatch kwarg exact_prior_factors: bool = True on optimise() threading through to disable auto-selection (safety valve for release notes).

Step 3: document the contract (candidate 2)

• autofit/graphical/README.md §3.2/“Exact path”: state the duck-typed protocol — a factor callable may provide has_exact_projection(*cavity_dists) -> bool and calc_exact_projection(*cavity_dists) -> Tuple[Message, ...] returning the projected tilted posterior including log_norm (cite ProbitModel.calc_exact_projection as the reference implementation); note the distinction from calc_exact_update (returns the factor's own message, multiplied into the cavity by ExactFactorFit).
• Optionally add a typing.Protocol (ExactlyProjectable) in expectation_propagation/factor_optimiser.py for reference/type-checking; no runtime enforcement.
• Merge-order note: PR docs: formal Bayesian specification of autofit.graphical (EP review phase 2) #1334 owns README.md — land WP1 after it merges, or rebase.

Step 4: tests + validation

• Unit (test_autofit/graphical/functionality/): PriorFactor(af.GaussianPrior(...)).has_exact_projection(...) true with Normal cavity; false (or true, if Step 1 extension done) for UniformPrior with mismatched stack.
• Declarative integration: build the shared-centre toy (pattern: autofit_workspace/scripts/features/expectation_propagation.py), assert every PriorFactor received an ExactFactorFit in opt.factor_optimisers, and posterior means/sigmas match a exact_prior_factors=False baseline within tight tolerance.
• Timing evidence (no assert): compare ep_history.csv (Phase 4 tooling, feat: EP review phase 4 — diagnostics & monitoring for EP fits #1335) wall-clock per cycle with/without; post numbers to this issue.
• F7 interplay: exact updates set log_norm via message multiplication — do not build evidence assertions until research: EP review phase 1 — statistics audit of autofit/graphical #1332 F7 (log_norm bookkeeping) is fixed; assert on means/sigmas only.

---

WP2 — First-class linear-Gaussian factor (≈1 week)

Phase 6 candidate 3. Depends on WP1 (contract documented).

• New autofit/graphical/factors/linear_gaussian.py: LinearGaussianFactor(x, A, b, sigma) representing log f = −½‖(A·x + b − y)/σ‖² (y observed; allow y to be a second variable for factor-to-factor comparison messages). Implements: __call__ (log density), has_exact_projection (all cavities Gaussian/NormalMessage), calc_exact_projection (standard linear-Gaussian conditioning: posterior precision Λ = Λ_cav + AᵀΣ⁻¹A, mean via solve; return per-variable marginals as NormalMessages with the correct joint log_norm), factor_jacobian fallback so non-Gaussian cavities still work numerically. Plate support for batched rows.
• Reference implementation to generalise: test_autofit/graphical/regression/test_linear_regression.py::LinearModel (verified working under EPOptimiser.from_meanfield in test_exact_updates).
• Validation: (i) reproduce test_exact_updates posteriors with the new class; (ii) an IC50-shaped integration script in autofit_workspace_test/scripts/graphical/ — the IC50 global factor (-0.5‖(latents @ coef_matrix + coef_mean − μ)/σ‖², see ic50_workspace/scripts/util.py) is exactly this form and is the highest-dimensional per-cycle fit on that graph; assert EP posterior parity vs the sampler-fitted baseline, report wall-clock delta.
• Export as af.LinearGaussianFactor (or autofit.graphical only — decide at review; lean autofit.graphical to keep af. surface small).

---

WP3 — Wider conjugate table (Gamma–Poisson, Beta–Bernoulli/Binomial) (≈1 week, GATED)

Phase 6 candidate 4. Do not start until the #1331 fix batch lands (Gamma from_mode invariant — decision 3) and #1332 F2 (Gamma/Beta KL direction) is fixed: both bugs sit directly under these updates.

• PoissonCountFactor (counts k, rate variable λ, Gamma cavity → Gamma tilted posterior α+Σk, β+n) and BernoulliFactor/BinomialFactor (Beta cavity → Beta posterior). Same protocol as WP2; each with property tests comparing calc_exact_projection moments against numerical integration of the tilted density (scipy.integrate.quad), the test pattern Phase 0 used.
• Requires a Poisson/Bernoulli likelihood story in the declarative layer only if demanded — ship as low-level factors first; no declarative sugar in this WP.

---

WP4 — Truncated-normal analytic moments (≈2–3 days, GATED)

Phase 6 candidate 5. Do not start until #1331-04 (TruncatedNormalMessage.log_partition missing the Gaussian term) is merged — the moments below feed the same normaliser.

• Implement calc_exact_update/has_exact_projection for truncated-normal messages against Normal cavities: truncated moments via the standard erf expressions (mean = μ + σ(φ(α)−φ(β))/Z etc.), projecting the truncated tilted distribution back to the message family. Also fixes the F6 approximation (untruncated KL) as a by-product if the exact truncated KL is added alongside (decide at review — same erf machinery).
• Property tests vs scipy.stats.truncnorm moments across central/near-bound/at-bound regimes (the Phase 1 F6 test grid: errors were 1.5%→140% — the analytic version should be exact to float precision).

---

Sequencing & bookkeeping

• Order: WP1 → WP2; WP3/WP4 unblock after the Priors & messages: 9 confirmed bugs — guidance wanted on 5 decisions #1331/research: EP review phase 1 — statistics audit of autofit/graphical #1332 fix batches. Each WP = its own PyAutoMind prompt → issue → PR when picked up (no bulk-issuing; file each as its predecessor nears shipping).
• Every WP: full test_autofit/ suite + the autofit_workspace_test/scripts/graphical/ scripts (incl. Phase 3's ep_exact.py, which pins the existing behaviour) must stay green; library unit tests numpy-only.
• Speed-up reporting: use Phase 4's ep_history.csv + graph_factors.png for before/after evidence on each WP; post to this issue.

Filed from the Phase 6 scoping session; may be picked up by a different model/session — everything needed is in this issue, #1337, and autofit/graphical/README.md.