Refactor scores

* rename scoring_rules module to scores
* mean and median wrappers for normed difference scores
* argument handling for quantile score
* scoring rule tests

Author: han-ol

bayesflow/scores/__init__.py

@@ -0,0 +1,9 @@
+from .scores import (
+    ScoringRule,
+    ParametricDistributionRule,
+    NormedDifferenceScore,
+    MedianScore,
+    MeanScore,
+    QuantileScore,
+    MultivariateNormalScore,
+)

bayesflow/scores/scores.py

@@ -0,0 +1,184 @@
+from collections.abc import Callable, Sequence
+
+from bayesflow.types import Shape, Tensor
+
+from bayesflow.links import OrderedQuantiles, PositiveSemiDefinite
+
+from bayesflow.utils import logging
+
+import keras
+
+import math
+
+PI = keras.ops.convert_to_tensor(math.pi)
+
+
+class ScoringRule:
+    def get_link(self):
+        return keras.layers.Activation("linear")
+
+    def build(self, reference_shape: Shape):
+        pass
+
+    def score(self, reference, target):
+        raise NotImplementedError
+
+
+class NormedDifferenceScore(ScoringRule):
+    def __init__(
+        self,
+        k: int,  # results in an estimator for the mean
+    ):
+        self.k = k
+        self.target_shape = (1,)
+
+    def score(self, reference: Tensor, target: Tensor) -> Tensor:
+        pointwise_differance = target - reference[:, None, :]
+        score = keras.ops.absolute(pointwise_differance) ** self.k
+        score = keras.ops.mean(score)
+        return score
+
+
+class MedianScore(NormedDifferenceScore):
+    def __init__(self):
+        super().__init__(k=1)
+
+
+class MeanScore(NormedDifferenceScore):
+    def __init__(self):
+        super().__init__(k=2)
+
+
+class WeightedNormedDifferenceScore(ScoringRule):
+    def __init__(
+        self,
+        weighting_function: Callable,
+        k: int = 2,
+    ):
+        if weighting_function:
+            self.weighting_function = weighting_function
+        else:
+            self.weighting_function = lambda input: 1
+        self.k = k
+        self.target_shape = (1,)
+
+    def score(self, reference: Tensor, target: Tensor) -> Tensor:
+        pointwise_differance = target - reference[:, None, :]
+        score = self.weighting_function(reference) * keras.ops.absolute(pointwise_differance) ** self.k
+        score = keras.ops.mean(score)
+        return score
+
+
+class QuantileScore(ScoringRule):
+    def __init__(
+        self,
+        q: Sequence[float] = None,
+    ):
+        if q is None:
+            q = [0.1, 0.5, 0.9]
+            logging.info(f"QuantileScore was not provided with argument `q`. Using the default quantile levels: {q}.")
+
+        self.q = keras.ops.convert_to_tensor(q)
+        self.target_shape = (len(self.q),)
+
+    def get_link(self):
+        if self.q is None:
+            raise AssertionError("Needs q to construct link")
+        else:
+            print(self.q)
+            return OrderedQuantiles(self.q)
+
+    def score(self, reference: Tensor, target: Tensor) -> Tensor:
+        pointwise_differance = target - reference[:, None, :]
+
+        score = pointwise_differance * (keras.ops.cast(pointwise_differance > 0, float) - self.q[None, :, None])
+        score = keras.ops.mean(score)
+        return score
+
+
+class ParametricDistributionRule(ScoringRule):
+    """
+    TODO
+    """
+
+    def __init__(self, target_mappings: dict[str, str] = None):
+        self.target_mappings = target_mappings
+        self.build(
+            (
+                None,
+                0,
+            )
+        )  # TODO: make less confusing: this currently needs to be called initially AND another time when the
+        # scoring rules attributes are updated instead of just once in the end. This is probably not the Keras way.
+
+    def build(self, reference_shape: Shape):
+        if self.target_mappings is None:
+            self.target_mappings = {key: key for key in self.compute_target_shape().keys()}
+
+        self.target_shape = self.map_target_keys(self.compute_target_shape(), inverse=True)
+
+    def map_target_keys(self, target_dict: dict[str, str], inverse=False):
+        if inverse:
+            map = {v: k for k, v in self.target_mappings.items()}
+        else:
+            map = self.target_mappings
+        return {map[key]: value for key, value in target_dict.items()}
+
+    def compute_target_shape(self):
+        raise NotImplementedError
+
+    def log_prob(self, x, **kwargs):
+        raise NotImplementedError
+
+    def score(self, reference: Tensor, target: dict[str, Tensor]) -> Tensor:
+        score = -self.log_prob(x=reference, **self.map_target_keys(target))
+        score = keras.ops.mean(score)
+        return score * 0.01
+
+
+class MultivariateNormalScore(ParametricDistributionRule):
+    def __init__(self, D: int = None, **kwargs):
+        super().__init__(**kwargs)
+        self.D = D
+
+    def build(self, reference_shape: Shape):
+        if reference_shape is None:
+            raise AssertionError("Cannot build before setting D.")
+        elif isinstance(reference_shape, tuple) and len(reference_shape) == 2:
+            self.D = reference_shape[1]
+        else:
+            raise AssertionError(f"Cannot build from reference_shape {reference_shape}")
+        super().build(reference_shape)
+
+    def compute_target_shape(self) -> dict[str, Shape]:
+        return dict(mean=(1,), covariance=(self.D,))
+
+    def get_link(self):
+        return self.map_target_keys(
+            dict(mean=keras.layers.Activation("linear"), covariance=PositiveSemiDefinite()), inverse=True
+        )
+
+    def log_prob(self, x: Tensor, mean: Tensor, covariance: Tensor) -> Tensor:
+        diff = x[:, None, :] - mean
+        inv_covariance = keras.ops.inv(covariance)
+        log_det_covariance = keras.ops.slogdet(covariance)[1]  # Only take the log of the determinant part
+
+        # Compute the quadratic term in the exponential of the multivariate Gaussian
+        quadratic_term = keras.ops.einsum("...i,...ij,...j->...", diff, inv_covariance, diff)
+
+        # Compute the log probability density
+        log_prob = -0.5 * (self.D * keras.ops.log(2 * PI) + log_det_covariance + quadratic_term)
+
+        return log_prob
+
+    def sample(self, sample_size, mean, covariance):
+        batch_size, D = mean.shape
+        # Ensure covariance is (batch_size, D, D)
+        assert covariance.shape == (batch_size, D, D)
+
+        # Use Cholesky decomposition to generate samples
+        chol = keras.ops.cholesky(covariance)
+        normal_samples = keras.random.normal((batch_size, D, sample_size))
+        samples = mean[:, :, None] + keras.ops.einsum("ijk,ikl->ijl", chol, normal_samples)
+
+        return samples

bayesflow/scoring_rules/__init__.py

@@ -0,0 +1,69 @@
+import keras
+import pytest
+
+
+@pytest.fixture()
+def batch_size():
+    return 16
+
+
+@pytest.fixture()
+def num_variables():
+    return 4
+
+
+@pytest.fixture()
+def reference(batch_size, num_variables):
+    return keras.random.uniform((batch_size, num_variables))
+
+
+@pytest.fixture()
+def median_score():
+    from bayesflow.scores import MedianScore
+
+    return MedianScore()
+
+
+@pytest.fixture()
+def mean_score():
+    from bayesflow.scores import MeanScore
+
+    return MeanScore()
+
+
+@pytest.fixture()
+def normed_diff_score():
+    from bayesflow.scores import NormedDifferenceScore
+
+    return NormedDifferenceScore(k=3)
+
+
+@pytest.fixture()
+def quantile_score():
+    from bayesflow.scores import QuantileScore
+
+    return QuantileScore()
+
+
+@pytest.fixture(params=["median_score", "mean_score", "normed_diff_score", "quantile_score"], scope="function")
+def basic_scoring_rule(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.fixture()
+def mvn_target(batch_size, num_variables):
+    mean_target = keras.ops.zeros((batch_size, 1, num_variables))
+    inputs = keras.random.normal((batch_size, num_variables, num_variables))
+    print(inputs.shape)
+    covariance_target = keras.ops.einsum("...ij,...kj->...ik", inputs, inputs)
+    return dict(
+        mean=mean_target,
+        covariance=covariance_target,
+    )
+
+
+@pytest.fixture()
+def multivariate_normal_score(num_variables):
+    from bayesflow.scores import MultivariateNormalScore
+
+    return MultivariateNormalScore(D=num_variables)

bayesflow/scoring_rules/scoring_rules.py

@@ -0,0 +1,36 @@
+import keras
+import pytest
+
+
+def test_require_argument_k():
+    from bayesflow.scores import NormedDifferenceScore
+
+    with pytest.raises(TypeError) as excinfo:
+        NormedDifferenceScore()
+
+    assert "missing 1 required positional argument: 'k'" in str(excinfo)
+
+
+def test_score_output(basic_scoring_rule, reference):
+    target_shape = (reference.shape[0], *basic_scoring_rule.target_shape, reference.shape[-1])
+    target = keras.ops.zeros(target_shape)
+    score = basic_scoring_rule.score(reference, target)
+
+    assert score.ndim == 0
+
+
+def test_mean_score_optimality(mean_score, reference):
+    suboptimal_target = keras.ops.expand_dims(keras.random.uniform(reference.shape), axis=1)
+    optimal_target = keras.ops.expand_dims(reference, axis=1)
+
+    suboptimal_score = mean_score.score(reference, suboptimal_target)
+    optimal_score = mean_score.score(reference, optimal_target)
+
+    assert suboptimal_score > optimal_score
+    assert keras.ops.isclose(optimal_score, 0)
+
+
+def test_multivariate_normal_score_output(multivariate_normal_score, reference, mvn_target):
+    score = multivariate_normal_score.score(reference, mvn_target)
+
+    assert score.ndim == 0