Add summary statistics to ProductMeasure

qmcpy/true_measure/product_measure.py

@@ -1,4 +1,5 @@
 import numpy as np
+from scipy import sparse
 
 from .abstract_true_measure import AbstractTrueMeasure
 from ..discrete_distribution.abstract_discrete_distribution import (
@@ -46,6 +47,9 @@ class ProductMeasure(AbstractTrueMeasure):
 
     Notes
     -----
+    For independent marginal blocks, means, variances, and standard deviations
+    are concatenated in marginal order, while covariance is block diagonal.
+
     Exact product weights are supported for direct marginal true measures. For
     recursively composed marginal measures, sampling is supported through
     QMCPy's recursive transform helper, but exact final-space product weights
@@ -177,6 +181,93 @@ def __init__(self, sampler, marginals):
 
         super(ProductMeasure, self).__init__()
 
+        for statistic in (
+            "mean",
+            "variance",
+            "standard_deviation",
+            "covariance",
+        ):
+            if all(hasattr(marginal, statistic) for marginal in self.marginals):
+                self.parameters.append(statistic)
+
+    def _marginal_statistic(self, marginal, marginal_index, statistic):
+        """Return a statistic or identify the marginal that does not provide it."""
+        try:
+            return getattr(marginal, statistic)
+        except AttributeError as error:
+            raise AttributeError(
+                f"ProductMeasure marginal {marginal_index} "
+                f"({type(marginal).__name__}) does not provide {statistic}."
+            ) from error
+
+    def _concatenate_marginal_statistic(self, statistic):
+        """Concatenate a coordinate-wise statistic in marginal order."""
+        values = []
+        for marginal_index, marginal in enumerate(self.marginals):
+            value = self._marginal_statistic(
+                marginal, marginal_index, statistic
+            )
+            value = np.atleast_1d(np.asarray(value))
+            if value.shape != (marginal.d,):
+                raise DimensionError(
+                    f"ProductMeasure marginal {marginal_index} "
+                    f"({type(marginal).__name__}) {statistic} must have shape "
+                    f"({marginal.d},), got {value.shape}."
+                )
+            values.append(value)
+
+        combined = self._read_only_array(np.concatenate(values))
+        return self._scalar_if_univariate(combined)
+
+    @property
+    def mean(self):
+        return self._concatenate_marginal_statistic("mean")
+
+    @property
+    def variance(self):
+        return self._concatenate_marginal_statistic("variance")
+
+    @property
+    def standard_deviation(self):
+        return self._concatenate_marginal_statistic("standard_deviation")
+
+    @property
+    def covariance(self):
+        blocks = []
+        for marginal_index, marginal in enumerate(self.marginals):
+            block = self._marginal_statistic(
+                marginal, marginal_index, "covariance"
+            )
+            if not sparse.issparse(block):
+                block = np.atleast_2d(np.asarray(block))
+            expected_shape = (marginal.d, marginal.d)
+            if block.shape != expected_shape:
+                raise DimensionError(
+                    f"ProductMeasure marginal {marginal_index} "
+                    f"({type(marginal).__name__}) covariance must have shape "
+                    f"{expected_shape}, got {block.shape}."
+                )
+            blocks.append(block)
+
+        if any(sparse.issparse(block) for block in blocks):
+            covariance = sparse.block_diag(blocks, format="dia")
+            data = covariance.data
+            data.setflags(write=False)
+            covariance.data = self._read_only_view(data)
+            return covariance
+
+        covariance = np.zeros(
+            (self.d, self.d),
+            dtype=np.result_type(*[block.dtype for block in blocks]),
+        )
+        start = 0
+        for block in blocks:
+            stop = start + block.shape[0]
+            covariance[start:stop, start:stop] = block
+            start = stop
+        covariance.setflags(write=False)
+        return self._read_only_view(covariance)
+
     @staticmethod
     def _expand_bounds(bounds, dimension, name):
         """

test/test_product_measure.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+import scipy.sparse as sp
 import scipy.stats as stats
 
 from qmcpy import (
@@ -48,6 +49,167 @@ def test_product_measure_replication_shape():
     assert x.shape == (r, n, 2)
 
 
+def test_product_measure_statistics_for_multiple_1d_marginals():
+    marginals = [
+        Uniform(DummySampler(1), lower_bound=8.0, upper_bound=12.0),
+        Uniform(DummySampler(1), lower_bound=-1.0, upper_bound=5.0),
+    ]
+    tm = ProductMeasure(sampler=DigitalNetB2(2, seed=23), marginals=marginals)
+
+    np.testing.assert_allclose(tm.mean, [10.0, 2.0])
+    np.testing.assert_allclose(tm.variance, [4.0 / 3.0, 3.0])
+    np.testing.assert_allclose(
+        tm.standard_deviation, [np.sqrt(4.0 / 3.0), np.sqrt(3.0)]
+    )
+    cov = tm.covariance
+    cov_dense = cov.toarray() if sp.issparse(cov) else cov
+    np.testing.assert_allclose(cov_dense, np.diag([4.0 / 3.0, 3.0]))
+
+    assert tm.mean.shape == (2,)
+    assert tm.variance.shape == (2,)
+    assert tm.standard_deviation.shape == (2,)
+    assert tm.covariance.shape == (2, 2)
+    for statistic in ("mean", "variance", "standard_deviation", "covariance"):
+        value = getattr(tm, statistic)
+        flags = value.data.flags if sp.issparse(value) else value.flags
+        assert not flags.writeable
+
+
+def test_product_measure_normalizes_scalar_1d_statistics():
+    marginals = [
+        Uniform(DummySampler(1), lower_bound=8.0, upper_bound=12.0),
+        Gaussian(DummySampler(1), mean=2.0, covariance=9.0),
+    ]
+    for marginal in marginals:
+        assert isinstance(marginal.mean, float)
+        assert isinstance(marginal.variance, float)
+        assert isinstance(marginal.standard_deviation, float)
+
+    tm = ProductMeasure(sampler=DigitalNetB2(2, seed=29), marginals=marginals)
+
+    np.testing.assert_allclose(tm.mean, [10.0, 2.0])
+    np.testing.assert_allclose(tm.variance, [4.0 / 3.0, 9.0])
+    np.testing.assert_allclose(
+        tm.standard_deviation, [np.sqrt(4.0 / 3.0), 3.0]
+    )
+    np.testing.assert_allclose(
+        tm.covariance.toarray(), np.diag([4.0 / 3.0, 9.0])
+    )
+    assert tm.mean.shape == (2,)
+    assert tm.variance.shape == (2,)
+    assert tm.standard_deviation.shape == (2,)
+    assert tm.covariance.shape == (2, 2)
+
+
+def test_product_measure_statistics_preserve_order_and_covariance_blocks():
+    marginals = [
+        Uniform(DummySampler(1), lower_bound=8.0, upper_bound=12.0),
+        Gaussian(
+            DummySampler(2),
+            mean=[2.0, 5.0],
+            covariance=[[2.0, 0.5], [0.5, 3.0]],
+        ),
+    ]
+    tm = ProductMeasure(sampler=DigitalNetB2(3, seed=31), marginals=marginals)
+    expected_covariance = np.array(
+        [
+            [4.0 / 3.0, 0.0, 0.0],
+            [0.0, 2.0, 0.5],
+            [0.0, 0.5, 3.0],
+        ]
+    )
+
+    np.testing.assert_allclose(tm.mean, [10.0, 2.0, 5.0])
+    np.testing.assert_allclose(tm.variance, [4.0 / 3.0, 2.0, 3.0])
+    np.testing.assert_allclose(
+        tm.standard_deviation,
+        [np.sqrt(4.0 / 3.0), np.sqrt(2.0), np.sqrt(3.0)],
+    )
+    covariance = tm.covariance.tocsr()
+    np.testing.assert_allclose(covariance.toarray(), expected_covariance)
+
+    assert tm.mean.shape == (3,)
+    assert tm.variance.shape == (3,)
+    assert tm.standard_deviation.shape == (3,)
+    assert covariance.shape == (3, 3)
+    assert covariance[:1, 1:].nnz == 0
+    assert covariance[1:, :1].nnz == 0
+
+
+def test_product_measure_mixed_covariance_blocks_remain_sparse():
+    d = 128
+    tm = ProductMeasure(
+        DummySampler(d + 2),
+        [
+            Uniform(DummySampler(d)),
+            Gaussian(
+                DummySampler(2),
+                covariance=np.array([[1.0, 0.5], [0.5, 1.0]]),
+            ),
+        ],
+    )
+
+    covariance = tm.covariance
+    expected = sp.block_diag(
+        [marginal.covariance for marginal in tm.marginals], format="dia"
+    )
+
+    assert sp.issparse(covariance)
+    assert covariance.format == "dia"
+    assert covariance.shape == (d + 2, d + 2)
+    difference = (covariance - expected).tocsr()
+    difference.eliminate_zeros()
+    assert difference.nnz == 0
+    covariance_csr = covariance.tocsr()
+    np.testing.assert_allclose(
+        covariance_csr[-2:, -2:].toarray(), [[1.0, 0.5], [0.5, 1.0]]
+    )
+    assert covariance_csr[:d, d:].nnz == 0
+    assert covariance_csr[d:, :d].nnz == 0
+    assert not covariance.data.flags.writeable
+    with pytest.raises(ValueError):
+        covariance.data.setflags(write=True)
+
+
+def test_product_measure_dense_covariance_cannot_be_made_writeable():
+    tm = ProductMeasure(
+        DummySampler(3),
+        [
+            Gaussian(DummySampler(1), covariance=2.0),
+            Gaussian(
+                DummySampler(2), covariance=[[3.0, 0.25], [0.25, 4.0]]
+            ),
+        ],
+    )
+
+    covariance = tm.covariance
+
+    assert isinstance(covariance, np.ndarray)
+    np.testing.assert_allclose(
+        covariance,
+        [[2.0, 0.0, 0.0], [0.0, 3.0, 0.25], [0.0, 0.25, 4.0]],
+    )
+    assert not covariance.flags.writeable
+    with pytest.raises(ValueError):
+        covariance.setflags(write=True)
+
+
+def test_product_measure_missing_marginal_statistic_is_identified():
+    marginals = [
+        ZeroInflatedExpUniform(DummySampler(1), p_zero=0.4, lam=1.5),
+        Uniform(DummySampler(1), lower_bound=2.0, upper_bound=5.0),
+    ]
+    tm = ProductMeasure(sampler=DigitalNetB2(2, seed=23), marginals=marginals)
+
+    np.testing.assert_allclose(tm.mean, [0.4, 3.5])
+    assert "covariance" not in tm.parameters
+    with pytest.raises(
+        AttributeError,
+        match=r"marginal 0 \(ZeroInflatedExpUniform\) does not provide covariance",
+    ):
+        _ = tm.covariance
+
+
 def test_product_measure_marginals_with_different_dimensions():
     n = 32
     marginals = [