From aa3bcd4d36d97f8b7b515fb873320f576484d3ca Mon Sep 17 00:00:00 2001
From: Mohit Arvind Khakharia <Mohit-Ak@users.noreply.github.com>
Date: Fri, 30 Oct 2026 02:29:04 -0400
Subject: [PATCH] Map sample node IDs to rows in pca

_expand_indices used node IDs as row indices into the weight matrix,
which has one row per sample in ts.samples() order, so pca raised
IndexError whenever the samples were not nodes 0..n-1.

Fixes #3443
Fixes #3365
---
 python/CHANGELOG.rst                    |  2 ++
 python/tests/test_relatedness_vector.py | 31 +++++++++++++++++++++++++
 python/tskit/trees.py                   |  7 +++++-
 3 files changed, 39 insertions(+), 1 deletion(-)

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