Repository: vnmabus/dcor

Issue:
Energy distance using medians?
Hi, thank you for your phenomenal work writing and documenting this library.

As I'm sure you're aware, there has been some literature suggesting that an energy statistic that is more robust to outliers can be calculated by taking the median rather than mean when calculating the average distance between samples. See: James, N. A., Kejariwal, A., & Matteson, D. S. (2016). Leveraging cloud data to mitigate user experience from ‘breaking bad.’ 2016 IEEE International Conference on Big Data (Big Data), 3499–3508. https://doi.org/10.1109/BigData.2016.7841013. Specifically section 3a of that article, "Robustness against Anomalies".

From looking at this library, it seems to me that this change would be as simple as allowing a configurable "average" function which would replace the use of `mean` in this code:

https://github.com/vnmabus/dcor/blob/e7351553fb277f271ede1bf3e7148b408185707a/dcor/_energy.py#L24-L28

Would you be interested in such an implementation?

Discussion:

Patch:
diff --git a/dcor/_energy.py b/dcor/_energy.py
index 7ff1788..1167721 100644
--- a/dcor/_energy.py
+++ b/dcor/_energy.py
@@ -22,13 +22,28 @@ def _check_valid_energy_exponent(exponent):
 
 
 def _energy_distance_from_distance_matrices(
-        distance_xx, distance_yy, distance_xy):
-    """Compute energy distance with precalculated distance matrices."""
-    return (2 * np.mean(distance_xy) - np.mean(distance_xx) -
-            np.mean(distance_yy))
+        distance_xx, distance_yy, distance_xy, average=None):
+    """
+    Compute energy distance with precalculated distance matrices.
+
+    Parameters
+    ----------
+    average: Callable[[ArrayLike], float]
+        A function that will be used to calculate an average of distances.
+        This defaults to np.mean.
+
+    """
+    if average is None:
+        average = np.mean
+
+    return (
+        2 * average(distance_xy) -
+        average(distance_xx) -
+        average(distance_yy)
+    )
 
 
-def _energy_distance_imp(x, y, exponent=1):
+def _energy_distance_imp(x, y, average=None, exponent=1):
     """
     Real implementation of :func:`energy_distance`.
 
@@ -47,12 +62,13 @@ def _energy_distance_imp(x, y, exponent=1):
 
     return _energy_distance_from_distance_matrices(distance_xx=distance_xx,
                                                    distance_yy=distance_yy,
-                                                   distance_xy=distance_xy)
+                                                   distance_xy=distance_xy,
+                                                   average=average)
 
 
 def energy_distance(x, y, **kwargs):
     """
-    energy_distance(x, y, *, exponent=1)
+    energy_distance(x, y, *, exponent=1, average=None)
 
     Computes the estimator for the energy distance of the
     random vectors corresponding to :math:`x` and :math:`y`.
@@ -68,6 +84,9 @@ def energy_distance(x, y, **kwargs):
         variables while the rows are individual instances of the random vector.
     exponent: float
         Exponent of the Euclidean distance, in the range :math:`(0, 2)`.
+    average: Callable[[ArrayLike], float]
+        A function that will be used to calculate an average of distances.
+        This defaults to np.mean.
 
     Returns
     -------
diff --git a/dcor/homogeneity.py b/dcor/homogeneity.py
index 0ab7a25..bb2c4d2 100644
--- a/dcor/homogeneity.py
+++ b/dcor/homogeneity.py
@@ -22,11 +22,11 @@ def _energy_test_statistic_coefficient(n, m):
 
 
 def _energy_test_statistic_from_distance_matrices(
-        distance_xx, distance_yy, distance_xy, n, m):
+        distance_xx, distance_yy, distance_xy, n, m, average=None):
     """Test statistic with precomputed distance matrices."""
     energy_distance = _energy._energy_distance_from_distance_matrices(
         distance_xx=distance_xx, distance_yy=distance_yy,
-        distance_xy=distance_xy
+        distance_xy=distance_xy, average=average
     )
 
     return _energy_test_statistic_coefficient(n, m) * energy_distance
@@ -34,7 +34,7 @@ def _energy_test_statistic_from_distance_matrices(
 
 def energy_test_statistic(x, y, **kwargs):
     """
-    energy_test_statistic(x, y, *, exponent=1)
+    energy_test_statistic(x, y, *, exponent=1, average=None)
 
     Homogeneity statistic.
 
@@ -51,6 +51,9 @@ def energy_test_statistic(x, y, **kwargs):
         variables while the rows are individual instances of the random vector.
     exponent: float
         Exponent of the Euclidean distance, in the range :math:`(0, 2)`.
+    average: Callable[[ArrayLike], float]
+        A function that will be used to calculate an average of distances.
+        This defaults to np.mean.
 
     Returns
     -------
@@ -92,7 +95,7 @@ def energy_test_statistic(x, y, **kwargs):
 
 
 def _energy_test_statistic_multivariate_from_distance_matrix(
-        distance, indexes, sizes):
+        distance, indexes, sizes, average=None):
     """Statistic for several random vectors given the distance matrix."""
     energy = 0.0
 
@@ -110,7 +113,7 @@ def _energy_test_statistic_multivariate_from_distance_matrix(
 
             pairwise_energy = _energy_test_statistic_from_distance_matrices(
                 distance_xx=distance_xx, distance_yy=distance_yy,
-                distance_xy=distance_xy, n=n, m=m)
+                distance_xy=distance_xy, n=n, m=m, average=average)
 
             energy += pairwise_energy
 
@@ -118,7 +121,7 @@ def _energy_test_statistic_multivariate_from_distance_matrix(
 
 
 def _energy_test_imp(samples, num_resamples=0,
-                     exponent=1, random_state=None):
+                     exponent=1, random_state=None, average=None):
     """
     Real implementation of :func:`energy_test`.
 
@@ -153,7 +156,8 @@ def _energy_test_imp(samples, num_resamples=0,
         return _energy_test_statistic_multivariate_from_distance_matrix(
             distance=distance_matrix,
             indexes=sample_indexes,
-            sizes=sample_sizes)
+            sizes=sample_sizes,
+            average=average)
 
     return _hypothesis._permutation_test_with_sym_matrix(
         sample_distances,
@@ -164,7 +168,8 @@ def _energy_test_imp(samples, num_resamples=0,
 
 def energy_test(*args, **kwargs):
     """
-    energy_test(*args, num_resamples=0, exponent=1, random_state=None)
+    energy_test(*args, num_resamples=0, exponent=1, random_state=None, \
+        average=None)
 
     Test of homogeneity based on the energy distance.
 
@@ -185,6 +190,9 @@ def energy_test(*args, **kwargs):
         Exponent of the Euclidean distance, in the range :math:`(0, 2)`.
     random_state: {None, int, array_like, numpy.random.RandomState}
         Random state to generate the permutations.
+    average: Callable[[ArrayLike], float]
+        A function that will be used to calculate an average of distances.
+        This defaults to np.mean.
 
     Returns
     -------


Test patch:
diff --git a/dcor/tests/test_homogeneity.py b/dcor/tests/test_homogeneity.py
index 2fbc952..254e6dd 100644
--- a/dcor/tests/test_homogeneity.py
+++ b/dcor/tests/test_homogeneity.py
@@ -119,3 +119,67 @@ class TestEnergyTest(unittest.TestCase):
             a, b, num_resamples=num_resamples, random_state=random_state)
 
         self.assertLess(result.p_value, significance)
+
+    def test_different_means_median(self):
+        """
+        Test that the test works on the same distribution with different means,
+        using the median average.
+        """
+        num_samples = 100
+
+        random_state = np.random.RandomState(0)
+
+        a = random_state.normal(loc=0, size=(num_samples, 1))
+        b = random_state.normal(loc=1, size=(num_samples, 1))
+
+        significance = 0.01
+        num_resamples = int(3 / significance + 1)
+
+        median_result = dcor.homogeneity.energy_test(
+            a,
+            b,
+            num_resamples=num_resamples,
+            random_state=random_state,
+            average=np.median
+        )
+
+        mean_result = dcor.homogeneity.energy_test(
+            a,
+            b,
+            num_resamples=num_resamples,
+            random_state=random_state,
+            average=np.mean
+        )
+
+        # Check that we are actually using a different average
+        self.assertNotAlmostEqual(
+            mean_result.statistic,
+            median_result.statistic
+        )
+
+        # Check that we detected the heterogeneity
+        self.assertLess(median_result.p_value, significance)
+
+    def test_different_distributions_median(self):
+        """
+        Test that the test works on different distributions using the median.
+        """
+        num_samples = 100
+
+        random_state = np.random.RandomState(0)
+
+        a = random_state.normal(loc=1, size=(num_samples, 1))
+        b = random_state.exponential(scale=1, size=(num_samples, 1))
+
+        significance = 0.01
+        num_resamples = int(3 / significance + 1)
+
+        result = dcor.homogeneity.energy_test(
+            a,
+            b,
+            average=np.median,
+            num_resamples=num_resamples,
+            random_state=random_state
+        )
+
+        self.assertLess(result.p_value, significance)
diff --git a/dcor/tests/test_independence.py b/dcor/tests/test_independence.py
index 5f0c9bc..97132a1 100644
--- a/dcor/tests/test_independence.py
+++ b/dcor/tests/test_independence.py
@@ -124,4 +124,4 @@ class TestDcovTest(unittest.TestCase):
         result = dcor.independence.distance_covariance_test(
             a, b, num_resamples=num_resamples, random_state=random_state)
 
-        self.assertLess(result.p_value, significance)
\ No newline at end of file
+        self.assertLess(result.p_value, significance)