From 833aea9861ef12be2c0067a83871e99de8cafc73 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 13 Aug 2025 09:13:00 +0000 Subject: [PATCH 1/2] feat: Add BatchCorr for batch-wise correlation calculation --- README.md | 1 + batchstats/__init__.py | 2 + batchstats/stats/__init__.py | 2 + batchstats/stats/corr.py | 150 +++++++++++++++++++++++++++++++++++ docs/source/core_classes.rst | 3 + tests/test_stats.py | 27 ++++++- 6 files changed, 184 insertions(+), 1 deletion(-) create mode 100644 batchstats/stats/corr.py diff --git a/README.md b/README.md index fc565c2..913d81e 100644 --- a/README.md +++ b/README.md @@ -107,5 +107,6 @@ print(f"NaN-aware mean shape: {nan_mean.shape}") * `BatchVar` * `BatchStd` * `BatchCov` +* `BatchCorr` For more details on each class, see the [API Reference](https://batchstats.readthedocs.io/en/latest/api.html). diff --git a/batchstats/__init__.py b/batchstats/__init__.py index 6af030a..432ccee 100644 --- a/batchstats/__init__.py +++ b/batchstats/__init__.py @@ -3,6 +3,7 @@ from .base import BatchNanStat, BatchStat from .nanstats import BatchNanMax, BatchNanMean, BatchNanMin, BatchNanPeakToPeak, BatchNanSum from .stats import ( + BatchCorr, BatchCov, BatchMax, BatchMean, @@ -14,6 +15,7 @@ ) __all__ = [ + "BatchCorr", "BatchCov", "BatchMax", "BatchMean", diff --git a/batchstats/stats/__init__.py b/batchstats/stats/__init__.py index c3fd678..950c76d 100644 --- a/batchstats/stats/__init__.py +++ b/batchstats/stats/__init__.py @@ -1,3 +1,4 @@ +from .corr import BatchCorr from .cov import BatchCov from .max import BatchMax from .mean import BatchMean @@ -8,6 +9,7 @@ from .var import BatchVar __all__ = [ + "BatchCorr", "BatchCov", "BatchMax", "BatchMean", diff --git a/batchstats/stats/corr.py b/batchstats/stats/corr.py new file mode 100644 index 0000000..e4fb553 --- /dev/null +++ b/batchstats/stats/corr.py @@ -0,0 +1,150 @@ +import numpy as np + +from .._misc import NoValidSamplesError, UnequalSamplesNumber, any_nan +from ..base import BatchStat +from .cov import BatchCov +from .var import BatchVar + + +class BatchCorr(BatchStat): + """ + Class for calculating the correlation of batches of data. + + The algorithm is an implementation of an online correlation calculation. + It is numerically stable and avoids a two-pass approach. + + Args: + ddof (int, optional): Means Delta Degrees of Freedom. The divisor used in calculations is N - ddof, where N represents the number of elements. By default ddof is zero. + + .. code:: python + + import numpy as np + from batchstats import BatchCorr + + # create some data + data1 = np.array([[1, 2], [3, 4]]) + data2 = np.array([[5, 6], [7, 8]]) + + # create a BatchCorr object + bc = BatchCorr() + + # update with the first batch + bc.update_batch(data1) + + # update with the second batch + bc.update_batch(data2) + + # get the correlation + total_corr = bc() + + """ + + def __init__(self, ddof=0): + super().__init__() + self.cov = BatchCov(ddof=ddof) + self.var1 = BatchVar(ddof=ddof) + self.var2 = BatchVar(ddof=ddof) + self.ddof = ddof + self._is_1d = None + + def update_batch(self, batch1, batch2=None, assume_valid=False): + """ + Update the correlation with new batches of data. + + Args: + batch1 (numpy.ndarray): Input batch 1. + batch2 (numpy.ndarray, optional): Input batch 2. Default is None. + assume_valid (bool, optional): If True, assumes all elements in the batches are valid. Default is False. + + Returns: + BatchCorr: Updated BatchCorr object. + + """ + is_1d = batch2 is None + if self._is_1d is None: + self._is_1d = is_1d + elif self._is_1d != is_1d: + raise ValueError("Inconsistent use of BatchCorr. Cannot mix calls with and without batch2.") + + _batch1 = batch1 + _batch2 = batch2 + if self._is_1d: + _batch2 = _batch1 + + if not assume_valid: + if _batch2 is None: + _batch1 = _batch2 = np.atleast_2d(np.asarray(_batch1)) + else: + _batch1, _batch2 = np.atleast_2d(np.asarray(_batch1)), np.atleast_2d(np.asarray(_batch2)) + + if len(_batch1) != len(_batch2): + raise UnequalSamplesNumber("batch1 and batch2 don't have the same lengths.") + + mask = ~any_nan(_batch1, axis=1) & ~any_nan(_batch2, axis=1) + + _batch1 = _batch1[mask] + _batch2 = _batch2[mask] + assume_valid = True + + if len(_batch1) > 0: + self.cov.update_batch(_batch1, _batch2, assume_valid=assume_valid) + self.var1.update_batch(_batch1, assume_valid=assume_valid) + if self._is_1d: + self.var2 = self.var1 + else: + self.var2.update_batch(_batch2, assume_valid=assume_valid) + self.n_samples = self.cov.n_samples + + return self + + def __call__(self) -> np.ndarray: + """ + Calculate the correlation. + + Returns: + numpy.ndarray: Correlation of the batches. + + Raises: + NoValidSamplesError: If no valid samples are available. + """ + if self.cov.cov is None: + raise NoValidSamplesError("No valid samples for calculating correlation.") + + cov = self.cov() + var1 = self.var1() + var2 = self.var2() + + std1 = np.sqrt(var1) + std2 = np.sqrt(var2) + + with np.errstate(divide='ignore', invalid='ignore'): + corr = cov / np.outer(std1, std2) + + return corr + + def __add__(self, other): + """ + Merge two BatchCorr objects. + """ + self.merge_test(other) + + if self.n_samples == 0: + return other + elif other.n_samples == 0: + return self + + if self._is_1d != other._is_1d: + raise ValueError("Cannot merge BatchCorr objects with different setups (1d vs 2d).") + + ret = BatchCorr(ddof=self.ddof) + ret.cov = self.cov + other.cov + ret.var1 = self.var1 + other.var1 + + if self._is_1d: + ret.var2 = ret.var1 + else: + ret.var2 = self.var2 + other.var2 + + ret.n_samples = ret.cov.n_samples + ret._is_1d = self._is_1d + return ret diff --git a/docs/source/core_classes.rst b/docs/source/core_classes.rst index a3a0b3c..0cd834d 100644 --- a/docs/source/core_classes.rst +++ b/docs/source/core_classes.rst @@ -19,6 +19,9 @@ These are the standard classes for computing statistics on datasets. If your dat .. autoclass:: batchstats.BatchCov :members: __init__, update_batch, __call__ +.. autoclass:: batchstats.BatchCorr + :members: __init__, update_batch, __call__ + .. autoclass:: batchstats.BatchMin :members: __init__, update_batch, __call__ diff --git a/tests/test_stats.py b/tests/test_stats.py index 5484e40..23086aa 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from batchstats import BatchCov, BatchMax, BatchMean, BatchMin, BatchPeakToPeak, BatchStd, BatchSum, BatchVar +from batchstats import BatchCorr, BatchCov, BatchMax, BatchMean, BatchMin, BatchPeakToPeak, BatchStd, BatchSum, BatchVar @pytest.fixture @@ -133,6 +133,31 @@ def test_cov_2(data, n_batches): assert np.allclose(cov, true_cov[:, index]) +def test_corr(data, n_batches): + true_corr = np.corrcoef(data.T) + + batchcorr = BatchCorr() + for batch_data in np.array_split(data, n_batches): + batchcorr.update_batch(batch_data) + + corr = batchcorr() + + assert np.allclose(corr, true_corr) + + +def test_corr_2(data, n_batches): + index = np.arange(25) + true_corr = np.corrcoef(data.T, data[:, index].T)[:50, 50:] + + batchcorr = BatchCorr() + for batch_data in np.array_split(data, n_batches): + batchcorr.update_batch(batch_data, batch_data[:, index]) + + corr = batchcorr() + + assert np.allclose(corr, true_corr) + + def test_mean_2d_features(data_2d_features, n_batches): true_stat = np.mean(data_2d_features, axis=0) From 32ed99dbf15e1c2d5251ed4226e34aaf2f694377 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 13 Aug 2025 09:32:17 +0000 Subject: [PATCH 2/2] feat: Add BatchCorr for batch-wise correlation calculation --- batchstats/stats/corr.py | 18 ++++++++++++++++-- tests/test_merge.py | 4 ++-- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/batchstats/stats/corr.py b/batchstats/stats/corr.py index e4fb553..74e8837 100644 --- a/batchstats/stats/corr.py +++ b/batchstats/stats/corr.py @@ -1,6 +1,13 @@ import numpy as np -from .._misc import NoValidSamplesError, UnequalSamplesNumber, any_nan +from .._misc import ( + DifferentAxisError, + DifferentShapesError, + DifferentStatsError, + NoValidSamplesError, + UnequalSamplesNumber, + any_nan, +) from ..base import BatchStat from .cov import BatchCov from .var import BatchVar @@ -126,7 +133,14 @@ def __add__(self, other): """ Merge two BatchCorr objects. """ - self.merge_test(other) + if type(self) != type(other): + raise DifferentStatsError() + if self.axis != other.axis: + raise DifferentAxisError() + + if self.cov.cov is not None and other.cov.cov is not None: + if self.cov.cov.shape != other.cov.cov.shape: + raise DifferentShapesError() if self.n_samples == 0: return other diff --git a/tests/test_merge.py b/tests/test_merge.py index 3c30988..f4c0c2d 100644 --- a/tests/test_merge.py +++ b/tests/test_merge.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from batchstats import BatchCov, BatchMax, BatchMean, BatchMin, BatchPeakToPeak, BatchStd, BatchSum, BatchVar +from batchstats import BatchCorr, BatchCov, BatchMax, BatchMean, BatchMin, BatchPeakToPeak, BatchStd, BatchSum, BatchVar @pytest.fixture @@ -13,7 +13,7 @@ def data(): def test_merge(data): data1, data2 = data data0 = np.concatenate([data1, data2]) - for stat in (BatchCov, BatchMax, BatchMean, BatchMin, BatchPeakToPeak, BatchStd, BatchSum, BatchVar): + for stat in (BatchCorr, BatchCov, BatchMax, BatchMean, BatchMin, BatchPeakToPeak, BatchStd, BatchSum, BatchVar): s0 = stat().update_batch(data0) s1 = stat().update_batch(data1) s2 = stat().update_batch(data2)