notebooks_and_scripts/probability_analysis/idprob_expanded_with_caching.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle \n",
    "import random\n",
    "import itertools\n",
    "import sqlite3\n",
    "\n",
    "import numpy as np\n",
    "from scipy import optimize\n",
    "\n",
    "from matplotlib import pyplot as plt, rcParams, patches as mplPatches\n",
    "from matplotlib.axes import Axes as mplAxes\n",
    "\n",
    "rcParams[\"font.size\"] = 8\n",
    "rcParams[\"font.family\"] = \"Roboto Condensed\"\n",
    "\n",
    "from idpp.db.util import IdPPdb\n",
    "from idpp.probability.trees import (\n",
    "    DatasetQueries,\n",
    "    construct_mz_tree,\n",
    "    construct_ccs_tree,\n",
    "    construct_rt_tree,\n",
    "    construct_ms2_tree_for_adduct_ids\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Constants"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Search Tolerances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MZ_PPMS = [0.6, 1.0, 1.8, 3.2, 5.6, 10.0, 17.8, 31.6]\n",
    "CCS_PERCENTS = [0.06, 0.10, 0.18, 0.32, 0.56, 1.00, 1.78, 3.16, 5.62]\n",
    "RT_TOLERANCES = [0.1, 0.2, 0.4, 0.8, 1.6]  # FOR IDPP RTP RTS ONLY\n",
    "MS2_TOLERANCES = [0.99, 0.95, 0.9, 0.8, 0.5, 0.65, 0.25, 0.1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# queries associated with the m/z results database\n",
    "MZ_RESULTS_QUERIES = {\n",
    "    \"create\": \"\"\"--sqlite3\n",
    "        CREATE TABLE IF NOT EXISTS Results (\n",
    "            ppm REAL NOT NULL,\n",
    "            query_cid INT NOT NULL,\n",
    "            matches BLOB NOT NULL,\n",
    "            n_matches INT NOT NULL\n",
    "        )\n",
    "    ;\"\"\",\n",
    "    \"clear\": \"\"\"--sqlite3\n",
    "        DELETE FROM Results\n",
    "    ;\"\"\",\n",
    "    \"insert\": \"\"\"--sqlite3\n",
    "        INSERT INTO Results VALUES (?,?,?,?)\n",
    "    ;\"\"\",\n",
    "    \"select\": \"\"\"--sqlite3\n",
    "        NULL   \n",
    "    ;\"\"\",\n",
    "}\n",
    "\n",
    "MZ_QRY = \"\"\"--sqlite3\n",
    "SELECT \n",
    "    cmpd_id,\n",
    "    adduct_mz\n",
    "FROM \n",
    "    Adducts\n",
    "    JOIN\n",
    "        Compounds USING(cmpd_id)\n",
    "WHERE \n",
    "    cmpd_id >= 0\n",
    "    AND adduct_id >= 0\n",
    "    AND adduct != \"none\"\n",
    ";\"\"\"\n",
    "\n",
    "RT_QRY = \"\"\"--sqlite3\n",
    "SELECT\n",
    "    cmpd_id,\n",
    "    rt\n",
    "FROM \n",
    "    RTs\n",
    "    JOIN\n",
    "        Adducts USING(adduct_id)\n",
    "    JOIN\n",
    "        Compounds USING(cmpd_id)\n",
    "WHERE\n",
    "    src_id=408  --> PREDICTED_idpp_rtp\n",
    ";\"\"\"\n",
    "\n",
    "CCS_QRY = \"\"\"--sqlite3\n",
    "SELECT\n",
    "    cmpd_id,\n",
    "    ccs\n",
    "FROM \n",
    "    CCSs\n",
    "    JOIN\n",
    "        Adducts USING(adduct_id)\n",
    "    JOIN\n",
    "        Compounds USING(cmpd_id)\n",
    ";\"\"\"\n",
    "\n",
    "MS2_QRY_A = \"\"\"--sqlite3\n",
    "SELECT \n",
    "    adduct_id,\n",
    "    COUNT(*) AS cnt \n",
    "FROM \n",
    "    MS2Spectra \n",
    "    JOIN \n",
    "        Adducts USING(adduct_id)\n",
    "GROUP BY \n",
    "    adduct_id\n",
    ";\"\"\"\n",
    "\n",
    "MS2_QRY_B = \"\"\"--sqlite3\n",
    "SELECT \n",
    "    cmpd_id, \n",
    "    adduct_id,\n",
    "\tfrag_imz, \n",
    "\tSUM(frag_ii)\n",
    "FROM \n",
    "    MS2Spectra \n",
    "    JOIN \n",
    "        Adducts USING(adduct_id)\n",
    "    JOIN\n",
    "        Compounds USING(cmpd_id)\n",
    "\tJOIN\n",
    "\t\tMS2Fragments USING(ms2_id)\n",
    "GROUP BY \n",
    "    adduct_id, \n",
    "\tfrag_imz\n",
    ";\"\"\"\n",
    "\n",
    "QUERIES = DatasetQueries(\n",
    "    mz_qry=MZ_QRY, \n",
    "    rt_qry=RT_QRY,\n",
    "    ccs_qry=CCS_QRY,\n",
    "    ms2_qry=(MS2_QRY_A, MS2_QRY_B)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fexp(x, beta, mu):\n",
    "    return (1. / beta) * np.exp(-(x - mu) / beta)\n",
    "\n",
    "def flin(x, a, b):\n",
    "    return x * a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# m/z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize and Cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the m/z results database\n",
    "mzres_con = sqlite3.connect(\"_cache/mz_results.db\")\n",
    "mzres_cur = mzres_con.cursor()\n",
    "_ = mzres_cur.execute(MZ_RESULTS_QUERIES[\"create\"])\n",
    "_ = mzres_cur.execute(MZ_RESULTS_QUERIES[\"clear\"])\n",
    "mzres_con.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = IdPPdb(\"idpp_cleaned_expanded.db\", read_only=True, enforce_idpp_ver=False)\n",
    "mzt = construct_mz_tree(db, QUERIES)\n",
    "print(\"constructed m/z tree\")\n",
    "\n",
    "print(\"performing m/z only queries and mapping compound/adduct IDs\")\n",
    "# at the same time, map compound IDs to adduct ids\n",
    "cmpd_id_to_adduct_id = {}\n",
    "for ppm in MZ_PPMS:\n",
    "    print(f\"{ppm=}\")\n",
    "    for qry_cid, matches in mzt.query_all_gen(ppm):\n",
    "        mzres_cur.execute(MZ_RESULTS_QUERIES[\"insert\"], \n",
    "                          (ppm, qry_cid, pickle.dumps(matches), len(matches)))\n",
    "        # add compound ID to adduct ID if mapping not already present\n",
    "        if qry_cid not in cmpd_id_to_adduct_id:\n",
    "            cmpd_id_to_adduct_id[qry_cid] = [\n",
    "                adduct_id for adduct_id in db.fetch_adduct_id_by_cmpd_id(qry_cid)\n",
    "            ]\n",
    "print(\"done\")\n",
    "\n",
    "with open(\"_cache/cmpd_id_to_adduct_id.pkl\", \"wb\") as pf:\n",
    "    pickle.dump(cmpd_id_to_adduct_id, pf)\n",
    "print(\"cached compound to adduct id map\")\n",
    "\n",
    "db.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mzres_con.commit()\n",
    "mzres_con.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Cached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"_cache/mz_results.pkl\", \"rb\") as pf:\n",
    "    mz_results = pickle.load(pf)\n",
    "\n",
    "with open(\"_cache/mz_only_probs.pkl\", \"rb\") as pf:\n",
    "    mz_only_probs = pickle.load(pf)\n",
    "\n",
    "with open(\"_cache/cmpd_id_to_adduct_id.pkl\", \"rb\") as pf:\n",
    "    cmpd_id_to_adduct_id = pickle.load(pf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "betas, mus, obs_meds = [], [], []\n",
    "for ppm in MZ_PPMS:\n",
    "    print(f\"{ppm=}\")\n",
    "    fig, ax = plt.subplots(figsize=(2.5, 3))\n",
    "    bin_max = max(mz_only_probs[ppm]) + 1\n",
    "    c, x = np.histogram(mz_only_probs[ppm], np.arange(1, bin_max))\n",
    "    cn = c / max(c)\n",
    "    ax.hist(mz_only_probs[ppm], bins=np.arange(1, bin_max), color=\"k\", histtype=\"step\", lw=1., \n",
    "            weights=[1 / max(c)] * len(mz_only_probs[ppm]), label=\"obs.\")\n",
    "    obs_med = np.median(mz_only_probs[ppm])\n",
    "    obs_med_y = cn[np.argmin(np.abs(x - obs_med))]\n",
    "    ax.axvline(obs_med, ymax=obs_med_y, ls=\"--\", lw=1., label=f\"obs. median={obs_med:.0f}\", c=\"k\")\n",
    "    (beta, mu), _ = optimize.curve_fit(fexp, x[:-1], cn)\n",
    "    ax.plot(x[:-1], fexp(x[:-1], beta, mu), 'b-', lw=1., label=f\"fit({beta=:.3f}, {mu=:.3f})\")\n",
    "    fit_med = beta * np.log(2.)\n",
    "    ax.axvline(fit_med, ymax=fexp(fit_med, beta, mu), ls=\"--\", lw=1., label=f\"fit median={fit_med:.3f}\", c=\"b\")\n",
    "    ax.legend(frameon=False)\n",
    "    for d in [\"top\", \"right\"]:\n",
    "        ax.spines[d].set_visible(False)\n",
    "    #ax.ticklabel_format(axis='y', style='sci', scilimits=(0, 0))\n",
    "    ax.set_xlabel(\"# matches\")\n",
    "    ax.set_ylabel(\"density\")\n",
    "    ax.set_xlim([1, 50 + 0.5])\n",
    "    plt.savefig(f\"_figures/mz/{ppm=}_dist_with_fit.png\", dpi=400, bbox_inches=\"tight\")\n",
    "    plt.show()\n",
    "    plt.close()\n",
    "    betas.append(beta)\n",
    "    mus.append(mu)\n",
    "    obs_meds.append(obs_med)\n",
    "fit_meds = [_ * np.log(2) for _ in betas]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "plt.plot(MZ_PPMS, obs_meds, \"ko\", fillstyle=\"none\", label=\"obs.\")\n",
    "p, _ = optimize.curve_fit(flin, MZ_PPMS, obs_meds)\n",
    "plt.plot(MZ_PPMS, flin(np.array(MZ_PPMS), *p), \"k--\", lw=1., label=f\"linear(a={p[0]:.3f}, b={p[1]:.3f})\")\n",
    "plt.plot(MZ_PPMS, fit_meds, \"bo\", fillstyle=\"none\", label=\"fit\")\n",
    "p, _ = optimize.curve_fit(flin, MZ_PPMS, fit_meds)\n",
    "plt.plot(MZ_PPMS, flin(np.array(MZ_PPMS), *p), \"b--\", lw=1., label=f\"linear(a={p[0]:.3f}, b={p[1]:.3f})\")\n",
    "ax.legend(frameon=False)\n",
    "ax.set_xscale(\"log\")\n",
    "#ax.set_yscale(\"log\")\n",
    "for d in [\"top\", \"right\"]:\n",
    "    ax.spines[d].set_visible(False)\n",
    "ax.set_xlabel(\"query tolerance (ppm)\")\n",
    "ax.set_ylabel(\"median # matches\")\n",
    "plt.savefig(\"_figures/mz/medians_obs_fit_vs_tol.png\", dpi=400, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# m/z + CCS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize and Cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = IdPPdb(\"idpp_cleaned_expanded.db\", read_only=True, enforce_idpp_ver=False)\n",
    "ccst = construct_ccs_tree(db, QUERIES)\n",
    "\n",
    "mz_ccs_probs = {}\n",
    "for ppm in MZ_PPMS:\n",
    "    for percent in CCS_PERCENTS:\n",
    "        counts = []\n",
    "        for id_A, matched_A in mz_results[ppm].items():\n",
    "            # id_A is a compound ID, need to convert to adduct ID(s) then use that(those)\n",
    "            # to fetch CCS values to query the CcsTree with\n",
    "            # treat each set of adduct CCS values as a separate addition to the counts list\n",
    "            for add_id_A in cmpd_id_to_adduct_id[id_A]:\n",
    "                # fetch CCS values (if any) then average them and query\n",
    "                if len(ccss := [_[1] for _ in db.fetch_ccs_by_adduct_id(add_id_A)]) > 0:\n",
    "                    matched_B = ccst.query_radius_single(np.mean(ccss), percent)\n",
    "                    counts.append(len(matched_A & matched_B))\n",
    "        if len(counts) > 0:\n",
    "            mz_ccs_probs[(ppm, percent)] = counts\n",
    "print(\"created m/z + CCS query results\")\n",
    "\n",
    "with open(\"_cache/mz_ccs_probs.pkl\", \"wb\") as pf:\n",
    "    pickle.dump(mz_ccs_probs, pf)\n",
    "print(\"cached m/z + CCS query results\")\n",
    "\n",
    "db.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Cached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"_cache/mz_ccs_probs.pkl\", \"rb\") as pf:\n",
    "    mz_ccs_probs = pickle.load(pf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep track of fit parameters and summary stats\n",
    "ppms, percents = [], []\n",
    "betas, mus, obs_meds = [], [], []\n",
    "for ppm in MZ_PPMS:\n",
    "    for percent in CCS_PERCENTS:\n",
    "        ppms.append(ppm)\n",
    "        percents.append(percent)\n",
    "        print(f\"{ppm=} {percent=}\")\n",
    "        fig, ax = plt.subplots(figsize=(2.5, 3))\n",
    "        bin_max = max(mz_ccs_probs[(ppm, percent)]) + 1\n",
    "        c, x = np.histogram(mz_ccs_probs[(ppm, percent)], np.arange(1, bin_max))\n",
    "        cn = c / max(c)\n",
    "        ax.hist(mz_ccs_probs[(ppm, percent)], bins=np.arange(1, bin_max), color=\"k\", histtype=\"step\", lw=1., \n",
    "                weights=[1 / max(c)] * len(mz_ccs_probs[(ppm, percent)]), label=\"obs.\") \n",
    "        obs_med = np.median(mz_ccs_probs[(ppm, percent)])\n",
    "        ax.axvline(obs_med, ls=\"--\", lw=1., label=f\"obs. median={obs_med:.0f}\", c=\"k\")\n",
    "        (beta, mu), _ = optimize.curve_fit(fexp, x[:-1], cn)\n",
    "        ax.plot(x[:-1], fexp(x[:-1], beta, mu), 'b-', lw=1., label=f\"fit({beta=:.3f}, {mu=:.3f})\")\n",
    "        fit_med = beta * np.log(2.)\n",
    "        ax.axvline(fit_med, ls=\"--\", lw=1., label=f\"fit median={fit_med:.3f}\", c=\"b\")\n",
    "        ax.legend(frameon=False)\n",
    "        for d in [\"top\", \"right\"]:\n",
    "            ax.spines[d].set_visible(False)\n",
    "        #ax.ticklabel_format(axis='y', style='sci', scilimits=(0, 0))\n",
    "        ax.set_xlabel(\"# matches\")\n",
    "        ax.set_ylabel(\"density\")\n",
    "        ax.set_xlim([1, 50 + 0.5])\n",
    "        plt.savefig(f\"_figures/mz_ccs/{ppm=}_{percent=}_dist_with_fit.png\", dpi=400, bbox_inches=\"tight\")\n",
    "        plt.show()\n",
    "        plt.close()\n",
    "        betas.append(beta)\n",
    "        mus.append(mu)\n",
    "        obs_meds.append(obs_med)\n",
    "fit_meds = [_ * np.log(2) for _ in betas]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax, axcb) = plt.subplots(ncols=2, width_ratios=(9, 1), figsize=(5, 4.5))\n",
    "mz_bins = MZ_PPMS + [56]\n",
    "ccs_bins = CCS_PERCENTS + [10]\n",
    "levels = np.arange(0, 4.1, 0.5)\n",
    "tcf = ax.tricontourf(ppms, percents, fit_meds, levels, cmap=\"binary_r\")\n",
    "ax.tricontour(ppms, percents, fit_meds, levels, colors=\"k\", linewidths=0.5)\n",
    "ax.plot(ppms, percents, \"wo\", ms=5, fillstyle=\"none\", mew=0.5, alpha=0.2)\n",
    "ax.set_xlabel(\"m/z tolerance (ppm)\")\n",
    "ax.set_ylabel(\"CCS tolerance (%)\")\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_yscale(\"log\")\n",
    "cb = fig.colorbar(tcf, cax=axcb)\n",
    "cb.set_ticks(levels)\n",
    "cb.set_ticklabels(levels)\n",
    "for l in levels[1:-1]:\n",
    "    cb.ax.axhline(l, lw=0.75, c=\"k\")\n",
    "cb.set_label(\"median # matches (fit)\")\n",
    "plt.savefig(\"_figures/mz_ccs/2D_tolerance_contour.png\", dpi=400, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# m/z + RT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize and Cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = IdPPdb(\"idpp_cleaned_expanded.db\", read_only=True, enforce_idpp_ver=False)\n",
    "rtt = construct_rt_tree(db, QUERIES)\n",
    "\n",
    "mz_rt_probs = {}\n",
    "for ppm in MZ_PPMS:\n",
    "    for tol in RT_TOLERANCES:\n",
    "        counts = []\n",
    "        for id_A, matched_A in mz_results[ppm].items():\n",
    "            # id_A is a compound ID, need to convert to adduct ID(s) then use that(those)\n",
    "            # to fetch CCS values to query the CcsTree with\n",
    "            # treat each set of adduct CCS values as a separate addition to the counts list\n",
    "            for add_id_A in cmpd_id_to_adduct_id[id_A]:\n",
    "                # fetch CCS values (if any) then average them and query\n",
    "                if len(rts := [_[1] for _ in db.fetch_rt_by_adduct_id(add_id_A, select_sources=[408])]) > 0:\n",
    "                    matched_B = rtt.query_radius_single(np.mean(rts), tol)\n",
    "                    counts.append(len(matched_A & matched_B))\n",
    "        if len(counts) > 0:\n",
    "            mz_rt_probs[(ppm, tol)] = counts\n",
    "print(\"created m/z + RT query results\")\n",
    "\n",
    "with open(\"_cache/mz_rt_probs.pkl\", \"wb\") as pf:\n",
    "    pickle.dump(mz_rt_probs, pf)\n",
    "print(\"cached m/z + RT query results\")\n",
    "\n",
    "db.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Cached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"_cache/mz_rt_probs.pkl\", \"rb\") as pf:\n",
    "    mz_rt_probs = pickle.load(pf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep track of fit parameters and summary stats\n",
    "ppms, tols = [], []\n",
    "betas, mus, obs_meds = [], [], []\n",
    "for ppm in MZ_PPMS:\n",
    "    for tol in RT_TOLERANCES:\n",
    "        print(f\"{ppm=} {tol=}\")\n",
    "        fig, ax = plt.subplots(figsize=(2.5, 3))\n",
    "        bin_max = max(mz_rt_probs[(ppm, tol)]) + 1\n",
    "        c, x = np.histogram(mz_rt_probs[(ppm, tol)], np.arange(1, bin_max))\n",
    "        cn = c / max(c)\n",
    "        ax.hist(mz_rt_probs[(ppm, tol)], bins=np.arange(1, bin_max), color=\"k\", histtype=\"step\", lw=1., \n",
    "                weights=[1 / max(c)] * len(mz_rt_probs[(ppm, tol)]), label=\"obs.\") \n",
    "        obs_med = np.median(mz_rt_probs[(ppm, tol)])\n",
    "        ax.axvline(obs_med, ls=\"--\", lw=1., label=f\"obs. median={obs_med:.0f}\", c=\"k\")\n",
    "        if len(cn) > 1:\n",
    "            ppms.append(ppm)\n",
    "            tols.append(tol)\n",
    "            (beta, mu), _ = optimize.curve_fit(fexp, x[:-1], cn)\n",
    "            ax.plot(x[:-1], fexp(x[:-1], beta, mu), 'b-', lw=1., label=f\"fit({beta=:.3f}, {mu=:.3f})\")\n",
    "            fit_med = beta * np.log(2.)\n",
    "            ax.axvline(fit_med, ls=\"--\", lw=1., label=f\"fit median={fit_med:.3f}\", c=\"b\")\n",
    "            ax.legend(frameon=False)\n",
    "            for d in [\"top\", \"right\"]:\n",
    "                ax.spines[d].set_visible(False)\n",
    "            #ax.ticklabel_format(axis='y', style='sci', scilimits=(0, 0))\n",
    "            ax.set_xlabel(\"# matches\")\n",
    "            ax.set_ylabel(\"density\")\n",
    "            ax.set_xlim([1, 50 + 0.5])\n",
    "            plt.savefig(f\"_figures/mz_rt/{ppm=}_{tol=}_dist_with_fit.png\", dpi=400, bbox_inches=\"tight\")\n",
    "            plt.show()\n",
    "            plt.close()\n",
    "            betas.append(beta)\n",
    "            mus.append(mu)\n",
    "            obs_meds.append(obs_med)\n",
    "fit_meds = [_ * np.log(2) for _ in betas]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax, axcb) = plt.subplots(ncols=2, width_ratios=(9, 1), figsize=(5, 4.5))\n",
    "mz_bins = MZ_PPMS + [56]\n",
    "ccs_bins = RT_TOLERANCES + [3.2]\n",
    "levels = np.arange(0, 1.1, 0.1)\n",
    "tcf = ax.tricontourf(ppms, tols, fit_meds, levels, cmap=\"binary_r\")\n",
    "ax.tricontour(ppms, tols, fit_meds, levels, colors=\"k\", linewidths=0.5)\n",
    "ax.plot(ppms, tols, \"wo\", ms=5, fillstyle=\"none\", mew=0.5, alpha=0.2)\n",
    "ax.set_xlabel(\"m/z tolerance (ppm)\")\n",
    "ax.set_ylabel(\"RT tolerance (min)\")\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_yscale(\"log\")\n",
    "cb = fig.colorbar(tcf, cax=axcb)\n",
    "cb.set_ticks(levels)\n",
    "cb.set_ticklabels(levels)\n",
    "for l in levels[1:-1]:\n",
    "    cb.ax.axhline(l, lw=0.75, c=\"k\")\n",
    "cb.set_label(\"median # matches (obs.)\")\n",
    "plt.savefig(\"_figures/mz_rt/2D_tolerance_contour.png\", dpi=400, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(min(fit_meds), max(fit_meds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(min(obs_meds), max(obs_meds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# m/z + MS2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize and Cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mz_ms2 = {\"mz_ms2_probs\": {}, \"mzs\": {}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = IdPPdb(\"idpp_cleaned_expanded.db\", read_only=True, enforce_idpp_ver=False)\n",
    "\n",
    "for i in range(8):\n",
    "    random.seed(i + 100)\n",
    "    for ppm in MZ_PPMS:\n",
    "        print(f\"{ppm=}\")\n",
    "        selected = 0\n",
    "        for i, (id_A, matched_A) in enumerate(mz_results[ppm].items()):\n",
    "            if random.random() > 0.05:\n",
    "                # most of the time just skip\n",
    "                continue\n",
    "            selected += 1\n",
    "            #print(\"---\")\n",
    "            #print(f\"{id_A=} {matched_A=}\")\n",
    "            aid_A = cmpd_id_to_adduct_id[id_A]\n",
    "            #print(f\"{aid_A=}\")\n",
    "            # id_A is a compound ID, need to convert to adduct ID(s) then use that(those) to\n",
    "            # create an MS2Tree on the fly for matches to this compound\n",
    "            matched_aids = set([aid for cid in matched_A for aid in cmpd_id_to_adduct_id[cid]])\n",
    "            #print(f\"{matched_aids=}\")\n",
    "            if (ms2t := construct_ms2_tree_for_adduct_ids(db, matched_aids)) is not None:\n",
    "                for tol in MS2_TOLERANCES:\n",
    "                    #print(f\"{tol=}\")\n",
    "                    qres = ms2t.query_all(tol)\n",
    "                    #print(f\"{qres=}\")\n",
    "                    q_aids = set()\n",
    "                    for aid_A in cmpd_id_to_adduct_id[id_A]:\n",
    "                        if aid_A in qres:\n",
    "                            q_aids |= qres[aid_A]\n",
    "                    if len(q_aids) > 0:\n",
    "                        #print(f\"{q_aids=}\")\n",
    "                        common_aids = matched_aids & q_aids\n",
    "                        common_cids = set([ms2t.adduct_to_cmpd_id[aid] for aid in common_aids])\n",
    "                        #print(f\"{common_aids=}\")\n",
    "                        #print(f\"{common_cids=}\")\n",
    "                        n = len(common_cids)\n",
    "                        m = len(matched_A)\n",
    "                        #print(f\"{n=} {m=}\")\n",
    "                        k = (ppm, tol)\n",
    "                        if k in mz_ms2[\"mz_ms2_probs\"]:\n",
    "                            mz_ms2[\"mz_ms2_probs\"][k].append(n)\n",
    "                        else:\n",
    "                            mz_ms2[\"mz_ms2_probs\"][k] = [n]\n",
    "                        if k in mz_ms2[\"mzs\"]:\n",
    "                            mz_ms2[\"mzs\"][k].append(m)\n",
    "                        else:\n",
    "                            mz_ms2[\"mzs\"][k] = [m]\n",
    "            print(f\"\\rsampled {selected} of {i + 1}  ({100 * (selected / (i + 1)):.1f} %)\", end=\"    \")\n",
    "        print()\n",
    "        # cache the results after each ppm completes\n",
    "        with open(\"_cache/mz_ms2.pkl\", \"wb\") as pf:\n",
    "            pickle.dump(mz_ms2, pf)\n",
    "\n",
    "db.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Cached"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"_cache/mz_ms2.pkl\", \"rb\") as pf:\n",
    "    mz_ms2 = pickle.load(pf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep track of fit parameters and summary stats\n",
    "ppms = []\n",
    "tols = []\n",
    "betas = []\n",
    "mus = []\n",
    "obs_meds = []\n",
    "for ppm in MZ_PPMS:\n",
    "    for tol in MS2_TOLERANCES:\n",
    "        ppms.append(ppm)\n",
    "        tols.append(tol)\n",
    "        print(f\"{ppm=} {tol=}\")\n",
    "        fig, ax = plt.subplots(figsize=(2.5, 3))\n",
    "        bin_max = max(mz_ms2[\"mz_ms2_probs\"][(ppm, tol)]) + 2\n",
    "        c, x = np.histogram(mz_ms2[\"mz_ms2_probs\"][(ppm, tol)], np.arange(1, bin_max))\n",
    "        cn = c / max(c)\n",
    "        ax.hist(mz_ms2[\"mz_ms2_probs\"][(ppm, tol)], bins=np.arange(1, bin_max), color=\"k\", histtype=\"step\", lw=1., \n",
    "                weights=[1 / max(c)] * len(mz_ms2[\"mz_ms2_probs\"][(ppm, tol)]), label=\"obs.\") \n",
    "        obs_med = np.median(mz_ms2[\"mz_ms2_probs\"][(ppm, tol)])\n",
    "        ax.axvline(obs_med, ls=\"--\", lw=1., label=f\"obs. median={obs_med:.0f}\", c=\"k\")\n",
    "        #print(x[:-1])\n",
    "        #print(cn)\n",
    "        (beta, mu), _ = optimize.curve_fit(fexp, x[:-1], cn)\n",
    "        ax.plot(x[:-1], fexp(x[:-1], beta, mu), 'b-', lw=1., label=f\"fit({beta=:.3f}, {mu=:.3f})\")\n",
    "        fit_med = beta * np.log(2.)\n",
    "        ax.axvline(fit_med, ls=\"--\", lw=1., label=f\"fit median={fit_med:.3f}\", c=\"b\")\n",
    "        ax.legend(frameon=False)\n",
    "        for d in [\"top\", \"right\"]:\n",
    "            ax.spines[d].set_visible(False)\n",
    "        #ax.ticklabel_format(axis='y', style='sci', scilimits=(0, 0))\n",
    "        ax.set_xlabel(\"# matches\")\n",
    "        ax.set_ylabel(\"density\")\n",
    "        ax.set_xlim([1, 40 + 0.5])\n",
    "        plt.savefig(f\"_figures/mz_ms2/{ppm=}_{tol=}_dist_with_fit.png\", \n",
    "                    dpi=400, bbox_inches=\"tight\")\n",
    "        plt.show()\n",
    "        plt.close()\n",
    "        betas.append(beta)\n",
    "        mus.append(mu)\n",
    "        obs_meds.append(obs_med)\n",
    "fit_meds = [_ * np.log(2) for _ in betas]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax, axcb) = plt.subplots(ncols=2, width_ratios=(9, 1), figsize=(5, 4.5))\n",
    "mz_bins = MZ_PPMS + [56]\n",
    "ms2_bins = [1.0] + MS2_TOLERANCES\n",
    "levels = np.arange(1.25, 3.1, 0.25)\n",
    "tcf = ax.tricontourf(ppms, tols, fit_meds, levels, cmap=\"binary_r\", )#norm=\"log\")\n",
    "ax.tricontour(ppms, tols, fit_meds, levels, colors=\"k\", linewidths=0.5)\n",
    "ax.plot(ppms, tols, \"wo\", ms=5, fillstyle=\"none\", mew=0.5, alpha=0.2)\n",
    "ax.set_xlabel(\"m/z tolerance (ppm)\")\n",
    "ax.set_ylabel(\"MS/MS similarity threshold\")\n",
    "ax.set_xscale(\"log\")\n",
    "#ax.set_yscale(\"log\")\n",
    "cb = fig.colorbar(tcf, cax=axcb)\n",
    "cb.set_ticks(levels)\n",
    "cb.set_ticklabels(levels)\n",
    "for l in levels[1:-1]:\n",
    "    cb.ax.axhline(l, lw=0.75, c=\"k\")\n",
    "cb.set_label(\"median # matches (fit)\")\n",
    "plt.savefig(f\"_figures/mz_ms2/2D_tolerance_contour.png\", \n",
    "            dpi=400, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(min(fit_meds), max(fit_meds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}