feat: Implement BatchWeightedSum and BatchWeightedMean with full axis and merge support

This commit introduces `BatchWeightedSum` and `BatchWeightedMean` classes to compute weighted statistics over batches of data.

This final version incorporates several rounds of feedback to improve robustness and design:

- **Core Implementation**: Added `BatchWeightedSum` and `BatchWeightedMean` classes.
- **Axis Support**: The classes now support summation over any axis, including non-batch axes, by correctly handling batch-wise results.
- **Weight Broadcasting**: The `update_batch` method accepts weights that can be broadcast to the data's shape.
- **Design Refactoring**: `BatchWeightedMean` is refactored to use a `BatchWeightedSum` instance for its `sum_of_weights` calculation, improving code reuse and simplifying the design.
- **Robustness**: Added a consistency check to prevent mixing different weight shapes. Division-by-zero in `BatchWeightedMean` is handled gracefully.
- **Testing**: The test suite is comprehensive, using 3D data and parameterizing for various axes, weight shapes, and testing the merge (`+`) operation.
- **Documentation**: All new classes are included in the documentation.

Author: google-labs-jules[bot]

batchstats/__init__.py

@@ -12,6 +12,8 @@
     BatchStd,
     BatchSum,
     BatchVar,
+    BatchWeightedMean,
+    BatchWeightedSum,
 )
 
 __all__ = [
@@ -25,6 +27,8 @@
     "BatchStd",
     "BatchSum",
     "BatchVar",
+    "BatchWeightedMean",
+    "BatchWeightedSum",
     "BatchNanMax",
     "BatchNanMean",
     "BatchNanMin",

batchstats/stats/__init__.py

@@ -7,6 +7,8 @@
 from .std import BatchStd
 from .sum import BatchSum
 from .var import BatchVar
+from .weighted_mean import BatchWeightedMean
+from .weighted_sum import BatchWeightedSum
 
 __all__ = [
     "BatchCorr",
@@ -18,4 +20,6 @@
     "BatchStd",
     "BatchSum",
     "BatchVar",
+    "BatchWeightedMean",
+    "BatchWeightedSum",
 ]

batchstats/stats/weighted_mean.py

@@ -0,0 +1,74 @@
+import numpy as np
+
+from .._misc import NoValidSamplesError
+from ..base.batch_stat import BatchStat
+from .weighted_sum import BatchWeightedSum
+
+
+class BatchWeightedMean(BatchStat):
+    """
+    Class for calculating the weighted mean of batches of data.
+    It computes `sum(w*x) / sum(w)` by using two `BatchWeightedSum` instances.
+    """
+
+    def __init__(self, axis=0):
+        super().__init__(axis=axis)
+        # For calculating sum(w*x)
+        self.weighted_sum = BatchWeightedSum(axis=axis)
+        # For calculating sum(w)
+        self.sum_of_weights = BatchWeightedSum(axis=axis)
+        self.n_samples = None # n_samples is not used in weighted stats
+
+    def update_batch(self, batch, weights):
+        """
+        Update the weighted mean with a new batch of data.
+        """
+        # Update sum(w*x)
+        self.weighted_sum.update_batch(batch, weights=weights)
+
+        # Update sum(w) by calculating sum(broadcasted_w * 1)
+        broadcasted_weights = np.broadcast_to(weights, batch.shape)
+        self.sum_of_weights.update_batch(batch=broadcasted_weights, weights=1)
+
+        return self
+
+    def __call__(self) -> np.ndarray:
+        """
+        Calculate the weighted mean.
+        """
+        try:
+            # Calculate sum(w*x) and sum(w)
+            w_sum = self.weighted_sum()
+            sow = self.sum_of_weights()
+
+            # To prevent division by zero, where sow is zero, we should return NaN or raise an error.
+            # Using np.errstate to handle division by zero gracefully.
+            with np.errstate(divide='ignore', invalid='ignore'):
+                mean_ = w_sum / sow
+                if hasattr(mean_, "dtype") and np.issubdtype(mean_.dtype, np.floating):
+                    mean_[sow == 0] = np.nan # Set to NaN where sum of weights is zero
+
+        except NoValidSamplesError:
+            # Re-raise with a more specific message if no samples were processed
+            raise NoValidSamplesError("No valid samples for calculating weighted mean.")
+
+        return mean_
+
+    def __add__(self, other):
+        """
+        Merge two BatchWeightedMean objects.
+        """
+        # Basic checks
+        if type(self) != type(other):
+            from .._misc import DifferentStatsError
+            raise DifferentStatsError()
+        if self.axis != other.axis:
+            from .._misc import DifferentAxisError
+            raise DifferentAxisError()
+
+        # Create a new object and merge the internal calculators
+        ret = BatchWeightedMean(axis=self.axis)
+        ret.weighted_sum = self.weighted_sum + other.weighted_sum
+        ret.sum_of_weights = self.sum_of_weights + other.sum_of_weights
+
+        return ret

batchstats/stats/weighted_sum.py

@@ -0,0 +1,133 @@
+import numpy as np
+
+from .._misc import NoValidSamplesError
+from ..base.batch_stat import BatchStat
+
+
+class BatchWeightedSum(BatchStat):
+    """
+    Class for calculating the weighted sum of batches of data.
+
+    The algorithm used is a simple cumulative sum. Each time a new batch is added,
+    the weighted sum of the new batch is added to the existing sum.
+
+    .. code:: python
+
+        import numpy as np
+        from batchstats.stats.weighted_sum import BatchWeightedSum
+
+        # create some data
+        data1 = np.array([[1, 2], [3, 4]])
+        weights1 = np.array([[0.1, 0.2], [0.3, 0.4]])
+        data2 = np.array([[5, 6], [7, 8]])
+        weights2 = np.array([[0.5, 0.6], [0.7, 0.8]])
+
+        # create a BatchWeightedSum object
+        bws = BatchWeightedSum()
+
+        # update with the first batch
+        bws.update_batch(data1, weights=weights1)
+
+        # update with the second batch
+        bws.update_batch(data2, weights=weights2)
+
+        # get the weighted sum
+        total_weighted_sum = bws()
+
+        # verify the result
+        expected_sum = np.array([8.4, 12.0])
+        np.testing.assert_allclose(total_weighted_sum, expected_sum)
+    """
+
+    def __init__(self, axis=0):
+        super().__init__(axis=axis)
+        self.sum = None
+        self.n_samples = None
+        self._weights_pattern = None
+
+        axis_tuple = self.axis if isinstance(self.axis, tuple) else (self.axis,)
+        # If batch axis 0 is not summed, we need to collect results in a list
+        self._is_list_mode = 0 not in axis_tuple
+        if self._is_list_mode:
+            self.sum = []
+
+    def update_batch(self, batch, weights):
+        """
+        Update the weighted sum with a new batch of data.
+        """
+        weights = np.asarray(weights)
+
+        axis_tuple = self.axis if isinstance(self.axis, tuple) else (self.axis,)
+        axis_tuple = tuple(ax if ax >= 0 else ax + batch.ndim for ax in axis_tuple)
+
+        # Check for consistent weight shapes
+        # The pattern ignores the batch axis (0) and any summed axes
+        axes_to_ignore = set(axis_tuple)
+        axes_to_ignore.add(0) # always ignore batch axis
+        current_pattern = tuple(
+            s for i, s in enumerate(weights.shape) if i not in axes_to_ignore
+        )
+        if self._weights_pattern is None:
+            self._weights_pattern = current_pattern
+        elif self._weights_pattern != current_pattern:
+            raise ValueError(
+                f"Inconsistent weights shape pattern. "
+                f"Expected pattern for non-summed axes: {self._weights_pattern}, "
+                f"but got {current_pattern}."
+            )
+
+        batch_sum = np.sum(a=batch * weights, axis=self.axis, keepdims=True)
+
+        if self._is_list_mode:
+            self.sum.append(batch_sum)
+        else:
+            if self.sum is None:
+                self.sum = batch_sum
+            else:
+                self.sum += batch_sum
+        return self
+
+    def __call__(self) -> np.ndarray:
+        """
+        Calculate the weighted sum.
+        """
+        if self.sum is None or (self._is_list_mode and not self.sum):
+            raise NoValidSamplesError("No valid samples for calculating weighted sum.")
+
+        if self._is_list_mode:
+            sum_ = np.concatenate(self.sum, axis=0)
+        else:
+            sum_ = self.sum.copy()
+
+        if self.axis is not None:
+            return sum_.squeeze(axis=self.axis)
+        return sum_
+
+    def __add__(self, other):
+        # Basic checks from merge_test
+        if type(self) != type(other):
+            from .._misc import DifferentStatsError
+            raise DifferentStatsError()
+        if self.axis != other.axis:
+            from .._misc import DifferentAxisError
+            raise DifferentAxisError()
+
+        if self._is_list_mode:
+            if not self.sum:
+                return other
+            if not other.sum:
+                return self
+            ret = BatchWeightedSum(axis=self.axis)
+            ret.sum = self.sum + other.sum
+            ret._weights_pattern = self._weights_pattern
+            return ret
+        else:
+            self.merge_test(other, field="sum") # ok to call here
+            if self.sum is None:
+                return other
+            elif other.sum is None:
+                return self
+            ret = BatchWeightedSum(axis=self.axis)
+            ret.sum = self.sum + other.sum
+            ret._weights_pattern = self._weights_pattern
+            return ret

docs/source/core_classes.rst

@@ -30,3 +30,14 @@ These are the standard classes for computing statistics on datasets. If your dat
 
 .. autoclass:: batchstats.BatchPeakToPeak
    :members: __init__, update_batch, __call__
+
+Weighted Statistics
+===================
+
+These classes are used for computing weighted statistics on datasets.
+
+.. autoclass:: batchstats.BatchWeightedSum
+   :members: __init__, update_batch, __call__
+
+.. autoclass:: batchstats.BatchWeightedMean
+   :members: __init__, update_batch, __call__

tests/test_weighted_stats.py

@@ -0,0 +1,130 @@
+import numpy as np
+import pytest
+
+from batchstats.stats.weighted_mean import BatchWeightedMean
+from batchstats.stats.weighted_sum import BatchWeightedSum
+
+
+@pytest.fixture
+def data_3d():
+    m, n, p = 100, 50, 10
+    return 1e1 * np.random.randn(m, n, p) + 1e3
+
+
+axis_weight_scenarios = [
+    (0, "full"),
+    (0, "broadcast_row"),
+    (1, "full"),
+    (1, "broadcast_col"),
+    (2, "full"),
+    ((1, 2), "full"),
+    ((1, 2), "broadcast_plane"),
+]
+
+
+@pytest.fixture(params=axis_weight_scenarios)
+def scenario(request):
+    return request.param
+
+
+@pytest.fixture
+def axis(scenario):
+    return scenario[0]
+
+
+@pytest.fixture
+def weights_type(scenario):
+    return scenario[1]
+
+
+@pytest.fixture
+def weights_3d(data_3d, axis, weights_type):
+    shape = data_3d.shape
+    axis_tuple = (axis,) if isinstance(axis, int) else axis
+
+    if weights_type == "full":
+        return np.random.rand(*shape)
+
+    w_shape = list(shape)
+    # This logic is a bit naive, but covers the test cases
+    if weights_type == "broadcast_row": # axis 0
+        w_shape = [shape[0], 1, 1]
+    elif weights_type == "broadcast_col": # axis 1
+        w_shape = [1, shape[1], 1]
+    elif weights_type == "broadcast_plane": # axis (1,2)
+        w_shape = [1, shape[1], shape[2]]
+
+    return np.random.rand(*w_shape)
+
+
+@pytest.fixture
+def n_batches():
+    return 13
+
+
+@pytest.mark.parametrize("klass", [BatchWeightedSum, BatchWeightedMean])
+def test_weighted_stats_3d(data_3d, n_batches, axis, weights_3d, klass):
+    if klass == BatchWeightedSum:
+        true_stat = np.sum(data_3d * weights_3d, axis=axis)
+    else:
+        broadcasted_weights = np.broadcast_to(weights_3d, data_3d.shape)
+        true_stat = np.sum(data_3d * weights_3d, axis=axis) / np.sum(broadcasted_weights, axis=axis)
+
+    batch_op = klass(axis=axis)
+
+    data_batches = np.array_split(data_3d, n_batches, axis=0)
+
+    if weights_3d.shape[0] > 1:
+        weights_batches = np.array_split(weights_3d, n_batches, axis=0)
+    else:
+        weights_batches = [weights_3d] * n_batches
+
+    for batch_data, batch_weights in zip(data_batches, weights_batches):
+        batch_op.update_batch(batch=batch_data, weights=batch_weights)
+
+    batch_stat = batch_op()
+    assert np.allclose(true_stat, batch_stat)
+
+
+@pytest.mark.parametrize("klass", [BatchWeightedSum, BatchWeightedMean])
+def test_weighted_merge_3d(data_3d, axis, weights_3d, klass):
+    if klass == BatchWeightedSum:
+        true_stat = np.sum(data_3d * weights_3d, axis=axis)
+    else:
+        broadcasted_weights = np.broadcast_to(weights_3d, data_3d.shape)
+        true_stat = np.sum(data_3d * weights_3d, axis=axis) / np.sum(broadcasted_weights, axis=axis)
+
+    # Split data and weights into two halves
+    d1, d2 = np.array_split(data_3d, 2, axis=0)
+    if weights_3d.shape[0] > 1:
+        w1, w2 = np.array_split(weights_3d, 2, axis=0)
+    else:
+        w1, w2 = weights_3d, weights_3d
+
+    # Create and update two separate objects
+    op1 = klass(axis=axis)
+    op1.update_batch(d1, w1)
+
+    op2 = klass(axis=axis)
+    op2.update_batch(d2, w2)
+
+    # Merge them
+    merged_op = op1 + op2
+
+    merged_stat = merged_op()
+    assert np.allclose(true_stat, merged_stat)
+
+
+def test_inconsistent_weights_shape_raises_error(data_3d):
+    bws = BatchWeightedSum(axis=1) # Sum over a non-batch axis
+
+    # First batch with per-column weights
+    batch1 = data_3d[:10]
+    weights1 = np.random.rand(1, data_3d.shape[1], 1)
+    bws.update_batch(batch1, weights1)
+
+    # Second batch with per-plane weights
+    batch2 = data_3d[10:20]
+    weights2 = np.random.rand(1, data_3d.shape[1], data_3d.shape[2])
+    with pytest.raises(ValueError, match="Inconsistent weights shape pattern"):
+        bws.update_batch(batch2, weights2)