Merge pull request #1345 from PyAutoLabs/feature/priors-messages-fixes

fix(priors): correctness batch — crashes, log-partition, projections

Author: Jammy2211

autofit/mapper/prior/abstract.py

@@ -37,6 +37,27 @@ class attribute.
 
         self.width_modifier = None
 
+    def log_normalisation(self, xp=np) -> float:
+        """
+        The additive constant dropped from :meth:`log_prior_from_value`.
+
+        By contract ``log_prior_from_value`` returns the log prior density *up to
+        an additive constant*: the value-independent normaliser is dropped because
+        it is irrelevant to posterior shape (it cancels in the Metropolis ratio,
+        and nested samplers use the unit-cube transform rather than the prior
+        density). This is harmless for sampling but means absolute ``log_prior``
+        values are not comparable across prior types.
+
+        For code that needs the *fully normalised* log density — e.g. evidence or
+        Bayes-factor arithmetic — the normalised value is::
+
+            log_prior_from_value(value) + log_normalisation()
+
+        The default is ``0.0`` (constant unknown / already normalised). Priors with
+        a known closed-form normaliser override this.
+        """
+        return 0.0
+
     @classmethod
     def tree_unflatten(cls, aux_data, children):
         """

autofit/mapper/prior/gaussian.py

@@ -114,6 +114,12 @@ def parameter_string(self) -> str:
         """
         return f"mean = {self.mean}, sigma = {self.sigma}"
 
+    def log_normalisation(self, xp=np) -> float:
+        """The constant ``-log(sigma) - 0.5*log(2*pi)`` dropped from the density-form
+        quadratic returned by ``NormalMessage.log_prior_from_value``. See
+        ``Prior.log_normalisation``."""
+        return -xp.log(self.sigma) - 0.5 * xp.log(2.0 * np.pi)
+
     def value_for(self, unit, xp=np):
         """
         Map a unit value in [0, 1] to a physical value drawn from this Gaussian prior.

autofit/mapper/prior/log_gaussian.py

@@ -95,8 +95,6 @@ def with_limits(cls, lower_limit: float, upper_limit: float) -> "LogGaussianPrio
         return cls(
             mean=(lower_limit + upper_limit) / 2,
             sigma=upper_limit - lower_limit,
-            lower_limit=lower_limit,
-            upper_limit=upper_limit,
         )
 
     def _new_for_base_message(self, message):
@@ -108,9 +106,7 @@ def _new_for_base_message(self, message):
         """
         return LogGaussianPrior(
             *message.parameters,
-            lower_limit=self.lower_limit,
-            upper_limit=self.upper_limit,
-            id_=self.instance().id,
+            id_=self.id,
         )
 
     def value_for(self, unit, xp=np):
@@ -141,6 +137,12 @@ def value_for(self, unit, xp=np):
     def parameter_string(self) -> str:
         return f"mean = {self.mean}, sigma = {self.sigma}"
 
+    def log_normalisation(self, xp=np) -> float:
+        """The constant ``-log(sigma) - 0.5*log(2*pi)`` dropped from the Gaussian-in-log
+        density in ``log_prior_from_value`` (the value-dependent ``-log(value)``
+        change-of-variables Jacobian is kept). See ``Prior.log_normalisation``."""
+        return -xp.log(self.sigma) - 0.5 * xp.log(2.0 * np.pi)
+
     def log_prior_from_value(self, value, xp=np):
         """
         Compute the log prior density of a given physical value under this log-Gaussian prior.

autofit/mapper/prior/log_uniform.py

@@ -153,6 +153,12 @@ def log_prior_from_value(self, value, xp=np):
         in_bounds = (value >= self.lower_limit) & (value <= self.upper_limit)
         return xp.where(in_bounds, -xp.log(value), -xp.inf)
 
+    def log_normalisation(self, xp=np) -> float:
+        """The constant ``-log(log(upper / lower))`` dropped from
+        ``log_prior_from_value`` (which returns ``-log(value)``). See
+        ``Prior.log_normalisation``."""
+        return -xp.log(xp.log(self.upper_limit / self.lower_limit))
+
     def value_for(self, unit, xp=np):
         """
         Returns a physical value from an input unit value according to the limits of the log10 uniform prior.

autofit/mapper/prior/truncated_gaussian.py

@@ -133,6 +133,21 @@ def parameter_string(self) -> str:
                 f"upper_limit = {self.upper_limit}"
                 )
 
+    def log_normalisation(self, xp=np) -> float:
+        """The constant ``-log(sigma) - 0.5*log(2*pi) - log(Z)`` dropped from
+        ``TruncatedNormalMessage.log_prior_from_value``, where
+        ``Z = Phi((upper - mean)/sigma) - Phi((lower - mean)/sigma)`` is the
+        truncation mass. See ``Prior.log_normalisation``."""
+        if xp.__name__.startswith("jax"):
+            import jax.scipy.stats as jstats
+            norm = jstats.norm
+        else:
+            from scipy.stats import norm
+        a = (self.lower_limit - self.mean) / self.sigma
+        b = (self.upper_limit - self.mean) / self.sigma
+        Z = norm.cdf(b) - norm.cdf(a)
+        return -xp.log(self.sigma) - 0.5 * xp.log(2.0 * np.pi) - xp.log(Z)
+
     def value_for(self, unit, xp=np):
         """
         Map a unit value in [0, 1] to a physical value drawn from this truncated Gaussian prior.

autofit/mapper/prior/uniform.py

@@ -106,13 +106,21 @@ def logpdf(self, x):
         Parameters
         ----------
         x
-            The value at which to evaluate the log PDF.
+            The value(s) at which to evaluate the log PDF. May be a scalar or a
+            numpy array.
         """
-        # TODO: handle x as a numpy array
-        if x == self.lower_limit:
-            x += epsilon
-        elif x == self.upper_limit:
-            x -= epsilon
+        # Nudge values sitting exactly on a boundary inwards by epsilon, where the
+        # PDF is otherwise undefined. The scalar path is kept bit-identical to the
+        # historical behaviour; the array path vectorises the same snap with
+        # ``np.where`` (the previous scalar-only ``==`` comparison raised on arrays).
+        if np.ndim(x) == 0:
+            if x == self.lower_limit:
+                x += epsilon
+            elif x == self.upper_limit:
+                x -= epsilon
+        else:
+            x = np.where(x == self.lower_limit, x + epsilon, x)
+            x = np.where(x == self.upper_limit, x - epsilon, x)
         return self.message.logpdf(x)
 
     def dict(self) -> dict:
@@ -178,6 +186,11 @@ def log_prior_from_value(self, value, xp=np):
         in_bounds = (value >= self.lower_limit) & (value <= self.upper_limit)
         return xp.where(in_bounds, xp.zeros_like(value), -xp.inf)
 
+    def log_normalisation(self, xp=np) -> float:
+        """The constant ``-log(upper - lower)`` dropped from ``log_prior_from_value``
+        (which returns ``0.0``). See ``Prior.log_normalisation``."""
+        return -xp.log(self.upper_limit - self.lower_limit)
+
     @property
     def limits(self) -> Tuple[float, float]:
         """The (lower_limit, upper_limit) bounds of this uniform prior."""

autofit/messages/beta.py

@@ -74,9 +74,12 @@ def inv_beta_suffstats(
     b
         Estimated beta parameter(s) of the Beta distribution.
 
-    Warnings
-    --------
-    Emits a RuntimeWarning if negative parameters are found, and clamps them to 0.5.
+    Raises
+    ------
+    ValueError
+        If the Newton-Raphson projection produces a negative alpha or beta, which
+        is not a valid Beta distribution. Previously this was silently warned and
+        clamped into a local that was immediately overwritten (a no-op).
     """
 
     _lnX, _ln1X = np.ravel(lnX), np.ravel(ln1X)
@@ -94,12 +97,14 @@ def inv_beta_suffstats(
         ab += np.linalg.solve(jac, - f)
 
     if np.any(ab < 0):
-        warnings.warn(
-            "invalid negative parameters found for inv_beta_suffstats, "
-            "clampling value to 0.5",
-            RuntimeWarning
+        raise ValueError(
+            "inv_beta_suffstats produced negative Beta parameters "
+            f"(alpha, beta):\n\n{ab}\n\n"
+            "A negative alpha or beta is not a valid Beta distribution, so the "
+            "moment-matching projection has failed. The previous behaviour "
+            "silently warned and clamped into a local that was immediately "
+            "overwritten (a no-op), letting the invalid parameters escape."
         )
-        b = np.clip(ab, 0.5, None)
 
     shape = np.shape(lnX)
     if shape:

autofit/messages/fixed.py

@@ -54,12 +54,12 @@ def sample(self, n_samples: Optional[int] = None) -> np.ndarray:
             return self.value
         return np.array([self.value])
 
-    logpdf_cache = {}
-
     def logpdf(self, x: np.ndarray) -> np.ndarray:
-        if x.shape not in FixedMessage.logpdf_cache:
-            FixedMessage.logpdf_cache[x.shape] = np.zeros_like(x)
-        return FixedMessage.logpdf_cache[x.shape]
+        # A fixed message contributes zero log-density everywhere. Return a
+        # fresh zero array each call: the previous class-level ``logpdf_cache``
+        # was an unbounded dict keyed on shape that also handed back an aliased
+        # mutable array, so mutating one result silently corrupted later calls.
+        return np.zeros_like(x)
 
     @cached_property
     def mean(self) -> np.ndarray:

autofit/messages/gamma.py

@@ -76,8 +76,13 @@ def sample(self, n_samples=None):
     def from_mode(cls, mode, covariance, **kwargs):
         m, V = cls._get_mean_variance(mode, covariance)
 
-        alpha = 1 + m ** 2 * V  # match variance
-        beta = alpha / m  # match mean
+        # Match the requested mean m and variance V, consistent with the Normal
+        # family. For a Gamma, mean = alpha / beta and variance = alpha / beta**2,
+        # which inverts to alpha = m**2 / V, beta = m / V. The previous
+        # ``alpha = 1 + m**2 * V`` had the variance the wrong way up (requesting
+        # variance 0.25 produced 2.0, and 4.0 produced 0.235).
+        alpha = m ** 2 / V
+        beta = m / V
         return cls(alpha, beta, **kwargs)
 
     def kl(self, dist):

autofit/messages/truncated_normal.py

@@ -30,11 +30,20 @@ def log_partition(self, xp=np) -> float:
         """
         Compute the log-partition function (normalizer) of the truncated Gaussian.
 
-        This is the log of the normalization constant Z of the truncated normal:
+        For the exponential-family interface this is the full cumulant, made of
+        two parts:
 
-            Z = Φ((b - μ)/σ) - Φ((a - μ)/σ)
+        1. the untruncated Gaussian log-partition ``A(η) = μ²/(2σ²) + log σ``
+           (identical to ``NormalMessage.log_partition``), and
+        2. the truncation-mass correction ``log Z`` with
+           ``Z = Φ((b - μ)/σ) - Φ((a - μ)/σ)``, where Φ is the standard normal
+           CDF and ``[a, b]`` are the truncation bounds.
 
-        where Φ is the standard normal CDF and [a, b] are the truncation bounds.
+        Previously only ``log Z`` was returned, dropping the Gaussian term. That
+        made the generic exponential-family pdf integrate to
+        ``σ·exp(μ²/2σ²)`` (e.g. 2.27 for the unit case) instead of 1.0 — the path
+        the EP machinery consumes. Sampling and ``log_prior_from_value`` use a
+        separate, correct path and were unaffected.
 
         Returns
         -------
@@ -43,10 +52,14 @@ def log_partition(self, xp=np) -> float:
         """
         from scipy.stats import norm
 
+        # Untruncated Gaussian log-partition — see NormalMessage.log_partition.
+        gaussian = (self.mean ** 2) / (2 * self.sigma ** 2) + xp.log(self.sigma)
+
         a = (self.lower_limit - self.mean) / self.sigma
         b = (self.upper_limit - self.mean) / self.sigma
         Z = norm.cdf(b) - norm.cdf(a)
-        return xp.log(Z) if Z > 0 else -xp.inf
+        log_Z = xp.log(Z) if Z > 0 else -xp.inf
+        return gaussian + log_Z
 
     log_base_measure = -0.5 * np.log(2 * np.pi)
 
@@ -472,9 +485,19 @@ def log_prior_from_value(self, value: float, xp=np) -> float:
         """
         Compute the log prior probability of a given physical value under this truncated Gaussian prior.
 
-        This accounts for truncation by normalizing the Gaussian density over the
-        interval [lower_limit, upper_limit], returning -inf if the value lies outside
-        these limits.
+        Returns ``log p(value)`` in density form, up to an additive constant, and
+        ``-inf`` for values outside ``[lower_limit, upper_limit]``.
+
+        The value-independent constants ``-log(sigma) - 0.5*log(2*pi)`` (the
+        Gaussian normaliser) and ``-log(Z)`` (the truncation mass) are dropped, so
+        this matches the constant-dropping convention already used by
+        ``NormalMessage.log_prior_from_value`` (and Uniform / LogUniform / LogGaussian).
+        Previously this method returned the *fully normalised* truncated density,
+        making it the odd one out — harmless to posterior shape (constants cancel
+        in the Metropolis ratio and nested samplers use the unit-cube transform),
+        but inconsistent for anyone reading absolute ``log_prior`` values or doing
+        evidence arithmetic. The dropped constant is recoverable via
+        ``TruncatedGaussianPrior.log_normalisation``.
 
         Parameters
         ----------
@@ -483,33 +506,18 @@ def log_prior_from_value(self, value: float, xp=np) -> float:
 
         Returns
         -------
-        The log prior probability of the given value, or -inf if outside truncation bounds.
+        The log prior density at the given value up to an additive constant, or
+        -inf if outside the truncation bounds.
         """
 
-        if xp.__name__.startswith("jax"):
-            import jax.scipy.stats as jstats
-            norm = jstats.norm
-        else:
-            from scipy.stats import norm
-
-        # Normalization term (truncation)
-        a = (self.lower_limit - self.mean) / self.sigma
-        b = (self.upper_limit - self.mean) / self.sigma
-        Z = norm.cdf(b) - norm.cdf(a)
-
-        # Log pdf
+        # Density-form quadratic, constants dropped (see docstring / NormalMessage).
         z = (value - self.mean) / self.sigma
-        log_pdf = (
-                -0.5 * z ** 2
-                - xp.log(self.sigma)
-                - 0.5 * xp.log(2.0 * xp.pi)
-        )
-        log_trunc_pdf = log_pdf - xp.log(Z)
+        log_pdf = -0.5 * z ** 2
 
-        # Truncation mask (must be xp.where for JAX)
+        # Truncation mask (must be xp.where for JAX).
         in_bounds = (self.lower_limit <= value) & (value <= self.upper_limit)
 
-        return xp.where(in_bounds, log_trunc_pdf, -xp.inf)
+        return xp.where(in_bounds, log_pdf, -xp.inf)
 
     def __str__(self):
         """