BC_Spatial_LDA_1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-16T16:27:39.516711Z",
     "start_time": "2020-07-16T16:27:39.495615Z"
    },
    "cellView": "both",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 923
    },
    "colab_type": "code",
    "id": "9dxxGAa9QhMd",
    "outputId": "581b0d94-1a72-471c-81a8-752fcbd67a9d"
   },
   "outputs": [],
   "source": [
    "#@title Install dependencies and setup paths\n",
    "import os\n",
    "import sys\n",
    "  \n",
    "PATH_TO_MODELS = f'./cycif_models/'\n",
    "\n",
    "paths = [\n",
    "  PATH_TO_MODELS,\n",
    "]\n",
    "\n",
    "paths_exist = {path: os.path.exists(path) for path in paths}\n",
    "display(paths_exist)\n",
    "assert all(paths_exist.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-16T16:27:41.458979Z",
     "start_time": "2020-07-16T16:27:39.518512Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "colab_type": "code",
    "id": "nDnN8F4I4wNS",
    "outputId": "20293984-cdf9-49ac-fd0a-bea67dc2d3f3"
   },
   "outputs": [],
   "source": [
    "#@title Import and set processes and set training params\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import functools\n",
    "import logging\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import os\n",
    "import pickle\n",
    "import scipy\n",
    "import seaborn as sns\n",
    "from sklearn.model_selection import train_test_split\n",
    "import time\n",
    "import tqdm\n",
    "# Spatial LDA imports\n",
    "from spatial_lda.featurization import neighborhood_to_cluster\n",
    "from spatial_lda.featurization import make_nearest_neighbor_graph\n",
    "from spatial_lda.featurization import make_merged_difference_matrices\n",
    "from spatial_lda.featurization import featurize_spleens\n",
    "#from spatial_lda.visualization import plot_samples_in_a_row\n",
    "#from spatial_lda.visualization import plot_bcell_topic_multicolor\n",
    "import spatial_lda.model\n",
    "#import scimap as sm\n",
    "import scanpy as sc\n",
    "\n",
    "N_PARALLEL_PROCESSES = 8 #@param\n",
    "TRAIN_SIZE_FRACTION = 0.99 #@param\n",
    "N_TOPICS_LIST = [8, 10, 6] #@param\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## functions\n",
    "%matplotlib inline\n",
    "import palettable.cartocolors.qualitative as qual_palettes\n",
    "import matplotlib.patches as mpatches\n",
    "import matplotlib.markers as mmark\n",
    "from matplotlib.lines import Line2D\n",
    "\n",
    "def plot_bcell_topic_multicolor(ax, sample_idx, topic_weights, spleen_dfs):\n",
    "    topic_weights_slide = topic_weights[topic_weights.index.map(lambda x: x[0])==s_slide]\n",
    "    color_palette = qual_palettes.Bold_10.mpl_colors\n",
    "    colors = color_palette[:topic_weights_slide.shape[1]]\n",
    "    d_color = dict(zip(sorted(set(np.argmax(np.array(topic_weights_slide),axis=1))),colors))\n",
    "    cell_coords = spleen_dfs[sample_idx]\n",
    "    non_b_coords = cell_coords[~cell_coords.isb]\n",
    "    ax.scatter(\n",
    "        non_b_coords['sample.X'],\n",
    "        non_b_coords['sample.Y'],\n",
    "        s=1,\n",
    "        c='k',\n",
    "        marker='x',\n",
    "        label='Non-tumor',\n",
    "        alpha=.2)\n",
    "\n",
    "    cell_indices = topic_weights_slide.index.map(lambda x: x[1])\n",
    "    coords = spleen_dfs[sample_idx].loc[spleen_dfs[sample_idx].index.isin(cell_indices)]\n",
    "\n",
    "    ax.scatter(coords['sample.X'], coords['sample.Y'], s=3,\n",
    "               c=np.array(colors)[np.argmax(np.array(topic_weights_slide), axis=1), :],\n",
    "               #label=d_color\n",
    "              )\n",
    "    handles = [mpatches.Circle((1,1), radius=5, color=color) for color in d_color.values()] #rectangles\n",
    "    #handles = [Line2D([], [], c=color, lw=0, marker=\"o\", markersize=8, label=species) for species, color in d_color.items()] #dots\n",
    "    ax.axes.get_yaxis().set_visible(False)\n",
    "    ax.axes.get_xaxis().set_visible(False)\n",
    "    ax.axes.set_title(f'Sample {sample_idx}')\n",
    "    ax.axis('equal')\n",
    "    ax.set_ylim(ax.get_ylim()[::-1])\n",
    "    ax.legend(handles, d_color.keys(),\n",
    "                    bbox_to_anchor=(1.05,.5), title=\"Topics\")\n",
    "    #return(fig, ax)\n",
    "\n",
    "\n",
    "def plot_samples_in_a_row(features_df, plot_fn, patient_dfs, tumor_set=None):\n",
    "    sns.set_style(\"white\")\n",
    "    tumor_idx = features_df.index.map(lambda x: x[0])\n",
    "    if tumor_set is None:\n",
    "        tumor_set = np.unique(tumor_idx)\n",
    "\n",
    "    n = len(tumor_set)\n",
    "    num_rows = 1\n",
    "    num_cols = (len(tumor_set) // num_rows)\n",
    "    fig, axes = plt.subplots(\n",
    "        num_rows, num_cols, figsize=(num_cols * 4, num_rows * 4), dpi=300)\n",
    "\n",
    "    for i, tumor in enumerate(tumor_set):\n",
    "        plot_fn(axes[i], tumor, features_df[tumor_idx == tumor], patient_dfs)\n",
    "\n",
    "    sns.despine(left=True, bottom=True)\n",
    "    return(fig, axes)\n",
    "    \n",
    "def plot_topics_heatmap(topics, features, ax, normalizer=None):\n",
    "    n_topics = topics.shape[0]\n",
    "    if normalizer is not None:\n",
    "        topics = normalizer(topics)\n",
    "    else:\n",
    "        topics = _standardize_topics(topics)\n",
    "\n",
    "    topics = pd.DataFrame(topics, index=features,\n",
    "                          columns=['Topic %d' % x for x in range(n_topics)])\n",
    "    sns.heatmap(topics, square=True, cmap='RdBu_r', ax=ax)\n",
    "    #return(fig, ax)\n",
    "from spatial_lda.visualization import plot_adjacency_graph\n",
    "\n",
    "def make_plot_fn(difference_matrices):  \n",
    "    def plot_fn(ax, tumor_idx, features_df, patient_dfs):\n",
    "        plot_adjacency_graph(ax, tumor_idx, features_df, patient_dfs, difference_matrices)\n",
    "    return plot_fn\n",
    "\n",
    "def plot_topic_multicolor(ax, sample_idx, topic_weights, spleen_dfs):\n",
    "    color_palette = qual_palettes.Bold_10.mpl_colors\n",
    "    colors = color_palette[:topic_weights.shape[1]]\n",
    "    d_color = dict(zip(sorted(set(np.argmax(np.array(topic_weights),axis=1))),colors))\n",
    "    cell_coords = spleen_dfs[sample_idx]\n",
    "    non_b_coords = cell_coords[~cell_coords.isb]\n",
    "    ax.scatter(\n",
    "        non_b_coords['sample.X'],\n",
    "        non_b_coords['sample.Y'],\n",
    "        s=1,\n",
    "        c='k',\n",
    "        marker='x',\n",
    "        label='Non-tumor',\n",
    "        alpha=.2)\n",
    "\n",
    "    cell_indices = topic_weights.index#.map(lambda x: x[1])\n",
    "    coords = spleen_dfs[sample_idx].loc[cell_indices]\n",
    "\n",
    "    ax.scatter(coords['sample.X'], coords['sample.Y'], s=3,\n",
    "               c=np.array(colors)[np.argmax(np.array(topic_weights), axis=1), :],\n",
    "               #label=d_color\n",
    "              )\n",
    "    handles = [mpatches.Circle((1,1), radius=5, color=color) for color in d_color.values()] #rectangles\n",
    "    #handles = [Line2D([], [], c=color, lw=0, marker=\"o\", markersize=8, label=species) for species, color in d_color.items()] #dots\n",
    "    ax.axes.get_yaxis().set_visible(False)\n",
    "    ax.axes.get_xaxis().set_visible(False)\n",
    "    ax.axes.set_title(f'Sample {sample_idx}')\n",
    "    ax.axis('equal')\n",
    "    ax.set_ylim(ax.get_ylim()[::-1])\n",
    "    ax.legend(handles, d_color.keys(),\n",
    "                    bbox_to_anchor=(1.05,.5), title=\"Topics\")\n",
    "    #return(fig, ax)\n",
    "    \n",
    "def plot_results(adata, s_col, x='X_centroid',y='Y_centroid'):\n",
    "    fig, ax = plt.subplots()\n",
    "    df = adata.to_df()\n",
    "    df[adata.obs.columns] = adata_lda.obs\n",
    "    color_palette = qual_palettes.Bold_10.mpl_colors\n",
    "    colors = color_palette[:len(df.loc[:,s_col].unique())]\n",
    "    d_color = dict(zip(sorted(set(df.loc[:,s_col])),colors))\n",
    "    if adata_lda.obs._get_numeric_data().columns.isin([s_col]).any():\n",
    "        sp = ax.scatter(df[x], df[y], s=3,c=df.loc[:,s_col],\n",
    "           cmap = 'plasma',\n",
    "              )\n",
    "        fig.colorbar(sp)\n",
    "    else: \n",
    "        ax.scatter(df[x], df[y], s=3,\n",
    "           c=df.loc[:,s_col].astype('str').map(d_color),\n",
    "              )\n",
    "        handles = [mpatches.Circle((1,1), radius=5, color=color) for color in d_color.values()]\n",
    "        ax.legend(handles, d_color.keys(),\n",
    "                    bbox_to_anchor=(1.05,.8), title=s_col)\n",
    "    ax.axis('equal')\n",
    "    ax.set_ylim(ax.get_ylim()[::-1])\n",
    "\n",
    "    return(fig, ax)\n",
    "\n",
    "def plot_one_tumor_all_topics(ax, tumor_idx, topic_weights, patient_dfs):\n",
    "    color_palette = qual_palettes.Bold_10.mpl_colors\n",
    "    colors = np.array(color_palette[:topic_weights.shape[1]])\n",
    "    d_color = dict(zip(sorted(set(np.argmax(np.array(topic_weights),axis=1))),colors))\n",
    "    cell_coords = patient_dfs[tumor_idx]\n",
    "    immune_coords = cell_coords[cell_coords.isimmune]\n",
    "    cell_indices = topic_weights.index.map(lambda x: x[1])\n",
    "    coords = patient_dfs[tumor_idx].loc[cell_indices]\n",
    "    ax.scatter(immune_coords['x'], -immune_coords['y'],\n",
    "               s=5, c='k', label='__Immune', alpha=0.1)\n",
    "    ax.scatter(coords['x'], -coords['y'], s=2,\n",
    "               c=colors[np.argmax(np.array(topic_weights), axis=1), :])\n",
    "    ax.set_title(f\"Tumor {tumor_idx}\")\n",
    "    ax.axes.get_yaxis().set_visible(False)\n",
    "    ax.axes.get_xaxis().set_visible(False)\n",
    "    ax.legend()\n",
    "    handles = [Line2D([], [], c=color, lw=0, marker=\"o\", markersize=8, label=species) for species, color in d_color.items()] #dots\n",
    "    ax.legend(handles, d_color.keys(),\n",
    "                    bbox_to_anchor=(1.05,.9), title=\"Topics\")\n",
    "\n",
    "def plot_one_tumor_topic(ax, tumor_idx, topic_weights, patient_dfs):\n",
    "    cell_coords = patient_dfs[tumor_idx]\n",
    "    immune_coords = cell_coords[cell_coords.isimmune]\n",
    "    cell_indices = topic_weights.index.map(lambda x: x[1])\n",
    "    coords = patient_dfs[tumor_idx].loc[cell_indices]\n",
    "    ax.scatter(immune_coords['x'], -immune_coords['y'],\n",
    "               s=5, c='k', label='__Immune', alpha=0.1)\n",
    "    ax.scatter(coords['x'], -coords['y'], s=2, c=topic_weights, cmap=\"plasma\",\n",
    "               label=topic_weights.name)\n",
    "    ax.set_title(f\"Tumor {tumor_idx}\")\n",
    "    ax.axes.get_yaxis().set_visible(False)\n",
    "    ax.axes.get_xaxis().set_visible(False)\n",
    "    ax.legend()\n",
    "    \n",
    "def plot_sm_topic(ax, s_slide, topic_weights, patient_dfs,topic_idx):\n",
    "    cell_coords = patient_dfs[patient_dfs.imageid==s_slide]\n",
    "    cell_indices = topic_weights[topic_weights.index.str.contains(s_slide)].index\n",
    "    coords = cell_coords.loc[cell_indices]\n",
    "    #ax.scatter(cell_coords['X_centroid'], -cell_coords['Y_centroid'],\n",
    "    #           s=5, c='k', label='__Immune', alpha=0.1)\n",
    "    ax.scatter(coords['X_centroid'], -coords['Y_centroid'], s=2, c=topic_weights.loc[cell_indices,topic_idx], cmap=\"plasma\",\n",
    "               label=topic_idx)\n",
    "    ax.set_title(f\"{s_slide}\")\n",
    "    ax.axes.get_yaxis().set_visible(False)\n",
    "    ax.axes.get_xaxis().set_visible(False)\n",
    "    ax.axis('equal')\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aLIaqZzV9Xcx"
   },
   "source": [
    "## Load data and featurize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir = '/home/groups/graylab_share/Chin_Lab/ChinData/engje/Data/20200000/20200406_JP-TMAs'#'/home/groups/graylab_share/data/engje/Data/20200000/20200406_JP-TMAs'\n",
    "#this data has been stadardized to 1 pixel = 1 um\n",
    "df_cyc = pd.read_csv(f'{datadir}/data/20220420_JP-TMAs_IMC-TMAs_MIBI_CombinedCelltypes_all.csv',index_col = 0,dtype='object')\n",
    "\n",
    "df_cyc['x'] = df_cyc.DAPI_X.astype('float64')\n",
    "df_cyc['y'] = df_cyc.DAPI_Y.astype('float64')\n",
    "\n",
    "df_cyc['sample.X'] = df_cyc.DAPI_X.astype('float64')\n",
    "df_cyc['sample.Y'] = df_cyc.DAPI_Y.astype('float64')\n",
    "\n",
    "df_cyc['isb'] = df_cyc.loc[:,'leidencelltype3']=='epithelial'\n",
    "df_cyc['isimmune'] = df_cyc.loc[:,'leidencelltype3']!='epithelial'\n",
    "#MIBI\n",
    "df_cyc.loc[df_cyc.Platform=='MIBI','isb'] = df_cyc.loc[df_cyc.Platform=='MIBI','leidencelltype2']=='epithelial'\n",
    "df_cyc.loc[df_cyc.Platform=='MIBI','isimmune'] = df_cyc.loc[df_cyc.Platform=='MIBI','leidencelltype2']!='epithelial'\n",
    "\n",
    "df_cyc['cluster'] = df_cyc.loc[:,'leiden']\n",
    "df_cyc['sample.Z'] = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load subtype annotation\n",
    "df_surv = pd.read_csv(f'{datadir}/annotation/BC-TMAs_clinical_data.csv',index_col=0)\n",
    "\n",
    "df_cyc_sub = pd.read_csv(f'{datadir}/annotation/20210403_JP-TMA1_Annotation+Clinical-Subtype.csv',index_col=0)\n",
    "df_cyc_sub.rename({'ClinicalSubtype':'ID'},axis=1,inplace=True)\n",
    "#collapse Subtypes\n",
    "d_replace = {'0':'other','?':'other', np.nan:'other'}\n",
    "df_cyc_sub.ID.replace(d_replace,inplace=True)\n",
    "df_cyc_sub = df_cyc_sub.append(df_surv[df_surv.index.str.contains('JP-TMA2-1')].rename({'subtype':'ID'},axis=1))\n",
    "df_cyc_sub.loc[df_cyc_sub.index.str.contains('JP-TMA2-1'),'Accession'] = df_cyc_sub.loc[df_cyc_sub.index.str.contains('JP-TMA2-1')].index\n",
    "#add subtype and patient\n",
    "#d_cyc_sub = dict(zip(df_cyc_sub.index.tolist(),df_cyc_sub.ID.tolist()))\n",
    "#d_patient = dict(zip(df_cyc_sub.index,df_cyc_sub.Accession))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_surv.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## specify platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for s_plat in ['IMC','cycIF']: #'cycIF','MIBI','IMC'\n",
    "    ls_slide = sorted(df_cyc[df_cyc.Platform==s_plat].slide_scene.unique())\n",
    "    ls_drop = df_cyc.columns[df_cyc[df_cyc.Platform==s_plat].isna().sum() > 0]\n",
    "    print(len(ls_slide))\n",
    "    ## load data\n",
    "    PATH_TO_SPLEEN_DF_PKL = f'./data/{s_plat}_df.pkl'\n",
    "    if os.path.exists(PATH_TO_SPLEEN_DF_PKL):\n",
    "        print('Loading saved features.')\n",
    "        with open(PATH_TO_SPLEEN_DF_PKL, 'rb') as f:\n",
    "            spleen_dfs = pickle.load(f)\n",
    "    else:\n",
    "        spleen_dfs = {}\n",
    "        for s_slide in ls_slide:\n",
    "            df_scene = df_cyc.loc[(df_cyc.slide_scene==s_slide),~df_cyc.columns.isin(ls_drop)]\n",
    "            if len(df_scene) > 0:\n",
    "                if s_slide == '50_Ay11x6-202': #no tumor\n",
    "                    print(s_slide)\n",
    "                elif s_slide == '59_Cy8x3-472':\n",
    "                    print(s_slide)\n",
    "                elif s_slide == '61_X4Y4-215':\n",
    "                    print(s_slide)\n",
    "                elif s_slide == '187_X3Y2-99':\n",
    "                    print(s_slide)\n",
    "                elif s_slide == '241_X8Y3-337':\n",
    "                    print(s_slide)\n",
    "                elif s_slide == '19_Cy8x8-509':\n",
    "                    print(s_slide)\n",
    "                elif s_slide == 'JP-TMA1-1_scene103':\n",
    "                    print(s_slide)\n",
    "                elif df_scene.isb.sum()<4:\n",
    "                    print(s_slide)\n",
    "                else:\n",
    "                    spleen_dfs.update({s_slide:df_scene})\n",
    "        print(len(spleen_dfs.keys()))\n",
    "        # save the pickle!\n",
    "        with open(PATH_TO_SPLEEN_DF_PKL, 'wb') as f:\n",
    "            pickle.dump(spleen_dfs, f)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#pretty fast\n",
    "PATH_TO_SPLEEN_FEATURES_PKL = f'./data/{s_plat}_cells_features.pkl'\n",
    "print(PATH_TO_SPLEEN_FEATURES_PKL)\n",
    "if os.path.exists(PATH_TO_SPLEEN_FEATURES_PKL):\n",
    "    print('Loading saved features.')\n",
    "    with open(PATH_TO_SPLEEN_FEATURES_PKL, 'rb') as f:\n",
    "        spleen_cells_features = pickle.load(f)\n",
    "else:\n",
    "    print('Featurizing samples ...')\n",
    "\n",
    "    spleen_cells_features = featurize_spleens(spleen_dfs, neighborhood_to_cluster, radius=100,\n",
    "                                            n_processes=N_PARALLEL_PROCESSES)\n",
    "    with open(PATH_TO_SPLEEN_FEATURES_PKL, 'wb') as f:\n",
    "        pickle.dump(spleen_cells_features, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_UzTiFBFQhMo"
   },
   "source": [
    "## Split training set and compute difference matrices\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_subtype ='ER+'#'TNBC'# \n",
    "\n",
    "if s_plat == 'cycIF':\n",
    "    ls_index = df_cyc_sub[df_cyc_sub.ID==s_subtype].index\n",
    "else:\n",
    "    ls_index = df_cyc[df_cyc.subtype==s_subtype].slide_scene\n",
    "ls_index_dict = ls_index[ls_index.isin(spleen_dfs.keys())].unique()\n",
    "spleen_cells_features_sub = spleen_cells_features[spleen_cells_features.index.map(lambda x: x[0]).isin(ls_index_dict)]\n",
    "if s_plat == 'IMC':\n",
    "    try:\n",
    "        spleen_cells_features_sub.drop('20',axis=1,inplace=True)\n",
    "    except:\n",
    "        print('no 20')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spleen_dfs_sub = {}\n",
    "for key in ls_index_dict:\n",
    "        spleen_dfs_sub.update({key:spleen_dfs[key]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-16T16:27:46.120069Z",
     "start_time": "2020-07-16T16:27:42.175104Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "colab_type": "code",
    "id": "KJTIrqcwQhMo",
    "outputId": "0e91fd0c-a13d-420b-cee7-28620b2ef863",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "#pretty fast\n",
    "\n",
    "spleen_difference_matrices = make_merged_difference_matrices(#spleen_cells_features, spleen_dfs,\n",
    "                                                             spleen_cells_features_sub, spleen_dfs_sub,\n",
    "                                                             'sample.X', 'sample.Y')\n",
    "#all_sample_idxs = spleen_cells_features.index.map(lambda x: x[0])\n",
    "all_sample_idxs = spleen_cells_features_sub.index.map(lambda x: x[0])\n",
    "_sets = train_test_split(spleen_cells_features_sub, #spleen_cells_features, \n",
    "                         test_size=1. - TRAIN_SIZE_FRACTION,\n",
    "                         stratify=all_sample_idxs)\n",
    "train_spleen_cells_features, test_spleen_cells__features = _sets\n",
    "train_difference_matrices = make_merged_difference_matrices(\n",
    "    train_spleen_cells_features, spleen_dfs_sub, #spleen_dfs,\n",
    "    'sample.X', 'sample.Y')\n",
    "spleen_idxs = train_spleen_cells_features.index.map(lambda x: x[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Yblqo1-MQhMr"
   },
   "source": [
    "## Visualize graphs used to tie neighboring environments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "mpl.rcParams['pdf.fonttype'] = 42\n",
    "#fig.savefig('LDA_example.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-16T16:30:08.986457Z",
     "start_time": "2020-07-16T16:27:46.121820Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 240
    },
    "colab_type": "code",
    "id": "_4pYoBDWQhMr",
    "outputId": "2000a82d-f66b-4d3d-a137-5f8a885a3cfa",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## \n",
    "s_date = '20241112'\n",
    "if s_plat == 'cycIF':\n",
    "    #os.mkdir(s_date)\n",
    "    ls_sllide =[ 'JP-TMA1-1_scene034','JP-TMA2-1_scene35']#'JP-TMA2-1_scene35',\n",
    "    if s_subtype == 'ER+':\n",
    "        ls_sllide =[ 'JP-TMA1-1_scene044','JP-TMA1-1_scene044']\n",
    "    _plot_fn = make_plot_fn(spleen_difference_matrices)\n",
    "    \n",
    "    #plot_samples_in_a_row(spleen_cells_features, _plot_fn, spleen_dfs, tumor_set=ls_slide[0:2])\n",
    "    plot_samples_in_a_row(spleen_cells_features_sub, _plot_fn, spleen_dfs_sub,tumor_set=ls_sllide)#ls_index_dict[0:5])\n",
    "    fig=plt.gcf()\n",
    "    plt.tight_layout()\n",
    "    fig.savefig(f\"{s_date}/graphs_LDA_{s_subtype}_{s_plat}.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bu2vNJML20gJ"
   },
   "source": [
    "# Spatial LDA results\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DyFs0CTqm0NB"
   },
   "source": [
    "## Parameter sweep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_topics = 8\n",
    "N_TOPICS_LIST = [8]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_plat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_subtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-16T16:30:09.512100Z",
     "start_time": "2020-07-16T16:30:09.063087Z"
    },
    "colab": {},
    "colab_type": "code",
    "id": "J_lPB6KkQhMw",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#this is the slow step\n",
    "from spatial_lda.model import order_topics_consistently\n",
    "spatial_lda_models = {}  \n",
    "difference_penalty = 0.25  \n",
    "for n_topics in N_TOPICS_LIST:\n",
    "    path_to_train_model = '_'.join((f'{PATH_TO_MODELS}{s_plat}_{s_subtype}_training',\n",
    "                                  f'penalty={difference_penalty}',\n",
    "                                  f'topics={n_topics}',\n",
    "                                  f'trainfrac={TRAIN_SIZE_FRACTION}')) + '.pkl'\n",
    "    print(path_to_train_model)\n",
    "    if not os.path.exists(path_to_train_model):\n",
    "        print(f'Running n_topics={n_topics}, d={difference_penalty}\\n')\n",
    "        spatial_lda_model = spatial_lda.model.train(sample_features=train_spleen_cells_features, \n",
    "                                                difference_matrices=train_difference_matrices,\n",
    "                                                difference_penalty=difference_penalty,\n",
    "                                                n_topics=n_topics,\n",
    "                                                n_parallel_processes=N_PARALLEL_PROCESSES,                                                                         \n",
    "                                                verbosity=1,\n",
    "                                                admm_rho=0.1,\n",
    "                                                primal_dual_mu=1e+5)\n",
    "        spatial_lda_models[n_topics] = spatial_lda_model\n",
    "        with open(path_to_train_model, 'wb') as f:\n",
    "            pickle.dump(spatial_lda_model, f)    \n",
    "    else:\n",
    "        with open(path_to_train_model, 'rb') as f:\n",
    "            spatial_lda_models[n_topics] = pickle.load(f)\n",
    "      \n",
    "order_topics_consistently(spatial_lda_models.values())     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Qv1EllIVmyNF"
   },
   "source": [
    "## Load models with different number of topics and the same difference penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-16T16:30:09.550391Z",
     "start_time": "2020-07-16T16:30:09.514070Z"
    },
    "colab": {},
    "colab_type": "code",
    "id": "sVY6Vc7b1DiR"
   },
   "outputs": [],
   "source": [
    "#can select more\n",
    "lda_5 = spatial_lda_models[num_topics]\n",
    "topic_weights_5 = lda_5.topic_weights\n",
    "\n",
    "samples = ls_slide[0:5] #[s_slide,s_slide]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "G7LAvf-EndIU"
   },
   "source": [
    "### Visualizing topics on regular cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if s_plat == 'cycIF':\n",
    "    s_slide= 'JP-TMA1-1_scene034'\n",
    "    if s_subtype == 'ER+':\n",
    "        s_slide ='JP-TMA1-1_scene044'\n",
    "    num_topics= 8 \n",
    "    fig, ax = plot_samples_in_a_row(spatial_lda_models[num_topics].topic_weights, plot_one_tumor_all_topics, spleen_dfs, tumor_set=[s_slide,s_slide])\n",
    "    ax[0].legend(frameon=False)\n",
    "    plt.tight_layout()\n",
    "    fig.savefig(f\"{s_date}/graphs_LDA_topics_{s_subtype}_{s_plat}.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if s_plat == 'cycIF':\n",
    "    for t in [0,1,2,3,4,5,6]:\n",
    "        print(t)\n",
    "        fig, ax = plot_samples_in_a_row(spatial_lda_models[num_topics].topic_weights.iloc[:, t], plot_one_tumor_topic, spleen_dfs, tumor_set=[s_slide,s_slide])\n",
    "        ax[0].legend(frameon=False)\n",
    "        plt.tight_layout()\n",
    "        fig.savefig(f\"{s_date}/graphs_LDA_topic_{t}_{s_subtype}_{s_plat}.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clinical annotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load Cyclic patients and subtypes\n",
    "df_cyc_sub = pd.read_csv(f'{datadir}/annotation/20210403_JP-TMA1_Annotation+Clinical-Subtype.csv',index_col=0)\n",
    "df_cyc_sub.rename({'ClinicalSubtype':'ID'},axis=1,inplace=True)\n",
    "#collapse Subtypes\n",
    "d_replace = {'0':'other','?':'other', np.nan:'other'}\n",
    "df_cyc_sub.ID.replace(d_replace,inplace=True)\n",
    "df_cyc_sub.loc[df_cyc_sub.index.str.contains('JP-TMA2-1'),'Accession'] = df_cyc_sub.loc[df_cyc_sub.index.str.contains('JP-TMA2-1'),'TMA_scene.1']\n",
    "#add subtype and patient\n",
    "d_cyc_sub = dict(zip(df_cyc_sub.index.tolist(),df_cyc_sub.ID.tolist()))\n",
    "d_patient = dict(zip(df_cyc_sub.index,df_cyc_sub.Accession))\n",
    "# JP-TMA2 subtypes\n",
    "\n",
    "#survival\n",
    "df_surv = pd.read_csv(f'{datadir}/data/cycIF_clinical_outcome.csv',index_col=0)\n",
    "df_surv['Platform'] = 'cycIF'\n",
    "df_surv.rename({'Recurence_time':'Recurrence_time'},axis=1,inplace=True)\n",
    "\n",
    "df_cyc_a = pd.read_csv(f'{datadir}/annotation/JP-TMAs_Clinical_Variables.csv',index_col=0)\n",
    "\n",
    "d_sub_both = dict(zip(df_cyc_sub.Accession,df_cyc_sub.ID))\n",
    "#d_sub_both.update(d_cyc_sub2)\n",
    "df_surv['subtype'] = df_surv.index.map(d_sub_both)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_surv.head()\n",
    "#df_cyc_a.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_weights_5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save tissue means\n",
    "if s_plat == 'cycIF':\n",
    "    s_sample = 'JP-TMA1'\n",
    "else:\n",
    "    s_sample = 'IMC'\n",
    "    if s_subtype == 'TNBC':\n",
    "        spleen_cells_features.drop('20',axis=1,inplace=True)\n",
    "s_grouper='Patient'\n",
    "s_cell= 'all'\n",
    "for num_topics in [8]:\n",
    "    lda_5 = spatial_lda_models[num_topics]\n",
    "    topic_weights_5 = lda_5.topic_weights\n",
    "    topic_weights_5['Patient'] = topic_weights_5.index.map(lambda x: x[0]).map(d_patient)\n",
    "\n",
    "    df_out = topic_weights_5.groupby(s_grouper).mean()\n",
    "    s_out = f'results_{s_sample}_SpatialLDA_by{s_grouper}_by{s_cell}_k{num_topics}.csv'\n",
    "    print(s_out)\n",
    "    df_out.to_csv(f'{datadir}/{s_out}')\n",
    "    break\n",
    "topic_weights_5.drop('Patient',axis=1,inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-16T16:30:20.188616Z",
     "start_time": "2020-07-16T16:30:18.913716Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "7XjTO5PmQhM3",
    "outputId": "ccf6c4a0-35fe-4a65-8a7e-b84586048931"
   },
   "outputs": [],
   "source": [
    "# plot heatmap\n",
    "normalizer = lambda ar: (ar/np.sum(ar, axis=1, keepdims=True)).T\n",
    "\n",
    "fig, ax = plt.subplots(dpi=300)\n",
    "plot_topics_heatmap(lda_5.components_, spleen_cells_features.columns, ax, normalizer)\n",
    "plt.tight_layout()\n",
    "fig.savefig(f'{datadir}/{s_date}/heatmap_spatial_lda_k{num_topics}_{s_plat}_{s_subtype}.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scimap\n",
    "\n",
    "not used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# for s_plat in ['cycIF','IMC','MIBI',]:\n",
    "#     results_file = f'{datadir}/{s_plat}_spatial_lda.h5ad'\n",
    "#     if not os.path.exists(results_file):\n",
    "#         print(f'generating {s_plat}')\n",
    "#         ls_slide = sorted(df_cyc[df_cyc.Platform==s_plat].slide_scene.unique())\n",
    "#         ls_drop = df_cyc.columns[df_cyc[df_cyc.Platform==s_plat].isna().sum() > 0]\n",
    "#         df_plat = df_cyc[df_cyc.Platform==s_plat]\n",
    "#         print(len(ls_slide))\n",
    "#         #create adata\n",
    "#         ls_drop_splda = ['leiden', 'leidencelltype3', 'leidencelltype4','leidencelltype2',\n",
    "#                'gatedcelltype3', 'gatedcelltype5', 'leidencelltype5', 'slide','core','corec',\n",
    "#                'celltype', 'DAPI_X', 'DAPI_Y', 'slide_scene', 'subtype', 'Patient',\n",
    "#                'Platform','sample.Y', 'sample.X','x', 'y', 'isb', 'isimmune', 'cluster', 'sample.Z']\n",
    "#         adata = sc.AnnData(df_plat.drop(set(ls_drop).union(set(ls_drop_splda)),axis=1).astype('float64')) #\n",
    "#         adata.obs['X_centroid'] = df_plat.loc[adata.obs.index,'DAPI_X'].astype('float64')#*.325 already corrected\n",
    "#         adata.obs['Y_centroid'] = df_plat.loc[adata.obs.index,'DAPI_Y'].astype('float64')#*.325\n",
    "#         adata.obs['phenotype'] = df_plat.loc[adata.obs.index,'leiden'].replace({'CD4 T cell': 'CD3 T cell','CD8 T cell': 'CD3 T cell'})\n",
    "#         adata.obs['imageid'] = df_plat.loc[adata.obs.index,'slide_scene']\n",
    "#         adata.write_h5ad(results_file)\n",
    "#         sm.tl.spatial_aggregate(adata, x_coordinate='X_centroid', y_coordinate='Y_centroid', purity=50,\n",
    "#                               phenotype='phenotype', method='radius', radius=30, knn=10, imageid='imageid', subset=None, label='spatial_aggregate')\n",
    "#         pd.DataFrame(adata.obs['spatial_aggregate']).to_csv(f'{datadir}/{s_plat}_spatial_aggregate_scimap.csv')\n",
    "        \n",
    "#         break\n",
    "#         adata_lda = sm.tl.spatial_lda (adata[adata.obs.imageid.isin(ls_slide)], x_coordinate='X_centroid', y_coordinate='Y_centroid',\n",
    "#                                phenotype='phenotype', method='radius', radius=30, knn=10, \n",
    "#                                imageid='imageid', num_motifs=10, random_state=0, subset=None, label='spatial_lda')\n",
    "#         adata_lda.write_h5ad(results_file)\n",
    "#         fig, ax = plt.subplots(figsize=(8,8),dpi=200)\n",
    "#         sns.heatmap(adata_lda.uns['spatial_lda_probability'],cmap='RdBu_r',ax=ax)\n",
    "#         plt.tight_layout()\n",
    "#         fig.savefig(f'{datadir}/{s_date}/heatmap_{s_plat}_sm_spatial_lda_k{len(adata_lda.uns[\"spatial_lda_probability\"])}.png')\n",
    "#     else:\n",
    "#         print(f'loading {s_plat}')\n",
    "#         adata = sc.read_h5ad(results_file)\n",
    "#         sm.tl.spatial_aggregate(adata, x_coordinate='X_centroid', y_coordinate='Y_centroid', purity=50,\n",
    "#                               phenotype='phenotype', method='radius', radius=30, knn=10, imageid='imageid', subset=None, label='spatial_aggregate')\n",
    "#         pd.DataFrame(adata.obs['spatial_aggregate']).to_csv(f'{datadir}/{s_plat}_spatial_aggregate_scimap.csv')\n",
    "      \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df_plat.drop(set(ls_drop).union(set(ls_drop_splda)),axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for idx in [1,3,5,6,7,8]:\n",
    "#     fig, ax = plt.subplots(figsize=(4,4),dpi=200)\n",
    "#     plot_sm_topic(ax, s_slide, topic_weights=adata_lda.uns['spatial_lda'],patient_dfs=adata_lda.obs,topic_idx=f'Topic-{idx}')\n",
    "#     #break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# adata_lda = sm.tl.spatial_cluster(adata_lda, df_name='spatial_lda', method='kmeans', k=10,\n",
    "#                    n_pcs=None, resolution=1, phenograph_clustering_metric='euclidean',\n",
    "#                    nearest_neighbors=30, random_state=0, label='spatial_kmeans_lda', output_dir=None)\n",
    "\n",
    "# adata_lda.uns['spatial_lda'].columns = [item.replace('Motif_','Topic-') for item in adata_lda.uns['spatial_lda'].columns]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df = adata_lda.to_df()\n",
    "# df[adata_lda.obs.columns] = adata_lda.obs\n",
    "# df_plot = df.loc[:,['spatial_kmeans_lda','phenotype']].groupby('spatial_kmeans_lda').value_counts(normalize=True).unstack().fillna(0)\n",
    "# fig, ax = plt.subplots(figsize=(8,8),dpi=200)\n",
    "# plot_topics_heatmap(np.array(df_plot), df_plot.columns,ax,normalizer)\n",
    "# ax.set_xticklabels(df_plot.index)\n",
    "# plt.tight_layout()\n",
    "# fig.savefig(f'{datadir}/{s_date}/heatmap_spatial_kmeans_lda_k{len(df_plot)}.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# #mean topics per tissue\n",
    "# s_sample = 'JP-TMA1'\n",
    "# s_grouper='Patient'\n",
    "# s_cell= 'all'\n",
    "# for num_topics in [10]:\n",
    "#     topic_weights_5 = adata_lda.uns['spatial_lda']\n",
    "#     topic_weights_5['Patient'] = topic_weights_5.index.map(lambda x: x.split('_cell')[0]).map(d_patient)\n",
    "#     df_out = topic_weights_5.groupby(s_grouper).mean()\n",
    "#     s_out = f'results_{s_sample}_smSpatialLDA_by{s_grouper}_by{s_cell}_k{num_topics}.csv'\n",
    "#     print(s_out)\n",
    "#     df_out.to_csv(f'{datadir}/{s_out}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # mean topic clusters per tissue\n",
    "\n",
    "# adata_lda.obs['Patient'] = adata_lda.obs.imageid.map(d_patient)\n",
    "# k = adata_lda.obs['spatial_kmeans_lda'].nunique()\n",
    "# s_sample = 'JP-TMA1'\n",
    "# s_grouper='Patient'\n",
    "# s_cell= 'all'\n",
    "\n",
    "# df_out = (adata_lda.obs.groupby(['spatial_kmeans_lda','Patient']).count().X_centroid/adata_lda.obs.groupby(['Patient']).count().X_centroid).unstack().fillna(0).T\n",
    "# s_out = f'results_{s_sample}_KmeansLDA_by{s_grouper}_by{s_cell}_k{k}.csv'\n",
    "# print(s_out)\n",
    "# df_out.to_csv(f'{datadir}/{s_out}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## the plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for s_slide in ls_slide[0:10]:\n",
    "#     s_slide = 'JP-TMA1-1_scene034'\n",
    "#     plot_results(adata_lda[adata_lda.obs.imageid==s_slide], s_col='spatial_kmeans_lda', x='X_centroid',y='Y_centroid')\n",
    "#     break\n",
    "# fig, ax = plt.subplots()\n",
    "# plot_topics_heatmap(np.array(df_plot), df_plot.columns, ax=ax, normalizer=normalizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# num_topics = 8\n",
    "# print(num_topics)\n",
    "# lda_5 = spatial_lda_models[num_topics]\n",
    "# topic_weights_5 = lda_5.topic_weights\n",
    "\n",
    "# for s_slide in spleen_dfs.keys():\n",
    "#     fig, ax = plt.subplots(1,2,figsize=(8,4))\n",
    "#     ax=ax.ravel()\n",
    "#     plot_bcell_topic_multicolor(ax[0], sample_idx=s_slide, topic_weights=topic_weights_5, spleen_dfs=spleen_dfs)\n",
    "#     plot_topics_heatmap(lda_5.components_, spleen_cells_features.columns, ax[1], normalizer)\n",
    "#     plt.tight_layout()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# survival analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## find best cutpoint \n",
    "import lifelines\n",
    "from lifelines import KaplanMeierFitter, CoxPHFitter\n",
    "from lifelines.statistics import multivariate_logrank_test\n",
    "from lifelines import exceptions\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\",category = exceptions.ApproximationWarning)\n",
    "\n",
    "def single_km(df_all,s_cell,s_subtype,s_plat,s_col,savedir,alpha=0.05,cutp=0.5,s_time='Survival_time',s_censor='Survival'):\n",
    "    df_all.index = df_all.index.astype('str')\n",
    "    df = df_all[(df_all.Platform==s_plat) & (df_all.subtype==s_subtype)].copy()\n",
    "    df = df.loc[:,[s_col,s_time,s_censor]].dropna()\n",
    "    if len(df) > 0:\n",
    "        #KM\n",
    "        i_cut = np.quantile(df.loc[:,s_col],cutp)\n",
    "        b_low = df.loc[:,s_col] <= i_cut\n",
    "        s_title1 = f'{s_subtype} {s_plat}'\n",
    "        s_title2 = f'{s_col} in {s_cell} cells n={len(df)}'\n",
    "        if i_cut == 0:\n",
    "            b_low = df.loc[:,s_col] <= 0\n",
    "        df.loc[b_low,'abundance'] = 'low'\n",
    "        df.loc[~b_low,'abundance'] = 'high'\n",
    "        #log rank\n",
    "        results = multivariate_logrank_test(event_durations=df.loc[:,s_time],\n",
    "                                            groups=df.abundance, event_observed=df.loc[:,s_censor])\n",
    "        #kaplan meier plotting\n",
    "        if results.summary.p[0] < alpha:\n",
    "            kmf = KaplanMeierFitter()\n",
    "            fig, ax = plt.subplots(figsize=(3,3),dpi=300)\n",
    "            for s_group in ['high','low']:\n",
    "                df_abun = df[df.abundance==s_group]\n",
    "                durations = df_abun.loc[:,s_time]\n",
    "                event_observed = df_abun.loc[:,s_censor]\n",
    "                try:\n",
    "                    kmf.fit(durations, event_observed,label=s_group)\n",
    "                    kmf.plot(ax=ax,ci_show=False,show_censors=True)\n",
    "                except:\n",
    "                    results.summary.p[0] = 1\n",
    "            ax.set_title(f'{s_title1}\\n{s_title2}\\np={results.summary.p[0]:.2}',fontsize=10)\n",
    "            ax.set_xlabel(s_censor)\n",
    "            ax.legend(loc='upper right',title=f'{cutp}({i_cut:.2})')\n",
    "            plt.tight_layout()\n",
    "            fig.savefig(f\"{savedir}/KM_{s_title1.replace(' ','_')}_{s_title2.replace(' ','_')}_{cutp}_{s_censor}.png\",dpi=300)\n",
    "        return(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#example file names \n",
    "s_sample = 'JP-TMA1'#'20220413_JP-TMAs_IMC-TMAs_MIBI'#'20220411_JP-TMAs_IMC-TMAs_MIBI' # '20220409_JP-TMAs_IMC-TMAs'\n",
    "df_file = pd.DataFrame()\n",
    "for s_file in os.listdir(datadir):\n",
    "    if s_file.find(f'results_{s_sample}') > -1:\n",
    "        s_type = s_file.split('.csv')[0].split('_')[2]\n",
    "        s_subtype = s_file.split('.csv')[0].split('_')[-1]\n",
    "        s_partition = 'spatial'\n",
    "        s_cell = s_file.split('.csv')[0].split('_')[-2].split('by')[1]\n",
    "    else:\n",
    "        continue\n",
    "    df_file.loc[s_file,'subtype'] = s_subtype\n",
    "    df_file.loc[s_file,'type'] = s_type\n",
    "    df_file.loc[s_file,'partition'] = s_partition\n",
    "    df_file.loc[s_file,'cell'] = s_cell\n",
    "    #break\n",
    "#df_file.to_csv(f'{s_sample}_results_files.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "ls_file = [#'results_JP-TMA1_KmeansLDA_byPatient_byall_k10.csv',\n",
    "       #'results_JP-TMA1_smSpatialLDA_byPatient_byall_k10.csv',\n",
    "       'results_JP-TMA1_SpatialLDA_byPatient_byall_k6.csv',\n",
    "       #'results_JP-TMA1_SpatialLDA_byPatient_byall_k8.csv'\n",
    "          ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 0.05\n",
    "s_date = '20220808'\n",
    "df_surv.index = df_surv.index.astype('str')\n",
    "savedir = f'{datadir}/{s_date}'\n",
    "s_plat = 'cycIF'\n",
    "\n",
    "\n",
    "for s_index in ls_file: #df_file.index:\n",
    "    s_type = df_file.loc[s_index,'type']\n",
    "    s_cell = df_file.loc[s_index,'type']\n",
    "    df_all=pd.read_csv(f'{datadir}/{s_index}',index_col=0)\n",
    "    df_all.index = df_all.index.astype('str')\n",
    "    ls_marker = df_all.columns\n",
    "    df_all = df_all.merge(df_surv,left_index=True,right_index=True)  \n",
    "    for s_subtype in ['ER+','TNBC']: #,'HER2+'\n",
    "        #for s_plat in ['cycIF']:\n",
    "        for s_time, s_censor in [('Survival_time','Survival'),('Recurrence_time','Recurrence')]:\n",
    "            for s_col in ls_marker:\n",
    "                for cutp in [.5,.66,0.33]:\n",
    "                    single_km(df_all,s_cell,s_subtype,s_plat,s_col,savedir,alpha,cutp,s_time,s_censor)\n",
    "                    #break\n",
    "                #break\n",
    "            #break\n",
    "        #break\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scimap testing (skipped output results)\n",
    "\n",
    "USE:\n",
    "- spatial_count: radius from center cell followed by spatial cluster\n",
    "\n",
    "- spatial aggregate; very straight forward. how does getis-ord compare? \n",
    "\n",
    "FUTURE:\n",
    "\n",
    "- spatial_interaction: permutation test for significant proximity: slow\n",
    "\n",
    "- spatial_expression: product  of  the  expression  matrix  and  a  weighted  proximity  matrix NOT working well\n",
    "\n",
    "- spatial_pscore: cell cell proximity between cells types of interest\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "- spatial_similarity_search: could be useful in the future\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_lda = sm.tl.spatial_aggregate(adata_lda, x_coordinate='X_centroid', y_coordinate='Y_centroid', purity=50,\n",
    "                              phenotype='phenotype', method='radius', radius=30, knn=10, imageid='imageid', subset=None, label='spatial_aggregate')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_lda.obs['spatial_aggregate']= adata_lda.obs['spatial_aggregate'].fillna('non-significant')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for s_slide in ls_slide[0:10]:\n",
    "    fig, ax =plot_results(adata_lda[adata_lda.obs.imageid==s_slide], s_col='spatial_aggregate', x='X_centroid',y='Y_centroid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.tl.spatial_count(adata_lda, x_coordinate='X_centroid', y_coordinate='Y_centroid', \n",
    "              phenotype='phenotype', method='radius', radius=30, knn=10, imageid='imageid', subset=None, label='spatial_count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.tl.spatial_cluster(adata_lda, df_name='spatial_count', method='kmeans', k=10, n_pcs=None, \n",
    "                resolution=1, phenograph_clustering_metric='euclidean', nearest_neighbors=30, random_state=0, label='spatial_count_kmeans', output_dir=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for s_slide in ls_slide[0:10]:\n",
    "    fig, ax =plot_results(adata_lda[adata_lda.obs.imageid==s_slide], s_col='spatial_count_kmeans', x='X_centroid',y='Y_centroid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.tl.spatial_distance(adata_lda, x_coordinate='X_centroid', y_coordinate='Y_centroid',\n",
    "                 z_coordinate=None, phenotype='phenotype', subset=None, imageid='imageid', label='spatial_distance')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_lda.obs[adata_lda.uns['spatial_distance'].columns] = adata_lda.uns['spatial_distance']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_col='CD8 T cell'\n",
    "fig, ax = plot_results(adata_lda, s_col='CD8 T cell',  x='X_centroid',y='Y_centroid')\n",
    "ax.set_title(f'distance to {s_col}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#something now working with clustering?\n",
    "adata_lda.raw = adata_lda\n",
    "sm.tl.spatial_expression(adata_lda, x_coordinate='X_centroid', y_coordinate='Y_centroid', method='radius', \n",
    "                   radius=30, knn=10, imageid='imageid', use_raw=True, log=True, subset=s_slide, label='spatial_expression', output_dir=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.tl.spatial_cluster(adata_lda, df_name='spatial_expression', method='kmeans', k=10, n_pcs=None, resolution=1,\n",
    "                phenograph_clustering_metric='euclidean', nearest_neighbors=30, random_state=0, label='spatial_kmeans_exp', output_dir=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for s_slide in ls_slide[0:10]:\n",
    "    fig, ax =plot_results(adata_lda[adata_lda.obs.imageid==s_slide], s_col='spatial_kmeans_exp', x='X_centroid',y='Y_centroid')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_col = 'spatial_kmeans_exp'\n",
    "df = adata_lda.to_df()\n",
    "df[adata_lda.obs.columns] = adata_lda.obs\n",
    "df_plot = df.loc[:,[s_col]+adata_lda.var.index.tolist()].groupby(s_col).mean().fillna(0)\n",
    "fig, ax = plt.subplots(figsize=(12,12))\n",
    "plot_topics_heatmap(np.array(df_plot), df_plot.columns,ax,normalizer)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.pp.rescale(adata_lda)\n",
    "sm.tl.spatial_expression(adata_lda, x_coordinate='X_centroid', y_coordinate='Y_centroid', method='radius', \n",
    "                   radius=30, knn=10, imageid='imageid', use_raw=False, log=True, subset=s_slide, label='spatial_expression', output_dir=None)\n",
    "sm.tl.spatial_cluster(adata_lda, df_name='spatial_expression', method='kmeans', k=10, n_pcs=None, resolution=1,\n",
    "                phenograph_clustering_metric='euclidean', nearest_neighbors=30, random_state=0, label='spatial_kmeans_rexp', output_dir=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_col = 'spatial_kmeans_exp'\n",
    "df = adata_lda.to_df()\n",
    "df[adata_lda.obs.columns] = adata_lda.obs\n",
    "df_plot = df.loc[:,[s_col]+adata_lda.var.index.tolist()].groupby(s_col).mean().fillna(0)\n",
    "fig, ax = plt.subplots(figsize=(12,12))\n",
    "plot_topics_heatmap(np.array(df_plot), df_plot.columns,ax,normalizer)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for s_slide in ls_slide[0:10]:\n",
    "    fig, ax =plot_results(adata_lda[adata_lda.obs.imageid==s_slide], s_col='spatial_kmeans_rexp', x='X_centroid',y='Y_centroid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#takes forever, but interesting\n",
    "sm.tl.spatial_interaction(adata_lda, x_coordinate='X_centroid', y_coordinate='Y_centroid', z_coordinate=None, phenotype='phenotype',\n",
    "                    method='radius', radius=30, knn=10, permutation=1000, imageid='imageid', subset=None, pval_method='zscore', label='spatial_interaction')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_data = sm.pl.spatial_interaction(adata_lda, spatial_interaction='spatial_interaction', summarize_plot=True, p_val=0.05, row_cluster=True, col_cluster=True, cmap='vlag', \n",
    "                    nonsig_color='grey', subset_phenotype=None, subset_neighbour_phenotype=None, binary_view=False, return_data=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.tl.spatial_pscore(adata, proximity=['Prolif. t.','CD8 T cell'], score_by='imageid', x_coordinate='X_centroid', y_coordinate='Y_centroid',\n",
    "               phenotype='phenotype', method='radius', radius=20, knn=3, imageid='imageid', subset=None, label='spatial_pscore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.uns['spatial_pscore']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.pl.spatial_pscore(adata, label='spatial_pscore', plot_score='Proximity Density', order_xaxis=None, color='grey')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#save it!\n",
    "results_file = f'scimap_test_{s_slide}.h5ad'\n",
    "if not os.path.exists(results_file):\n",
    "    adata_lda.write(results_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_file = f'scimap_test_{s_slide}.h5ad'\n",
    "if os.path.exists(results_file):\n",
    "                print('loading data')\n",
    "                adata = sc.read_h5ad(results_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "CODEX_Spatial_LDA.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "scimap",
   "language": "python",
   "name": "scimap"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}