point estimate of covariance compatible with standardization

Author: han-ol

bayesflow/approximators/point_approximator.py

@@ -61,6 +61,15 @@ def estimate(
         conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
 
         estimates = self._estimate(**conditions, **kwargs)
+
+        if "inference_variables" in self.standardize:
+            for score_key, score in self.inference_network.scores.items():
+                for head_key in estimates[score_key].keys():
+                    trafo_type = score.TRANSFORMATION_TYPE.get(head_key, "rank1+shift")
+                    estimates[score_key][head_key] = self.standardize_layers["inference_variables"](
+                        estimates[score_key][head_key], forward=False, transformation_type=trafo_type
+                    )
+
         estimates = self._apply_inverse_adapter_to_estimates(estimates, **kwargs)
 
         # Optionally split the arrays along the last axis.

bayesflow/networks/standardization/standardization.py

@@ -53,6 +53,7 @@ def call(
         stage: str = "inference",
         forward: bool = True,
         log_det_jac: bool = False,
+        transformation_type: str = "rank1+shift",
         **kwargs,
     ) -> Tensor | Sequence[Tensor]:
         """
@@ -75,6 +76,14 @@ def call(
         Tensor or Sequence[Tensor]
             Transformed tensor, and optionally the log-determinant if `log_det_jac=True`.
         """
+        msg = """
+        Non-default transformation (i.e. transformation_type != "rank1+shift")
+        is not supported for forward or log_det_jac.
+        """
+        if forward or log_det_jac:
+            if transformation_type != "rank1+shift":  # non default transformation
+                raise ValueError(msg)
+
         flattened = keras.tree.flatten(x)
         outputs, log_det_jacs = [], []
 
@@ -91,7 +100,13 @@ def call(
                 # we can just replace them with zeros.
                 out = keras.ops.nan_to_num(out, nan=0.0)
             else:
-                out = mean + std * val
+                match transformation_type:
+                    case "rank1+shift":
+                        # x_i = x_i' * std + mean
+                        out = val * std + mean
+                    case "rank02":
+                        # x_ij = x_ij * sigma_i * sigma_j
+                        out = val * std * keras.ops.moveaxis(std, -1, -2)
 
             outputs.append(out)
 

bayesflow/scores/multivariate_normal_score.py

@@ -26,16 +26,16 @@ class MultivariateNormalScore(ParametricDistributionScore):
     For more information see :py:class:`ScoringRule`.
     """
 
-    RANK: dict[str, int] = {"covariance": 2}
+    TRANSFORMATION_TYPE: dict[str, str] = {"covariance": "rank02"}
     """
-    The covariance matrix is a rank 2 tensor and as such the inverse of the standardization operation is
+    Marks covariance head to handle de-standardization as for covariant rank-(0,2) tensors.
 
-    x = x' * sigma ^ 2
+    The appropriate inverse of the standardization operation is
 
-    Accordingly, covariance is also included in :py:attr:`NO_SHIFT`.
-    """
+    x_ij = x_ij * sigma_i * sigma_j.
 
-    NO_SHIFT: tuple[str] = ("covariance",)
+    For the mean head the default ("rank1+shift") is not overridden.
+    """
 
     def __init__(self, dim: int = None, links: dict = None, **kwargs):
         super().__init__(links=links, **kwargs)

bayesflow/scores/scoring_rule.py

@@ -37,31 +37,15 @@ class ScoringRule:
     with a type of estimate whenever the adapter is applied to them in inverse direction.
     """
 
-    RANK: dict[str, int] = {}
+    TRANSFORMATION_TYPE: dict[str, str] = {"covariance": "rank02"}
     """
-    Mapping of prediction head names to their tensor rank for inverse standardization.
+    Defines nonstandard transformation behaviour for de-standardization.
 
-    The rank indicates the power to which the standard deviation is raised before being multiplied to some estimate
-    in standardized space.
+    The standard transformation
 
-    x = x' * sigma ^ rank [ + mean ]
+    x_i = x_i' * std + mean
 
-    If a head is not present in this mapping, a default rank of 1 is assumed.
-
-    Typically, if :py:attr:`RANK` is modified for an estimate, it is also included in :py:attr:`NO_SHIFT`.
-    """
-
-    NO_SHIFT: tuple[str] = tuple()
-    """
-    Names of prediction heads whose estimates should not be shifted when applying inverse standardization.
-
-    During inverse standardization, point estimates are typically shifted by the stored mean vector. Any head
-    listed in this tuple will skip the shift step and only be scaled. By default, this tuple is empty,
-    meaning all heads will be shifted to undo standardization.
-
-    x = x' * sigma ^ rank + mean
-
-    See also :py:attr:`RANK`.
+    is referred to as "rank1+shift". Keys not specified here will fallback to that default.
     """
 
     def __init__(