Fix pca when the sample nodes are not 0..n-1

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`).
 
 
 --------------------

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

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