Tskit: Remove arbitrary sites from a tree sequence

[hyanwong]

I keep using something like the function below. Should we tidy it up and add it to the API? It seems useful to me (it's also coded up semi-independently as remove sites in https://github.com/tskit-dev/tskit/pull/292)

def prune_sites(ts, keep_sites):
    # Remove sites and mutations not marked as `True` in the keep_sites boolean array
    def keep_with_offset(keep, data, offset):
        """Copied from keep_intervals"""
        # We need the astype here for 32 bit machines
        lens = np.diff(offset).astype(np.int32)
        return (
            data[np.repeat(keep, lens)],
            np.concatenate([
                np.array([0], dtype=offset.dtype),
                np.cumsum(lens[keep], dtype=offset.dtype)]))

    tables = ts.dump_tables()
    sites = ts.tables.sites
    mutations = ts.tables.mutations

    new_as, new_as_offset = keep_with_offset(
        keep_sites, sites.ancestral_state, sites.ancestral_state_offset)
    new_md, new_md_offset = keep_with_offset(
        keep_sites, sites.metadata, sites.metadata_offset)
    keep_mutations = keep_sites[mutations.site]
    tables.sites.set_columns( #replaces existing
        position=sites.position[keep_sites],
        ancestral_state=new_as,
        ancestral_state_offset=new_as_offset,
        metadata=new_md,
        metadata_offset=new_md_offset)

    new_ds, new_ds_offset = keep_with_offset(
        keep_mutations, mutations.derived_state, mutations.derived_state_offset)
    new_md, new_md_offset = keep_with_offset(
        keep_mutations, mutations.metadata, mutations.metadata_offset)

    site_map = np.cumsum(keep_sites, dtype=mutations.site.dtype) - 1
    parent_map = np.cumsum(keep_mutations, dtype=mutations.parent.dtype) - 1
    # parent -1 always maps to parent -1
    parent_map = np.append(parent_map, np.array([-1], dtype=mutations.parent.dtype))
    # Don't bother dealing with connected mutations as we are removing the whole site
    tables.mutations.set_columns(
        site=site_map[mutations.site[keep_mutations]],
        node=mutations.node[keep_mutations],
        derived_state=new_ds,
        derived_state_offset=new_ds_offset,
        parent=parent_map[mutations.parent[keep_mutations]],
        metadata=new_md,
        metadata_offset=new_md_offset)

    return tables.tree_sequence()

[petrelharp]

Perhaps as a method of the TableCollection? If you have use for it, then sure?

[jeromekelleher]

Other people have used similar thing to this all right (say, to get rid of sites with low-frequency mutations). I'm not overly keen on the name though. Maybe filter_sites or something?

[hyanwong]

Right - the name was just picked from thin air. Happy with whatever seems sensible (I previously used remove_sites). If we use filter_sites(to_remove=None) that allows us to keep the option of adding an additional to_keep parameter later, if that seems useful.

Presumably to_remove could either be a boolean array of the same length as num_sites, or a list of integer indices, used to index into the numpy arrays?

[petrelharp]

I vote for remove rather than filter, since it's more clear what's happening to them.

[jeromekelleher]

Lets use delete_sites --- this mirrors the terminology for delete_intervals.

[hyanwong]

We should put keep_with_offset somewhere for reuse, either at the top of tables.py or in util.py

[hyanwong]

This is done. Anyone want to review https://github.com/hyanwong/tskit/commit/6be14754b80f70366e3905697e5748af848961ed