removed restriction on having >1 sample set for relatedness

Author: petrelharp

c/tskit/trees.c

@@ -4470,12 +4470,12 @@ tsk_treeseq_pi2_unbiased(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
  ***********************************/
 
 static int
-check_sample_stat_inputs(tsk_size_t num_sample_sets, tsk_size_t tuple_size,
-    tsk_size_t num_index_tuples, const tsk_id_t *index_tuples)
+check_sample_stat_inputs(tsk_size_t num_sample_sets, tsk_size_t min_num_sample_sets,
+    tsk_size_t tuple_size, tsk_size_t num_index_tuples, const tsk_id_t *index_tuples)
 {
     int ret = 0;
 
-    if (num_sample_sets < tuple_size) {
+    if (num_sample_sets < min_num_sample_sets) {
         ret = tsk_trace_error(TSK_ERR_INSUFFICIENT_SAMPLE_SETS);
         goto out;
     }
@@ -4520,7 +4520,7 @@ tsk_treeseq_divergence(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
     const double *windows, tsk_flags_t options, double *result)
 {
     int ret = 0;
-    ret = check_sample_stat_inputs(num_sample_sets, 2, num_index_tuples, index_tuples);
+    ret = check_sample_stat_inputs(num_sample_sets, 2, 2, num_index_tuples, index_tuples);
     if (ret != 0) {
         goto out;
     }
@@ -4584,7 +4584,8 @@ tsk_treeseq_genetic_relatedness(const tsk_treeseq_t *self, tsk_size_t num_sample
     const double *windows, tsk_flags_t options, double *result)
 {
     int ret = 0;
-    ret = check_sample_stat_inputs(num_sample_sets, 2, num_index_tuples, index_tuples);
+    // we want to allow self comparisons
+    ret = check_sample_stat_inputs(num_sample_sets, 1, 2, num_index_tuples, index_tuples);
     if (ret != 0) {
         goto out;
     }
@@ -4729,7 +4730,7 @@ tsk_treeseq_Y2(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
     const double *windows, tsk_flags_t options, double *result)
 {
     int ret = 0;
-    ret = check_sample_stat_inputs(num_sample_sets, 2, num_index_tuples, index_tuples);
+    ret = check_sample_stat_inputs(num_sample_sets, 2, 2, num_index_tuples, index_tuples);
     if (ret != 0) {
         goto out;
     }
@@ -4770,7 +4771,7 @@ tsk_treeseq_f2(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
     const double *windows, tsk_flags_t options, double *result)
 {
     int ret = 0;
-    ret = check_sample_stat_inputs(num_sample_sets, 2, num_index_tuples, index_tuples);
+    ret = check_sample_stat_inputs(num_sample_sets, 2, 2, num_index_tuples, index_tuples);
     if (ret != 0) {
         goto out;
     }
@@ -4816,7 +4817,7 @@ tsk_treeseq_Y3(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
     const double *windows, tsk_flags_t options, double *result)
 {
     int ret = 0;
-    ret = check_sample_stat_inputs(num_sample_sets, 3, num_index_tuples, index_tuples);
+    ret = check_sample_stat_inputs(num_sample_sets, 3, 3, num_index_tuples, index_tuples);
     if (ret != 0) {
         goto out;
     }
@@ -4859,7 +4860,7 @@ tsk_treeseq_f3(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
     const double *windows, tsk_flags_t options, double *result)
 {
     int ret = 0;
-    ret = check_sample_stat_inputs(num_sample_sets, 3, num_index_tuples, index_tuples);
+    ret = check_sample_stat_inputs(num_sample_sets, 3, 3, num_index_tuples, index_tuples);
     if (ret != 0) {
         goto out;
     }
@@ -4908,7 +4909,7 @@ tsk_treeseq_f4(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
     const double *windows, tsk_flags_t options, double *result)
 {
     int ret = 0;
-    ret = check_sample_stat_inputs(num_sample_sets, 4, num_index_tuples, index_tuples);
+    ret = check_sample_stat_inputs(num_sample_sets, 4, 4, num_index_tuples, index_tuples);
     if (ret != 0) {
         goto out;
     }

python/tests/test_tree_stats.py

@@ -2307,32 +2307,49 @@ def test_shapes(self, proportion):
         else:
             assert x.shape == (2,)
 
-
-class TestBranchGeneticRelatedness(TestGeneticRelatedness, TopologyExamplesMixin):
-    mode = "branch"
-
     def test_single_sample_set_self_comparison(self, ts_12_highrecomb_fixture):
+        if self.mode is None:
+            return
         # Test for issue #3055 - self-comparisons with single sample set
         ts = ts_12_highrecomb_fixture
         # Single sample set with self-comparison
-        result = ts.genetic_relatedness([[0]], indexes=[(0, 0)], mode="branch")
-        assert result.shape == (1,)
+        result = ts.genetic_relatedness([[0]], indexes=[(0, 0)], mode=self.mode)
+        result_shape = (ts.num_nodes, 1) if self.mode == "node" else (1,)
+        assert result.shape == result_shape
         # Should work for multiple samples in single set too
-        result = ts.genetic_relatedness([[0, 1, 2]], indexes=[(0, 0)], mode="branch")
-        assert result.shape == (1,)
+        result = ts.genetic_relatedness([[0, 1, 2]], indexes=[(0, 0)], mode=self.mode)
+        assert result.shape == result_shape
+        # Test with multiple self-comparisons
+        result = ts.genetic_relatedness(
+            [[0, 1], [2, 3]], indexes=[(0, 0), (1, 1)], mode=self.mode
+        )
+        result_shape = (ts.num_nodes, 2) if self.mode == "node" else (2,)
+        assert result.shape == result_shape
 
     def test_single_sample_set_invalid_indexes(self, ts_12_highrecomb_fixture):
+        if self.mode is None:
+            return
         # Test that invalid indexes raise ValueError with single sample set
         ts = ts_12_highrecomb_fixture
         # Index out of bounds (only have 1 sample set, but trying to access index 1)
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0]], indexes=[(0, 1)], mode="branch")
+        with pytest.raises(
+            exceptions.LibraryError, match="TSK_ERR_BAD_SAMPLE_SET_INDEX"
+        ):
+            ts.genetic_relatedness([[0]], indexes=[(0, 1)], mode=self.mode)
         # Negative index
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0]], indexes=[(-1, 0)], mode="branch")
+        with pytest.raises(
+            exceptions.LibraryError, match="TSK_ERR_BAD_SAMPLE_SET_INDEX"
+        ):
+            ts.genetic_relatedness([[0]], indexes=[(-1, 0)], mode=self.mode)
         # Both indexes out of bounds
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0, 1]], indexes=[(2, 2)], mode="branch")
+        with pytest.raises(
+            exceptions.LibraryError, match="TSK_ERR_BAD_SAMPLE_SET_INDEX"
+        ):
+            ts.genetic_relatedness([[0, 1]], indexes=[(2, 2)], mode=self.mode)
+
+
+class TestBranchGeneticRelatedness(TestGeneticRelatedness, TopologyExamplesMixin):
+    mode = "branch"
 
     @pytest.mark.parametrize("polarised", [True, False])
     def test_simple_tree_noncentred(self, polarised):
@@ -2388,61 +2405,10 @@ def test_simple_tree_noncentred(self, polarised):
 class TestNodeGeneticRelatedness(TestGeneticRelatedness, TopologyExamplesMixin):
     mode = "node"
 
-    def test_single_sample_set_self_comparison(self, ts_12_highrecomb_fixture):
-        # Test for issue #3055 - self-comparisons with single sample set
-        ts = ts_12_highrecomb_fixture
-        # Single sample set with self-comparison
-        result = ts.genetic_relatedness([[0]], indexes=[(0, 0)], mode="node")
-        assert result.shape == (ts.num_nodes, 1)
-        # Should work for multiple samples in single set too
-        result = ts.genetic_relatedness([[0, 1, 2]], indexes=[(0, 0)], mode="node")
-        assert result.shape == (ts.num_nodes, 1)
-
-    def test_single_sample_set_invalid_indexes(self, ts_12_highrecomb_fixture):
-        # Test that invalid indexes raise ValueError with single sample set
-        ts = ts_12_highrecomb_fixture
-        # Index out of bounds (only have 1 sample set, but trying to access index 1)
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0]], indexes=[(0, 1)], mode="node")
-        # Negative index
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0]], indexes=[(-1, 0)], mode="node")
-        # Both indexes out of bounds
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0, 1]], indexes=[(2, 2)], mode="node")
-
 
 class TestSiteGeneticRelatedness(TestGeneticRelatedness, MutatedTopologyExamplesMixin):
     mode = "site"
 
-    def test_single_sample_set_self_comparison(self, ts_12_highrecomb_fixture):
-        # Test for issue #3055 - self-comparisons with single sample set
-        ts = ts_12_highrecomb_fixture
-        # Single sample set with self-comparison
-        result = ts.genetic_relatedness([[0]], indexes=[(0, 0)], mode="site")
-        assert result.shape == (1,)
-        # Should work for multiple samples in single set too
-        result = ts.genetic_relatedness([[0, 1, 2]], indexes=[(0, 0)], mode="site")
-        assert result.shape == (1,)
-        # Test with multiple self-comparisons
-        result = ts.genetic_relatedness(
-            [[0], [1]], indexes=[(0, 0), (1, 1)], mode="site"
-        )
-        assert result.shape == (2,)
-
-    def test_single_sample_set_invalid_indexes(self, ts_12_highrecomb_fixture):
-        # Test that invalid indexes raise ValueError with single sample set
-        ts = ts_12_highrecomb_fixture
-        # Index out of bounds (only have 1 sample set, but trying to access index 1)
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0]], indexes=[(0, 1)], mode="site")
-        # Negative index
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0]], indexes=[(-1, 0)], mode="site")
-        # Both indexes out of bounds
-        with pytest.raises(ValueError, match="Index out of bounds"):
-            ts.genetic_relatedness([[0, 1]], indexes=[(2, 2)], mode="site")
-
     def test_match_K_c0(self):
         # This test checks that ts.genetic_relatedness() matches K_c0
         # from Speed & Balding (2014) https://www.nature.com/articles/nrg3821

python/tskit/trees.py

@@ -8106,7 +8106,6 @@ def __k_way_sample_set_stat(
         span_normalise=True,
         polarised=False,
         centre=True,
-        allow_self_comparisons=False,
     ):
         sample_set_sizes = np.array(
             [len(sample_set) for sample_set in sample_sets], dtype=np.uint32
@@ -8133,24 +8132,6 @@ def __k_way_sample_set_stat(
                 "Indexes must be convertable to a 2D numpy array with {} "
                 "columns".format(k)
             )
-        # For genetic_relatedness, we allow self-comparisons with a single sample set
-        if allow_self_comparisons and len(sample_sets) < k:
-            # Check that all indexes are valid
-            if np.any(indexes >= len(sample_sets)) or np.any(indexes < 0):
-                raise ValueError("Index out of bounds")
-            # Find which sample sets we actually need
-            unique_indexes = np.unique(indexes)
-            if np.max(unique_indexes) < len(sample_sets):
-                # we need to pad with dummy sets to satisfy the C side
-                # requirement of having at least k sets
-                sample_sets = list(sample_sets)
-                sample_set_sizes = list(sample_set_sizes)
-                while len(sample_sets) < k:
-                    # Add a dummy sample set that won't be used
-                    sample_sets.append(np.array([sample_sets[0][0]], dtype=np.int32))
-                    sample_set_sizes.append(1)
-                sample_set_sizes = np.array(sample_set_sizes, dtype=np.uint32)
-                flattened = util.safe_np_int_cast(np.hstack(sample_sets), np.int32)
         stat = self.__run_windowed_stat(
             windows,
             ll_method,
@@ -8721,7 +8702,6 @@ def genetic_relatedness(
             span_normalise=span_normalise,
             polarised=polarised,
             centre=centre,
-            allow_self_comparisons=True,
         )
         if proportion:
             # TODO this should be done in C also