---
jupytext:
  formats: md:myst
  text_representation:
    extension: .md
    format_name: myst
kernelspec:
  display_name: Python 3
  language: python
  name: python3
---

(ts-data-types)=

# Data structures related to tree sequences

## Low-level data types

* {class}`fwdpy11.Node` defines nodes
* {class}`fwdpy11.Edge` defines edges
* {class}`fwdpy11.Site` defines a site (genomic location) where a mutation is present.
* {class}`fwdpy11.MutationRecord` defines mutations associated with {class}`fwdpy11.Site`
  objects on trees and in {class}`fwdpy11.PopulationBase` objects.
* {class}`fwdpy11.NodeTable` represents a node table
* {class}`fwdpy11.EdgeTable` represents an edge table
* {class}`fwdpy11.MutationTable` represents a mutation table
* {class}`fwdpy11.SiteTable` represents a site table

These types are all analogous to the corresponding [tskit][tskit] types.
Here, {class}`fwdpy11.MutationRecord` plays the role of `tskit.Mutation`,
but with additional data fields that are useful for the forward-time simulations.

## Table collections

The above data types are encapsulated into the Python class {class}`fwdpy11.TableCollection`.  Instances of this
class are data fields of populations, via {attr}`fwdpy11.PopulationBase.tables`.

(representing-trees)=

## Representing trees

An edge table and a node table contain all the data required to generate what we may call "marginal trees", which are the trees
corresponding to genomic intervals {math}`[left, right)`. In fwdpy11, {class}`fwdpy11.TreeIterator` allows
left-to-tright traversal across marginal trees as well as access to their data.

The intent of this section is to describe the structure of the marginal trees and get some terminology out of the way.  I will use Jerome Kelleher's
[tskit][tskit] package to illustrate the concepts.  fwdpy11 uses the same basic data model, and it is more convenient to use trees
generated via a coalescent simulation that to run forward simulations.  Plus, we can use the great "draw" methods from
[tskit][tskit]!

:::{note}

Do not mistake the [tskit][tskit] operations used below for the operations needed in fwdpy11.  They differ in some
important ways.  Namely, the [tskit][tskit] API is currently richer and the fwdpy11 API is currently more low-level.
The main goal here is to define terms and visualize them.

:::

We are going to look at a single tree, but the concepts apply to any of the trees generated by a simulation with nonzero
recombination rates.  First, let's get a tree sequence using the linear-time algorithm of Hudson (1990),
which involves repeated swapping of the IDs of potential ancestors until the MRCA is reached.

```{code-cell} python
import tskit
import numpy as np

# The following makes use of the Kirk Lohmueller seed.
np.random.seed(101 * 405 * 10 * 110)

tc = tskit.TableCollection(1)

nsam = 10

for i in range(nsam):
    tc.nodes.add_row(time=0, flags=1)

    nodes = np.arange(2 * nsam - 1, dtype=np.int32)
    time = 0.0
```

```{code-cell} python
# The citation for this algorithm is
# Hudson, Richard R. 1990.
# “Gene Genealogies and the Coalescent Process.”
# Oxford Surveys in Evolutionary Biology 7 (1): 44.
n = nsam
while n > 1:
    rcoal = (n * (n - 1)) / 2.0
    tcoal = np.random.exponential(1.0 / rcoal)
    time += tcoal
    tc.nodes.add_row(time=time)
    ancestor = 2 * nsam - n
    p = np.random.choice(n, 1)[0]
    tc.edges.add_row(left=0.0, right=1.0, parent=ancestor, child=nodes[p])
    nodes[p] = nodes[n - 1]
    p = np.random.choice(n - 1, 1)[0]
    tc.edges.add_row(left=0.0, right=1.0, parent=ancestor, child=nodes[p])
    nodes[p] = nodes[2 * nsam - n]
    n -= 1
```

```{code-cell} python
tc.sort()
ts = tc.tree_sequence()
print(ts.first().draw(format="unicode"))
```

The above tree corresponds to a sample size of 10 haplotypes.  The tree shows the node labels.  The nodes labelled 0
through 9 correspond to the present time point--these are "alive nodes" or the "current generation" if we are thinking
about a forward simulation.  Further, we can describe the branches leading to these sample nodes as the "tips" or
"leaves" of a tree.

### "Leaves" of a tree

One thing that we often want to know is "how many samples descend from node `i`?" To do that, we
may look at the leaf counts attribute of a marginal tree:

```{code-cell} python
# Let's store our
# tree in a variable now
t = ts.first()
```

```{code-cell} python
def get_leaf_counts(tree, i):
    return len([j for j in tree.leaves(i)])
```

```{code-cell} python
# Map node ids to their leaf counts using a dict
lcmap = {i: "".format(get_leaf_counts(t, i)) for i in range(len(ts.tables.nodes))}
print(t.draw(format="unicode", node_labels=lcmap))
```

### Children and siblings

Let's take another look at our tree, labelled with node ids:

```{code-cell} python
print(t.draw(format="unicode"))
```

Let's ponder node 15 for a moment.  It has two immediate descendants, nodes 11 and 14.  We may consider
these the left and right children, respectively, of node 15.

```{code-cell} python
def get_children(tree, i):
    lc = tree.left_child(i)
    rc = tree.right_child(i)
    if lc == tskit.NULL and rc == tskit.NULL:
        return "->NULL".format(i)
    return "->".format(i) + str((lc, rc))


cmap = {i: get_children(t, i) for i in range(len(ts.tables.nodes))}
print(t.draw(format="unicode", node_labels=cmap))
```

Likewise, we may look at the sibling relationships amongst nodes:

```{code-cell} python
def get_sibs(tree, i):
    ls = tree.left_sib(i)
    rs = tree.right_sib(i)
    if ls == tskit.NULL and rs == tskit.NULL:
        return "->NULL".format(i)
    return "->".format(i) + str((ls, rs))
```

```{code-cell} python
smap = {i: get_sibs(t, i) for i in range(len(ts.tables.nodes))}
print(t.draw(format="unicode", node_labels=smap))
```

### Multiply-linked lists

Under the hood, the data structures representing marginal trees consist of several arrays
represening a multiply-linked list allowing traversal up/down/left/right along a marginal tree.
By convention, the value -1 is taken as a NULL value, signifying that there are no more nodes in
"that" direction along the tree.

We can look directly at what theses arrays look like in our tree:

```{code-cell} python
nnodes = len(ts.tables.nodes)
# First, let's write down our node ids:
print([i for i in range(nnodes)])
# Now, get the parents of each node, moving "up" the tree
print([t.parent(i) for i in range(nnodes)])
# The left child list allows moving "down left" along a tree
print([t.left_child(i) for i in range(nnodes)])
# The right child list allows moving "down right" along a tree
print([t.right_child(i) for i in range(nnodes)])
# The left sib list allows moving "left" along a tree
print([t.left_sib(i) for i in range(nnodes)])
# The right sib list allows moving "right" along a tree
print([t.right_sib(i) for i in range(nnodes)])
```

These lists are interpreted a lists of nodes referring to other nodes.  For example, the value 10 in position 0 of the
parents list means, "The parent of the node with index 0 has index 10", where the indexes refer to the node table.

Given the above lists, you may start at any valid node id (*e.g.* a value not equal to -1 in the first array),
and then "walk" in any direction you choose along the tree until you hit a value of -1, meaning that you cannot proceed
any further.

:::{note}

The trees generated by `msprime.simulate` are bifurcating, which is a consequence of simulating from the Kingman
coalescent.  In forward-time simulations, it is not uncommon to have more than two descendants of a node.  When that
happens, left_child and right_child refer to the left-most and right-most children, respectively. Thus, to "walk"
along the descendants of a node, you proceed to left_child, and then march along right_sib until a value of -1 is
seen.  This "walking" method is the same as what you would do for a bifurcating tree, but I want to point out that
counting the number of immediate descendants of a node requires counting the number of steps that the walk requires,
and that it may be longer than two steps.

:::

The last five arrays show above correspond to the following attributes of {class}`fwdpy11.TreeIterator`:

* {attr}`fwdpy11.TreeIterator.parent`
* {attr}`fwdpy11.TreeIterator.left_child`
* {attr}`fwdpy11.TreeIterator.right_child`
* {attr}`fwdpy11.TreeIterator.left_sib`
* {attr}`fwdpy11.TreeIterator.right_sib`

The key to efficiency is how these linked lists are updated as you move from tree {math}`i` to tree {math}`i+1`.  Given
an approprate set of indexes, these lists are only updated at the positions that differ between the two trees.  In
practice, adjacent trees are highly-correlated, meaning that very few values need updating. The
construction of these indexes is described immediatlely above the description of "Algorithm T" in the Kelleher et al. (2016) paper describing [msprime][msprime].

The building of those indexes is a key determinant of performance for algorithms on tree sequences.  Fortunately, you
don't need to worry about that, as that machinery is hidden in the internals of fwdpp (and tskit).

[tskit]: http://tskit.readthedocs.io

[msprime]: http://msprime.readthedocs.io


