{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Local Effects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the more important aspects of random forest nodes, and by extension node clusters, is that they describe what we would call \"Local Effects\"\n",
    "\n",
    "While a conventional linear regression might describe a linear relationship between the behavior of a feature and a target that is true across the entire dataset, a node in a random forest may just as easily be a child of another node, and thus only trained on a small part of the dataset. Therefore a relationship that it describes between a feature and a target may be true across the entire dataset, or it may only be true conditionally on the predictions made by the parents of the node."
   ]
  },
  {
   "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 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",
    "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",
    "forest = tr.Forest.load(data_location + 'cv_forest_trimmed_extra')\n",
    "forest.arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_feature_mask = np.zeros(2000,dtype=bool)\n",
    "\n",
    "for feature in forest.output_features:\n",
    "    f_i = list(young.var_names).index(feature)\n",
    "    filtered_feature_mask[f_i] = True\n",
    "    \n",
    "young_filtered = young[:,filtered_feature_mask]\n",
    "young_filtered.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forest.reset_split_clusters()\n",
    "\n",
    "# forest.interpret_splits(\n",
    "#     depth=8,\n",
    "#     mode='additive_mean',\n",
    "#     metric='cosine',\n",
    "#     pca=100,\n",
    "#     relatives=True,\n",
    "#     k=50,\n",
    "#     resolution=2,\n",
    "# )\n",
    "\n",
    "# print(len(forest.split_clusters))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forest.maximum_spanning_tree(mode='samples')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forest.html_tree_summary(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forest.backup(data_location + \"full_clustering\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We now would like to see if there are any local associations that are dramatically different\n",
    "# from global ones, to the degree that it is impossible to recapture them using PCA-based analysis. \n",
    "\n",
    "# We will need to perform a PCA analysis first. \n",
    "\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "model = PCA(n_components=40).fit(young.X)\n",
    "transformed = model.transform(young.X)\n",
    "recovered = model.inverse_transform(transformed)\n",
    "\n",
    "centered = young.X - np.mean(young.X,axis=0)\n",
    "null_squared_residual = np.power(centered,2)\n",
    "\n",
    "recovered_residual = young.X - recovered\n",
    "recovered_squared_residual = np.power(recovered_residual,2)\n",
    "\n",
    "pca_recovered_per_sample = np.sum(recovered_squared_residual,axis=1)\n",
    "pca_recovered_fraction_per_sample = np.sum(recovered_squared_residual,axis=1) / np.sum(null_squared_residual,axis=1)\n",
    "print(np.sum(null_squared_residual))\n",
    "print(np.sum(recovered_squared_residual))\n",
    "\n",
    "print(f\"Remaining variance:{(np.sum(recovered_squared_residual) / np.sum(null_squared_residual))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Here we specify two interesting features and see what the weights for them are in each PC\n",
    "\n",
    "f1 = \"Tmem119\"\n",
    "f2 = \"Cd74\"\n",
    "\n",
    "f1_index = forest.truth_dictionary.feature_dictionary[f1]\n",
    "f2_index = forest.truth_dictionary.feature_dictionary[f2]\n",
    "\n",
    "f1_loadings = model.components_[:,f1_index]\n",
    "f2_loadings = model.components_[:,f2_index]\n",
    "\n",
    "# plt.figure()\n",
    "# plt.scatter(forest.output[:,f1_index],forest.output[:,f2_index],s=1)\n",
    "# plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.title(f\"PC Loadings of {f1} vs {f2}\")\n",
    "plt.scatter(f1_loadings,f2_loadings)\n",
    "plt.xlabel(f\"{f1} weight\")\n",
    "plt.ylabel(f\"{f2} weight\")\n",
    "plt.plot([.2,-.2],np.array([-.2,.2])*.55,color='red',label=\"Slope of -.55\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "for i,pc in enumerate(model.components_):\n",
    "    print(f\"PC:{i}, {f1}:{pc[f1_index]}, {f2}:{pc[f2_index]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Here we visualize the loadings of each PC to get a sense for where the PC is making meaningful predictions. \n",
    "# (It may give us a hint as to whether or not it specifies a particular cell type)\n",
    "\n",
    "for i,pc in enumerate(transformed.T):\n",
    "    plt.figure()\n",
    "    plt.title(f\"PC {i} Loadings\")\n",
    "    ab_max = np.max(np.abs(pc))\n",
    "    plt.scatter(*forest.tsne_coordinates.T,c=pc,s=3,alpha=.4,cmap='bwr',vmin=-ab_max,vmax=ab_max)\n",
    "    plt.xlabel(\"UMAP Embedding, (AU)\")\n",
    "    plt.ylabel(\"UMAP Embedding, (AU)\")\n",
    "    plt.colorbar()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we will look for features that have an especially large discrepancy in the local \n",
    "# correlation compared to the global correlation for each factor. \n",
    "\n",
    "for factor in forest.split_clusters:\n",
    "    print(\"=====================================\")\n",
    "    print(factor.name())\n",
    "    print(\"=====================================\")\n",
    "    fi_pairs = factor.most_local_correlations()\n",
    "    features = forest.output_features\n",
    "    f_names = [(features[i],features[j]) for (i,j) in fi_pairs]\n",
    "    local_correlations = factor.local_correlations()\n",
    "    global_correlations = forest.global_correlations()\n",
    "    discrepancy = [(local_correlations[i,j],global_correlations[i,j]) for (i,j) in fi_pairs]\n",
    "    print(f_names)\n",
    "    print(discrepancy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "interesting_pairs = []\n",
    "\n",
    "for factor in forest.split_clusters:\n",
    "    interesting_pairs.extend(factor.most_local_correlations(n=1))\n",
    "    \n",
    "uniques = list(set([y for x in interesting_pairs for y in x]))\n",
    "  \n",
    "factor_correlation_table = np.zeros((len(interesting_pairs),len(forest.split_clusters)))\n",
    "\n",
    "for i,factor in enumerate(forest.split_clusters):\n",
    "    local_correlations = factor.local_correlations(indices=uniques)\n",
    "    for j,(f1,f2) in enumerate(interesting_pairs):\n",
    "        f1_u = uniques.index(f1)\n",
    "        f2_u = uniques.index(f2)\n",
    "        factor_correlation_table[j,i] = local_correlations[f1_u,f2_u]\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(factor_correlation_table,interpolation='none',aspect='auto',cmap='bwr',vmin=-1,vmax=1)\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n",
    "from scipy.cluster.hierarchy import linkage,dendrogram\n",
    "\n",
    "# factor_agglomeration = dendrogram(linkage(factor_correlation_table, metric='cosine', method='average'), no_plot=True)['leaves']\n",
    "\n",
    "# plt.figure()\n",
    "# plt.imshow(factor_correlation_table.T[factor_agglomeration].T,interpolation='none',aspect='auto',cmap='bwr',vmin=-1,vmax=1)\n",
    "# plt.colorbar()\n",
    "# plt.show()\n",
    "\n",
    "print([(x,y) for x,y in enumerate(interesting_pairs)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "forest.output_features[1639]\n",
    "\n",
    "print(forest.split_clusters[23].local_correlations(indices=[717,1639]))\n",
    "print(forest.split_clusters[20].local_correlations(indices=[717,1639]))\n",
    "\n",
    "# cluster 23, Rrares2 (717), Meg3 (1639)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we check the naive linear fit between two features (eg a simple correlation among all cells)\n",
    "\n",
    "from scipy.stats import linregress\n",
    "\n",
    "f1 = \"Ctsd\"\n",
    "f2 = \"H2-Ab1\"\n",
    "\n",
    "f1_index = forest.truth_dictionary.feature_dictionary[f1]\n",
    "f2_index = forest.truth_dictionary.feature_dictionary[f2]\n",
    "\n",
    "f1_values = forest.output[:,f1_index]\n",
    "f2_values = forest.output[:,f2_index]\n",
    "\n",
    "slope,intercept,r_fit,_,_ = linregress(f1_values,f2_values)\n",
    "\n",
    "plt.figure(figsize=(3,2.5))\n",
    "plt.title(f\"Linar Fit, {f1}, {f2}, Naive\")\n",
    "plt.scatter(f1_values,f2_values,s=3)\n",
    "plt.plot(np.arange(7), intercept + (np.arange(7) * slope),c='red',label=f\"Slope:{np.around(slope,3)},R2:{np.around(r_fit,3)}\")\n",
    "plt.legend()\n",
    "plt.xlabel(f\"{f1}\")\n",
    "plt.ylabel(f\"{f2}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we filter only for cells that have a high or low sister score for a particular factor\n",
    "# and linearly regress two genes to check for a \"local\" association. \n",
    "\n",
    "\n",
    "from scipy.stats import linregress\n",
    "\n",
    "factor = forest.split_clusters[34]\n",
    "factor_threshold = .2\n",
    "factor_mask = np.abs(factor.sister_scores() > factor_threshold)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(f\"Sister scores, {factor.name()}\")\n",
    "plt.hist(factor.sister_scores(),bins=50)\n",
    "plt.plot([factor_threshold,factor_threshold],[-100,100],color='red')\n",
    "plt.plot([-factor_threshold,-factor_threshold],[-100,100],color='red',label=\"Sister score threshold\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "f1 = \"Tmem119\"\n",
    "f2 = \"Cd74\"\n",
    "\n",
    "f1_index = forest.truth_dictionary.feature_dictionary[f1]\n",
    "f2_index = forest.truth_dictionary.feature_dictionary[f2]\n",
    "\n",
    "f1_values = forest.output[:,f1_index][factor_mask]\n",
    "f2_values = forest.output[:,f2_index][factor_mask]\n",
    "\n",
    "slope,intercept,r_fit,_,_ = linregress(f1_values,f2_values)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(f\"Linar Fit, {f1}, {f2}, Factor {factor.name()}, Filtered\")\n",
    "plt.scatter(f1_values,f2_values,s=3)\n",
    "plt.plot(np.arange(7), intercept + (np.arange(7) * slope),c='red',label=f\"Slope:{np.around(slope,3)},R2:{np.around(r_fit,3)}\")\n",
    "plt.xlabel(f\"{f1}\")\n",
    "plt.ylabel(f\"{f2}\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Here we filter only for cells that have a high or low sister score for a particular factor\n",
    "# and linearly regress two genes to check for a \"local\" association. \n",
    "\n",
    "\n",
    "from scipy.stats import linregress\n",
    "\n",
    "factor = forest.split_clusters[41]\n",
    "factor_threshold = .05\n",
    "factor_mask = np.abs(factor.sister_scores() > factor_threshold)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(f\"Sister scores, {factor.name()}\")\n",
    "plt.hist(factor.sister_scores(),bins=50)\n",
    "plt.plot([factor_threshold,factor_threshold],[-100,100],color='red')\n",
    "plt.plot([-factor_threshold,-factor_threshold],[-100,100],color='red',label=\"Sister score threshold\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "f1 = \"Tmem119\"\n",
    "f2 = \"Cd74\"\n",
    "\n",
    "f1_index = forest.truth_dictionary.feature_dictionary[f1]\n",
    "f2_index = forest.truth_dictionary.feature_dictionary[f2]\n",
    "\n",
    "f1_values = forest.output[:,f1_index][factor_mask]\n",
    "f2_values = forest.output[:,f2_index][factor_mask]\n",
    "\n",
    "slope,intercept,r_fit,_,_ = linregress(f1_values,f2_values)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(f\"Linar Fit, {f1}, {f2}, Factor {factor.name()}, Filtered\")\n",
    "plt.scatter(f1_values,f2_values,s=3)\n",
    "plt.plot(np.arange(7), intercept + (np.arange(7) * slope),c='red',label=f\"Slope:{np.around(slope,3)},R2:{np.around(r_fit,3)}\")\n",
    "plt.xlabel(f\"{f1}\")\n",
    "plt.ylabel(f\"{f2}\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we find highly weighted genes for a particular PC, as well as the rankings of particular features of interest\n",
    "# Our objective is to see if the two featurs represent an important part of the variance captured by the PC\n",
    "pc = 8\n",
    "\n",
    "f1 = \"Tmem119\"\n",
    "f2 = \"Cd74\"\n",
    "\n",
    "f1_index = forest.truth_dictionary.feature_dictionary[f1]\n",
    "f2_index = forest.truth_dictionary.feature_dictionary[f2]\n",
    "\n",
    "weights = model.components_[pc]\n",
    "\n",
    "weight_sort = np.argsort(np.abs(weights))\n",
    "\n",
    "print(list(forest.output_features[weight_sort[:-20:-1]]))\n",
    "print(list(weights[weight_sort[:-20:-1]]))\n",
    "\n",
    "print(f\"{f1}: {len(weights) - list(weight_sort).index(f1_index)}\")\n",
    "print(f\"{f2}: {len(weights) - list(weight_sort).index(f2_index)}\")\n",
    "\n",
    "print(weights[f1_index])\n",
    "print(weights[f2_index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for s_c in forest.split_clusters:\n",
    "    scores = s_c.sister_scores()\n",
    "    log_scores = s_c.log_sister_scores(prior=10)\n",
    "\n",
    "    abmax=np.max(np.abs(scores))\n",
    "\n",
    "    plt.figure()\n",
    "    plt.title(\"Regular\")\n",
    "    plt.scatter(*forest.tsne_coordinates.T,c=scores,cmap='bwr',s=1,vmin=-abmax,vmax=abmax)\n",
    "    plt.colorbar()\n",
    "    plt.show()\n",
    "\n",
    "    abmax=np.max(np.abs(log_scores))\n",
    "\n",
    "    plt.figure()\n",
    "    plt.title(\"Log\")\n",
    "    plt.scatter(*forest.tsne_coordinates.T,c=log_scores,cmap='bwr',s=1,vmin=-abmax,vmax=abmax)\n",
    "    plt.colorbar()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "factor = forest.split_clusters[34]\n",
    "\n",
    "samples = factor.sample_scores()\n",
    "sisters = factor.sister_scores()\n",
    "log_sisters = factor.log_sister_scores()\n",
    "\n",
    "plt.figure(figsize=(3,2.5))\n",
    "plt.title(f\"Distribution of Sample Scores In {factor.name()}\")\n",
    "plt.hist(samples,bins=50)\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.xlabel(\"Sample Scores\")\n",
    "plt.show()\n",
    "\n",
    "plt.figure(figsize=(3,2.5))\n",
    "plt.title(f\"Distribution of Sister Scores In {factor.name()}\")\n",
    "plt.hist(sisters,bins=50)\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.xlabel(\"Sister Scores\")\n",
    "plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.hist(log_sisters,bins=50)\n",
    "plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we test whether or not a particular factor over-expresses a gene of interest\n",
    "# (Used as a statistical test for cell type identity, eg \"is factor 34 immune cells?\")\n",
    "from scipy.stats import ttest_ind\n",
    "\n",
    "factor = forest.split_clusters[34]\n",
    "factor_threshold = .2\n",
    "mask = factor.sister_scores() > factor_threshold\n",
    "\n",
    "feature = \"Cd45\"\n",
    "\n",
    "f_index = forest.truth_dictionary.feature_dictionary[f1]\n",
    "\n",
    "test = ttest_ind(young.X[mask][:,f_index],young.X[~mask][:,f_index],equal_var=False)\n",
    "\n",
    "print(f\"{feature} in {factor.name()} vs all other: {test}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import ttest_ind\n",
    "\n",
    "factor = forest.split_clusters[34]\n",
    "factor_threshold = .2\n",
    "mask = factor.sister_scores() > factor_threshold\n",
    "\n",
    "feature = \"C1qa\"\n",
    "\n",
    "f_index = forest.truth_dictionary.feature_dictionary[f1]\n",
    "\n",
    "test = ttest_ind(young.X[mask][:,f_index],young.X[~mask][:,f_index],equal_var=False)\n",
    "\n",
    "print(f\"{feature} in {factor.name()} vs all other: {test}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature = \"Tmem119\"\n",
    "\n",
    "f_index = forest.truth_dictionary.feature_dictionary[feature]\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(*forest.tsne_coordinates.T,c=forest.output[:,f_index],s=1)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature = \"Cd74\"\n",
    "\n",
    "f_index = forest.truth_dictionary.feature_dictionary[feature]\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(forest.split_clusters[24].sister_scores(),forest.output[:,f_index],s=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(3,2.5))\n",
    "plt.title(\"Cd45 Mean Expression\")\n",
    "plt.bar([0,1],[0.0030044015,0.3989004],yerr=[0.0006448814299982911,0.017244088969228192],width=.5,tick_label=[\"Rest\",\"NC 34\"])\n",
    "plt.ylabel(\"Mean Expression (Log TPM)\")\n",
    "plt.show()"
   ]
  },
  {
   "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
}

