feat: expose opt-in EP updater through declarative optimise

autofit/__init__.py

@@ -18,6 +18,12 @@
 from .graphical.declarative.factor.analysis import EPAnalysisFactor
 from .graphical.declarative.collection import FactorGraphModel
 from .graphical.declarative.factor.hierarchical import HierarchicalFactor
+from .graphical.expectation_propagation.optimiser import (
+    ApproxUpdater,
+    DynamicUpdater,
+    FactorUpdater,
+    SimplerUpdater,
+)
 from .graphical.laplace import LaplaceOptimiser
 from .non_linear.grid.grid_list import GridList
 from .non_linear.samples.summary import SamplesSummary

autofit/graphical/README.md

@@ -147,6 +147,19 @@ equivalently, an exponential moving average on natural parameters:
 factors update more slowly). `δ` may therefore be a scalar or a
 per-variable `MeanField` of scalars.
 
+The declarative API accepts this policy explicitly:
+
+    factor_graph.optimise(
+        optimiser,
+        updater=af.SimplerUpdater(delta=0.5),
+    )
+
+`updater=None` is the default and preserves the existing undamped
+`SimplerUpdater(delta=1.0)` behaviour. Damping is problem-dependent rather
+than a universal convergence fix; in particular, it has worsened hierarchical
+scale collapse in repeated-run diagnostics, so a damped configuration should
+be validated across repeated fits.
+
 **Invalid-projection fallback**: if the division produces an invalid
 message (e.g. negative variance — possible because Eq. (10)'s
 subtraction of natural parameters is not closed in the family),
@@ -278,6 +291,7 @@ re-verify (see the seam tests in
 | `AnalysisFactor(prior_model, analysis, optimiser)` | one `Factor` whose value is `analysis.log_likelihood_function` on the instance built from its variables | carries its own tilted-fit optimiser (§3.2) |
 | each free `Prior` | one graph `Variable` **and** one `PriorFactor` | priors are ordinary factors (§1); `PriorFactor` currently wraps the message's bound `factor` method, which strips the exact-update hooks — the conjugate update of §3.2 is *not* auto-selected declaratively (tracked: #1337 / plan #1338 WP1) |
 | the *same prior object* assigned to several models | one shared `Variable` connecting those factors | this is how information flows between datasets |
+| `optimise(..., updater=...)` | `EPOptimiser(updater=...)` | the supplied update policy survives lowering unchanged; omitting it preserves the undamped `delta=1.0` default |
 | compound prior (`prior_a * x + prior_b`, `mapper/prior/arithmetic/`) | **no graph variable** — the arithmetic is evaluated at instance-creation inside every factor that references it; only its component priors are variables | the relation is enforced *exactly* inside each tilted fit (no extra approximation — cf. §6); consequently the compound quantity has no message, no marginal, no evidence contribution of its own. (A `model.<property>` sugar for building these was deliberately reverted in `be6411755`.) |
 | `HierarchicalFactor` | one `Factor` per drawn variable (plus the distribution's parameter variables) | deliberate dimensionality choice |
 | — (no declarative expression) | `Factor(..., factor_out=v)` graph-level deterministic variables (§6.1) | **not reachable** from the declarative layer, by design as of the 2026-07 review (Phase 5, #1336): it trades the exact in-factor relation for a factorised q(v) with messages |

autofit/graphical/__init__.py

@@ -5,7 +5,13 @@
 from .declarative.factor.hierarchical import _HierarchicalFactor, HierarchicalFactor
 from .expectation_propagation.diagnostics import EPDiagnostics, check_sigma_collapse, mean_field_summary
 from .expectation_propagation.ep_mean_field import EPMeanField
-from .expectation_propagation.optimiser import EPOptimiser
+from .expectation_propagation.optimiser import (
+    ApproxUpdater,
+    DynamicUpdater,
+    EPOptimiser,
+    FactorUpdater,
+    SimplerUpdater,
+)
 from .expectation_propagation import StochasticEPOptimiser
 from .factor_graphs import FactorGraph
 from .factor_graphs.factor import Factor

autofit/graphical/declarative/abstract.py

@@ -5,7 +5,10 @@
 
 from autofit.graphical.declarative.factor.prior import PriorFactor
 from autofit.graphical.declarative.graph import DeclarativeFactorGraph
-from autofit.graphical.expectation_propagation import AbstractFactorOptimiser
+from autofit.graphical.expectation_propagation import (
+    AbstractFactorOptimiser,
+    ApproxUpdater,
+)
 from autofit.graphical.expectation_propagation import EPMeanField, EPOptimiser
 from autofit.mapper.prior.abstract import Prior
 from autofit.mapper.prior_model.collection import Collection
@@ -151,6 +154,7 @@ def _make_ep_optimiser(
         optimiser: AbstractFactorOptimiser,
         paths: Optional[AbstractPaths] = None,
         ep_history: Optional = None,
+        updater: Optional[ApproxUpdater] = None,
     ) -> EPOptimiser:
         return EPOptimiser(
             self.graph,
@@ -162,13 +166,15 @@ def _make_ep_optimiser(
             },
             ep_history=ep_history,
             paths=paths,
+            updater=updater,
         )
 
     def optimise(
         self,
         optimiser: AbstractFactorOptimiser,
         paths: Optional[AbstractPaths] = None,
         ep_history: Optional = None,
+        updater: Optional[ApproxUpdater] = None,
         **kwargs,
     ):
         """
@@ -182,6 +188,9 @@ def optimise(
             object is copied to every optimiser.
         optimiser
             An optimiser that acts on graphs
+        updater
+            An optional policy controlling how strongly EP factor messages are
+            updated. If omitted, EP keeps its existing undamped update policy.
 
         Returns
         -------
@@ -190,7 +199,12 @@ def optimise(
         """
         from autofit.graphical.declarative.result import EPResult
 
-        opt = self._make_ep_optimiser(optimiser, paths=paths, ep_history=ep_history)
+        opt = self._make_ep_optimiser(
+            optimiser,
+            paths=paths,
+            ep_history=ep_history,
+            updater=updater,
+        )
         updated_ep_mean_field = opt.run(self.mean_field_approximation(), **kwargs)
 
         return EPResult(

autofit/graphical/expectation_propagation/__init__.py

@@ -1,6 +1,12 @@
 from .diagnostics import EPDiagnostics, check_sigma_collapse, mean_field_summary
 from .ep_mean_field import EPMeanField
 from .history import FactorHistory, EPHistory
-from .optimiser import AbstractFactorOptimiser
-from .optimiser import EPOptimiser
+from .optimiser import (
+    AbstractFactorOptimiser,
+    ApproxUpdater,
+    DynamicUpdater,
+    EPOptimiser,
+    FactorUpdater,
+    SimplerUpdater,
+)
 from .stochastic import StochasticEPOptimiser

autofit/graphical/expectation_propagation/diagnostics.py

@@ -245,8 +245,11 @@ def check_sigma_collapse(
             warnings_list.append(
                 f"sigma-collapse: variable '{name}' has std {stds[-1]:.3g} "
                 f"below the floor {std_floor:.1g} — the fit has likely "
-                f"collapsed to a point (see PyAutoFit #1332 F10; consider "
-                f"damping, e.g. delta < 1, or per-factor sampler optimisers)."
+                f"collapsed to a point (see PyAutoFit #1332 F10). Mitigations "
+                f"are problem-dependent: optional damping is available via "
+                f"updater=af.SimplerUpdater(delta=0.5), but has worsened "
+                f"hierarchical scale collapse in repeated-run diagnostics; "
+                f"validate it across repeated fits."
             )
             continue
 

autofit/graphical/expectation_propagation/optimiser.py

@@ -208,6 +208,10 @@ def __init__(
             be optimised
         paths
             Optionally define how data should be output
+        updater
+            An optional policy controlling the strength of factor-message
+            updates. If omitted, full updates are used via
+            ``SimplerUpdater(delta=1.0)``.
         """
         factor_optimisers = factor_optimisers or {}
         self.factor_graph = factor_graph

test_autofit/graphical/functionality/test_diagnostics.py

@@ -131,6 +131,8 @@ def test_sigma_collapse_floor():
     assert len(warnings_list) == 1
     assert "collapsed" in warnings_list[0]
     assert "floor" in warnings_list[0]
+    assert "updater=af.SimplerUpdater(delta=0.5)" in warnings_list[0]
+    assert "problem-dependent" in warnings_list[0]
 
 
 def test_sigma_collapse_monotone():

test_autofit/graphical/gaussian/test_declarative.py

@@ -92,6 +92,54 @@ def test_custom_optimiser(make_model_factor):
     assert factor_optimisers[factor_2] is default_optimiser
 
 
+def test_custom_updater_reaches_ep_optimiser(factor_model):
+    updater = af.SimplerUpdater(delta=0.5)
+
+    ep_optimiser = factor_model._make_ep_optimiser(
+        ep.LaplaceOptimiser(),
+        updater=updater,
+    )
+
+    assert ep_optimiser.updater is updater
+
+
+def test_default_updater_is_undamped(factor_model):
+    ep_optimiser = factor_model._make_ep_optimiser(ep.LaplaceOptimiser())
+
+    assert isinstance(ep_optimiser.updater, af.SimplerUpdater)
+    assert ep_optimiser.updater.delta(factor=None, model_approx=None) == 1.0
+
+
+@pytest.mark.parametrize(
+    "updater",
+    [None, af.SimplerUpdater(delta=0.5)],
+)
+def test_optimise_forwards_optional_updater(factor_model, updater, monkeypatch):
+    captured = {}
+
+    class StubEPOptimiser:
+        ep_history = None
+
+        @staticmethod
+        def run(model_approx, **kwargs):
+            return model_approx
+
+    def make_ep_optimiser(
+        optimiser,
+        paths=None,
+        ep_history=None,
+        updater=None,
+    ):
+        captured["updater"] = updater
+        return StubEPOptimiser()
+
+    monkeypatch.setattr(factor_model, "_make_ep_optimiser", make_ep_optimiser)
+
+    factor_model.optimise(ep.LaplaceOptimiser(), updater=updater)
+
+    assert captured["updater"] is updater
+
+
 def test_factor_model_attributes(factor_model):
     """
     There are: