{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scanpy as sc\n",
    "\n",
    "import pickle \n",
    "\n",
    "data_location = \"../data/aging_brain/\"\n",
    "\n",
    "young = pickle.load(open(data_location + \"aging_brain_young.pickle\",mode='rb'))\n",
    "old = pickle.load(open(data_location + \"aging_brain_old.pickle\",mode='rb'))\n",
    "\n",
    "filtered = pickle.load(open(data_location + \"aging_brain_filtered.pickle\",mode='rb'))\n",
    "\n",
    "batch_encoding = np.loadtxt(data_location + 'aging_batch_encoding.tsv')\n",
    "batch_encoding = batch_encoding.astype(dtype=bool)\n",
    "\n",
    "young_mask = np.zeros(37069,dtype=bool)\n",
    "old_mask = np.zeros(37069,dtype=bool)\n",
    "\n",
    "young_mask[:young.shape[0]] = True\n",
    "old_mask[young.shape[0]:] = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sc.pp.neighbors(filtered)\n",
    "# sc.tl.umap(filtered)\n",
    "# sc.tl.louvain(filtered)\n",
    "# sc.pl.umap(filtered,color='louvain')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_1 = np.zeros(young.X.shape[0],dtype=bool)\n",
    "\n",
    "mask_1[batch_encoding[young_mask][:,0]] = True\n",
    "mask_1[batch_encoding[young_mask][:,2]] = True\n",
    "mask_1[batch_encoding[young_mask][:,4]] = True\n",
    "mask_1[batch_encoding[young_mask][:,6]] = True\n",
    "\n",
    "mask_2 = np.zeros(young.X.shape[0],dtype=bool)\n",
    "\n",
    "mask_2[batch_encoding[young_mask][:,1]] = True\n",
    "mask_2[batch_encoding[young_mask][:,3]] = True\n",
    "mask_2[batch_encoding[young_mask][:,5]] = True\n",
    "mask_2[batch_encoding[young_mask][:,7]] = True\n",
    "\n",
    "young_1 = young[mask_1]\n",
    "young_2 = young[mask_2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_mask_1 = np.zeros(old.X.shape[0],dtype=bool)\n",
    "\n",
    "old_mask_1[batch_encoding[old_mask][:,8]] = True\n",
    "old_mask_1[batch_encoding[old_mask][:,10]] = True\n",
    "old_mask_1[batch_encoding[old_mask][:,12]] = True\n",
    "old_mask_1[batch_encoding[old_mask][:,14]] = True\n",
    "\n",
    "old_1 = old[old_mask_1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "# sys.path.append('/localscratch/bbrener1/rusty_forest_v3/src')\n",
    "sys.path.append('../src/python/')\n",
    "import tree_reader as tr \n",
    "import lumberjack\n",
    "\n",
    "data_location = \"../data/aging_brain/\"\n",
    "\n",
    "first_forest = tr.Forest.load(data_location + 'scanpy_aging_brain_restricted_1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "# sys.path.append('/localscratch/bbrener1/rusty_forest_v3/src')\n",
    "sys.path.append('../src/python')\n",
    "import tree_reader as tr \n",
    "import lumberjack\n",
    "\n",
    "first_forest = lumberjack.fit(\n",
    "    young_1.X,\n",
    "    header=filtered.var_names,\n",
    "    trees=100,\n",
    "    ifs=700,\n",
    "    ofs=700,\n",
    "    ss=500,\n",
    "    depth=8,\n",
    "    leaves=10,\n",
    "    dispersion_mode='ssme',\n",
    "    sfr=0,\n",
    "    norm='l1',\n",
    "    reduction=2,\n",
    "    reduce_input='true',\n",
    "    reduce_output='false'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_forest.set_cache(True)\n",
    "first_forest.backup(data_location + \"scanpy_aging_brain_restricted_1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_forest.self_prediction = first_forest.predict(young_1.X)\n",
    "first_forest.self_prediction.prediction_report()"
   ]
  },
  {
   "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": [
    "feature_residuals = first_forest.self_prediction.residuals()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "centered = young_1.X - np.mean(young_1.X,axis=0)\n",
    "null_squared_residual = np.power(centered,2)\n",
    "feature_null = np.sum(null_squared_residual,axis=0) + 1\n",
    "\n",
    "forest_squared_residuals = np.power(feature_residuals,2)\n",
    "\n",
    "forest_feature_error = np.sum(forest_squared_residuals,axis=0) + 1\n",
    "forest_feature_remaining = forest_feature_error/feature_null\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.hist(forest_feature_remaining.copy(),bins=50)\n",
    "plt.show()\n",
    "\n",
    "filtered_feature_mask = forest_feature_remaining < .6\n",
    "\n",
    "print(np.sum(filtered_feature_mask))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filtered_feature_mask = np.zeros(2000,dtype=bool)\n",
    "# for feature in second_forest.output_features:\n",
    "#     f_i = list(young.var_names).index(feature)\n",
    "#     filtered_feature_mask[f_i] = True\n",
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "young_1_filtered = young_1.copy().T[filtered_feature_mask].copy().T\n",
    "young_2_filtered = young_2.copy().T[filtered_feature_mask].copy().T\n",
    "old_1_filetered = old_1.copy().T[filtered_feature_mask].copy().T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest = lumberjack.fit(\n",
    "    young_1_filtered.X,\n",
    "    header=young_1_filtered.var_names,\n",
    "    trees=100,\n",
    "    ifs=150,\n",
    "    ofs=150,\n",
    "    ss=500,\n",
    "    depth=8,\n",
    "    leaves=10,\n",
    "    dispersion_mode='ssme',\n",
    "    sfr=0,\n",
    "    norm='l1',\n",
    "    reduction=2,\n",
    "    reduce_input='true',\n",
    "    reduce_output='false'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.set_cache(True)\n",
    "second_forest.backup(\"scanpy_aging_brain_restricted_2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.trim(.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.self_prediction = second_forest.predict(young_1_filtered.X)\n",
    "second_forest.self_prediction.node_sample_encoding()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.self_prediction.prediction_report()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.other_prediction = second_forest.predict(young_2_filtered.X)\n",
    "second_forest.other_prediction.node_sample_encoding()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "# sys.path.append('/localscratch/bbrener1/rusty_forest_v3/src')\n",
    "sys.path.append('../src')\n",
    "import tree_reader as tr \n",
    "import lumberjack\n",
    " \n",
    "second_forest = tr.Forest.load('scanpy_aging_brain_restricted_2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.reset_split_clusters()\n",
    "second_forest.interpret_splits(\n",
    "    k=10,\n",
    "    pca=False, \n",
    "    depth=6,\n",
    "    metric=\"cosine\",  \n",
    "    mode='additive_mean', \n",
    "    relatives=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.maximum_spanning_tree(mode='samples')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "second_forest.tsne_coordinates = filtered.obsm['X_umap'][young_mask][mask_1]\n",
    "second_forest.html_tree_summary(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.split_clusters[20].sisters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.old_prediction = second_forest.predict(old_1_filetered.X)\n",
    "second_forest.old_prediction.node_sample_encoding()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "self_fvu = []\n",
    "old_fvu = []\n",
    "other_fvu = []\n",
    "\n",
    "other_mwu = []\n",
    "old_mwu = []\n",
    "\n",
    "self_fraction = []\n",
    "other_fraction = []\n",
    "old_fraction = []\n",
    "\n",
    "for factor in second_forest.split_clusters[1:]:\n",
    "    print(\"++++++++++++++++++++\")\n",
    "    print(factor.name())\n",
    "    print(\"++++++++++++++++++++\")\n",
    "#     second_forest.self_prediction.compare_factor_means(second_forest.old_prediction,factor,metric='euclidean',plot=['distance'])\n",
    "#     second_forest.self_prediction.compare_factor_residuals(second_forest.old_prediction,factor)\n",
    "    (sfvu,olfvu,olmwu) = second_forest.self_prediction.compare_factor_fvu(second_forest.old_prediction,factor,plot=True)\n",
    "    (_,otfvu,otmwu) = second_forest.self_prediction.compare_factor_fvu(second_forest.other_prediction,factor,plot=True)\n",
    "    self_fvu.append(sfvu)\n",
    "    old_fvu.append(olfvu)\n",
    "    other_fvu.append(otfvu)\n",
    "    old_mwu.append(olmwu)\n",
    "    other_mwu.append(otmwu)\n",
    "#     (s_f,ot_f,_) = second_forest.self_prediction.compare_factor_fractions(second_forest.old_prediction,factor,plot=True)\n",
    "#     (s_f,ol_f,_) = second_forest.self_prediction.compare_factor_fractions(second_forest.other_prediction,factor,plot=True)\n",
    "#     self_fraction.append(s_f)\n",
    "#     other_fraction.append(ot_f)\n",
    "#     old_fraction.append(ol_f)\n",
    "    #     second_forest.self_prediction.compare_factor_fractions(second_forest.other_prediction,factor)\n",
    "    \n",
    "#     print(factor.name())\n",
    "#     sm,om = second_forest.self_prediction.compare_factor_marginals(second_forest.old_prediction,factor)\n",
    "\n",
    "# print(fraction_return)\n",
    "\n",
    "self_cod = 1 - np.array(self_fvu) \n",
    "old_cod = 1 - np.array(old_fvu)\n",
    "other_cod = 1 - np.array(other_fvu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Coefficient of Determination Ratio\")\n",
    "plt.plot([0,.5],[0,.5],label=\"Ideal Fit\", color='red')\n",
    "for i,(p1,p2) in enumerate(zip(zip(self_cod,old_cod),zip(self_cod,other_cod))):\n",
    "    plt.plot([p1[0],p2[0]],[p1[1],p2[1]],linewidth=.5)\n",
    "    plt.text(p1[0]+.001,p1[1]-.005,str(i+1),fontsize=5)\n",
    "plt.scatter(self_cod,old_cod,label=\"Old Mice\",s=2)\n",
    "plt.scatter(self_cod,other_cod,label=\"Test Set\",s=2)\n",
    "plt.legend()\n",
    "plt.xlabel(\"Trained COD\")\n",
    "plt.ylabel(\"Observed COD\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Volcano plot based on MWU results for COD discrepancy. P values seem a bit jumpy, but then again we don't have a\n",
    "# ton of data points. \n",
    "\n",
    "old_lpv = [-np.log10(m[1]) for m in old_mwu]\n",
    "other_lpv = [-np.log10(m[1]) for m in other_mwu]\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(self_cod - other_cod, other_lpv)\n",
    "plt.xlabel(\"Difference in COD\")\n",
    "plt.ylabel(\"-1og 10 P value\")\n",
    "plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(self_cod - old_cod, old_lpv)\n",
    "for i,(x,y) in enumerate(zip(self_cod-old_cod,old_lpv)):\n",
    "    plt.text(x+.001,y-.005,str(i+1),fontsize=5)\n",
    "plt.xlabel(\"Difference in COD\")\n",
    "plt.ylabel(\"-1og 10 P value\")\n",
    "plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Change in COD vs P Value\")\n",
    "plt.scatter(self_cod - other_cod, other_lpv,label='Test',s=2)\n",
    "plt.scatter(self_cod - old_cod, old_lpv,label='Old',s=2)\n",
    "plt.xlabel(\"Difference in COD\")\n",
    "plt.ylabel(\"-1og 10 P value\")\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Fraction Ratio\")\n",
    "# for i,(p1,p2) in enumerate(zip(zip(self_fraction,old_fraction),zip(self_fraction,other_fraction))):\n",
    "#     plt.plot([p1[0],p2[0]],[p1[1],p2[1]],linewidth=.5)\n",
    "#     plt.text(p1[0]+.001,p1[1]-.005,str(i+1),fontsize=5)\n",
    "plt.scatter(self_fraction,other_fraction,label=\"Test Set\",s=2)\n",
    "# plt.scatter(self_fraction,old_fraction,label=\"Old\",s=2)\n",
    "plt.plot([.4,.6],[.4,.6])\n",
    "plt.legend()\n",
    "plt.xlabel(\"Trained Fraction\")\n",
    "plt.ylabel(\"Observed Fraction\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = second_forest.split_clusters[5]\n",
    "\n",
    "cods = []\n",
    "\n",
    "for node in cluster.nodes:\n",
    "    (n,p) = second_forest.self_prediction.node_r2_doublet(node)\n",
    "    fvu = n/p\n",
    "    cod = 1 - fvu\n",
    "    cods.append(cod)\n",
    "\n",
    "plt.figure()\n",
    "plt.hist(cods,bins=np.arange(-1,1,.01),alpha=.5)\n",
    "\n",
    "cods = []\n",
    "\n",
    "for node in cluster.nodes:\n",
    "    (n,p) = second_forest.old_prediction.node_r2_doublet(node)\n",
    "    fvu = n/p\n",
    "    cod = 1 - fvu\n",
    "    cods.append(cod)\n",
    "\n",
    "plt.hist(cods,bins=np.arange(-1,1,.01),alpha=.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fraction_return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for factor in second_forest.split_clusters:\n",
    "    pops = [len(n.samples()) for n in factor.nodes]\n",
    "    plt.figure()\n",
    "    plt.title(factor.name())\n",
    "    plt.hist(pops,bins=np.arange(0,8000,100))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.self_prediction.prediction_report()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_forest.other_prediction.prediction_report()"
   ]
  },
  {
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}

