tests/monotonicsplineepistasisbottlenecklikelihood_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `MonotonicSplineEpistasisBottleneckLikelihood` model\n",
    "Simulates some variants and analyzes them under the `globalepistasis.MonotonicSplineEpistasisBottleneckLikelihood` model.\n",
    "\n",
    "This notebook is designed as a test, not for user documentation--and it uses some private methods that should not be used in real analyses."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up\n",
    "Import Python modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:40.857077Z",
     "iopub.status.busy": "2024-02-06T21:52:40.855312Z",
     "iopub.status.idle": "2024-02-06T21:52:44.027523Z",
     "shell.execute_reply": "2024-02-06T21:52:44.026392Z",
     "shell.execute_reply.started": "2024-02-06T21:52:40.857022Z"
    }
   },
   "outputs": [],
   "source": [
    "import cProfile\n",
    "import math\n",
    "import pstats\n",
    "import random\n",
    "import tempfile\n",
    "import time\n",
    "import warnings\n",
    "\n",
    "import binarymap\n",
    "\n",
    "import numpy\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from plotnine import *\n",
    "\n",
    "import scipy\n",
    "\n",
    "import dms_variants.globalepistasis\n",
    "import dms_variants.codonvarianttable\n",
    "import dms_variants.simulate\n",
    "from dms_variants.constants import CBPALETTE, CODONS_NOSTOP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameters for simulations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:44.032639Z",
     "iopub.status.busy": "2024-02-06T21:52:44.032162Z",
     "iopub.status.idle": "2024-02-06T21:52:44.038591Z",
     "shell.execute_reply": "2024-02-06T21:52:44.037876Z",
     "shell.execute_reply.started": "2024-02-06T21:52:44.032588Z"
    }
   },
   "outputs": [],
   "source": [
    "seed = 1  # random number seed\n",
    "genelength = 30  # gene length in codons\n",
    "bclen = 16  # length of barcodes\n",
    "libs = [\"lib_1\"]  # distinct libraries of gene\n",
    "variants_per_lib = 1000 * genelength  # variants per library\n",
    "avgmuts = 2.5  # average codon mutations per variant\n",
    "variant_error_rate = 0  # rate variants are mis-sequenced\n",
    "depth = 5000 * variants_per_lib  # per-sample sequencing depth\n",
    "lib_uniformity = 3  # uniformity of library pre-selection\n",
    "noise = 0.0  # non-statistical noise in selections\n",
    "bottleneck = 5 * variants_per_lib\n",
    "bottlenecks = {  # bottlenecks from pre- to post-selection\n",
    "    \"loose_bottle\": bottleneck,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seed random number generators:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:44.042635Z",
     "iopub.status.busy": "2024-02-06T21:52:44.042162Z",
     "iopub.status.idle": "2024-02-06T21:52:44.046542Z",
     "shell.execute_reply": "2024-02-06T21:52:44.045658Z",
     "shell.execute_reply.started": "2024-02-06T21:52:44.042595Z"
    }
   },
   "outputs": [],
   "source": [
    "random.seed(seed)\n",
    "numpy.random.seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppress warnings that clutter output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:44.051013Z",
     "iopub.status.busy": "2024-02-06T21:52:44.050640Z",
     "iopub.status.idle": "2024-02-06T21:52:44.055314Z",
     "shell.execute_reply": "2024-02-06T21:52:44.054468Z",
     "shell.execute_reply.started": "2024-02-06T21:52:44.050976Z"
    }
   },
   "outputs": [],
   "source": [
    "warnings.simplefilter(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate variants\n",
    "Simulate wildtype gene sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:44.059285Z",
     "iopub.status.busy": "2024-02-06T21:52:44.058938Z",
     "iopub.status.idle": "2024-02-06T21:52:44.064059Z",
     "shell.execute_reply": "2024-02-06T21:52:44.063294Z",
     "shell.execute_reply.started": "2024-02-06T21:52:44.059254Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wildtype gene of 30 codons:\n",
      "AGATCCGTGATTCTGCGTGCTTACACCAACTCACGGGTGAAACGTGTAATCTTATGCAACAACGACTTACCTATCCGCAACATCCGGCTG\n"
     ]
    }
   ],
   "source": [
    "geneseq = \"\".join(random.choices(CODONS_NOSTOP, k=genelength))\n",
    "print(f\"Wildtype gene of {genelength} codons:\\n{geneseq}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulate a `CodonVariantTable`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:44.068296Z",
     "iopub.status.busy": "2024-02-06T21:52:44.067795Z",
     "iopub.status.idle": "2024-02-06T21:52:46.137225Z",
     "shell.execute_reply": "2024-02-06T21:52:46.136393Z",
     "shell.execute_reply.started": "2024-02-06T21:52:44.068247Z"
    }
   },
   "outputs": [],
   "source": [
    "variants = dms_variants.simulate.simulate_CodonVariantTable(\n",
    "    geneseq=geneseq,\n",
    "    bclen=bclen,\n",
    "    library_specs={\n",
    "        lib: {\"avgmuts\": avgmuts, \"nvariants\": variants_per_lib} for lib in libs\n",
    "    },\n",
    "    seed=seed,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Number of amino-acid mutations per variant in library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:46.140743Z",
     "iopub.status.busy": "2024-02-06T21:52:46.140454Z",
     "iopub.status.idle": "2024-02-06T21:52:47.199942Z",
     "shell.execute_reply": "2024-02-06T21:52:47.199188Z",
     "shell.execute_reply.started": "2024-02-06T21:52:46.140716Z"
    }
   },
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "p = variants.plotNumMutsHistogram(\"aa\", samples=None, max_muts=8)\n",
    "p = p + theme(panel_grid_major_x=element_blank())  # no vertical grid lines\n",
    "_ = p.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate variant counts\n",
    "We will simulate using a `SigmoidPhenotypeSimulator`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:47.201787Z",
     "iopub.status.busy": "2024-02-06T21:52:47.201362Z",
     "iopub.status.idle": "2024-02-06T21:52:47.214353Z",
     "shell.execute_reply": "2024-02-06T21:52:47.213216Z",
     "shell.execute_reply.started": "2024-02-06T21:52:47.201755Z"
    }
   },
   "outputs": [],
   "source": [
    "phenosimulator = dms_variants.simulate.SigmoidPhenotypeSimulator(\n",
    "    geneseq,\n",
    "    seed=seed,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot relationship between simulated and observed phenotypes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:47.216146Z",
     "iopub.status.busy": "2024-02-06T21:52:47.215702Z",
     "iopub.status.idle": "2024-02-06T21:52:47.639830Z",
     "shell.execute_reply": "2024-02-06T21:52:47.639075Z",
     "shell.execute_reply.started": "2024-02-06T21:52:47.216105Z"
    }
   },
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "for value in [\"enrichment\", \"phenotype\"]:\n",
    "    p = phenosimulator.plotLatentVsObserved(value)\n",
    "    _ = p.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot latent and observed phenotypes of all single amino-acid mutants:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:47.641454Z",
     "iopub.status.busy": "2024-02-06T21:52:47.641122Z",
     "iopub.status.idle": "2024-02-06T21:52:48.374533Z",
     "shell.execute_reply": "2024-02-06T21:52:48.373733Z",
     "shell.execute_reply.started": "2024-02-06T21:52:47.641429Z"
    }
   },
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "for value in [\"latentPhenotype\", \"observedPhenotype\", \"observedEnrichment\"]:\n",
    "    p = phenosimulator.plotMutsHistogram(value)\n",
    "    _ = p.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we simulate the counts according to the observed phenotype:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:48.377400Z",
     "iopub.status.busy": "2024-02-06T21:52:48.376758Z",
     "iopub.status.idle": "2024-02-06T21:52:49.131784Z",
     "shell.execute_reply": "2024-02-06T21:52:49.130667Z",
     "shell.execute_reply.started": "2024-02-06T21:52:48.377345Z"
    }
   },
   "outputs": [],
   "source": [
    "counts = dms_variants.simulate.simulateSampleCounts(\n",
    "    variants=variants,\n",
    "    phenotype_func=phenosimulator.observedEnrichment,\n",
    "    variant_error_rate=variant_error_rate,\n",
    "    pre_sample={\"total_count\": depth, \"uniformity\": lib_uniformity},\n",
    "    pre_sample_name=\"pre-selection\",\n",
    "    post_samples={\n",
    "        name: {\"noise\": noise, \"total_count\": depth, \"bottleneck\": bottle}\n",
    "        for name, bottle in bottlenecks.items()\n",
    "    },\n",
    "    seed=seed,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add the counts to the variant table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:49.133341Z",
     "iopub.status.busy": "2024-02-06T21:52:49.133014Z",
     "iopub.status.idle": "2024-02-06T21:52:49.334874Z",
     "shell.execute_reply": "2024-02-06T21:52:49.333828Z",
     "shell.execute_reply.started": "2024-02-06T21:52:49.133318Z"
    }
   },
   "outputs": [],
   "source": [
    "variants.add_sample_counts_df(counts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Confirm expected selection on mutations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:49.336417Z",
     "iopub.status.busy": "2024-02-06T21:52:49.336066Z",
     "iopub.status.idle": "2024-02-06T21:52:50.127784Z",
     "shell.execute_reply": "2024-02-06T21:52:50.126795Z",
     "shell.execute_reply.started": "2024-02-06T21:52:49.336391Z"
    }
   },
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "p = variants.plotNumCodonMutsByType(\"all\")\n",
    "p = p + theme(panel_grid_major_x=element_blank())  # no vertical grid lines\n",
    "_ = p.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get functional scores\n",
    "Get the functional scores:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:50.129468Z",
     "iopub.status.busy": "2024-02-06T21:52:50.129097Z",
     "iopub.status.idle": "2024-02-06T21:52:50.609853Z",
     "shell.execute_reply": "2024-02-06T21:52:50.609113Z",
     "shell.execute_reply.started": "2024-02-06T21:52:50.129443Z"
    }
   },
   "outputs": [],
   "source": [
    "func_scores = variants.func_scores(\"pre-selection\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Classify the functional scores:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:50.611089Z",
     "iopub.status.busy": "2024-02-06T21:52:50.610878Z",
     "iopub.status.idle": "2024-02-06T21:52:51.151166Z",
     "shell.execute_reply": "2024-02-06T21:52:51.150422Z",
     "shell.execute_reply.started": "2024-02-06T21:52:50.611068Z"
    }
   },
   "outputs": [],
   "source": [
    "func_scores = variants.classifyVariants(func_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `plotnine` to plot distribution of functional scores:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:51.152364Z",
     "iopub.status.busy": "2024-02-06T21:52:51.152091Z",
     "iopub.status.idle": "2024-02-06T21:52:54.171690Z",
     "shell.execute_reply": "2024-02-06T21:52:54.169485Z",
     "shell.execute_reply.started": "2024-02-06T21:52:51.152343Z"
    }
   },
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "p = (\n",
    "    ggplot(func_scores, aes(\"variant_class\", \"func_score\"))\n",
    "    + geom_violin(aes(fill=\"variant_class\"))\n",
    "    + ylab(\"functional score\")\n",
    "    + xlab(\"\")\n",
    "    + facet_grid(\"post_sample ~ library\")\n",
    "    + theme(\n",
    "        figure_size=(2.75 * len(libs), 2 * len(bottlenecks)),\n",
    "        axis_text_x=element_text(angle=90),\n",
    "        panel_grid_major_x=element_blank(),  # no vertical grid lines\n",
    "    )\n",
    "    + scale_fill_manual(values=CBPALETTE[1:])\n",
    ")\n",
    "_ = p.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a binary map for the variants:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.175467Z",
     "iopub.status.busy": "2024-02-06T21:52:54.174520Z",
     "iopub.status.idle": "2024-02-06T21:52:54.611530Z",
     "shell.execute_reply": "2024-02-06T21:52:54.610415Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.175408Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of binary map is 600\n"
     ]
    }
   ],
   "source": [
    "binarymap = binarymap.BinaryMap(func_scores)\n",
    "print(f\"Length of binary map is {binarymap.binarylength}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now initialize the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.619692Z",
     "iopub.status.busy": "2024-02-06T21:52:54.619274Z",
     "iopub.status.idle": "2024-02-06T21:52:54.629954Z",
     "shell.execute_reply": "2024-02-06T21:52:54.629219Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.619662Z"
    }
   },
   "outputs": [],
   "source": [
    "model = dms_variants.globalepistasis.MonotonicSplineEpistasisBottleneckLikelihood(\n",
    "    binarymap, bottleneck=bottleneck\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at some model attributes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.631272Z",
     "iopub.status.busy": "2024-02-06T21:52:54.631021Z",
     "iopub.status.idle": "2024-02-06T21:52:54.639015Z",
     "shell.execute_reply": "2024-02-06T21:52:54.638106Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.631244Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "607"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.nparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.640419Z",
     "iopub.status.busy": "2024-02-06T21:52:54.640073Z",
     "iopub.status.idle": "2024-02-06T21:52:54.645798Z",
     "shell.execute_reply": "2024-02-06T21:52:54.644925Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.640391Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.likelihood_calc_params_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.647341Z",
     "iopub.status.busy": "2024-02-06T21:52:54.646972Z",
     "iopub.status.idle": "2024-02-06T21:52:54.677151Z",
     "shell.execute_reply": "2024-02-06T21:52:54.676383Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.647308Z"
    }
   },
   "outputs": [],
   "source": [
    "assert numpy.allclose(model.loglik, -193491, rtol=0.01, atol=3), model.loglik"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.678822Z",
     "iopub.status.busy": "2024-02-06T21:52:54.678210Z",
     "iopub.status.idle": "2024-02-06T21:52:54.683782Z",
     "shell.execute_reply": "2024-02-06T21:52:54.683093Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.678797Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.latent_phenotype_wt()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit model\n",
    "Perform initial least squares optimization followed by rescaling latent effects and epistasis function parameters (normally this is actually performed by the `fit` method):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.685124Z",
     "iopub.status.busy": "2024-02-06T21:52:54.684757Z",
     "iopub.status.idle": "2024-02-06T21:52:54.741571Z",
     "shell.execute_reply": "2024-02-06T21:52:54.740288Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.685095Z"
    }
   },
   "outputs": [],
   "source": [
    "model._fit_latent_leastsquares()\n",
    "model._prescale_params(\n",
    "    k=1, g_k_range=(binarymap.func_scores.min(), binarymap.func_scores.max())\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at model parameters again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.743628Z",
     "iopub.status.busy": "2024-02-06T21:52:54.743157Z",
     "iopub.status.idle": "2024-02-06T21:52:54.751517Z",
     "shell.execute_reply": "2024-02-06T21:52:54.750484Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.743589Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "c_alpha: -15.59\n",
      "alpha_1: 3.72\n",
      "alpha_2: 3.72\n",
      "alpha_3: 3.72\n",
      "alpha_4: 3.72\n",
      "alpha_5: 3.72\n"
     ]
    }
   ],
   "source": [
    "for key, val in model.epistasis_func_params_dict.items():\n",
    "    print(f\"{key}: {round(val, 2)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the gradients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:52:54.754904Z",
     "iopub.status.busy": "2024-02-06T21:52:54.753662Z",
     "iopub.status.idle": "2024-02-06T21:53:04.817017Z",
     "shell.execute_reply": "2024-02-06T21:53:04.815326Z",
     "shell.execute_reply.started": "2024-02-06T21:52:54.754860Z"
    }
   },
   "outputs": [],
   "source": [
    "err = scipy.optimize.check_grad(\n",
    "    model._loglik_by_allparams, model._dloglik_by_allparams, model._allparams\n",
    ")\n",
    "tot = math.sqrt(sum(model._dloglik_by_allparams(model._allparams) ** 2))\n",
    "if err / tot > 1e-2:\n",
    "    raise ValueError(\n",
    "        f\"Numerical and analytical gradients differ: {err}, {tot}, {err / tot}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print model parameters after this \"partial\" fitting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:04.819440Z",
     "iopub.status.busy": "2024-02-06T21:53:04.818973Z",
     "iopub.status.idle": "2024-02-06T21:53:04.825126Z",
     "shell.execute_reply": "2024-02-06T21:53:04.824108Z",
     "shell.execute_reply.started": "2024-02-06T21:53:04.819410Z"
    }
   },
   "outputs": [],
   "source": [
    "assert numpy.allclose(model.loglik, -120885, rtol=0.01, atol=3), model.loglik"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:04.826604Z",
     "iopub.status.busy": "2024-02-06T21:53:04.826125Z",
     "iopub.status.idle": "2024-02-06T21:53:04.835548Z",
     "shell.execute_reply": "2024-02-06T21:53:04.834486Z",
     "shell.execute_reply.started": "2024-02-06T21:53:04.826572Z"
    }
   },
   "outputs": [],
   "source": [
    "for param, val in model.likelihood_calc_params_dict.items():\n",
    "    print(f\"{param}: {val:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:04.837674Z",
     "iopub.status.busy": "2024-02-06T21:53:04.837024Z",
     "iopub.status.idle": "2024-02-06T21:53:04.845724Z",
     "shell.execute_reply": "2024-02-06T21:53:04.844759Z",
     "shell.execute_reply.started": "2024-02-06T21:53:04.837634Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "round(model.latent_phenotype_wt(), 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now fit the model by maximum likelihood, also profiling [as here](https://stackoverflow.com/a/1584468).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:04.848250Z",
     "iopub.status.busy": "2024-02-06T21:53:04.847410Z",
     "iopub.status.idle": "2024-02-06T21:53:51.409438Z",
     "shell.execute_reply": "2024-02-06T21:53:51.407734Z",
     "shell.execute_reply.started": "2024-02-06T21:53:04.848206Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fitting took 46.6 seconds\n",
      "Iterations: 589.\n",
      "Function evaluations: 705\n",
      "Termination message: CONVERGENCE: REL_REDUCTION_OF_F_<=_FACTR*EPSMCH\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "start = time.time()\n",
    "profiler = cProfile.Profile()\n",
    "optres = profiler.runcall(model.fit)\n",
    "print(f\"Fitting took {time.time() - start:.1f} seconds\")\n",
    "print(\n",
    "    f\"Iterations: {optres.nit}.\\n\"\n",
    "    f\"Function evaluations: {optres.nfev}\\n\"\n",
    "    f\"Termination message: {optres.message}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at some of the model parameters after optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:51.412372Z",
     "iopub.status.busy": "2024-02-06T21:53:51.411682Z",
     "iopub.status.idle": "2024-02-06T21:53:51.436132Z",
     "shell.execute_reply": "2024-02-06T21:53:51.435343Z",
     "shell.execute_reply.started": "2024-02-06T21:53:51.412343Z"
    }
   },
   "outputs": [],
   "source": [
    "assert numpy.allclose(model.loglik, -63335, rtol=0.01, atol=3), model.loglik"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:51.437399Z",
     "iopub.status.busy": "2024-02-06T21:53:51.437106Z",
     "iopub.status.idle": "2024-02-06T21:53:51.441260Z",
     "shell.execute_reply": "2024-02-06T21:53:51.440591Z",
     "shell.execute_reply.started": "2024-02-06T21:53:51.437374Z"
    }
   },
   "outputs": [],
   "source": [
    "assert numpy.allclose(model.latent_phenotype_wt(), 0), model.latent_phenotype_wt()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot correlation of measured functional scores with latent and observed phenotypes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:51.442633Z",
     "iopub.status.busy": "2024-02-06T21:53:51.442336Z",
     "iopub.status.idle": "2024-02-06T21:53:51.909812Z",
     "shell.execute_reply": "2024-02-06T21:53:51.907873Z",
     "shell.execute_reply.started": "2024-02-06T21:53:51.442600Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Correlation for observed:\n",
      "R = 0.9, P = 0\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "df = pd.DataFrame(\n",
    "    {\n",
    "        \"measured\": binarymap.func_scores,\n",
    "        \"latent\": model.phenotypes_frombinary(binarymap.binary_variants, \"latent\"),\n",
    "        \"observed\": model.phenotypes_frombinary(binarymap.binary_variants, \"observed\"),\n",
    "    }\n",
    ")\n",
    "\n",
    "for phenotype in [\"latent\", \"observed\"]:\n",
    "    _ = (ggplot(df, aes(phenotype, \"measured\")) + geom_bin2d()).draw()\n",
    "\n",
    "    if phenotype == \"observed\":\n",
    "        r, p = scipy.stats.pearsonr(\n",
    "            binarymap.func_scores,\n",
    "            model.phenotypes_frombinary(binarymap.binary_variants, phenotype),\n",
    "        )\n",
    "        print(f\"\\nCorrelation for {phenotype}:\\nR = {r:.1f}, P = {p:.2g}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot latent versus observed phenotypes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:51.913037Z",
     "iopub.status.busy": "2024-02-06T21:53:51.912075Z",
     "iopub.status.idle": "2024-02-06T21:53:52.191439Z",
     "shell.execute_reply": "2024-02-06T21:53:52.190597Z",
     "shell.execute_reply.started": "2024-02-06T21:53:51.912989Z"
    }
   },
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "_ = (ggplot(df, aes(\"latent\", \"observed\")) + geom_bin2d()).draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at profiling results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.192787Z",
     "iopub.status.busy": "2024-02-06T21:53:52.192557Z",
     "iopub.status.idle": "2024-02-06T21:53:52.228598Z",
     "shell.execute_reply": "2024-02-06T21:53:52.227561Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.192761Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tue Feb  6 13:53:52 2024    /tmp/tmp7vbri1ir\n",
      "\n",
      "         3227834 function calls (2977364 primitive calls) in 46.550 seconds\n",
      "\n",
      "   Ordered by: cumulative time\n",
      "   List reduced from 355 to 20 due to restriction <20>\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "        1    0.001    0.001   46.550   46.550 globalepistasis.py:1446(fit)\n",
      "        1    0.000    0.000   46.329   46.329 _minimize.py:51(minimize)\n",
      "        1    0.190    0.190   46.326   46.326 _lbfgsb_py.py:212(_minimize_lbfgsb)\n",
      "      705    0.013    0.000   46.047    0.065 _differentiable_functions.py:293(fun_and_grad)\n",
      "     2119    1.698    0.001   31.580    0.015 ispline.py:251(_calculate_Itotal_or_dItotal)\n",
      "    13215    4.041    0.000   28.173    0.002 ispline.py:706(_calculate_I_or_dI)\n",
      "      706    0.003    0.000   23.920    0.034 _differentiable_functions.py:260(_update_fun)\n",
      "      705    0.004    0.000   23.917    0.034 _differentiable_functions.py:162(update_fun)\n",
      "      705    0.016    0.000   23.913    0.034 _differentiable_functions.py:140(fun_wrapped)\n",
      "      705    0.006    0.000   23.890    0.034 globalepistasis.py:1521(_loglik_by_allparams)\n",
      "      706    0.004    0.000   22.122    0.031 _differentiable_functions.py:265(_update_grad)\n",
      "      705    0.005    0.000   22.118    0.031 _differentiable_functions.py:174(update_grad)\n",
      "      705    0.014    0.000   22.113    0.031 _differentiable_functions.py:170(grad_wrapped)\n",
      "      705    0.040    0.000   22.086    0.031 globalepistasis.py:1549(_dloglik_by_allparams)\n",
      "      708    2.032    0.003   21.824    0.031 globalepistasis.py:2265(loglik)\n",
      "     2824    0.386    0.000   19.779    0.007 globalepistasis.py:2357(_n_v_bottle)\n",
      "     2824    0.285    0.000   19.402    0.007 globalepistasis.py:2326(_base_to_observed_pheno)\n",
      "      707    0.026    0.000   19.075    0.027 globalepistasis.py:1685(_observed_phenotypes)\n",
      "      709    0.030    0.000   18.464    0.026 globalepistasis.py:2696(epistasis_func)\n",
      "      709    0.013    0.000   17.202    0.024 ispline.py:226(Itotal)\n",
      "\n",
      "\n",
      "<pstats.Stats object at 0x7f9345242a50>\n",
      "Tue Feb  6 13:53:52 2024    /tmp/tmp7vbri1ir\n",
      "\n",
      "         3227834 function calls (2977364 primitive calls) in 46.550 seconds\n",
      "\n",
      "   Ordered by: internal time\n",
      "   List reduced from 355 to 20 due to restriction <20>\n",
      "\n",
      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
      "69696/11616    8.246    0.000   10.297    0.001 ispline.py:1122(_calculate_dM_dx)\n",
      "62382/41172    7.111    0.000    7.699    0.000 ispline.py:1029(_calculate_M)\n",
      "    13215    4.041    0.000   28.173    0.002 ispline.py:706(_calculate_I_or_dI)\n",
      "    17788    2.511    0.000    2.943    0.000 shape_base.py:219(vstack)\n",
      "     1477    2.370    0.002    2.370    0.002 {method 'dot' of 'numpy.ndarray' objects}\n",
      "    72928    2.127    0.000    2.127    0.000 {method 'reduce' of 'numpy.ufunc' objects}\n",
      "      708    2.032    0.003   21.824    0.031 globalepistasis.py:2265(loglik)\n",
      "     2119    1.698    0.001   31.580    0.015 ispline.py:251(_calculate_Itotal_or_dItotal)\n",
      "     1410    1.441    0.001    1.546    0.001 globalepistasis.py:2372(_digamma_n_v_bottle_1)\n",
      "    13215    1.045    0.000    2.451    0.000 ispline.py:744(<listcomp>)\n",
      "     1436    1.012    0.001    1.012    0.001 {built-in method scipy.sparse._sparsetools.csr_matvec}\n",
      "      705    0.571    0.001    1.501    0.002 _compressed.py:367(multiply)\n",
      "    42352    0.523    0.000    0.523    0.000 {method 'astype' of 'numpy.ndarray' objects}\n",
      "    97218    0.521    0.000    0.533    0.000 ispline.py:948(_ti_le_x_lt_tiplusk)\n",
      "     2115    0.482    0.000    2.347    0.001 globalepistasis.py:2296(_dloglik_dobserved_phenotype)\n",
      "175422/97170    0.467    0.000    7.998    0.000 ispline.py:980(M)\n",
      "     4242    0.459    0.000    0.459    0.000 {method 'nonzero' of 'numpy.ndarray' objects}\n",
      "      705    0.387    0.001    0.387    0.001 {built-in method scipy.sparse._sparsetools.expandptr}\n",
      "     2824    0.386    0.000   19.779    0.007 globalepistasis.py:2357(_n_v_bottle)\n",
      "    41822    0.371    0.000    0.371    0.000 {built-in method numpy.zeros}\n",
      "\n",
      "\n",
      "<pstats.Stats object at 0x7f9345242a50>\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "with tempfile.NamedTemporaryFile(\"w\") as f:\n",
    "    profiler.dump_stats(f.name)\n",
    "    f.flush()\n",
    "    stats = pstats.Stats(f.name).strip_dirs()\n",
    "for stat in [\"cumtime\", \"tottime\"]:\n",
    "    print(stats.sort_stats(stat).print_stats(20))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examine single mutants\n",
    "Correlation of predicted and actual predicted latent effects.\n",
    "Get a data frame of the latent effects of all single mutations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.229929Z",
     "iopub.status.busy": "2024-02-06T21:53:52.229588Z",
     "iopub.status.idle": "2024-02-06T21:53:52.234955Z",
     "shell.execute_reply": "2024-02-06T21:53:52.234115Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.229899Z"
    }
   },
   "outputs": [],
   "source": [
    "singles_df = model.latent_effects_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add the true latent effects from the simulator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.236290Z",
     "iopub.status.busy": "2024-02-06T21:53:52.235939Z",
     "iopub.status.idle": "2024-02-06T21:53:52.243467Z",
     "shell.execute_reply": "2024-02-06T21:53:52.242355Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.236261Z"
    }
   },
   "outputs": [],
   "source": [
    "singles_df = singles_df.rename(columns={\"mutation\": \"aa_substitutions\"}).assign(\n",
    "    true_latent_effect=lambda x: x[\"aa_substitutions\"].map(phenosimulator.muteffects)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now add the predicted and true latent and observed phenotypes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.244973Z",
     "iopub.status.busy": "2024-02-06T21:53:52.244600Z",
     "iopub.status.idle": "2024-02-06T21:53:52.283316Z",
     "shell.execute_reply": "2024-02-06T21:53:52.282394Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.244946Z"
    }
   },
   "outputs": [],
   "source": [
    "singles_df = model.add_phenotypes_to_df(singles_df)\n",
    "\n",
    "singles_df = singles_df.assign(\n",
    "    true_latent_phenotype=lambda x: x[\"aa_substitutions\"].map(\n",
    "        phenosimulator.latentPhenotype\n",
    "    ),\n",
    "    true_enrichment=lambda x: x[\"aa_substitutions\"].map(\n",
    "        phenosimulator.observedEnrichment\n",
    "    ),\n",
    "    true_observed_phenotype=lambda x: x[\"aa_substitutions\"].map(\n",
    "        phenosimulator.observedPhenotype\n",
    "    ),\n",
    "    enrichment=lambda x: model.enrichments(x[\"observed_phenotype\"]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Correlations among the observed quantities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.284628Z",
     "iopub.status.busy": "2024-02-06T21:53:52.284305Z",
     "iopub.status.idle": "2024-02-06T21:53:52.291925Z",
     "shell.execute_reply": "2024-02-06T21:53:52.291134Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.284605Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ",observed_phenotype,true_enrichment,true_observed_phenotype,enrichment\n",
      "observed_phenotype,1.0,0.9,1.0,0.9\n",
      "true_enrichment,0.9,1.0,0.9,1.0\n",
      "true_observed_phenotype,1.0,0.9,1.0,0.9\n",
      "enrichment,0.9,1.0,0.9,1.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    singles_df[[c for c in singles_df.columns if \"latent\" not in c]]\n",
    "    .corr(numeric_only=True)\n",
    "    .to_csv(float_format=\"%.1f\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examine all variants used for fitting\n",
    "Look at correlation among all the variants used to fit the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.293218Z",
     "iopub.status.busy": "2024-02-06T21:53:52.292889Z",
     "iopub.status.idle": "2024-02-06T21:53:52.301041Z",
     "shell.execute_reply": "2024-02-06T21:53:52.300089Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.293191Z"
    }
   },
   "outputs": [],
   "source": [
    "variants_df = model.phenotypes_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.302636Z",
     "iopub.status.busy": "2024-02-06T21:53:52.302235Z",
     "iopub.status.idle": "2024-02-06T21:53:52.548347Z",
     "shell.execute_reply": "2024-02-06T21:53:52.547681Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.302588Z"
    }
   },
   "outputs": [],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "\n",
    "_ = (ggplot(variants_df, aes(\"func_score\", \"observed_phenotype\")) + geom_bin2d()).draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-06T21:53:52.550729Z",
     "iopub.status.busy": "2024-02-06T21:53:52.550102Z",
     "iopub.status.idle": "2024-02-06T21:53:52.559240Z",
     "shell.execute_reply": "2024-02-06T21:53:52.558232Z",
     "shell.execute_reply.started": "2024-02-06T21:53:52.550691Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ",func_score,latent_phenotype,observed_phenotype\n",
      "func_score,1.0,0.7,0.9\n",
      "latent_phenotype,0.7,1.0,0.8\n",
      "observed_phenotype,0.9,0.8,1.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    variants_df[[\"func_score\", \"latent_phenotype\", \"observed_phenotype\"]]\n",
    "    .corr()\n",
    "    .to_csv(float_format=\"%.1f\")\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {
    "height": "1033.71px",
    "left": "0px",
    "right": "1423.47px",
    "top": "110.284px",
    "width": "211.989px"
   },
   "toc_section_display": "block",
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}