Add summary statistics to ProductMeasure

qmcpy/true_measure/product_measure.py

@@ -1,10 +1,11 @@
 import numpy as np
+from scipy import sparse
 
 from .abstract_true_measure import AbstractTrueMeasure
 from ..discrete_distribution.abstract_discrete_distribution import (
     AbstractDiscreteDistribution,
 )
-from ..util import DimensionError, ParameterError
+from ..util import DimensionError, ParameterError, _univ_repr
 
 
 class ProductMeasure(AbstractTrueMeasure):
@@ -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,137 @@ def __init__(self, sampler, marginals):
 
         super(ProductMeasure, self).__init__()
 
+        self._mean_cache = None
+        self._variance_cache = None
+        self._standard_deviation_cache = None
+        self._covariance_cache = None
+
+        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):
+        if self._mean_cache is None:
+            self._mean_cache = self._concatenate_marginal_statistic("mean")
+        return self._mean_cache
+
+    @property
+    def variance(self):
+        if self._variance_cache is None:
+            self._variance_cache = self._concatenate_marginal_statistic("variance")
+        return self._variance_cache
+
+    @property
+    def standard_deviation(self):
+        if self._standard_deviation_cache is None:
+            self._standard_deviation_cache = self._concatenate_marginal_statistic(
+                "standard_deviation"
+            )
+        return self._standard_deviation_cache
+
+    def _compute_covariance(self):
+        """Build and protect the block-diagonal marginal covariance."""
+        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="csr")
+            # Rebuild from a read-only base so writes cannot be re-enabled.
+            data = self._read_only_array(covariance.data)
+            return sparse.csr_matrix(
+                (data, covariance.indices, covariance.indptr),
+                shape=covariance.shape,
+                copy=False,
+            )
+
+        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)
+
+    @property
+    def covariance(self):
+        if self._covariance_cache is None:
+            self._covariance_cache = self._compute_covariance()
+        return self._covariance_cache
+
+    def __repr__(self):
+        """Represent ProductMeasure without expanding marginal sparse matrices."""
+        lines = [f"{type(self).__name__} (AbstractTrueMeasure)"]
+        for parameter in dict.fromkeys(self.parameters):
+            if parameter == "marginals":
+                marginals = ", ".join(
+                    f"{type(marginal).__name__}(d={marginal.d})"
+                    for marginal in self.marginals
+                )
+                lines.append(f"    {parameter:<15} [{marginals}]")
+            elif parameter == "covariance" and sparse.issparse(self.covariance):
+                covariance = self.covariance
+                summary = (
+                    f"sparse {covariance.format.upper()}, "
+                    f"shape={covariance.shape}, nnz={covariance.nnz}"
+                )
+                lines.append(f"    {parameter:<15} {summary}")
+            else:
+                formatted = _univ_repr(self, "AbstractTrueMeasure", [parameter])
+                lines.extend(formatted.splitlines()[1:])
+        return "\n".join(lines)
+
     @staticmethod
     def _expand_bounds(bounds, dimension, name):
         """