preserve native dimensions instead of expanding at the end

Author: lkirk

python/tests/test_ld_matrix.py

@@ -2511,34 +2511,30 @@ def r2_ij(X, n):
         pAB, pAb, paB = X / n
         pA = pAb + pAB
         pB = paB + pAB
-        D2_ij = np.prod(pAB - (pA * pB))
-        denom = np.prod(np.sqrt(pA * pB * (1 - pA) * (1 - pB)))
+        D2_ij = np.prod(pAB - (pA * pB), keepdims=True)
+        denom = np.prod(np.sqrt(pA * pB * (1 - pA) * (1 - pB)), keepdims=True)
         with suppress_overflow_div0_warning():
-            return np.expand_dims(D2_ij / denom, axis=0)
+            return D2_ij / denom
 
     @staticmethod
     def D2_ij(X, n):
         pAB, pAb, paB = X / n
         pA = pAb + pAB
         pB = paB + pAB
-        return np.expand_dims(np.prod(pAB - (pA * pB)), axis=0)
+        return np.prod(pAB - (pA * pB), keepdims=True)
 
     @staticmethod
     def D2_ij_unbiased(X, n):
-        """
-        NB: the two sample sets must be disjoint
-            we have no way for testing equality
-        """
+        """NB: We use double brackets here to preserve the output shape of (1,)"""
         AB, Ab, aB = X
         ab = n - X.sum(0)
-        return np.expand_dims(
-            (Ab[0] * aB[0] - AB[0] * ab[0])
-            * (Ab[1] * aB[1] - AB[1] * ab[1])
-            / n[0]
-            / (n[0] - 1)
-            / n[1]
-            / (n[1] - 1),
-            axis=0,
+        return (
+            (Ab[[0]] * aB[[0]] - AB[[0]] * ab[[0]])
+            * (Ab[[1]] * aB[[1]] - AB[[1]] * ab[[1]])
+            / n[[0]]
+            / (n[[0]] - 1)
+            / n[[1]]
+            / (n[[1]] - 1)
         )
 
 
@@ -2624,7 +2620,7 @@ def test_general_one_way_two_locus_stat_multiallelic(stat):
     elif stat in {"D", "r", "D_prime"}:
         result = ts.two_locus_count_stat([ts.samples()], func, 1, polarised=True)
     else:
-        # default norm func is lambda X, n, nA, nB: np.expand_dims(1 / (nA * nB), axis=0)
+        # default norm func is `lambda X, n, nA, nB: 1 / (nA * nB)[np.newaxis,]`
         result = ts.two_locus_count_stat([ts.samples()], func, 1)
     np.testing.assert_array_almost_equal(ts.ld_matrix(stat=stat), result)
 
@@ -2641,13 +2637,14 @@ def test_general_two_way_two_locus_stat_multiallelic(stat):
     ts = tsutil.all_fields_ts()
     func = getattr(GeneralStatFuncs, stat)
     if stat == "r2_ij":
-
-        def norm_f(X, n, nA, nB):
-            return np.expand_dims(X[0].sum() / n.sum(), axis=0)
-
-        result = ts.two_locus_count_stat([ts.samples(), ts.samples()], func, 1, norm_f)
+        result = ts.two_locus_count_stat(
+            [ts.samples(), ts.samples()],
+            func,
+            1,
+            lambda X, n, nA, nB: X[0].sum(keepdims=True) / n.sum(),
+        )
     else:
-        # default norm func is lambda X, n, nA, nB: np.expand_dims(1 / (nA * nB), axis=0)
+        # default norm func is `lambda X, n, nA, nB: 1 / (nA * nB)[np.newaxis,]`
         result = ts.two_locus_count_stat([ts.samples(), ts.samples()], func, 1)
     np.testing.assert_array_almost_equal(
         ts.ld_matrix(

python/tskit/trees.py

@@ -10936,7 +10936,7 @@ def two_locus_count_stat(
         sample_sets,
         f,
         result_dim,
-        norm_f=lambda X, n, nA, nB: np.expand_dims(1 / (nA * nB), axis=0),
+        norm_f=lambda X, n, nA, nB: 1 / (nA * nB)[np.newaxis,],
         polarised=False,
         sites=None,
         positions=None,