feat: Add BatchCorr for batch-wise correlation calculation

Author: google-labs-jules[bot]

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

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)