Segregating variants w/ Xarray

sgkit_is_segregating.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculating segregating variants using Xarray w/o first computing allele counts.\n",
    "\n",
    "See: https://scikit-allel.readthedocs.io/en/stable/model/ndarray.html?highlight=is_segregating#allel.AlleleCountsArray.is_segregating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([False,  True,  True, False, False,  True, False, False])"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "\n",
    "calls = np.array([\n",
    "    [[0, 0], [0, 0]],\n",
    "    [[0, 0], [0, 1]],\n",
    "    [[0, 2], [1, 1]],\n",
    "    [[2, 2], [-1, -1]],\n",
    "    [[2, -1], [2, -1]],\n",
    "    [[2, -1], [1, -1]],\n",
    "    [[2, -1], [-1, -1]],\n",
    "    [[-1, -1], [-1, -1]],\n",
    "])\n",
    "dims = ('variants', 'samples', 'ploidy')\n",
    "ds = xr.Dataset(dict(call_genotype=xr.DataArray(calls, dims=dims)))\n",
    "\n",
    "is_segregating = (\n",
    "    ds\n",
    "    # This looks for any calls that are different from\n",
    "    # a nan-aware mean (which would be the same w/ no segregation)\n",
    "    .assign(cgo=ds.call_genotype, cgf=ds.call_genotype >= 0)\n",
    "    .assign(cgm=lambda ds: ds.cgo.weighted(ds.cgf).mean(dim=('samples', 'ploidy')))\n",
    "    .pipe(lambda ds: ds.cgf & (ds.cgo != ds.cgm))\n",
    "    .any(dim=('samples', 'ploidy'))\n",
    ")\n",
    "is_segregating.values"
   ]
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}