From a98a0195d49b9fce6442c7dfb7214d422a067678 Mon Sep 17 00:00:00 2001
From: Yan Wong <yan.wong@bdi.ox.ac.uk>
Date: Wed, 28 Aug 2024 15:45:09 +0100
Subject: [PATCH] Add example of excluding chromosomes using site masks

And switch to the preferred `zarr.open` instead of `zarr.load`.
---
 .../example_data.vcz/contig_id/.zarray        |   4 +-
 docs/_static/example_data.vcz/contig_id/0     | Bin 20 -> 32 bytes
 docs/usage.md                                 |  48 ++++++++++++++----
 3 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/docs/_static/example_data.vcz/contig_id/.zarray b/docs/_static/example_data.vcz/contig_id/.zarray
index ce072443..011f6ab3 100644
--- a/docs/_static/example_data.vcz/contig_id/.zarray
+++ b/docs/_static/example_data.vcz/contig_id/.zarray
@@ -9,8 +9,8 @@
         "id": "blosc",
         "shuffle": 1
     },
-    "dtype": "<U1",
-    "fill_value": null,
+    "dtype": "<U4",
+    "fill_value": "",
     "filters": null,
     "order": "C",
     "shape": [
diff --git a/docs/_static/example_data.vcz/contig_id/0 b/docs/_static/example_data.vcz/contig_id/0
index 113105d5dbd5e8e0105d6cac25faca56be3dc137..016433bdd1678fbc84c1db61da6bdf1b563664f1 100644
GIT binary patch
literal 32
dcmZQ#G!_tGU|;}Y1t3la;tU`z0%AiT1^^vx0rda?

literal 20
VcmZQ#G-hF8U|;}Y5g;}IVgLpC0EYko

diff --git a/docs/usage.md b/docs/usage.md
index 7250acaa..baaf0d4c 100644
--- a/docs/usage.md
+++ b/docs/usage.md
@@ -32,7 +32,7 @@ document. However, for the moment we'll just use a pre-generated dataset:
 
 ```{code-cell} ipython3
 import zarr
-ds = zarr.load("_static/example_data.vcz")
+ds = zarr.open("_static/example_data.vcz")
 ```
 
 This is what the genotypes stored in that datafile look like:
@@ -44,7 +44,7 @@ assert all(len(np.unique(a)) == len(a) for a in ds['variant_allele'])
 assert any([np.sum(g) == 1 for g in ds['call_genotype']]) # at least one singleton
 assert any([np.sum(g) == 0 for g in ds['call_genotype']]) # at least one non-variable
 
-alleles = ds['variant_allele'].astype(str)
+alleles = ds['variant_allele'][:].astype(str)
 sites = np.arange(ds['call_genotype'].shape[0])
 print(" " * 22, "Site:", " ".join(str(x) for x in range(8)), "\n")
 for sample in range(ds['call_genotype'].shape[1]):
@@ -76,11 +76,11 @@ and not used for inference (with a warning given).
 import tsinfer
 
 # For this example take the REF allele (index 0) as ancestral
-ancestral_alleles = ds['variant_allele'][:,0].astype(str)
+ancestral_allele = ds['variant_allele'][:,0].astype(str)
 # This is just a numpy array, set the last site to an unknown value, for demo purposes
-ancestral_alleles[-1] = "."
+ancestral_allele[-1] = "."
 
-vdata = tsinfer.VariantData("_static/example_data.vcz", ancestral_alleles)
+vdata = tsinfer.VariantData("_static/example_data.vcz", ancestral_allele)
 ```
 
 The `VariantData` object is a lightweight wrapper around the .vcz file.
@@ -98,14 +98,40 @@ Additionally, during the inference step, additional sites can be flagged as not
 inference, for example if they are deemed unreliable (this is done
 via the `exclude_positions` parameter).
 
+Sites which are not used for inference will
+still be included in the final tree sequence, with mutations at those sites being placed
+onto branches by {meth}`parsimony<tskit.Tree.map_mutations>`. 
+
 ### Masks
 
-Sites which are not used for inference will still be included in the final tree sequence, with
-mutations at those sites being placed onto branches by parsimony. However, it is also possible
-to completely exclude sites and samples from the final tree sequence, by specifing a `site_mask`
-and/or a `sample_mask` when creating the `VariantData` object. Such sites or samples will be
-completely omitted from both inference and the final tree sequence. This can be useful, for
-example, to reduce the amount of computation required for an inference.
+It is also possible to *completely* exclude sites and samples, by specifing a boolean
+`site_mask` and/or a `sample_mask` when creating the `VariantData` object. Sites or samples with
+a mask value of `True` will be completely omitted both from inference and the final tree sequence.
+This can be useful, for example, if your VCF file contains multiple chromosomes (in which case
+`tsinfer` will need to be run separately on each chromosome) or if you wish to select only a subset
+of the chromosome for inference (e.g. to reduce computational load). If a `site_mask` is provided,
+note that the ancestral alleles array only specifies alleles for the unmasked sites.
+
+Below, for instance, is an example of including only sites up to position six in the contig
+labelled "chr1" in the `example_data.vcz` file:
+
+```{code-cell}
+import numpy as np
+
+# mask out any sites not associated with the contig named "chr1"
+# (for demonstration: all sites in this .vcz file are from "chr1" anyway)
+chr1_index = np.where(ds.contig_id[:] == "chr1")[0]
+site_mask = ds.variant_contig[:] != chr1_index
+# also mask out any sites with a position >= 6
+site_mask[ds.variant_position[:] >= 6] = True
+
+smaller_vdata = tsinfer.VariantData(
+    "_static/example_data.vcz",
+    ancestral_allele=ancestral_allele[site_mask == False],
+    site_mask=site_mask,
+)
+print(f"The `smaller_vdata` object returns data for only {smaller_vdata.num_sites} sites")
+```
 
 ### Topology inference