Merge pull request #1383 from PyAutoLabs/feature/ep-projection-weights

fix: EP projection — transform samples to base space + use log weights

Author: Jammy2211

autofit/mapper/prior/abstract.py

@@ -198,7 +198,7 @@ def instance_for_arguments(
         _ = ignore_assertions
         return arguments[self]
 
-    def project(self, samples, weights):
+    def project(self, samples, log_weight_list):
         """Project this prior given samples and log weights from a search.
 
         Returns a copy of this prior whose message has been updated to
@@ -208,13 +208,17 @@ def project(self, samples, weights):
         ----------
         samples
             Array of sample values for this parameter.
-        weights
-            Log weights for each sample.
+        log_weight_list
+            Log weights for each sample. These must be log (not linear)
+            importance weights: the projection exponentiates them for its
+            weighted moment match. Passing linear weights makes every sample
+            count almost equally, which drags the projected message towards
+            the prior bounds (see PyAutoFit#1382).
         """
         result = copy(self)
         result.message = self.message.project(
             samples=samples,
-            log_weight_list=weights,
+            log_weight_list=log_weight_list,
             id_=self.id,
         )
         return result

autofit/messages/composed_transform.py

@@ -190,8 +190,20 @@ def __sub__(self, other):
     def project(
         self, samples, log_weight_list, **_,
     ):
+        """
+        Moment-matching projection of weighted samples.
+
+        ``samples`` arrive in the transformed (physical) space, so they are
+        mapped through the transform stack into the base message's space
+        before its moment match — exactly as ``@transform`` does for pdf and
+        friends. The weights need no Jacobian correction: a weighted
+        empirical distribution transforms its atoms, not its weights.
+        (Previously the physical samples were projected directly in base
+        space, pinning the projected message near a prior bound —
+        PyAutoFit#1382.)
+        """
         return TransformedMessage(
-            self.base_message.project(samples, log_weight_list),
+            self.base_message.project(self._transform(samples), log_weight_list),
             *self.transforms,
             id_=self.id,
         )

autofit/non_linear/result.py

@@ -343,12 +343,23 @@ def projected_model(self) -> AbstractPriorModel:
         Create a new model with the same structure as the previous model,
         replacing each prior with a new prior created by calculating sufficient
         statistics from samples and corresponding weights for that prior.
+
+        `Samples.weight_list` holds *linear* importance weights; the message
+        projection is an importance-weighted moment match over *log* weights,
+        so they are converted here. Zero-weight samples map to -inf and drop
+        out of the moments.
         """
-        weights = self.samples.weight_list
+        weights = np.asarray(self.samples.weight_list)
+        if not np.any(weights > 0.0):
+            raise ValueError(
+                "Cannot project a model from samples whose weights are all zero."
+            )
+        with np.errstate(divide="ignore"):
+            log_weight_list = np.log(weights)
         arguments = {
             prior: prior.project(
                 samples=np.array(self.samples.values_for_path(path)),
-                weights=weights,
+                log_weight_list=log_weight_list,
             )
             for path, prior in self.samples.model.path_priors_tuples
         }

test_autofit/graphical/test_unification.py

@@ -61,6 +61,68 @@ def test_projected_model():
     assert isinstance(projected_model.centre, af.UniformPrior)
 
 
+def test_projected_model_moments():
+    """
+    Regression test for PyAutoFit#1382: `projected_model` must convert the
+    samples' linear importance weights to log weights before the message
+    projection's weighted moment match.
+
+    The sample set mimics a nested-sampling run: a broad prior-exploration
+    phase with negligible weights plus a concentrated posterior with almost
+    all the weight. The projected prior's moments must match the posterior,
+    not the prior-phase spread — with linear weights fed in as log weights,
+    every sample counts almost equally and the projection lands near a
+    prior bound.
+    """
+    rng = np.random.default_rng(1)
+
+    posterior_mean = 2.05
+    posterior_sigma = 0.03
+
+    prior_phase = rng.uniform(1.5, 3.0, size=2000)
+    posterior_phase = rng.normal(posterior_mean, posterior_sigma, size=500)
+    values = np.concatenate([prior_phase, posterior_phase])
+
+    weights = np.concatenate(
+        [
+            np.full(prior_phase.size, 1.0e-9),
+            np.full(posterior_phase.size, (1.0 - prior_phase.size * 1.0e-9) / 500),
+        ]
+    )
+
+    model = af.Model(
+        af.ex.Gaussian,
+        centre=af.UniformPrior(lower_limit=1.5, upper_limit=3.0),
+    )
+    # normalization and sigma have UniformPrior(0, 1) under the test config, so
+    # they get their own in-support draws; the moment assertions are on centre.
+    unit_values = rng.uniform(0.05, 0.95, size=values.size)
+    samples = af.Samples(
+        model,
+        [
+            af.Sample(
+                -1.0,
+                -1.0,
+                weight=weight,
+                kwargs={
+                    ("centre",): value,
+                    ("normalization",): unit_value,
+                    ("sigma",): unit_value,
+                },
+            )
+            for value, unit_value, weight in zip(values, unit_values, weights)
+        ],
+    )
+    result = af.mock.MockResult(samples=samples)
+
+    projected_centre = result.projected_model.centre
+
+    assert projected_centre.mean == pytest.approx(posterior_mean, abs=0.02)
+    assert float(np.sqrt(projected_centre.variance)) == pytest.approx(
+        posterior_sigma, rel=0.5
+    )
+
+
 def test_uniform_normal(x):
     message = TransformedMessage(
         UniformNormalMessage,