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

(tskitconvert_vignette)=

# Converting data to tskit format

:::{note}
The examples in this vignette involve population objects with no ancestry.
This is intentional, as it keeps the execution times to a minimum while still illustrating the concepts.
:::

```{code-cell} python
---
tags: [ "hide-input" ]
---
import demes
import fwdpy11
```

At the end of a simulation, {func}`fwdpy11.DiploidPopulation.dump_tables_to_tskit` will generate a {class}`tskit.TreeSequence` object.
This object can be saved to a file using {meth}`tskit.TreeSequence.dump`, resulting in a "trees file".

Saving data in `tskit`'s format gives you access to a huge array of methods for downstream analysis of your simulated data.
See the [tskit docs](https://tskit.dev/tskit/docs/stable/index.html) for much more on what you can do.
Much of what you need is found in the docs for {class}`tskit.TreeSequence`.

What happens when we call {func}`fwdpy11.DiploidPopulation.dump_tables_to_tskit`?

1. All tables are converted from `fwdpy11` format to `tskit` format in a {class}`tskit.TableCollection`.
2. For individuals, populations, and mutations, row-level metadata are encoded.
   The {ref}`next vignette <tskit_metadata_vignette>` covers how to interact with this information.
3. Optionally, top-level metadata are encoded, which is controlled by keyword arguments to {func}`fwdpy11.DiploidPopulation.dump_tables_to_tskit`.

## Obtaining a `tskit` tree sequence

```{code-cell} python
pop = fwdpy11.DiploidPopulation(1000, 1.0)
```

By default, export to `tskit` returns a {class}`tskit.TreeSequence`:

```{code-cell} python
ts = pop.dump_tables_to_tskit()
type(ts)
```

In general, we want this type in order to dump straight to a file:

```python
# not executed
ts.dump("treefile.trees")
```

We can also do a lot of analysis using this type.
However, certain tasks would be easier if this type knew about the metadata generated by `fwdpy11`.

## Metadata

The `tskit` data model includes [metadata](https://tskit.dev/tskit/docs/stable/metadata.html).
Metadata means anything not strictly needed for operations on trees/tree sequences.
Examples may include population names, properties of mutations such as selection coefficients, geographic locations of individuals, etc..
These sorts of metadata are included as optional columns in tables.
See the `tskit` [data model](https://tskit.dev/tskit/docs/stable/data-model.html) documentation for more detail.

Metadata can also include "top-level" information that is stored as metadata for the tree sequence (or {class}`tskit.TableCollection`) itself.
Much of this vignette concerns top-level metadata.

### Metadata schema

The `tskit` metadata require a specification, or *schema*, that is used to validate the metadata contents.

The top-level metadata schema is:

```{code-cell} python
print(fwdpy11.tskit_tools.metadata_schema.TopLevelMetadata)
```

We see that the metadata are encoded as a [JSON schema](https://json-schema.org/).
Only one field is required, `generation`.
This required field is automatically populated by {func}`fwdpy11.DiploidPopulation.dump_tables_to_tskit`.
The remaining fields are filled in by keyword arguments passed to that function.

Much of what follows concerns filling in this top-level metadata.
To do this, we use incomplete data (populations with no ancestry) to illustrate the concepts.

## Populating top-level metadata

Top-level metadata are an attempt to improve reproducibility.
However, all of these features are optional.
Because a simulation project using Python may be very complex, with many local and non-local imports, etc., fully documenting the top-level metadata and/or "provenance" may be impractical.
In some cases, reproducibility may instead by best managed by a `git` repository where script and `Makefile`/`Snakefile` work flows coexist.

### Adding random number seeds

You may add a random number seed to the top-level metadata.
Doing so is optional because we cannot guarantee that you used this seed to construct a {class}`fwdpy11.GSLrng` and then immediately run the simulation.
For example, if you built the `rng` object and then simulated two replicates, the seed is only valid to reproduce the first.
You do not know the state of the `rng` when the second simulation started.

```{code-cell} python
ts = pop.dump_tables_to_tskit()
assert "seed" not in ts.metadata
```

```{code-cell} python
ts = pop.dump_tables_to_tskit(seed=54321)
assert ts.metadata["seed"] == 54321
```

```{code-cell} python
---
tags: [raises-exception]
---
ts = pop.dump_tables_to_tskit(seed=-4616)
```

### Adding `ModelParams` objects

Another optional top-level metadata item are instances of {class}`fwdpy11.ModelParams`.
Including this is optional for several reasons, but a big one is that this class can be arbitrarily hard to serialize.
The examples listed below will work with types provided by `fwdpy11` (although we haven't tested them all!), but user-defined types may be more challenging.

```{code-cell} python
ts = pop.dump_tables_to_tskit()
assert "model_params" not in ts.metadata
```

Here is an example of Gaussian stabilizing selection around an optimum trait value of `0.0`:

```{code-cell} python
optimum = fwdpy11.Optimum(optimum=0.0, VS=10.0, when=0)
gss = fwdpy11.GaussianStabilizingSelection.single_trait([optimum])
pdict = {
    'nregions': [],
    'sregions': [fwdpy11.GaussianS(0, 1, 1, 0.10)],
    'recregions': [fwdpy11.PoissonInterval(0, 1, 1e-3)],
    'gvalue': fwdpy11.Additive(2., gss),
    'rates': (0.0, 1e-3, None),
    'simlen': 10 * pop.N,
    'prune_selected': False,
    'demography': fwdpy11.ForwardDemesGraph.tubes([pop.N], 1)
}

params = fwdpy11.ModelParams(**pdict)
```

We can pass the `params` object on when exporting the data to `tskit`:

```{code-cell} python
ts = pop.dump_tables_to_tskit(model_params=params)
recovered_params = fwdpy11.ModelParams(**eval(ts.metadata["model_params"])) 
assert recovered_params == params
```
Note that we can recover our demographic model from the restored parameters:

```{code-cell} python
print(recovered_params.demography)
```

```{code-cell} python
print(recovered_params.demography.demes_graph)
```

### User-defined metadata

It may be useful to store arbitrary data in the output.
For example, this could be a {class}`dict` containing information obtained from a {ref}`time series <recorders_vignette>` analysis of a simulation.

The simplest situation is to use a `dict` containing simple objects:

```{code-cell} python
ts = pop.dump_tables_to_tskit(data={'x': 3})
ts.metadata["data"]
```

Any use of user-defined types will require conversion to string, however:

```{code-cell} python
class MyData(object):
    def __init__(self, x):
        self.x = x

    def __repr__(self):
        return f"MyData(x={self.x})"

# Send in the str representation:
ts = pop.dump_tables_to_tskit(data=str(MyData(x=77)))
```

Getting the data back out requires an `eval`:

```{code-cell} python
assert eval(ts.metadata["data"]).x == 77
```

:::{warning}
Recording data like this can become arbitrary difficult.
The Python import system may make it more difficult to recover data than the examples here would indicate.
Further, if the data are very large, then other output formats are likely more appropriate.
:::

## Setting population table metadata

This example involves simulating a multi-deme model.

```{code-cell} python
---
tags: ["hide-input"]
---

# Example 07 from the demes tutorial
yaml = """
time_units: generations
demes:
  - name: X
    epochs:
      - end_time: 1000
        start_size: 2000
  - name: A
    ancestors:
      - X
    epochs:
      - start_size: 2000
  - name: B
    ancestors:
      - X
    epochs:
      - start_size: 2000
"""
demography = fwdpy11.ForwardDemesGraph.from_demes(
    yaml,
    burnin=100,
    burnin_is_exact=True,
)
Opt = fwdpy11.Optimum
GSSmo_ancestor = fwdpy11.GaussianStabilizingSelection.single_trait(
    [Opt(when=0, optimum=0.0, VS=1.0), Opt(when=100, optimum=1.0, VS=1.0)]
)
gvalue_ancestor = fwdpy11.Additive(2.0, GSSmo_ancestor)
GSSmo_daughter_1 = fwdpy11.GaussianStabilizingSelection.single_trait(
    [Opt(when=100, optimum=1.0, VS=1.0)]
)
gvalue_daughter_1 = fwdpy11.Additive(2.0, GSSmo_daughter_1)

GSSmo_daughter_2 = fwdpy11.GaussianStabilizingSelection.single_trait(
    [Opt(when=100, optimum=1.0, VS=1.0)]
)
gvalue_daughter_2 = fwdpy11.Additive(2.0, GSSmo_daughter_2)
p = {
    "nregions": [],
    "sregions": [fwdpy11.GaussianS(0, 1, 1, 0.25)],
    "recregions": [fwdpy11.PoissonInterval(0, 1, 1e-3)],
    "rates": (0.0, 0.025, None),
    "gvalue": [gvalue_ancestor, gvalue_daughter_1, gvalue_daughter_2],
    "prune_selected": False,
    "demography": demography,
    "simlen": demography.final_generation,
}

rng = fwdpy11.GSLrng(12351235)
params = fwdpy11.ModelParams(**p)
pop = fwdpy11.DiploidPopulation(demography.initial_sizes, 1)

fwdpy11.evolvets(rng, pop, params, simplification_interval=100)
```

By default, deme names are placed in the metadata:

```{code-cell} python
ts = pop.dump_tables_to_tskit()
for p in ts.populations():
    print(p.metadata)
```

We can override this behavior by providing a dictionary
mapping a deme's integer id (index) to data.

```{code-cell} python
md = {}
for i, deme in enumerate(demography.demes_graph.demes):
    md[i] = {'name': deme.name, 'end_time': deme.end_time}
ts = pop.dump_tables_to_tskit(population_metadata=md)
for p in ts.populations():
    print(p.metadata)
```
