VariantData.from_tree_sequence

[hyanwong]

In many cases we want to infer from an existing tree sequence. At the moment we use SampleData.from_tree_sequence(). Following on from VariantData.from_arrays() in #924 , I think it would also be useful to have the from_tree_sequence() method. This is in addition to sgkit-dev/bio2zarr#232 (which would provide essentially the same functionality).

I think this would be useful for writing the unit tests too.

[hyanwong]

In sgkit-dev/bio2zarr#395 (comment), Jerome suggested a wrapper for bio2zarr, like:

with tsinfer.VariantData.from_tree_sequence(ts) as data:
    tsinfer.infer(data)

Perhaps this should be implemented once sgkit-dev/bio2zarr#387 is fixed?

[hyanwong]

Something like

import bio2zarr.tskit as ts2z
import tskit
import tsinfer
import tempfile
import os.path
import numpy as np
from contextlib import contextmanager

@contextmanager
def from_tree_sequence(ts, ancestral_state=None, **kwargs):
    if ancestral_state is None:
        sites = ts.tables.sites
        ancestral_state = np.array(tskit.unpack_strings(
            sites.ancestral_state, sites.ancestral_state_offset))
    with tempfile.TemporaryDirectory() as zarrdir:
        zarrfile = os.path.join(zarrdir, "tmp.vcz")
        ts2z.convert(ts, zarrfile)
        yield tsinfer.VariantData(zarrfile, ancestral_state=ancestral_state, **kwargs)

[hyanwong]

I forgot that we also need to implement use_sites_time and use_individuals_time for this method, which we use extensively in testing. For instance, we need the replacement from_tree_sequence to work here:

tsinfer/tests/test_formats.py

Line 1927 in 20788d3

sample_data = formats.SampleData.from_tree_sequence(

We should test that this all works.