Repository: tskit-dev/msprime (language: python)

Issue:
Add parse_ped function 
This should take a pedigree file as input and return a tskit TableCollection ready for simulating with. This will be the primary interface for people setting up the input to running a pedigree simulation, so we need to take some care about how we think about populations, etc, and whether we do things like supporting inserting additional samples.

This can be a more specialised version of the code that we added to tskit (probably we can just move the function over as a starting point).


Discussion / hints:
Closed in #1866 

Not clear that there's enough information in the FAM file format for us to actually build a simulation though.

Fix patch:
diff --git a/msprime/__init__.py b/msprime/__init__.py
index 655c5fa86..2702a7284 100644
--- a/msprime/__init__.py
+++ b/msprime/__init__.py
@@ -56,6 +56,7 @@
     Population,
 )
 
+from msprime.pedigrees import PedigreeBuilder
 from msprime.intervals import RateMap
 from msprime.likelihood import log_arg_likelihood, log_mutation_likelihood
 
@@ -130,6 +131,7 @@
     "NODE_IS_GC_EVENT",
     "NUCLEOTIDES",
     "PAM",
+    "PedigreeBuilder",
     "Population",
     "PopulationConfiguration",
     "PopulationParametersChange",
diff --git a/msprime/pedigrees.py b/msprime/pedigrees.py
index e8e3af6f0..2b38639fc 100644
--- a/msprime/pedigrees.py
+++ b/msprime/pedigrees.py
@@ -26,6 +26,64 @@
 import numpy as np
 import tskit
 
+from msprime import demography as demog_mod
+
+
+class PedigreeBuilder:
+    def __init__(self, demography=None):
+        if demography is None:
+            demography = demog_mod.Demography.isolated_model([1])
+        self.demography = demography
+        self.tables = tskit.TableCollection(0)
+        demography.insert_populations(self.tables)
+
+        assert len(self.tables.individuals) == 0
+        self.tables.individuals.metadata_schema = tskit.MetadataSchema(
+            {
+                "codec": "json",
+                "type": "object",
+                "properties": {},
+                "additionalProperties": True,
+            }
+        )
+
+    def add_individual(
+        self, *, time, is_sample=None, parents=None, population=None, metadata=None
+    ):
+        if is_sample is None:
+            # By default, individuals at time 0 are marked as samples
+            is_sample = time == 0
+
+        if parents is None:
+            parents = [-1, -1]
+        if len(parents) != 2:
+            raise ValueError("Must have exactly two parents")
+
+        ind_id = self.tables.individuals.add_row(parents=parents, metadata=metadata)
+        if population is None:
+            if self.demography.num_populations == 1:
+                population = 0
+            else:
+                raise ValueError(
+                    "Must specify the population for each individual in a "
+                    "multi-population demography"
+                )
+        population = self.demography[population].id
+        flags = tskit.NODE_IS_SAMPLE if is_sample else 0
+        for _ in range(2):
+            self.tables.nodes.add_row(
+                flags=flags, time=time, individual=ind_id, population=population
+            )
+        return ind_id
+
+    def finalise(self, sequence_length):
+        copy = self.tables.copy()
+        copy.sequence_length = sequence_length
+        # Not strictly necessary, but it's useful when the tables are written
+        # to file and are loaded as as a tree sequence.
+        copy.build_index()
+        return copy
+
 
 def sim_pedigree(
     *,
@@ -39,22 +97,10 @@ def sim_pedigree(
     # part of the public API and subject to arbitrary changes/removal
     # in the future.
 
-    tables = tskit.TableCollection(
-        sequence_length=-1 if sequence_length is None else sequence_length
-    )
-    # TODO we should be using the demography class here to set up the
-    # population metadata
-    permissive_json = tskit.MetadataSchema(
-        {
-            "codec": "json",
-            "type": "object",
-            "properties": {},
-            "additionalProperties": True,
-        }
-    )
-    tables.populations.metadata_schema = permissive_json
-    # tables.individuals.metadata_schema = permissive_json
-    tables.populations.add_row()
+    # We use the builder to get a standard set of tables, but bypass
+    # it so that we can use numpy methods add to it more efficiently.
+    tables = PedigreeBuilder().tables
+    tables.sequence_length = -1 if sequence_length is None else sequence_length
 
     population = np.array([tskit.NULL], dtype=np.int32)
     rng = np.random.RandomState(random_seed)


Test patch:
diff --git a/tests/test_pedigree.py b/tests/test_pedigree.py
index 8f53486c1..14335074c 100644
--- a/tests/test_pedigree.py
+++ b/tests/test_pedigree.py
@@ -16,20 +16,8 @@
 
 
 def get_base_tables(sequence_length):
-    tables = tskit.TableCollection(sequence_length=sequence_length)
-    # Add single population for simplicity. We put in a schema to
-    # avoid warnings.
-    permissive_json = tskit.MetadataSchema(
-        {
-            "codec": "json",
-            "type": "object",
-            "properties": {},
-            "additionalProperties": True,
-        }
-    )
-    tables.populations.metadata_schema = permissive_json
-    tables.individuals.metadata_schema = permissive_json
-    tables.populations.add_row()
+    tables = pedigrees.PedigreeBuilder().tables
+    tables.sequence_length = sequence_length
     return tables
 
 
@@ -259,23 +247,10 @@ def simulate_pedigree(
     random_seed: Random seed.
     """
     rng = np.random.RandomState(random_seed)
-    tables = get_base_tables(sequence_length)
-
-    def add_individual(generation, parents=(-1, -1)):
-        ind_id = tables.individuals.add_row(
-            parents=parents, metadata={"original_id": len(tables.individuals)}
-        )
-        time = num_generations - generation - 1
-        for _ in parents:
-            tables.nodes.add_row(
-                flags=tskit.NODE_IS_SAMPLE if time == 0 else 0,
-                time=time,
-                population=0,
-                individual=ind_id,
-            )
-        return ind_id
+    builder = msprime.PedigreeBuilder()
 
-    curr_gen = [add_individual(0) for _ in range(num_founders)]
+    time = num_generations - 1
+    curr_gen = [builder.add_individual(time=time) for _ in range(num_founders)]
     for generation in range(1, num_generations):
         num_pairs = len(curr_gen) // 2
         if num_pairs == 0 and num_children_prob[0] != 1:
@@ -283,15 +258,15 @@ def add_individual(generation, parents=(-1, -1)):
                 f"Not enough people to make children in generation {generation}"
             )
         all_parents = rng.choice(curr_gen, size=(num_pairs, 2), replace=False)
+        time -= 1
         curr_gen = []
         for parents in all_parents:
             num_children = rng.choice(len(num_children_prob), p=num_children_prob)
             for _ in range(num_children):
                 parents = np.sort(parents).astype(np.int32)
-                ind_id = add_individual(generation, parents=parents.astype(np.int32))
+                ind_id = builder.add_individual(time=time, parents=parents)
                 curr_gen.append(ind_id)
-    tables.build_index()
-    return tables
+    return builder.finalise(sequence_length)
 
 
 class TestSimPedigree:
@@ -1065,7 +1040,7 @@ def test_node_times_disagree(self):
 
     def test_node_populations_disagree(self):
         tables = get_base_tables(100)
-        tables.populations.add_row()
+        tables.populations.add_row({"name": "pop_1", "description": "Y"})
         ind = tables.individuals.add_row(parents=[-1, -1])
         tables.nodes.add_row(
             flags=tskit.NODE_IS_SAMPLE, time=0, individual=ind, population=0
@@ -1327,6 +1302,167 @@ def test_children_before_parents(self, tmp_path):
         assert np.array_equal(tb[2].parents, [1, 0])
 
 
+class TestPedigreeBuilder:
+    def test_is_sample_default(self):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0)
+        pb.add_individual(time=0.0001)
+        pb.add_individual(time=1)
+        tables = pb.finalise(1)
+        zero_nodes = tables.nodes.time == 0
+        assert np.all(tables.nodes.flags[zero_nodes] == tskit.NODE_IS_SAMPLE)
+        non_zero_nodes = tables.nodes.time != 0
+        assert np.all(tables.nodes.flags[non_zero_nodes] == 0)
+
+    def test_is_sample_default_override(self):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0, is_sample=True)
+        pb.add_individual(time=0.0001, is_sample=True)
+        pb.add_individual(time=1, is_sample=True)
+        tables = pb.finalise(1)
+        assert np.all(tables.nodes.flags == tskit.NODE_IS_SAMPLE)
+
+    @pytest.mark.parametrize("parents", [[], [0], [0, 1, 2]])
+    def test_bad_parents_length(self, parents):
+        pb = pedigrees.PedigreeBuilder()
+        with pytest.raises(ValueError, match="exactly two parents"):
+            pb.add_individual(time=0, parents=parents)
+
+    def test_default_parents(self):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0)
+        tables = pb.finalise(1)
+        np.testing.assert_array_equal(tables.individuals[0].parents, [-1, -1])
+
+    @pytest.mark.parametrize(
+        "parents", [[0, 1], [0.0, 1.0], np.array([0, 1], dtype=np.int8)]
+    )
+    def test_parents(self, parents):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=1)
+        pb.add_individual(time=1)
+        pb.add_individual(time=0, parents=parents)
+        tables = pb.finalise(1)
+        np.testing.assert_array_equal(tables.individuals[2].parents, parents)
+
+    def test_default_demography(self):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0)
+        ts = pb.finalise(1).tree_sequence()
+        assert ts.num_populations == 1
+        assert ts.population(0).metadata["name"] == "pop_0"
+        assert all(node.population == 0 for node in ts.nodes())
+
+    def test_one_pop_demography(self):
+        demography = msprime.Demography()
+        demography.add_population(name="A", initial_size=1)
+        pb = pedigrees.PedigreeBuilder(demography=demography)
+        pb.add_individual(time=0)
+        ts = pb.finalise(1).tree_sequence()
+        assert ts.num_populations == 1
+        assert ts.population(0).metadata["name"] == "A"
+        assert all(node.population == 0 for node in ts.nodes())
+
+    def test_two_pop_demography(self):
+        demography = msprime.Demography()
+        demography.add_population(name="A", initial_size=1)
+        demography.add_population(name="B", initial_size=1)
+        pb = pedigrees.PedigreeBuilder(demography=demography)
+        ts = pb.finalise(1).tree_sequence()
+        assert ts.num_populations == 2
+        assert ts.population(0).metadata["name"] == "A"
+        assert ts.population(1).metadata["name"] == "B"
+
+    def test_two_pop_demography_population_required(self):
+        demography = msprime.Demography()
+        demography.add_population(name="A", initial_size=1)
+        demography.add_population(name="B", initial_size=1)
+        pb = pedigrees.PedigreeBuilder(demography=demography)
+        with pytest.raises(ValueError, match="multi-population demography"):
+            pb.add_individual(time=0)
+
+    def test_two_pop_demography_populations_by_index(self):
+        demography = msprime.Demography()
+        demography.add_population(name="A", initial_size=1)
+        demography.add_population(name="B", initial_size=1)
+        pb = pedigrees.PedigreeBuilder(demography=demography)
+        pb.add_individual(time=0, population=0)
+        pb.add_individual(time=0, population=1)
+        ts = pb.finalise(1).tree_sequence()
+        assert all(ts.node(u).population == 0 for u in [0, 1])
+        assert all(ts.node(u).population == 1 for u in [2, 3])
+
+    def test_two_pop_demography_populations_by_name(self):
+        demography = msprime.Demography()
+        demography.add_population(name="A", initial_size=1)
+        demography.add_population(name="B", initial_size=1)
+        pb = pedigrees.PedigreeBuilder(demography=demography)
+        pb.add_individual(time=0, population="A")
+        pb.add_individual(time=0, population="B")
+        ts = pb.finalise(1).tree_sequence()
+        assert all(ts.node(u).population == 0 for u in [0, 1])
+        assert all(ts.node(u).population == 1 for u in [2, 3])
+
+    def test_node_individual_linkage(self):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0)
+        pb.add_individual(time=0)
+        tables = pb.finalise(1)
+        assert len(tables.individuals) == 2
+        assert len(tables.nodes) == 4
+        assert tables.nodes.individual[0] == 0
+        assert tables.nodes.individual[1] == 0
+        assert tables.nodes.individual[2] == 1
+        assert tables.nodes.individual[3] == 1
+
+    @pytest.mark.parametrize("sequence_length", [0.1, 1, 10.11])
+    def test_finalise_sequence_length(self, sequence_length):
+        pb = pedigrees.PedigreeBuilder()
+        tables = pb.finalise(sequence_length)
+        assert tables.sequence_length == sequence_length
+
+    def test_finalise_copy(self):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0)
+        t1 = pb.finalise(1)
+        t2 = pb.finalise(1)
+        assert t1 == t2
+        assert t1 is not t2
+
+    @pytest.mark.parametrize("time", [0, 1, 10.11, np.array([12345.6])[0]])
+    def test_time(self, time):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=time)
+        t1 = pb.finalise(1)
+        assert np.all(t1.nodes.time == time)
+
+    @pytest.mark.parametrize(
+        "metadata", [{}, {"A": "B"}, {"A": [0, 1, 2]}, {"A": {"B": "C"}}]
+    )
+    def test_metadata(self, metadata):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0, metadata=metadata)
+        ts = pb.finalise(1).tree_sequence()
+        assert ts.individual(0).metadata == metadata
+
+    def test_add_individual_return_value(self):
+        pb = pedigrees.PedigreeBuilder()
+        n = 5
+        for j in range(n):
+            ind_id = pb.add_individual(time=0, metadata={"id": j})
+            assert ind_id == j
+        ts = pb.finalise(1).tree_sequence()
+        # double check that these are the IDs in the tables
+        for j in range(n):
+            assert ts.individual(j).metadata == {"id": j}
+
+    def test_indexed(self):
+        pb = pedigrees.PedigreeBuilder()
+        pb.add_individual(time=0)
+        t1 = pb.finalise(1)
+        assert t1.has_index()
+
+
 if __name__ == "__main__":
     # Easy way to output a simulated pedigree for command line testing.
     # Run with python3 tests/test_pedigree.py NUM_FOUNDERS NUM_GENERATIONS > out.trees