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

Issue:
Document PedigreeBuilder
#1913 added the PedigreeBuilder class, but didn't add any documentation.


Discussion / hints:

Fix patch:
diff --git a/docs/ancestry.md b/docs/ancestry.md
index 5b0b5a312..13427f39d 100644
--- a/docs/ancestry.md
+++ b/docs/ancestry.md
@@ -2394,8 +2394,8 @@ definition of the ``draw_pedigree`` function used here.
 
 We then use this pedigree information as the
 {ref}`initial state<sec_ancestry_initial_state>` for the simulation
-(note that we must set the sequence length on the pedigree tables
-before calling ``sim_ancestry``):
+(note that we set the sequence length on the pedigree tables
+when we called {func}`.parse_pedigree`):
 
 ```{code-cell}
 ped_ts = msprime.sim_ancestry(
@@ -2591,8 +2591,6 @@ pedigree = msprime.parse_pedigree(
 draw_pedigree(pedigree.tree_sequence())
 ```
 
-
-
 ```{code-cell}
 ped_ts = msprime.sim_ancestry(
     initial_state=pedigree, model="fixed_pedigree", random_seed=41)
diff --git a/docs/api.md b/docs/api.md
index 1287dd8c8..85a0bde8b 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -272,6 +272,3 @@ The node was created by a census event. Please see the
 ```{eval-rst}
 .. autofunction:: msprime.parse_pedigree
 ```
-
-```{eval-rst}
-.. autofunction:: msprime.write_pedigree
diff --git a/docs/pedigrees.md b/docs/pedigrees.md
index a71537322..25701536f 100644
--- a/docs/pedigrees.md
+++ b/docs/pedigrees.md
@@ -36,11 +36,6 @@ and can be provided as input to constrain simulations of genetic ancestry
 In this section we describe the data structures used to encode pedigrees
 in msprime, and the utilities used to create input pedigrees.
 
-:::{todo}
-This page is incomplete. Needs another pass through.
-https://github.com/tskit-dev/msprime/issues/2008
-:::
-
 (sec_pedigrees_encoding)=
 
 ## Pedigree encoding
@@ -64,13 +59,13 @@ section for details), which is an artefact of tskit's node-centric design.
 
 It is possible to specify a pedigree model directly using tskit APIs,
 but it is simpler to use the {class}`.PedigreeBuilder` utility
-class (or the simplified text {ref}`sec_pedigrees_file_format`).
+class (or the {ref}`sec_pedigrees_file_format`).
 In the following example we build a simple trio:
 
 ```{code-cell}
 pb = msprime.PedigreeBuilder()
-mom_id  = pb.add_individual(time=1)
-dad_id  = pb.add_individual(time=1)
+mom_id = pb.add_individual(time=1)
+dad_id = pb.add_individual(time=1)
 pb.add_individual(time=0, parents=[mom_id, dad_id], is_sample=True)
 pedigree = pb.finalise()
 # TODO replace with display(pedigree) when its implemented in tskit
@@ -84,10 +79,55 @@ contains the pedigree information defined by the calls to
 the parents, and because they are founders, we don't provide any information
 about their parents. Each call to {meth}`~.PedigreeBuilder.add_individual`
 returns the integer ID of newly added individual.
-
 ### Requirements
 
+This section lists the detailed requirements of the low-level encoding
+for the tskit node and individual tables.
+
+:::{warning}
+Most of this information is only
+of interest if you are constructing these tables by hand. For most
+purposes it's best to use either the {class}`.PedigreeBuilder`
+API or the {ref}`sec_pedigrees_file_format`.
+:::
+
+- Each individual must have exactly two parents. Unknown parents are
+  indicated using ``-1`` (thus, a founder individual will have
+  ``parents=[-1, -1]``.
+- Each individual must be associated with exactly two nodes.
+- All nodes associated with an individual must have the same ``time``
+  and ``population`` values. (These are referred to as the individual's
+  time and population, as a shorthand.)
+- An individual's time must be less than all its parent's times.
+
+
+(sec_pedigrees_metadata)=
+
+## Metadata
+
+It is often useful to associate metadata with individuals when building
+a pedigree. This allows us to either directly store information about the
+individuals for later analysis, or to store an identifier to facilitate
+matching with existing records.
 
+In this example we update our simple trio by associating a string
+identifier with each individual. We use the
+{meth}`tskit.MetadataSchema.permissive_json` function here to create
+a schema, as this is an easy and quick way to get a functional
+metadata schema. Much more sophisticated approaches are possible, however: please
+see the {ref}`tskit metadata documentation<tskit:sec_metadata>` for
+more information.
+
+```{code-cell}
+pb = msprime.PedigreeBuilder(
+    individuals_metadata_schema=tskit.MetadataSchema.permissive_json())
+mom_id = pb.add_individual(time=1, metadata={"name": "mom"})
+dad_id = pb.add_individual(time=1, metadata={"name": "dad"})
+pb.add_individual(
+    time=0, parents=[mom_id, dad_id], is_sample=True, metadata={"name": "child"})
+pedigree = pb.finalise()
+display(pedigree.individuals)
+```
 
 (sec_pedigrees_file_format)=
 
@@ -97,8 +137,75 @@ The methods described in {ref}`sec_pedigrees_encoding` are general and
 allow for arbitrary metadata to be associated with individuals. It
 is often convenient to work with a text based representation of
 the pedigree, which is supported by the {func}`.parse_pedigree` function.
+The {ref}`sec_pedigrees_file_format_definition` section defines the columns and
+requirements for this file format.
+
+:::{seealso}
+See the
+ {ref}`sec_pedigrees_file_format_basic_example` and subsequent sections
+for examples.
+:::
+
+
+(sec_pedigrees_file_format_definition)=
+
+### Format definition
+
+This section describes the detailed rules for the pedigree text file format.
+
+The first line of the file must be a header starting with a ``#`` character,
+listing the included columns separated by white space.
+
+Columns can be listed in any order. Columns from the required columns list
+below must be included, and only recognised columns can be included in the
+file.
+
+All other lines in the file are rows defining a particular individual.
+Each row must contain data for all columns defined in the header,
+separated by whitespace.
+
+#### Required columns
+
+id
+: A unique string identifier for this individual.
+
+parent0
+: The parent that will be assigned to ``parents[0]`` in the
+  {ref}`tskit encoding<sec_pedigrees_encoding>`. This must
+  be either a identifier defined in the ``id`` column, or
+  the value ``.`` (representing missing data).
+
+parent1
+: The parent that will be assigned to ``parents[1]`` in the
+  {ref}`tskit encoding<sec_pedigrees_encoding>`. The value
+  requirements are identical to the ``parent0`` column.
+
+time
+: The time value to associate with the individual.
+
+:::{note}
+   The ``time`` column may become optional in later releases.
+:::
+
+#### Optional columns
+
+is_sample
+: The sample status of the individual. If equal to ``1``, the
+  nodes associated with this individual will be marked as samples.
+  If ``0`` they will not. No other values are supported.
+
+population
+: The population to associate with the individual. Values must
+  correspond to
+  {ref}`population identifiers<sec_demography_populations_identifiers>`
+  in a supplied {class}`.Demography` object
+
+
+(sec_pedigrees_file_format_basic_example)=
 
-### Basic structure
+### Basic example
+
+In this example we encode a trio, ``child``, ``mom`` and ``dad``:
 
 ```{code-cell}
 txt = """\
@@ -112,53 +219,131 @@ display(pedigree.individuals)
 display(pedigree.nodes)
 ```
 
+We have three individuals in our input pedigree, and there is therefore
+three rows in the individual table. In the tskit encoding, individuals
+are referred to by their integer ID (the corresponding row in the
+individual table). Individuals are added to the table in the same order
+they appear in the file, and therefore ``child`` corresponds to ID 0,
+``mom`` ID 1 and ``dad`` ID 2. These ID values are then used in the
+``parents`` column to describe the pedigree relationships.
 
-(sec_pedigrees_visualisation)=
+Time information about individuals is stored in the node table,
+where the time of an individual is associated with its two nodes.
+For example, nodes 0 and 1 correspond to the two genomes of
+individual 0 (``child``), and these are both at time 0.
 
-## Visualising pedigrees
+The individual table also includes the original string ``file_id`` as
+metadata to faciliate joining with existing data sources. We can
+access these IDs as follows:
+
+```{code-cell}
+for id, ind in enumerate(pedigree.individuals):
+    print(id, "->", ind.metadata["file_id"])
+```
+
+```{seealso}
+See the tskit {ref}`metadata documentation<tskit:sec_metadata>` for more information
+on how to use metadata.
+```
+
+### Specifying samples
+
+In the previous example we did not specify which of our individuals were
+the samples ("probands"). In this case, {func}`.parse_pedigree` assumes
+that any individual at time 0 is a sample. This is show in the example
+above in the node table, where we can see that the flags value for
+both of ``child``'s nodes is 1 (corresponding to {data}`tskit.NODE_IS_SAMPLE`)
+
+We can override this behaviour by providing an ``is_sample`` column:
 
 ```{code-cell}
 import io
 txt = """\
-# id parent0 parent1 time
-child mom dad 0
-mom    .   .  1
-dad    .   .  1
+# id parent0 parent1 time is_sample
+child1 mom1 dad1 0 1
+mom1     .   .   1 0
+dad1     .   .   1 0
+child2 mom2 dad2 1 1
+mom2     .   .   2 0
+dad2     .   .   2 0
 """
 pedigree = msprime.parse_pedigree(io.StringIO(txt))
 display(pedigree.individuals)
-
 display(pedigree.nodes)
 ```
 
+Here we have two trios, where the child in the second trio is from the
+same generation as the parents in the second. We use the ``is_sample``
+column to specify that ``child2`` is a sample as well as ``child1``.
+
+
+### Demography information
+
+If the founder individuals of the pedigree belong to different populations
+and we wish to simulate ancestry
+{ref}`beyond the pedigree<sec_ancestry_models_fixed_pedigree_completing>`,
+we must define a {ref}`demographic model<sec_demography>` and define
+the populations that the individuals belong to.
+
+:::{seealso}
+See the {ref}`sec_ancestry_models_fixed_pedigree_demography` section for
+a full discussion of how demography interacts with fixed pedigree simulations
+and important caveats.
+:::
+
+To do this we must pass the {class}`.Demography` instance defining the
+demographic model to {func}`.parse_pedigree`, and include
+a ``population`` column in the file.
+
 ```{code-cell}
-pedigree.sequence_length = 100
-ts = msprime.sim_ancestry(initial_state=pedigree, model="fixed_pedigree",
-    random_seed=1)
+import io
+txt = """\
+# id parent0 parent1 time is_sample population
+child1 mom1 dad1 0 1 A
+mom1     .   .   1 0 A
+dad1     .   .   1 0 A
+child2 mom2 dad2 1 1 B
+mom2     .   .   2 0 B
+dad2     .   .   2 0 B
+"""
+
+demography = msprime.Demography()
+demography.add_population(name="A", initial_size=10)
+demography.add_population(name="B", initial_size=20)
+demography.add_population(name="C", initial_size=100)
+demography.add_population_split(time=10, derived=["A", "B"], ancestral="C");
 
-SVG(ts.draw_svg())
+pedigree = msprime.parse_pedigree(io.StringIO(txt), demography=demography)
+display(pedigree.populations)
+display(pedigree.individuals)
+display(pedigree.nodes)
 ```
 
+(sec_pedigrees_visualisation)=
+
 ## Visualising pedigrees
 
+It is often useful to visualise a (small) pedigree. The following function
+is used in this documentation, which may be a useful recipe for others:
+
 ```{code-cell}
 def draw_pedigree(ped_ts):
+
     G = nx.DiGraph()
     for ind in ped_ts.individuals():
         time = ped_ts.node(ind.nodes[0]).time
-        G.add_node(ind.id, time=time)
+        pop = ped_ts.node(ind.nodes[0]).population
+        G.add_node(ind.id, time=time, population=pop)
         for p in ind.parents:
             if p != tskit.NULL:
                 G.add_edge(ind.id, p)
     pos = nx.multipartite_layout(G, subset_key="time", align="horizontal")
-    nx.draw_networkx(G, pos, with_labels=True)
-
-tables = msprime.pedigrees.sim_pedigree(
-    population_size=20, end_time=3, num_samples=5, direction="backward", random_seed=1)
-tables.individuals
-tables.sequence_length = 1
-
-draw_pedigree(tables.tree_sequence())
+    colours = plt.rcParams['axes.prop_cycle'].by_key()['color']
+    node_colours = [colours[node_attr["population"]] for node_attr in G.nodes.values()]
+    nx.draw_networkx(G, pos, with_labels=True, node_color=node_colours)
+    plt.show()
 ```
 
-
+See the {ref}`sec_ancestry_models_fixed_pedigree_demography` section
+for an example of this function drawing pedigrees in a multi-population
+model.
diff --git a/msprime/ancestry.py b/msprime/ancestry.py
index c07e7fe14..d9c61a28a 100644
--- a/msprime/ancestry.py
+++ b/msprime/ancestry.py
@@ -1720,16 +1720,28 @@ class DiscreteTimeWrightFisher(AncestryModel):
 
 
 class FixedPedigree(AncestryModel):
-    # TODO Complete documentation.
-    # TODO Since the pedigree is a necessary parameter for this simulation
-    # model and it cannot be used with any other model we should make it a
-    # parametric model where the parameter is the pedigree. This would
-    # streamline a bunch of logic.
     """
     Backwards-time simulations through a pre-specified pedigree, with diploid
     individuals and back-and-forth recombination. The string ``"fixed_pedigree"`` can
     be used to refer to this model.
+
+    The model has no parameters, but requires that the
+    :ref:`initial_state<sec_ancestry_initial_state>` parameter is provided to
+    :func:`.sim_ancestry` and that it contains a valid pedigree. See the
+    :ref:`sec_pedigrees` section for more information on the required
+    encoding and methods for translating existing pedigrees into this format.
+
+    Unlike other :ref:`ancestry models<sec_ancestry_models>`, the
+    :class:`FixedPedigree` cannot be combined with other simulation models, and
+    special attention must be paid to the role of founders when
+    :ref:`recapitating these
+    simulations<sec_ancestry_models_fixed_pedigree_completing>`.
+
+    See the :ref:`sec_ancestry_models_fixed_pedigree` section for more information
+    for more information and examples of running simulations within a
+    fixed pedigree.
     """
+
     name = "fixed_pedigree"
 
 
diff --git a/msprime/pedigrees.py b/msprime/pedigrees.py
index 64af1c9d6..21d989d76 100644
--- a/msprime/pedigrees.py
+++ b/msprime/pedigrees.py
@@ -48,7 +48,11 @@ def parse_value(self, value: str) -> float:
 
 class BooleanColumn(Column):
     def parse_value(self, value: str) -> bool:
-        return bool(int(value))
+        if value == "0":
+            return False
+        if value == "1":
+            return True
+        raise ValueError("Only 0 or 1 values are supported")
 
 
 class StringOrIntColumn(Column):
@@ -107,10 +111,51 @@ def _parse_pedigree_header(
     return columns
 
 
-def parse_pedigree(text_file, demography=None, sequence_length=None):
+def parse_pedigree(
+    text_file,
+    *,
+    demography: None | demog_mod.Demography = None,
+    sequence_length: None | float = None,
+) -> tskit.TableCollection:
     """
-    Parse a text describing a pedigree used for input to the
-    :class:`.FixedPedigree` ancestry model.
+    Parse a text file describing a pedigree used for input to the
+    :class:`.FixedPedigree` ancestry model. See the
+    :ref:`sec_pedigrees_encoding` section for more information
+    on the data encoding used for pedigrees.
+
+    .. seealso::
+        See the :ref:`sec_pedigrees_file_format_definition` section for
+        a detailed description of the columns and formatting requirements
+        for this file format.
+
+    The returned :class:`tskit.TableCollection` will contain an
+    individual and two nodes for each data row in the input file.
+    The individual will have a metadata field ``file_id`` containing
+    the value of the ``id`` column in the input. Individuals
+    (and their corresponding nodes) are added to the tables in the
+    order seen in the file. There is no ordering requirement for
+    parents and children.
+
+    If a :class:`.Demography` instance is provided to the ``demography``
+    parameter, this is used to translate and validate
+    :ref:`population identifiers<sec_demography_populations_identifiers>`
+    in the ``population`` column, and is also used to
+    fill the population table in the output :class:`tskit.TableCollection`.
+    See the :ref:`sec_ancestry_models_fixed_pedigree_demography` section for
+    more information on the interaction between demography and
+    :class:`.FixedPedigree` simulations.
+
+    :param text_file: A file-like object to read from.
+    :param demography: The :class:`.Demography` defining populations
+        referred to in the ``populations`` column, if specified. If
+        None (the default) a demography consisting of one
+        population is used (and only population 0 can be referred to).
+    :param float sequence_length: If specified, set the
+        ``sequence_length`` property of the returned TableCollection
+        to this value.
+    :return: The :class:`tskit.TableCollection` object containing the
+        corresponding pedigree data.
+    :rtype: :class:`tskit.TableCollection`
     """
     # First line must be the header and contain the column headers
     header = next(text_file, None)
@@ -146,6 +191,8 @@ def parse_pedigree(text_file, demography=None, sequence_length=None):
 
     # Convert the list of individuals to a tskit pedigree.
     builder = PedigreeBuilder(
+        # TODO the schema should document the file_id attribute. Once we have
+        # better API support in tskit we can add this.
         individuals_metadata_schema=tskit.MetadataSchema.permissive_json(),
         demography=demography,
     )
@@ -204,6 +251,33 @@ def write_pedigree(ts, out):
 
 
 class PedigreeBuilder:
+    """
+    Utility for building pedigrees in the format required for input to
+    the :class:`.FixedPedigree` ancestry model.
+
+    .. seealso::
+        See the :ref:`sec_pedigrees` section for more information on
+        how pedigrees are described and imported in msprime.
+
+    Example::
+
+        pb = msprime.PedigreeBuilder()
+        mom_id = pb.add_individual(time=1)
+        dad_id = pb.add_individual(time=1)
+        pb.add_individual(time=0, parents=[mom_id, dad_id], is_sample=True)
+        pedigree_tables = pb.finalise()
+
+    :param demography: The :class:`.Demography` defining populations
+        referred to in the ``populations`` column, if specified. If
+        None (the default) a demography consisting of one
+        population is used (and only population 0 can be referred to).
+    :param tskit.MetadataSchema individuals_metadata_schema: If specified,
+        set the ``metadata_schema`` for the individuals table in the
+        final table collection. Must be an instance of
+        :class:`tskit.MetadataSchema` See the :ref:`sec_pedigrees_metadata`
+        section for more information.
+    """
+
     def __init__(self, demography=None, individuals_metadata_schema=None):
         if demography is None:
             demography = demog_mod.Demography.isolated_model([1])
@@ -226,9 +300,34 @@ def add_individual(
         parents=None,
         population=None,
         metadata=None,
-    ):
+    ) -> int:
         """
-        Adds the specified individual, returning its ID.
+        Adds an individual with the specified properties, returning its ID.
+
+        :param float time: The time for this individual measured in generations ago.
+        :param bool is_sample: If True, the new individual is marked as a sample;
+            if False, the individual is not marked as a sample. If None (the default)
+            the individual is marked as a sample if its ``time`` is zero.
+            Parent IDs are not checked, and may refer to individuals not yet
+            added to the tables.
+        :param list(int) parents: The integer IDs of the specified individual's
+            parents. Exactly two parents must be specified. If None (the default),
+            the individual is treated as a founder by setting its parents to
+            ``[-1, -1]``.
+        :param str|int population: The population to associated with this
+            individual. The value can be a string or integer, following the
+            usual
+            :ref:`population identifier<sec_demography_populations_identifiers>`
+            rules. If None (the default), a population ID of 0 will be assigned
+            if the :class:`.Demography` associated with this PedigreeBuilder has a
+            single population (the default). If the demography has more than
+            one population, then a population must be explicitly specified for
+            each individual.
+        :param dict|bytes metadata: Any metadata to associate with the
+            new individual. See the :ref:`sec_pedigrees_metadata` section
+            for more information.
+        :return: The ID of the newly added individual.
+        :rtype: int
         """
         if is_sample is None:
             # By default, individuals at time 0 are marked as samples
@@ -278,9 +377,18 @@ def add_individuals(self, *, parents, time):
         )
         return ind_ids
 
-    def finalise(self, sequence_length=None):
+    def finalise(self, sequence_length=None) -> tskit.TableCollection:
         """
-        Returns the table collection representing the defined pedigree.
+        Returns the :class:`tskit.TableCollection` describing the pedigree
+        defined by calls to :meth:`.add_individual`.
+
+        The state of the pedigree builder is not modified by this method.
+
+        :param float sequence_length: If specified, set the ``sequence_length``
+            property of the returned TableCollection to this value. If
+            ``None`` (the default) the ``sequence_length`` is ``-1``.
+        :return: The TableCollection defining the pedigree.
+        :rtype: tskit.TableCollection
         """
         copy = self.tables.copy()
         copy.sequence_length = 1


Test patch:
diff --git a/tests/test_pedigree.py b/tests/test_pedigree.py
index a06936736..2d2d7f96f 100644
--- a/tests/test_pedigree.py
+++ b/tests/test_pedigree.py
@@ -1378,7 +1378,7 @@ def test_unknown_column(self):
         with pytest.raises(ValueError, match="'unknown_col' not supported"):
             parse_indented(text)
 
-    def test_parital_row_min(self):
+    def test_partial_row_min(self):
         text = """\
         # id parent0 parent1 time
         child
@@ -1386,7 +1386,7 @@ def test_parital_row_min(self):
         with pytest.raises(ValueError, match="Incorrect number of columns"):
             parse_indented(text)
 
-    def test_parital_row_missing_optional(self):
+    def test_partial_row_missing_optional(self):
         text = """\
         # id parent0 parent1 time is_sample
         child . . 0
@@ -1427,6 +1427,14 @@ def test_bad_is_sample(self):
         with pytest.raises(ValueError, match="value for column 'is_sample' on line 2"):
             parse_indented(text)
 
+    def test_is_sample_non_01(self):
+        text = """\
+        # id parent0 parent1 time is_sample
+        child . . 0 2
+        """
+        with pytest.raises(ValueError, match="value for column 'is_sample' on line 2"):
+            parse_indented(text)
+
     def test_id_is_dot(self):
         text = """\
         # id parent0 parent1 time
@@ -1452,7 +1460,7 @@ def verify(self, pedigree, demography=None):
         pedigrees.write_pedigree(pedigree.tree_sequence(), buff)
         buff.seek(0)
         parsed = pedigrees.parse_pedigree(
-            buff, demography, sequence_length=pedigree.sequence_length
+            buff, demography=demography, sequence_length=pedigree.sequence_length
         )
         pedigree.assert_equals(parsed, ignore_metadata=True)