4-tree-sequence-inference-and-statistics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Tree sequence inference and statistics\n",
    "\n",
    "So far, we've focused on *simulating* data with a tree sequence structure, and we've talked about the benefits of doing so. We've looked at some basic summary statistics of the tree sequences to try and interpret them, but you might be wondering,\n",
    " - So far we've been doing all this on simulated datasets. Can I use any of these techniques on my real dataset?\n",
    "  - what if I want to perform 'conventional' analyses on these? Do I need to convert my tree sequences out to VCFs and then apply all my regular tools, or can I do any of these analyses directly on the tree sequence files?\n",
    "\n",
    "In both cases, the answer is Yes! \n",
    "\n",
    "1. The `tsinfer` and `tsdate` packages let you estimate a tree sequence for your dataset. (Note: these are *estimates* -- they won't be exactly correct. See Brandt et al for some discussion of these, and what types of information may or may not be legit for these)\n",
    "2. The `tskit` package (which we've already seen and used) has a number of utility functions that let you explore, manipulate and analyse data stored in a tree sequence format. \n",
    "\n",
    "In this notebook, we'll touch on all of these to give you a flavour of what is possible with these packages.\n",
    "\n",
    " - [4.1 An overview of tsinfer](#5.1Overview)\n",
    " - [4.2 Hands on with tsinfer](#5.2HandsOn)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cyvcf2 # For reading VCF files into Python\n",
    "\n",
    "import tskit\n",
    "import tsinfer, tsdate\n",
    "from IPython.display import SVG\n",
    "import json\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='5.1Overview'></a>\n",
    "## 4.1 An overview of `tsinfer`\n",
    "\n",
    "Simulating a tree sequence is relatively simple compared to *inferring* a tree sequence from existing data.\n",
    "The [tsinfer software](https://tsinfer.readthedocs.io/en/stable/) implements a heuristic algorithm which does this in a scalable manner.\n",
    "\n",
    "`Tsinfer` (pronounced t-s-infer) is comparable in some ways to other ancestral inference software such as [ARGweaver](https://doi.org/10.1371/journal.pgen.1004342), [Relate](https://myersgroup.github.io/relate/), and [Rent+](https://doi.org/10.1093/bioinformatics/btw735). However, it differs considerably in approach and scalability.\n",
    "Note that none of these other software packages produce tree sequences as output, although is possible to convert their output to tree sequences.\n",
    "Also note that `tsinfer` produces trees with a relatively accurate topology, but unlike other ancestral inference tools, it makes no attempt at the moment to produce precise branch length estimates -- for this we need another tool like `tsdate`.\n",
    "\n",
    "An important restriction is that `tsinfer` requires phased sample sequences with known ancestral states for each variant. It also works better with full sequence data than with data from scattered target SNPs (e.g. as obtained from SNP chips)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "We'll be looking at a dataset of variants from sparrow chromosome 24 kindly shared by Mark Ravinet. There are 10 individuals represented in the samples: 5 from Norway and 5 from France.\n",
    "\n",
    "### 5.1.1 Preparing the data for `tsinfer`\n",
    "\n",
    "To prepare our data for `tsinfer`,\n",
    "we need to make a `SampleData` object with information about the individuals and phased, bi-allelic sites we wish to use for inference.\n",
    "\n",
    "These things are typically done with the `add_individual()` and `add_sites()` methods.\n",
    "To save us some time, let's define some helper functions (taken from the tsinfer website) that will make this object for us from a `cyvcf2` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some functions for importing data from VCF --> tsinfer format.\n",
    "# Properties these functions preserve:\n",
    "# - Our data is diploid\n",
    "# - Our individuals belong to two distinct populations, and we want our sample data to include that information.\n",
    "# - we only want the inference to use alleles at SNVs and deletions for now (others are possible!)\n",
    "\n",
    "def add_diploid_sites(vcf, samples):\n",
    "    \"\"\"\n",
    "    Read the sites in the vcf and add them to the samples object.\n",
    "    \"\"\"\n",
    "    # You may want to change the following line, e.g. here we allow\n",
    "    # \"*\" (a spanning deletion) to be a valid allele state\n",
    "    allele_chars = set(\"ATGCatgc*\")\n",
    "    pos = 0\n",
    "    # progressbar = tqdm(total=samples.sequence_length, desc=\"Read VCF\", unit='bp')\n",
    "    for variant in vcf:  # Loop over variants, each assumed at a unique site\n",
    "        # progressbar.update(variant.POS - pos)\n",
    "        if pos == variant.POS:\n",
    "            print(f\"Duplicate entries at position {pos}, ignoring all but the first\")\n",
    "            continue\n",
    "        else:\n",
    "            pos = variant.POS\n",
    "        if any([not phased for _, _, phased in variant.genotypes]):\n",
    "            raise ValueError(\"Unphased genotypes for variant at position\", pos)\n",
    "        alleles = [variant.REF.upper()] + [v.upper() for v in variant.ALT]\n",
    "        ancestral = variant.INFO.get(\"AA\", \".\")  # \".\" means unknown\n",
    "        # some VCFs (e.g. from 1000G) have many values in the AA field: take the 1st\n",
    "        ancestral = ancestral.split(\"|\")[0].upper()\n",
    "        if ancestral == \".\" or ancestral == \"\":\n",
    "            ancestral_allele = tskit.MISSING_DATA\n",
    "            # alternatively, you could specify `ancestral = variant.REF.upper()`\n",
    "        else:\n",
    "            ancestral_allele = alleles.index(ancestral)\n",
    "        # Check we have ATCG alleles\n",
    "        for a in alleles:\n",
    "            if len(set(a) - allele_chars) > 0:\n",
    "                print(f\"Ignoring site at pos {pos}: allele {a} not in {allele_chars}\")\n",
    "                continue\n",
    "        # Map original allele indexes to their indexes in the new alleles list.\n",
    "        genotypes = [g for row in variant.genotypes for g in row[0:2]]\n",
    "        samples.add_site(pos, genotypes, alleles, ancestral_allele=ancestral_allele)\n",
    "\n",
    "\n",
    "def chromosome_length(vcf):\n",
    "    assert len(vcf.seqlens) == 1\n",
    "    return vcf.seqlens[0]\n",
    "\n",
    "def add_populations(vcf, samples):\n",
    "    \"\"\"\n",
    "    Add tsinfer Population objects and returns a list of IDs corresponding to the VCF samples.\n",
    "    \"\"\"\n",
    "    # In this VCF, the first letter of the sample name refers to the population\n",
    "    samples_first_letter = [sample_name[0] for sample_name in vcf.samples]\n",
    "    pop_lookup = {}\n",
    "    pop_lookup[\"8\"] = samples.add_population(metadata={\"country\": \"Norway\"})\n",
    "    pop_lookup[\"F\"] = samples.add_population(metadata={\"country\": \"France\"})\n",
    "    return [pop_lookup[first_letter] for first_letter in samples_first_letter]\n",
    "\n",
    "\n",
    "def add_diploid_individuals(vcf, samples, populations):\n",
    "    for name, population in zip(vcf.samples, populations):\n",
    "        samples.add_individual(ploidy=2, metadata={\"name\": name}, population=population)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use these helper functions to read in information about our sparrow sample from the VCF file and save it into a `tsinfer.SampleData` object.\n",
    "First, we add in information about our populations.\n",
    "Then we add information about the sampled individuals in those populations,\n",
    "and finally then add the genotypes held by those individuals at each site."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vcf_location = \"P_dom_chr24_phased.vcf.gz\"\n",
    "vcf = cyvcf2.VCF(vcf_location)\n",
    "with tsinfer.SampleData(\n",
    "    path=\"P_dom_chr24_phased.samples\", sequence_length=chromosome_length(vcf)\n",
    ") as samples:\n",
    "    populations = add_populations(vcf, samples)\n",
    "    add_diploid_individuals(vcf, samples, populations)\n",
    "    add_diploid_sites(vcf, samples)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    \"Sample file created for {} samples \".format(samples.num_samples)\n",
    "    + \"({} individuals) \".format(samples.num_individuals)\n",
    "    + \"with {} variable sites.\".format(samples.num_sites),\n",
    "    flush=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1.2 Inferring a tree sequence with `tsinfer` and `tsdate`\n",
    "\n",
    "We're now ready to run `tsinfer`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_undated = tsinfer.infer(samples)\n",
    "ts_undated"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As usual, there's some useful information about the fact that we generated this with tsinfer in the provenance section of the table.\n",
    "\n",
    "Do you notice anything different about this tree sequence compared with the (simulated) ones that we've been working with up to this point? \n",
    "1. \"Time Units: uncalibrated\". That's because tsinfer infers the topologies for these trees, but doesn't infer times for them. For this, we'll apply `tsdate`.\n",
    "2. More sites than mutations. In the process of generating the trees, `tsinfer` also generated a bunch of ancestral haplotypes that *didn't* contribute any info to the final genotypes we see.\n",
    "\n",
    "We'll fix (2) by applying `simplify()` to remove some 'stray' unary nodes and sites outputted by tsinfer:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mutation rate is given as the probability of mutations per base per generation\n",
    "tss = ts_undated.simplify()\n",
    "tss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Then we'll run `tsdate` to obtain a time-calibrated version of this inferred tree sequence using an estimated mutation rate of `1e-8` per base per generation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_dated = tsdate.date(tss, mutation_rate=1e-8)\n",
    "ts_dated"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that each of these steps has added an extra row into the provenance table too.\n",
    "\n",
    "Let's print some trees to see what this output looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colours = {\"Norway\": \"red\", \"France\": \"blue\"}\n",
    "colours_for_node = {}\n",
    "for n in ts_dated.samples():\n",
    "    population_data = ts_dated.population(ts_dated.node(n).population)\n",
    "    colours_for_node[n] = colours[json.loads(population_data.metadata)[\"country\"]]\n",
    "\n",
    "individual_for_node = {}\n",
    "for n in ts_dated.samples():\n",
    "    individual_data = ts_dated.individual(ts_dated.node(n).individual)\n",
    "    individual_for_node[n] = json.loads(individual_data.metadata)[\"name\"]\n",
    "\n",
    "tree = ts_dated.at(4e6)\n",
    "tree.draw(\n",
    "    height=700,\n",
    "    width=1500,\n",
    "    node_labels=individual_for_node,\n",
    "    node_colours=colours_for_node\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Many of these samples are clustering by population label, but not all of them are. Let's look at another tree at a different position where there even fewer sample nodes are clustering by population:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = ts_dated.at(2e6)\n",
    "tree.draw(\n",
    "    path=\"tree_at_1Mb.svg\",\n",
    "    height=700,\n",
    "    width=1500,\n",
    "    node_labels=individual_for_node,\n",
    "    node_colours=colours_for_node,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The amount we can learn from any individual tree is perhaps limited, because of the stochasticity of the trees across the genome.\n",
    "\n",
    "Can we summarise whether the Norwegian and French sparrows exhibit population structure, and learn more about the nature of that structure?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: I know that we're getting to a point where printing individual trees might be a little overwhelming and hard to interpret.\n",
    "The `tsbrowse` application provides a number of useful functions that might help you visualise and summarise these larger tree sequences.\n",
    "See [here](https://tskit.dev/software/tsbrowse.html) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.2 Summary statistics with tskit\n",
    "\n",
    "Suppose you have $n$ sequences typed at $m$ different sites...\n",
    "\n",
    "```\n",
    "   ...GTAACGCGATAAGAGATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGAGATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGTAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGTAATAGCGTA...\n",
    "```\n",
    "\n",
    "...and you want to calculate mean pairwise diversity on these samples, i.e.\n",
    "\n",
    "$$ \\pi = \\dfrac{1}{n(n-1)/2}\\sum_{i=1}^{n-1} \\sum_{j=i+1}^n k_{ij}, $$\n",
    "\n",
    "where $k_{ij}$ is the number of sites at which sequences $i$ and $j$ carry a different allele.\n",
    "The scaling of this procedure is\n",
    "\n",
    "$$ O\\left( n^2 m \\right) $$\n",
    "\n",
    "ie. quadratic in the number of samples $n$, and linear in the number of sites $m$.\n",
    "However, there is an equivalent way of performing this calculation by assigning weights to the sample nodes, and propagating these values further up the tree using a 'summary function' at each mutation.\n",
    "\n",
    " <img src=\"pics/worksheet4-node-weights.jpeg\" width=\"350\" height=\"350\">\n",
    " \n",
    " This is what `tskit` calls a *site statistic* calculation, and because the operation is of order\n",
    " \n",
    " $$ O\\left( n + \\rho m (\\log(n))^2 \\right)  << O\\left( n^2 m \\right), $$\n",
    " \n",
    "the calculation is quick to run, especially on large datasets:\n",
    "\n",
    " <img src=\"pics/worksheet4-stat-speed.jpeg\" width=\"500\" height=\"500\">\n",
    "\n",
    "See the following paper for more details.\n",
    "\n",
    "Peter Ralph, Kevin Thornton, Jerome Kelleher, Efficiently Summarizing Relationships in Large Samples: A General Duality Between Statistics of Genealogies and Genomes, Genetics, Volume 215, Issue 3, 1 July 2020, Pages 779–797, https://doi.org/10.1534/genetics.120.303253"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The basic syntax: (nucleotide diversity)\n",
    "\n",
    "`tskit` uses very similar syntax for all of its inbuilt statistics, so we'll explore the options using `diversity()` as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_dated.diversity()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, `tskit` presents a normalised version of the statistic  scaled by the length of the region represented in `ts`. This allows you to make comparisons between different tree sequences that may be of different lengths. However, this isn't how all other genetic software computes diversity -- if you wish to disable `tskit`'s default behaviour, use the `span_normalise` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div = ts_dated.diversity(span_normalise=False)\n",
    "div[()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating statistics on subsets of the samples\n",
    "\n",
    "Remember that this dataset consists of samples from two contemporary populations here, of different sizes. We’d expect them to each have different diversity levels, and for these to differ from the overall (sample-wide) diversity rate. We can get this information out by specifying each of these with the `sample_nodes` argument.\n",
    "\n",
    "A quick and simple way to get all of the sample node IDs from a particular population is to use the `samples()` method. For instance, the following code returns a numpy array holding all of the sample node IDs from 'population 0':"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_of_interest=[ts_dated.samples(population=0),\n",
    "                          ts_dated.samples(population=1),\n",
    "                          ts_dated.samples()]\n",
    "\n",
    "ts_dated.diversity(sample_sets=samples_of_interest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output is a 1-dimensional numpy array, where each element is a diversity statistic value for one of the sample sets we specified.\n",
    "Nucleotide diversity is lowest in the set of samples from the small population, and largest in the pooled set of samples, as you'd expect.\n",
    "\n",
    "Note that you can use any list of node IDs as inputs to `sample_sets`. This may be useful if your samples of interest correspond to something other than populations (for instance, samples that hold some phenotype of interest)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_dated.diversity(sample_sets=[[12, 14, 18, 7, 8],\n",
    "                          [3, 19, 10]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Genome scans\n",
    "\n",
    "So far, we’ve just been calculating statistics summarising diversity values along the entire simulated genome. However, in many cases, we might be more interested in how diversity varies along the genome. We can do this using the `windows` argument.\n",
    "\n",
    "We specify the start and end points of the sequence, and the locations of the breakpoints between each window. For instance, suppose we wanted to specify some windows of length 1Mb covering our 50Mb chromosome:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "breakpoints = [i*1e5 for i in range(0, int(ts_dated.sequence_length/1e5))] + [ts_dated.sequence_length] \n",
    "print(breakpoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div = ts_dated.diversity(sample_sets=samples_of_interest, windows=breakpoints)\n",
    "print(\"Dimensions of the output:\", div.shape, \"\\n\")\n",
    "print(\"Diversity values over the first 10 windows in each sample set:\")\n",
    "print(div[:10,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output holds one row for each of our specified windows,\n",
    "and each element of the row holds diversity value in some particular window amongst one of our sample sets.\n",
    "Let’s plot these:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "names_to_plot = ['Norway', 'France', 'ALL']\n",
    "lines = plt.plot(breakpoints[:-1], div)\n",
    "plt.grid(alpha=0.5)\n",
    "plt.legend(lines, names_to_plot);\n",
    "plt.xticks()\n",
    "plt.xlabel(\"Position on chromosome\")\n",
    "plt.title(\"Windowed diversity values\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='4.2BranchStatistics'></a>\n",
    "### 4.2.2 Branch statistics and the 'duality' of  tree-based statistics\n",
    "\n",
    "There are several different types of randomness in genetic models that interact with each other in complex ways.\n",
    "In addition to randomness in the genealogical trees that are produced in a given demographic scenario,\n",
    "there is also randomness caused by the mutational process.\n",
    "\n",
    "When you calculate site statistics, or anything based on allele frequencies, *both* of these processes contribute to the statistical noisiness you see.\n",
    "\n",
    "However in tree sequences, you have information about branches which allows you to bypass this latter type of mutation.\n",
    "Instead of moving upwards along the trees and updating the statistic every time you come across a mutation, you can update the statistic based on the lengths of the branches.\n",
    "This should have some correspondence with the *number of mutations we may expect*. (This should certainly be true in simulated datasets, where we are certain of the correctness of the underlying trees).\n",
    "\n",
    "This is the basic idea behind the *branch statistics* in `tskit`.\n",
    "\n",
    "Here are the diversity stats we looked at before, this time with the branch versions included. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div_branch = ts_dated.diversity(\n",
    "    sample_sets=samples_of_interest,\n",
    "    windows=breakpoints,\n",
    "    mode='branch')\n",
    "names_to_plot = ['Norway', 'France', 'ALL']\n",
    "lines = plt.plot(breakpoints[:-1], div_branch)\n",
    "plt.grid(alpha=0.5)\n",
    "plt.legend(lines, names_to_plot);\n",
    "plt.title(\"Windowed diversity (branch)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extra material -- GNNs\n",
    "\n",
    "Consider removing, or turning into a supplementary notebook -- I think we just won't have time for this!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def find_nearest_neighbour_populations(t, focal_sample):\n",
    "    \"\"\"\n",
    "    Find the nearest neighbour populations for a given focal sample.\n",
    "    \"\"\"\n",
    "    p = t.parent(focal_sample)\n",
    "    if p == -1:\n",
    "        return (0, 0)\n",
    "    children = [c for c in t.leaves(u = p) if c != focal_sample]\n",
    "    relative_populations = [ts_dated.node(c).population for c in children]  \n",
    "    # Calculate the proportion of children in each population\n",
    "    pop_proportions = {}\n",
    "    for pop in relative_populations:\n",
    "        if pop not in pop_proportions:\n",
    "            pop_proportions[pop] = 0\n",
    "        pop_proportions[pop] += 1\n",
    "    # Count how many of the items in relative_populations are 0\n",
    "    count_pop_0 = sum(1 for pop in relative_populations if pop == 0)\n",
    "    count_pop_1 = len(relative_populations) - count_pop_0\n",
    "    # Calculate proportion of children in each population\n",
    "    prop_pop_0 = count_pop_0 / len(relative_populations)\n",
    "    prop_pop_1 = count_pop_1 / len(relative_populations)\n",
    "\n",
    "    return prop_pop_0, prop_pop_1\n",
    "\n",
    "\n",
    "# test\n",
    "t = ts_dated.at(3.49e6)\n",
    "focal_sample = 0\n",
    "\n",
    "print(find_nearest_neighbour_populations(t, focal_sample))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_interval(t, focal_sample, seq_length, ax):\n",
    "    nns = find_nearest_neighbour_populations(t, focal_sample)\n",
    "    i = t.interval\n",
    "    ax.add_patch(plt.Rectangle((i[0]/seq_length, 0), (i[1] - i[0])/seq_length, nns[0], color=\"red\"))\n",
    "    ax.add_patch(plt.Rectangle((i[0]/seq_length, nns[0]), (i[1] - i[0])/seq_length, nns[0] + nns[1], color=\"blue\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in ts_dated.individuals():\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def plot_proportions_for_individual(tree_seq, individual_id):\n",
    "\n",
    "    for i in tree_seq.individuals():\n",
    "        if i.id == individual_id:\n",
    "            node1, node2 = i.nodes\n",
    "            break\n",
    "\n",
    "    seq_length = tree_seq.sequence_length\n",
    "\n",
    "    # Plot it all\n",
    "    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 2.5))\n",
    "    # for ax in [ax1, ax2]:\n",
    "        # ax.set_xticks([])\n",
    "        # ax.set_xticklabels([])\n",
    "    for t in tree_seq.trees(sample_lists=True):\n",
    "        plot_interval(t, node1, seq_length, ax1) # CHANGE focal sample\n",
    "    ax1.set_yticks([])\n",
    "    ax1.set_yticklabels([])\n",
    "    ax1.set_ylabel(\"Hap 1\")\n",
    "    # Repeat the process for the second focal node\n",
    "    for t in tree_seq.trees(sample_lists=True):\n",
    "        plot_interval(t, node2, seq_length, ax2) # CHANGE focal sample\n",
    "    ax2.set_yticks([])\n",
    "    ax2.set_yticklabels([])\n",
    "    ax2.set_ylabel(\"Hap 2\")\n",
    "\n",
    "    # Add legend\n",
    "    colors = ['red', 'blue']\n",
    "    labels = ['Norway', 'France']\n",
    "    handles = [plt.Rectangle((0, 0), 1, 1, color=color) for color in colors]\n",
    "    ax2.legend(handles, labels, loc='upper right', bbox_to_anchor=(1.2, 1.5))\n",
    "\n",
    "    # Add overall title\n",
    "    fig.suptitle(f\"Nearest neighbour populations, sample {individual_id}\", fontsize=16)\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# test\n",
    "plot_proportions_for_individual(ts_dated, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 6)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 7)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_proportions_for_individual(ts_dated, 9)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}