3-simulating-mutations-with-msprime.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import msprime, tskit\n",
    "import numpy as np\n",
    "from IPython.display import SVG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Simulating mutations with msprime\n",
    "\n",
    "In this worksheet, we're going to focus on how to add simulated mutations to our existing tree sequences, should you need to look at actual genome **sequences**....\n",
    "\n",
    " - [3.0 Do you really need mutations?](#3.0DoYouReallyNeedMutations)\n",
    " - [3.1 The basic syntax](#3.1TheBasicSyntax)\n",
    " - [3.2 Mutation models](#3.2MutationModels)\n",
    " - [3.3 Mutation rates in different epochs](#3.3MutationRatesInDifferentEpochs)\n",
    " - [3.4 Exporting into other formats](#3.4ExportingIntoOtherFormats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='13.0DoYouReallyNeedMutations'></a>\n",
    "       \n",
    "## 3.0 Do you really need mutations?\n",
    "\n",
    "In tree sequences, the genetic genealogy exists independently of the mutations that\n",
    "generate the observed genetic variation.\n",
    "Our analyses may not actually require mutations: the genealogy on its own may be sufficent.\n",
    "This applies particularly to cases where you generate tree sequences by simulation,\n",
    "when you can be certain that your branch lengths are correct.\n",
    "\n",
    "Think carefully. Do you really need to analyse\n",
    "information about alleles and mutations? If not, here's why you could consider omitting\n",
    "them: \n",
    "\n",
    "1. Neutral mutations and sites can be added to a genealogy later\n",
    "2. Simulating sites and mutations increases memory requirements and computational time\n",
    "3. Storing mutations increases the tree sequence file size, and can slow some downstream analyses\n",
    "\n",
    "Consider the following simulation of samples from a 2-population [island model](https://tskit.dev/msprime/docs/stable/api.html#msprime.Demography.island_model):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = 1e7\n",
    "rho = 1e-8 \n",
    "n_subpops = 2\n",
    "subpop_size = 1e4\n",
    "migration_rate = 1e-4\n",
    "\n",
    "# Create a mutationless diploid tree sequence of n_subpops demes\n",
    "ts_orig = msprime.sim_ancestry(\n",
    "    samples={f\"pop_{i}\": 10 for i in range(n_subpops)},  # 10 samples from each subpop\n",
    "    demography=msprime.Demography.island_model([subpop_size] * n_subpops, migration_rate),\n",
    "    ploidy=2,\n",
    "    recombination_rate=rho,\n",
    "    sequence_length=L,\n",
    "    random_seed=123,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the `sim_mutations()` function to add neutral sites and mutations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = 1e-8\n",
    "\n",
    "ts_mutated = msprime.sim_mutations(ts_orig, rate=mu, random_seed=456)\n",
    "print(\n",
    "    \"Adding mutations has increased the tree sequence file size by \"\n",
    "    f\"{ts_mutated.nbytes / ts_orig.nbytes * 100:.0f}%\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analysis without variable sites\n",
    "\n",
    "Some genetic analyses are primarily focused on patterns of descent or ancestry. For instance, have you ever needed to study\n",
    "* local ancestry, global ancestry and identity-by-descent?\n",
    "* identification of most recent common ancestors and their descendants (including e.g. genealogical nearest neighbour analysis)?\n",
    "\n",
    "In cases like these, having the genealogy is usually sufficient to perform the analysis.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although many genetic analyses are based on patterns of genetic variation, for many purposes the genetic variation can be thought of as a measure of the relative length of branches on the local trees in a tree sequence. So while mutations are necessary to generate realistically variable genetic sequences, some statistical analyses do not necessarily require them to be present in a tree sequence. We'll talk about this further in the next notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A basic tree sequence to work with\n",
    "\n",
    "To emphasise that mutations and genealogy are truly separate in `msprime`, we will use the same simulated genetic genealogies (node and edge tables) throughout this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(2, sequence_length=100, random_seed=1987, recombination_rate=0.001)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_big = msprime.sim_ancestry(\n",
    "    100,\n",
    "    population_size = 1e4,\n",
    "    sequence_length=5e7,\n",
    "    random_seed=1982,\n",
    "    recombination_rate=1e-8\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3.1TheBasicSyntax'></a>\n",
    "## 3.1 The basic syntax\n",
    "\n",
    "To simulate mutations, apply the `sim_mutations()` method to an existing tree sequence object.\n",
    "At minimum, you must supply a per-base, per-generation mutation rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts = msprime.sim_mutations(ts, rate=0.01, random_seed=2016)\n",
    "SVG(mts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, the mutations are simulated under a discrete *Jukes-Cantor* model.\n",
    "(Under this model, there is an equal probability of each ancestral state (`A`, `C`, `G`, `T`), and an equal probability of each possible transition between these states (`A<->C`, `A<->G` etc.\n",
    "See [this](https://tskit.dev/msprime/docs/stable/api.html#msprime.JC69) for more information.)\n",
    "\n",
    "The output is yet another tree sequence, but this time with mutations on particular edges of the trees.\n",
    "There is now an additional provenance record to show that this dataset was obtained by applying `sim_mutations()` to the tree sequence initially simulated with `sim_ancestry()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Under the hood, `sim_mutations()` has added a *mutations table* to the original data:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts.tables.mutations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And a sites table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts.tables.sites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that some sites have experienced multiple mutations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To view the sequence information at each successive site, we can use the `variants()` iterator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for var in mts.variants():\n",
    "    print(var.site.position, var.alleles, var.genotypes, sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get all alleles at once, we can use `genotype_matrix()` (but **only** if the tree sequence is small!!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts.genotype_matrix()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1.1 Mutation rate variation\n",
    "\n",
    "You can specify mutation hotspots using the `RateMap` method as an argument to `rate`, instead of the single number:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ratemap = msprime.RateMap(position=[0, 40, 60, 100], rate=[0.01, 0.1, 0.01])\n",
    "mts = msprime.sim_mutations(ts, rate=ratemap, random_seed=104)\n",
    "SVG(mts.draw_svg(mutation_labels={}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Note: discrete vs continuous coordinates\n",
    " \n",
    " By default, the mutations generated by `sim_mutations()` will be at discrete positions.\n",
    " As with `sim_ancestry()`, we can also choose to place mutations at continuous (floating point) positions if wish.\n",
    " See the `discrete_genome` option [here](https://tskit.dev/msprime/docs/stable/api.html#msprime.sim_mutations),\n",
    " and note that the mutations produced with this option will conform to the classical `infinite sites` assumption`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3.2MutationModels'></a>\n",
    "## 3.2 Mutation models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Under the default Jukes-Cantor mutation model, there is an equal probability of each ancestral state (`A`, `C`, `G`, `T`), and an equal probability of each possible transition between these states (`A<->C`, `A<->G` etc).\n",
    "These properties are defined in the model's `root_distribution` and `transition_matrix`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msprime.JC69().alleles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msprime.JC69().root_distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "msprime.JC69().transition_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can summarise the mutations in this tree sequence with some extra code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_transitions(ts, alleles):\n",
    "    counts = np.zeros((len(alleles), len(alleles)), dtype='int')\n",
    "    for s in ts.sites():\n",
    "        aa = s.ancestral_state\n",
    "        for m in s.mutations:\n",
    "            pa = aa\n",
    "            da = m.derived_state\n",
    "            if m.parent != tskit.NULL:\n",
    "                pa = ts.mutation(m.parent).derived_state\n",
    "            counts[alleles.index(pa), alleles.index(da)] += 1\n",
    "    print(\"\\t\", \"\\t\".join(alleles))\n",
    "    for j, a in enumerate(alleles):\n",
    "        print(f\"{a}\\t\", \"\\t\".join(map(str, counts[j])))\n",
    "        \n",
    "count_transitions(mts, msprime.JC69().alleles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, we have a reasonably small number of mutations, so it's not necessarily obvious that our mutations satisfy the Jukes-Cantor model. This should be clearer looking at a larger dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts_big = msprime.sim_mutations(ts_big, rate=1e-8, random_seed=2016)\n",
    "count_transitions(mts_big, msprime.JC69().alleles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are *many* other mutations models you can choose from in `msprime`:\n",
    "\n",
    " - [BinaryMutationModel](https://tskit.dev/msprime/docs/stable/api.html#msprime.BinaryMutationModel): Basic binary mutation model with two flip-flopping alleles: “0” and “1”.\n",
    "\n",
    " - [JC69](https://tskit.dev/msprime/docs/stable/api.html#msprime.JC69): Jukes & Cantor model (‘69), equal probability of transitions between nucleotides\n",
    "\n",
    " - [HKY](https://tskit.dev/msprime/docs/stable/api.html#msprime.HKY): Hasegawa, Kishino & Yano model (‘85), different probabilities for transitions and transversions\n",
    "\n",
    " - [F84](https://tskit.dev/msprime/docs/stable/api.html#msprime.F84): Felsenstein model (‘84), different probabilities for transitions and transversions\n",
    "\n",
    " - [GTR](https://tskit.dev/msprime/docs/stable/api.html#msprime.GTR): The Generalised Time-Reversible nucleotide mutation model, a general parameterisation of a time-reversible mutation process\n",
    "\n",
    " - [BLOSUM62](https://tskit.dev/msprime/docs/stable/api.html#msprime.BLOSUM62): The BLOSUM62 model of time-reversible amino acid mutation\n",
    "\n",
    " - [PAM](https://tskit.dev/msprime/docs/stable/api.html#msprime.PAM): The PAM model of time-reversible amino acid mutation\n",
    "\n",
    " - [MatrixMutationModel](https://tskit.dev/msprime/docs/stable/api.html#msprime.MatrixMutationModel): Superclass of the specific mutation models with a finite set of states\n",
    "\n",
    " - [InfiniteAlleles](https://tskit.dev/msprime/docs/stable/api.html#msprime.InfiniteAlleles): A generic infinite-alleles mutation model\n",
    "\n",
    " - [SLiMMutationModel](https://tskit.dev/msprime/docs/stable/api.html#msprime.SLiMMutationModel): An infinite-alleles model of mutation producing SLiM-style mutations\n",
    "\n",
    "Any of these alternative models can be specified using the `model` argument of `sim_mutations()`.\n",
    "We'll practice with one of them:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2.1 An alternative example: Transitions and transversions\n",
    "\n",
    "In some situations we want nucleotide transitions (`A<->G`), (`C<->T`) (ie. exchanges of bases with a similar shape) to be more likely than nucleotide tranversions (all others).\n",
    "We can do this with the [Hasegawa, Kishino & Yano (HKY) model](https://tskit.dev/msprime/docs/stable/api.html#msprime.HKY).\n",
    "In addition to an overall mutation `rate`, you also specify `kappa`, a constant scaling parameter such that the probability of each type of transition is `kappa` times the probability of each type of transversion."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most of our simulated mutations are now transitions. (Also, note that the site with multiple mutations now experiences a back mutation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts_big = msprime.sim_mutations(\n",
    "    ts_big,\n",
    "    rate=1e-8,\n",
    "    random_seed=2016,\n",
    "    model=msprime.HKY(kappa=20))\n",
    "\n",
    "count_transitions(mts_big, msprime.HKY(kappa=20).alleles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is now a very small (non-0) probability of *silent mutations*, transitions that do not change the type of the allele. This behaviour is an artifact of the model adjustments that `msprime` makes to enable mutation 'stackability'. It usually doesn't matter too much, unless your mutation model has a large number of possible mutation allele types -- for more information on this, see [this discussion in the documentation](https://tskit.dev/msprime/docs/stable/mutations.html#sec-mutations-adjusting-for-silent)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "msprime.HKY(kappa=20).transition_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3.3MutationRatesInDifferentEpochs'></a>\n",
    "### 3.2.2 Mutation rates in different epochs\n",
    "\n",
    "Use the `start_time` and `end_time` arguments to simulate mutations only within specific timeframes.\n",
    "This can be useful if you want to simulate a change in mutation rates over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts = msprime.sim_mutations(ts, rate=0.01, random_seed=1714, start_time=2)\n",
    "SVG(mts.draw_svg(mutation_labels={}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mts.tables.mutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmts = msprime.sim_mutations(mts, rate=0.1, random_seed=851, start_time=1,\n",
    "                            end_time=2)\n",
    "SVG(mmts.draw_svg(mutation_labels={}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmts.tables.mutations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the original mutation has been re-numbered according to its age."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2.3 Stacking mutations\n",
    "\n",
    "As the previous example shows,\n",
    "we can apply `sim_mutations()` to *any* tree sequence, including one that already has mutations on it.\n",
    "This allows us to 'stack' mutations, which can be useful if you wish to simulate several different types of mutations from different models, or perhaps over different time periods. \n",
    "\n",
    "For instance, here's a tree sequence with some mutations that we made earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SVG(mts.draw_svg(mutation_labels={})) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can run `sim_mutations()` once more to overlay more mutations on top of this tree sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmts = msprime.sim_mutations(mts, rate=0.01, random_seed=1959)\n",
    "SVG(mmts.draw_svg(mutation_labels={}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's also no reason why we can't apply `sim_mutations()` to a tree sequence generated outside of `msprime`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2.4 Make-your-own mutation model!\n",
    "\n",
    "Mutation models consist of a few key ingredients:\n",
    "\n",
    " - **alleles** (the possible 'choices' you can see at each unit)\n",
    " - **root distribution** (what's the distribution of ancestral alleles?)\n",
    " - **transition matrix** (what's the probability of mutating from one allele to another?)\n",
    " \n",
    "You can use these to define your own (finite-sites) mutation model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = msprime.MatrixMutationModel(\n",
    "    alleles = [\"💩\", \"🎄\", \"🔥\"],\n",
    "    root_distribution = [1.0, 0.0, 0.0],\n",
    "    transition_matrix = [[0.0, 1.0, 0.0],\n",
    "                         [0.0, 0.8, 0.2],\n",
    "                         [1.0, 0.0, 0.0]]\n",
    ")\n",
    "mts = msprime.sim_mutations(\n",
    "    ts, rate=0.01, random_seed=1215112, model=model)\n",
    "\n",
    "for var in mts.variants():\n",
    "    print(var.site.position, var.alleles, var.genotypes, sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Exercise*. For most of their evolutionary history, mutations in your study organisms were well described by the HKY model. \n",
    "However, they were all were exposed to an environmental mutagen 50 generations ago, and from then on all new mutations were of type 👽. All of the nucleotide bases have an equal probability of mutating to an allele of type 👽. Once a nucleotide mutates to a 👽, it cannot mutate back to a normal nucleotide. Modify the code below to simulate this scenario."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The underlying genealogy\n",
    "ts_ex = msprime.sim_ancestry(\n",
    "    10, sequence_length=1000, random_seed=1987, recombination_rate=0.001,\n",
    "    population_size=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify code here \n",
    "alien_model = msprime.MatrixMutationModel(\n",
    "    alleles = ['A', 'C', 'G', 'T'],\n",
    "    root_distribution =  [0.25, 0.25, 0.25, 0.25],\n",
    "    transition_matrix = [[0, 1/3, 1/3, 1/3],\n",
    "       [1/3, 0, 1/3, 1/3],\n",
    "       [1/3, 1/3, 0, 1/3],\n",
    "       [1/3, 1/3, 1/3, 0]]\n",
    ")\n",
    "\n",
    "mts_ex = msprime.sim_mutations(\n",
    "    ts_ex, rate=1e-5, random_seed=2016, model=msprime.HKY(kappa=20),\n",
    "     start_time=0, end_time=None\n",
    ")\n",
    "mts_ex = msprime.sim_mutations(\n",
    "    mts_ex, rate=1e-5, random_seed=2299, model=alien_model,\n",
    "     start_time=0, end_time=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check here\n",
    "print(mts_ex.tables.mutations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check here\n",
    "for var in mts_ex.variants():\n",
    "    print(var.site.position, var.alleles, var.genotypes, sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3.4ExportingSequenceDataIntoOtherFormats'></a>\n",
    "## 3.3 Exporting sequence data into other formats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3.1 To VCF\n",
    "\n",
    "We'll use the `write_vcf()` method to export our new simulated sequences to a VCF file for use with external software:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"worksheet3-output.vcf\", \"w\") as vcf_file:\n",
    "    mts.write_vcf(vcf_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"worksheet3-output.vcf\", \"r\") as f:\n",
    "    print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and we may even want to specify our own individual names:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indv_names = [\"platypus_1\", \"platypus_2\"]\n",
    "with open(\"worksheet3-output.vcf\", \"w\") as vcf_file:\n",
    "    mts.write_vcf(vcf_file, individual_names=indv_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"worksheet3-output.vcf\", \"r\") as f:\n",
    "    print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3.2 To other Python objects\n",
    "\n",
    "Many of the underlying methods in `tskit` produce `numpy` objects as output.\n",
    "Because of this, tree sequences play well with other Python libraries, including plotting libraries like `matplotlib` and scientific libraries like `scikit-allel`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For instance, the `genotype_matrix()` method will return an object that works as a HaplotypeArray in `scikit-allel`, giving you easy access to its summary functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import allel\n",
    "\n",
    "gens = mts_big.genotype_matrix()\n",
    "haps = allel.HaplotypeArray(gens)\n",
    "allele_counts = haps.count_alleles()\n",
    "allel.mean_pairwise_difference(allele_counts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, many of these calculations can also be done within `tskit` itself.\n",
    "We'll see this in the next notebook!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ts2025",
   "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": 4
}