Remap individual ids after simplification

c/tests/test_trees.c

@@ -1076,9 +1076,9 @@ test_simplest_unary_with_individuals(void)
                         "1  2   6   4\n";
     const char *individuals = "0    0.5     -1,-1\n"
                               "0    1.5,3.1 -1,-1\n"
-                              "0    2.1     -1,-1\n"
-                              "0    3.2     -1,-1\n"
-                              "0    4.2     -1,-1\n";
+                              "0    2.1     0,1\n"
+                              "0    3.2     1,2\n"
+                              "0    4.2     2,3\n";
     const char *nodes_expect = "1  0   0   -1\n"
                                "1  0   0   0\n"
                                "0  1   0   1\n"
@@ -1091,8 +1091,8 @@ test_simplest_unary_with_individuals(void)
                                "1  2   5   4\n";
     const char *individuals_expect = "0    0.5     -1,-1\n"
                                      "0    1.5,3.1 -1,-1\n"
-                                     "0    2.1     -1,-1\n"
-                                     "0    3.2     -1,-1\n";
+                                     "0    2.1     0,1\n"
+                                     "0    3.2     1,2\n";
     tsk_treeseq_t ts, simplified, expected;
     tsk_id_t sample_ids[] = { 0, 1 };
 

c/tskit/tables.c

@@ -7352,6 +7352,14 @@ simplifier_finalise_references(simplifier_t *self)
         }
     }
 
+    /* Remap parent IDs */
+    for (j = 0; j < self->tables->individuals.parents_length; j++) {
+        self->tables->individuals.parents[j]
+            = self->tables->individuals.parents[j] == TSK_NULL
+                  ? TSK_NULL
+                  : individual_id_map[self->tables->individuals.parents[j]];
+    }
+
     /* Remap node IDs referencing the above */
     for (j = 0; j < num_nodes; j++) {
         pop_id = node_population[j];

python/tests/simplify.py

@@ -453,7 +453,10 @@ def finalise_references(self):
             if count > 0:
                 row = input_individuals[input_id]
                 output_id = self.tables.individuals.add_row(
-                    flags=row.flags, location=row.location, metadata=row.metadata
+                    flags=row.flags,
+                    location=row.location,
+                    parents=row.parents,
+                    metadata=row.metadata,
                 )
                 individual_id_map[input_id] = output_id
 
@@ -468,6 +471,23 @@ def finalise_references(self):
             population=population_id_map[nodes.population],
         )
 
+        # Remap the parent ids of individuals
+        individuals_copy = self.tables.individuals.copy()
+        self.tables.individuals.clear()
+        for row in individuals_copy:
+            mapped_parents = []
+            for p in row.parents:
+                if p == -1:
+                    mapped_parents.append(-1)
+                else:
+                    mapped_parents.append(individual_id_map[p])
+            self.tables.individuals.add_row(
+                flags=row.flags,
+                location=row.location,
+                parents=mapped_parents,
+                metadata=row.metadata,
+            )
+
         # We don't support migrations for now. We'll need to remap these as well.
         assert self.ts.num_migrations == 0
 

python/tests/test_tables.py

@@ -31,6 +31,7 @@
 import pickle
 import platform
 import random
+import struct
 import time
 import unittest
 import warnings
@@ -2361,6 +2362,109 @@ def test_samples_interface(self):
             with pytest.raises(OverflowError):
                 tables.simplify(samples=np.array([0, bad_node]))
 
+    @pytest.fixture(scope="session")
+    def wf_sim_with_individual_metadata(self):
+        tables = wf.wf_sim(
+            9,
+            10,
+            seed=1,
+            deep_history=True,
+            initial_generation_samples=False,
+            num_loci=5,
+            record_individuals=True,
+        )
+        assert tables.individuals.num_rows > 50
+        individuals_copy = tables.copy().individuals
+        tables.individuals.clear()
+        tables.individuals.metadata_schema = tskit.MetadataSchema({"codec": "json"})
+        for i, individual in enumerate(individuals_copy):
+            tables.individuals.add_row(
+                flags=individual.flags,
+                location=individual.location,
+                parents=individual.parents,
+                metadata={
+                    "original_id": i,
+                    "original_parents": [int(p) for p in individual.parents],
+                },
+            )
+        tables.sort()
+        return tables
+
+    def test_individual_parent_mapping(self, wf_sim_with_individual_metadata):
+        tables = wf_sim_with_individual_metadata.copy()
+        tables.simplify()
+        ts = tables.tree_sequence()
+        for individual in tables.individuals:
+            for parent, original_parent in zip(
+                individual.parents, individual.metadata["original_parents"]
+            ):
+                if parent != tskit.NULL:
+                    assert (
+                        ts.individual(parent).metadata["original_id"] == original_parent
+                    )
+        assert set(tables.individuals.parents) != {tskit.NULL}
+
+    def test_shuffled_individual_parent_mapping(self, wf_sim_with_individual_metadata):
+        tables = wf_sim_with_individual_metadata.copy()
+        tsutil.shuffle_tables(
+            tables,
+            42,
+            shuffle_edges=False,
+            shuffle_populations=False,
+            shuffle_individuals=True,
+            shuffle_sites=False,
+            shuffle_mutations=False,
+            shuffle_migrations=False,
+        )
+        # Check we have a mixed up order
+        with pytest.raises(
+            tskit.LibraryError,
+            match="Individuals must be provided in an order where"
+            " children are after their parent individuals",
+        ):
+            tables.tree_sequence()
+
+        tables.simplify()
+        metadata = [
+            tables.individuals.metadata_schema.decode_row(m)
+            for m in tskit.unpack_bytes(
+                tables.individuals.metadata, tables.individuals.metadata_offset
+            )
+        ]
+        for individual in tables.individuals:
+            for parent, original_parent in zip(
+                individual.parents, individual.metadata["original_parents"]
+            ):
+                if parent != tskit.NULL:
+                    assert metadata[parent]["original_id"] == original_parent
+        assert set(tables.individuals.parents) != {tskit.NULL}
+
+    def test_individual_mapping(self):
+        tables = wf.wf_sim(
+            9,
+            10,
+            seed=1,
+            deep_history=True,
+            initial_generation_samples=False,
+            num_loci=5,
+            record_individuals=True,
+        )
+        assert tables.individuals.num_rows > 50
+        node_md = []
+        individual_md = [b""] * tables.individuals.num_rows
+        for i, node in enumerate(tables.nodes):
+            node_md.append(struct.pack("i", i))
+            individual_md[node.individual] = struct.pack("i", i)
+        tables.nodes.packset_metadata(node_md)
+        tables.individuals.packset_metadata(individual_md)
+        tables.sort()
+        tables.simplify()
+        ts = tables.tree_sequence()
+        for node in tables.nodes:
+            if node.individual != tskit.NULL:
+                assert ts.individual(node.individual).metadata == node.metadata
+        assert set(tables.individuals.parents) != {tskit.NULL}
+
     def test_bad_individuals(self, simple_ts_fixture):
         tables = simple_ts_fixture.dump_tables()
         tables.individuals.clear()