diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index ee9bf8a2cf..a455a8d962 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -30,6 +30,8 @@ - ``TableCollection.delete_older`` now checks table integrity first, raising an error for invalid node or mutation references instead of reading out of bounds (:user:`jeromekelleher`, :issue:`3486`). +- ``TreeSequence.pca`` no longer fails when the sample nodes are not + ``0, ..., n - 1`` (:user:`Mohit-Ak`, :issue:`3443`, :issue:`3365`). -------------------- diff --git a/python/tests/test_relatedness_vector.py b/python/tests/test_relatedness_vector.py index 722e168759..d78b1d69b5 100644 --- a/python/tests/test_relatedness_vector.py +++ b/python/tests/test_relatedness_vector.py @@ -1329,6 +1329,37 @@ def test_err_samples(self, centre, num_windows): time_windows=[time_low, time_high], ) + @pytest.mark.parametrize("centre", (True, False)) + def test_samples_not_first_nodes(self, centre): + # Renumbering nodes so that samples are not 0..n-1 must not change the PCA + ts = msprime.sim_ancestry( + 10, + ploidy=2, + population_size=20, + sequence_length=100, + recombination_rate=0.01, + random_seed=12345, + ) + is_sample = np.isin(np.arange(ts.num_nodes), ts.samples()) + non_samples = np.where(~is_sample)[0] + order = np.concatenate([non_samples[:5], ts.samples(), non_samples[5:]]) + new_ts = ts.subset(order) + new_id = np.argsort(order) + assert np.all(new_ts.samples() == new_id[ts.samples()]) + assert new_ts.samples()[0] == 5 + samples = ts.samples()[[3, 0, 7, 12, 5, 9, 15]] + for kwargs, new_kwargs in [ + ({}, {}), + ({"samples": samples}, {"samples": new_id[samples]}), + ({"individuals": [4, 1, 7, 2]}, {"individuals": [4, 1, 7, 2]}), + ]: + pc = ts.pca(num_components=3, centre=centre, random_seed=1, **kwargs) + new_pc = new_ts.pca( + num_components=3, centre=centre, random_seed=1, **new_kwargs + ) + np.testing.assert_allclose(new_pc.factors, pc.factors) + np.testing.assert_allclose(new_pc.eigenvalues, pc.eigenvalues) + @pytest.mark.parametrize("centre", (True, False)) def test_individuals_matches_samples(self, centre): # ploidy 1 individuals should be the same as samples diff --git a/python/tskit/trees.py b/python/tskit/trees.py index f79b35a11d..3ac5b858ea 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -9233,8 +9233,13 @@ def genetic_relatedness_vector( return out def _expand_indices(self, x: np.ndarray, indices: np.ndarray) -> np.ndarray: + sample_index = np.full(self.num_nodes, -1) + sample_index[self.samples()] = np.arange(self.num_samples) + rows = sample_index[indices] + if np.any(rows < 0): + raise ValueError("All nodes must be samples.") y = np.zeros((self.num_samples, x.shape[1])) - y[indices] = x + y[rows] = x return y