figure5_ibd.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8256ad61-b10b-45e8-af56-822ef432c141",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0,'/home/cane/Documents/yoseflab/can/resolVI')\n",
    "from scvi.external import RESOLVI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9321c37a-5693-4030-a8ae-51c30b644291",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scanpy as sc\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scvi\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e1fb0e4-2c21-4cb2-9b3c-ec33649558eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "import pyro.poutine as poutine\n",
    "import torch\n",
    "from pyro.infer import Importance, EmpiricalMarginal, Trace_ELBO, SVI\n",
    "from pyro.infer.autoguide import AutoDiagonalNormal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d20721-911f-40d0-9afd-42a9085c6cad",
   "metadata": {},
   "outputs": [],
   "source": [
    "scvi.settings.seed = 0\n",
    "sc.set_figure_params(dpi=100, dpi_save=300, format='png', frameon=False, vector_friendly=True, fontsize=14, color_map='viridis', figsize=None)\n",
    "sc.settings.figdir = 'figure4_new/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee03098e-9ac9-4ae1-abe3-7d37e2e0afd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append('..')\n",
    "import _utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9dc6db-7fe7-46c6-88a1-b69ae09d2a24",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Cleanup data and initialize supervised dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692cb912-d721-44e8-b38e-8e8bf338b8a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvae = RESOLVI.load('/external_data/other/resolvi_final_other_files/mouse_colitis/original_q25_semisupervised_tier3_d0-d21/resolvae')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45892230-9712-49c4-aadb-d03e009befac",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvae.adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0bfb031-038b-4691-a4a1-f38d2f55bed5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a4f23ae-54db-45a4-be3f-eb3a18fe688d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "adata = resolvae.adata\n",
    "adata.obsm[\"X_resolVI\"] = resolvae.get_latent_representation()\n",
    "sc.pp.neighbors(adata, use_rep='X_resolVI', n_neighbors=15, method='rapids')\n",
    "sc.tl.umap(adata, min_dist=0.3, init_pos='spectral')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92453f19-ba15-4905-9dbe-51b72f5dda9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['total_counts'] = adata.layers['raw_counts'].sum(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11a2da8e-5e8b-423c-bd6e-f6c6d603da06",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['timepoint'] = [i.split('_')[1] for i in adata.obs['Mouse_ID']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35c00177-d165-440a-b4d6-355c13c12739",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(adata, color=[\"Tier1\", \"Tier3\", \"timepoint\"], save='unsupervised_d0_d21.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84154555-2a5c-417d-aa49-a3716a2df70d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "adata.obsm['X_tsne'] = adata.obs[['Tier1_umap_x', 'Tier1_umap_y']].values\n",
    "sc.pl.tsne(adata, color=[\"Tier1\", \"Tier3\", \"timepoint\"], save='original_d0_d21.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b97a864-e111-4f3e-b723-2a86f1c33f04",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_corr = resolvae.sample_posterior_predictive(\n",
    "    model=resolvae.module.model_corrected,\n",
    "    return_sites=['obs'],\n",
    "    num_samples=5, return_samples=False, batch_size=1000, batch_steps=10, summary_fun={\"q50\": np.median})\n",
    "samples_corr = pd.DataFrame(samples_corr).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bee26093-6ac8-4123-a4ab-abf4bc984757",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.layers['generated_expression'] = samples_corr.loc['post_sample_q50', 'obs']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38c9cb2c-d9f3-487e-8470-197ee65c7722",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "samples = resolvae.sample_posterior_predictive(\n",
    "    model=resolvae.module.model_residuals,\n",
    "    return_sites=[\n",
    "        'mixture_proportions', 'per_gene_background', \n",
    "        'diffusion_mixture_proportion', 'per_neighbor_diffusion', 'px_r_inv'\n",
    "        ],\n",
    "    num_samples=5, return_samples=False, batch_size=1000, batch_steps=10)\n",
    "samples = pd.DataFrame(samples).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd05e795-3176-45c7-b62c-fd70dc624ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['true_proportion'] = samples.loc['post_sample_means', 'mixture_proportions'][:, 0]\n",
    "adata.obs['diffusion_proportion'] = samples.loc['post_sample_means', 'mixture_proportions'][:, 1]\n",
    "adata.obs['background_proportion'] = samples.loc['post_sample_means', 'mixture_proportions'][:, 2]\n",
    "adata.varm['background'] = pd.DataFrame(samples.loc['post_sample_means', 'per_gene_background'][0, ...].squeeze().T, index=adata.var_names)\n",
    "adata.var['px_r'] = 1/(1e-6 + samples.loc['post_sample_means', 'px_r_inv'][0, :])\n",
    "\n",
    "_ = plt.hist(adata.obs['background_proportion'], bins=30, range=(0,1))\n",
    "_ = plt.hist(adata.obs['true_proportion'], bins=30, range=(0,1))\n",
    "_ = plt.hist(adata.obs['diffusion_proportion'], bins=30, range=(0,1))\n",
    "plt.legend(['Background_Proportion', 'True_proportions', 'Diffusion_Proportion'])\n",
    "plt.savefig(f'figure4_new/histogram_proportions.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ace0d696-0377-4ca6-b593-787f5c6d8764",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['true_counts'] = adata.obs['true_proportion'] * adata.obs['total_counts']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99aec695-c390-4d32-95ec-08179abbcc27",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(adata[adata.obs['true_counts']<15], color=['true_counts', 'total_counts', 'background_proportion', 'Tier3'], vmax='p95')\n",
    "sc.pl.umap(adata[adata.obs['true_counts']>15], color=['true_counts', 'total_counts', 'background_proportion', 'Tier3'], vmax='p95')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bfa7785-3d22-4443-9148-e507fd6b4cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(adata, groupby='Tier1')\n",
    "sc.pl.rank_genes_groups_dotplot(adata, standard_scale='var')\n",
    "sc.pl.rank_genes_groups_dotplot(adata[adata.obs['true_counts']>15], standard_scale='var', save='_marker_genes_hq_coarse.pdf')\n",
    "sc.pl.rank_genes_groups_dotplot(adata[adata.obs['true_counts']<15], standard_scale='var', save='_marker_genes_lq_coarse.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15d33038-56e2-441d-ac82-2db43318024a",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['fov_batch'] = adata.obs['Slice_ID'].astype(str) + '__' + adata.obs['FOV'].astype(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b72ec78-69f4-4a32-8bb5-56f550f8d5be",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['low_quality'] = 0\n",
    "adata.obs.loc[np.logical_and(adata.obs['background_proportion']>0.5, adata.obs['true_counts']<20), 'low_quality'] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faf487dd-ea0a-4e83-952b-90698a49c7f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = adata.obs.groupby('Slice_ID')['low_quality'].sum()\n",
    "b = pd.DataFrame(b)\n",
    "b.rename_axis(columns={'low_quality': 'counts'})\n",
    "b['mean'] = adata.obs.groupby('Slice_ID')['low_quality'].mean()\n",
    "list(b.sort_values(by='mean', ascending=False).head(10).index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "139eab41-06ea-4a8a-9a43-39f153c73f61",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sc.read_h5ad('mouse_colitis/cleaned_up_labels_colitis_d0-d21.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b80c4f73-a26b-44d6-9712-cad1c2c96de1",
   "metadata": {},
   "outputs": [],
   "source": [
    "slice = '082421_D21_m2_1_slice_1'\n",
    "sc.pl.spatial(adata[(adata.obs['Slice_ID']==slice) & (adata.obs['low_quality']==1)], spot_size=20, color='Tier1', groups=('EntericNervous', 'Smooth Muscle Cells', 'Epithelial'),\n",
    "             save='low_quality_spatial_cells.pdf')\n",
    "sc.pl.spatial(adata[(adata.obs['Slice_ID']==slice) & (adata.obs['low_quality']==0)], spot_size=10, color='Tier1', groups=('EntericNervous', 'Smooth Muscle Cells', 'Epithelial'),\n",
    "             save='high_quality_spatial_cells.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa6eeb1-4786-44f0-86f6-42b0167d6c61",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for slice in ['100221_D9_m5_2_slice_1', '062921_D9_m5_2_slice_3', '082421_D21_m2_1_slice_1']:\n",
    "    print(slice)\n",
    "    sc.pl.spatial(adata[(adata.obs['Slice_ID']==slice) & (adata.obs['low_quality']==1)], spot_size=20, color='Tier1', groups=('EntericNervous', 'Smooth Muscle Cells', 'Epithelial'))\n",
    "    sc.pl.spatial(adata[(adata.obs['Slice_ID']==slice) & (adata.obs['low_quality']==0)], spot_size=10, color='Tier1', groups=('EntericNervous', 'Smooth Muscle Cells', 'Epithelial'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcce7ab1-d5b7-4d7f-b7a5-86e83c023a8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "df_fov = adata.obs.groupby('fov_batch')['low_quality'].agg(['sum', 'mean'])\n",
    "\n",
    "# Extract unique timepoints; ensure it results in a single value per group\n",
    "df_fov['timepoint'] = adata.obs.groupby('fov_batch')['timepoint'].agg(lambda x: x.unique()[0])\n",
    "\n",
    "# Plotting\n",
    "# Prepare colors - map each unique timepoint to a color\n",
    "unique_timepoints = df_fov['timepoint'].unique()\n",
    "color_map = {tp: plt.cm.tab10(i) for i, tp in enumerate(unique_timepoints)}\n",
    "colors = df_fov['timepoint'].map(color_map)\n",
    "\n",
    "# Create scatter plot\n",
    "plt.scatter(df_fov['sum'], df_fov['mean'], s=3, c=colors, vmax=10)  # vmax used here if needed for normalization\n",
    "\n",
    "# Adding a colorbar with timepoint labels requires more complex handling,\n",
    "# because plt.scatter does not support automatic colorbar for categorical data.\n",
    "# We will create a custom legend:\n",
    "handles = [plt.Line2D([0], [0], marker='o', color=color_map[tp], linestyle='', label=tp) for tp in unique_timepoints]\n",
    "plt.legend(handles=handles, title='Timepoint', loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "\n",
    "plt.xlabel('Sum of Low Quality')\n",
    "plt.ylabel('Mean of Low Quality')\n",
    "plt.title('Scatter Plot of Low Quality by FOV Batch')\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c11d45c5-8efb-4742-8112-8912673429d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub = adata[adata.obs['low_quality']==0].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a287d107-324f-4cab-b2e0-28b27d19384c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pp.neighbors(sub, use_rep='X_resolVI', n_neighbors=100, method='rapids')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38c433db-e5a8-4674-9273-3eeed9fc08d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84a1c2b2-c367-4210-b9c0-029365f471c6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cell_types = sub.obs['Tier1'].cat.codes.values\n",
    "most_common_type = []\n",
    "connectivities = sub.obsp['connectivities']\n",
    "weighted_ratios = []\n",
    "\n",
    "for i in tqdm(range(connectivities.shape[0])):\n",
    "    neighbors = connectivities[i].nonzero()[1]\n",
    "    neighbor_types = cell_types[neighbors]\n",
    "    neighbor_weights = connectivities[i, neighbors].toarray().flatten()\n",
    "    unique_types, counts = np.unique(neighbor_types, return_counts=True)\n",
    "    weighted_counts = np.zeros(len(sub.obs['Tier1'].cat.categories))\n",
    "\n",
    "    for j, type_index in enumerate(unique_types):\n",
    "        weighted_counts[type_index] = neighbor_weights[neighbor_types == type_index].sum()\n",
    "\n",
    "    most_common = weighted_counts.argmax()\n",
    "    most_common_type.append(most_common)\n",
    "    total_weight = neighbor_weights.sum()\n",
    "    if total_weight > 0:\n",
    "        weighted_ratios.append(weighted_counts[most_common] / total_weight)\n",
    "    else:\n",
    "        weighted_ratios.append(0)\n",
    "\n",
    "sub.obs['smoothed_tier1'] = sub.obs['Tier1'].cat.categories[most_common_type]\n",
    "sub.obs['uncertainty_tier1'] = 1 - np.array(weighted_ratios)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8769c26-5ffc-4bdf-9def-da9bb209814f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(sub, color=['smoothed_tier1', 'uncertainty_tier1'], vmin=0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef23f0d-3a9a-42de-a310-b80afeeab34b",
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_types = sub.obs['Tier3'].cat.codes.values\n",
    "most_common_type = []\n",
    "connectivities = sub.obsp['connectivities']\n",
    "weighted_ratios = []\n",
    "\n",
    "for i in tqdm(range(connectivities.shape[0])):\n",
    "    neighbors = connectivities[i].nonzero()[1]\n",
    "    neighbor_types = cell_types[neighbors]\n",
    "    neighbor_weights = connectivities[i, neighbors].toarray().flatten()\n",
    "    unique_types, counts = np.unique(neighbor_types, return_counts=True)\n",
    "    weighted_counts = np.zeros(len(sub.obs['Tier3'].cat.categories))\n",
    "\n",
    "    for j, type_index in enumerate(unique_types):\n",
    "        weighted_counts[type_index] = neighbor_weights[neighbor_types == type_index].sum()\n",
    "\n",
    "    most_common = weighted_counts.argmax()\n",
    "    most_common_type.append(most_common)\n",
    "    total_weight = neighbor_weights.sum()\n",
    "    if total_weight > 0:\n",
    "        weighted_ratios.append(weighted_counts[most_common] / total_weight)\n",
    "    else:\n",
    "        weighted_ratios.append(0)\n",
    "\n",
    "sub.obs['smoothed_tier3'] = sub.obs['Tier3'].cat.categories[most_common_type]\n",
    "sub.obs['uncertainty_tier3'] = 1 - np.array(weighted_ratios)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85046b64-6047-4bc8-8dd6-e93bfe9b2907",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(sub, color=['smoothed_tier3', 'uncertainty_tier3'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e06dcc66-22b2-4b07-9f59-573d4fd929e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fine_coarse_dictionary = sub.obs[['Tier3', 'Tier1']].drop_duplicates().set_index('Tier3').to_dict()['Tier1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1806b44-394a-4490-9b89-f98a46627533",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(sub, groupby='Tier1')\n",
    "sc.pl.rank_genes_groups_dotplot(sub, n_genes=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df0a12d5-c28a-4f57-8b3c-17a742123024",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['redo_celltyping'] = 'unassigned'\n",
    "adata.obs.loc[sub[(sub.obs['uncertainty_tier1']<0.05) & (sub.obs['uncertainty_tier3']<0.4)].obs_names, 'redo_celltyping'] = sub.obs.loc[\n",
    "    (sub.obs['uncertainty_tier1']<0.05) & (sub.obs['uncertainty_tier3']<0.4), 'smoothed_tier3']\n",
    "adata.obs['redo_celltyping_coarse'] = 'unassigned'\n",
    "adata.obs.loc[sub[(sub.obs['uncertainty_tier1']<0.05)].obs_names, 'redo_celltyping_coarse'] = sub.obs.loc[\n",
    "    sub.obs['uncertainty_tier1']<0.05, 'smoothed_tier1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd86b744-f16e-4bed-8d98-421f9b756822",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_cell_types = list(sub.obs['Tier3'].unique())\n",
    "for i in range(0, len(unique_cell_types), 3):\n",
    "    groups = unique_cell_types[i:i+3]\n",
    "    print(groups)\n",
    "    sc.pl.umap(adata, color=['redo_celltyping', 'Tier3'], groups=groups, size=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8ac5656-d886-490c-be53-def644f0f1b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['undone_celltyping_coarse'] = adata.obs['Tier1'].astype(str)\n",
    "adata.obs.loc[sub[sub.obs['uncertainty_tier1']<0.1].obs_names, 'undone_celltyping_coarse'] = 'assigned'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f1f9fb7-cdee-4cef-836e-c947bb27c4b8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pl.umap(adata, color=['redo_celltyping', 'redo_celltyping_coarse'], ncols=1, save='redid_celltyping_unsupervised.pdf')\n",
    "sc.pl.umap(adata, color=['redo_celltyping', 'redo_celltyping_coarse'], ncols=1, groups=['unassigned'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b47f490b-8f3e-4a47-bb29-0622e1014e14",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.write_h5ad('mouse_colitis/cleaned_up_labels_colitis_d0-d21.h5ad')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cacfb113-f1a6-413f-a53c-fe4a8b447bd8",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Transfer learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98790246-fae9-4125-b85b-dd9435bf0002",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse import csr_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83080c6a-9805-4887-b662-76a788a4d7d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_data = sc.read_h5ad('/external_data/other/resolvi_final_other_files/ibd_moffitt/adata_day35.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3a6111c-f458-4662-82b1-f24bce1e6938",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_data.layers['raw_counts'] = csr_matrix(pd.read_csv('/external_data/other/resolvi_final_other_files/ibd_moffitt/X_raw_day35.csv', header=None).values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3be6fbfe-5c04-4f60-8865-be81dcf497b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_data.obs['timepoint'] = 'D35'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f567977c-704a-484a-838c-cc9c432cb7c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_data.obs_names = 'query_' + query_data.obs_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b48de5f7-5c71-402a-bc67-058fe12ab308",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query_data.obsm['spatial'] = query_data.obs[['x', 'y']].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6be27005-7847-41a7-b986-b808cd986510",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "resolvae = RESOLVI.load('mouse_colitis/original_q25_semisupervised_redo3/resolvae')"
   ]
  },
  {
   "cell_type": "raw",
   "id": "ebc40914-5efb-4a85-9fd7-86d1b8e990d7",
   "metadata": {},
   "source": [
    "resolvae_new = resolvae.load_query_data(adata, reference_model=resolvae)\n",
    "resolvae_new.train(max_epochs=50, weight_decay=0., batch_size=1024, n_epochs_kl_warmup=0)\n",
    "resolvae_new.save('mouse_colitis/original_q25_semisupervised_redo3/resolvae_transfer', save_anndata=True, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93c67637-cc2c-4f54-8f71-9eeb46ba6b44",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvae.prepare_query_anndata(query_data, resolvae)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b02f5001-1199-49d3-b8dc-276adb502896",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "retrain = False\n",
    "if retrain:\n",
    "    resolvae_query = resolvae.load_query_data(query_data, reference_model=resolvae)\n",
    "    resolvae_query.train(max_epochs=100, weight_decay=0., batch_size=1024, n_epochs_kl_warmup=0)\n",
    "    resolvae_query.save('mouse_colitis/original_q25_semisupervised_redo3/resolvae_query', save_anndata=True, overwrite=True)\n",
    "else:\n",
    "    resolvae_query = RESOLVI.load('mouse_colitis/original_q25_semisupervised_redo3/resolvae_query')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f7b17bc-d8a1-4042-9a47-bd8b04c25b0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(resolvae_query.history['elbo_train'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ddf9d4-368e-438d-8b42-930024497223",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sc.read('mouse_colitis/cleaned_up_labels_colitis_d0-d21.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8d9179-4667-4ad4-b97d-c8713eab495a",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data = adata.concatenate(query_data, join='outer', batch_key='source', batch_categories=['reference', 'query'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59e0dc82-c628-4654-850e-276aff4cafaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.obs['cluster'] = all_data.obs['cluster'].fillna(\"unassigned\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c0641b0-300c-4b4e-98c8-64963cd0baca",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data = all_data[all_data.obs['cluster']!='Mast cell']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4aaca630-94ad-43c0-9a2a-d306e7aac98b",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.obs_names = [i.split('-')[0] for i in all_data.obs_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f2a6ab5-e1f4-474d-93c6-c3a72d1704dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvae_both = resolvae_query.load_query_data(all_data, reference_model=resolvae_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2fd174b-f5d4-4ad5-9e82-58098b96870a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_data.obsm['celltype_predicted'] = resolvae_both.predict(soft=True, num_samples=5, batch_size=3000)\n",
    "all_data.obs['predicted_celltype'] = all_data.obsm['celltype_predicted'].idxmax(axis=1)\n",
    "all_data.obs['predicted_celltype_prob'] = all_data.obsm['celltype_predicted'].max(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1032ec06-8971-4f64-a464-e26bf53c04ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "fine_coarse_dictionary = adata.obs[['Tier3', 'Tier1']].drop_duplicates().set_index('Tier3').to_dict()['Tier1']\n",
    "all_data.obs['predicted_celltype_coarse'] = all_data.obs['predicted_celltype'].map(fine_coarse_dictionary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1081720a-b0fa-4175-a78f-bf66582b8442",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "all_data.obsm[\"X_resolVI\"] = resolvae_both.get_latent_representation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd599538-b28a-4723-beed-46c5a18bcadf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pp.neighbors(all_data, use_rep='X_resolVI', method='rapids')\n",
    "sc.tl.umap(all_data, min_dist=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b40081-1935-4464-afd3-1e5d4563a55b",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data = sc.read_h5ad('/external_data/other/resolvi_final_other_files/figure4/query_trained_data_d0_35.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e568d96-2e5a-49f7-aac1-d2e07667a547",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.obs.loc[all_data.obs['redo_celltyping'] == 'unassigned', 'redo_celltyping'] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8d5f2a5-7e48-426f-a54c-8da0acf8427a",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(all_data[all_data.obs['timepoint'].isin(['D0', 'D3', 'D21', 'D9'])],\n",
    "           color=[\"predicted_celltype\", \"predicted_celltype_prob\", \"redo_celltyping\"], ncols=1, save='semisupervised_celltypes_predicted.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f57516ef-920b-4492-b25f-ef004b0e173d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fibroblast = sub[(sub.obs['Tier1']=='Fibroblast') & (sub.obs['predicted_celltype_coarse']=='Fibroblast')].copy()\n",
    "sc.tl.rank_genes_groups(fibroblast, groupby='Tier3')\n",
    "sc.pl.rank_genes_groups_dotplot(fibroblast, standard_scale='var', dendrogram=True, n_genes=5)\n",
    "sc.pl.rank_genes_groups_dotplot(fibroblast[fibroblast.obs['redo_celltyping']!='unassigned'], standard_scale='var', dendrogram=True, n_genes=5, save='marker_genes_seed_fine.pdf')\n",
    "sc.pl.rank_genes_groups_dotplot(fibroblast[fibroblast.obs['redo_celltyping']=='unassigned'], standard_scale='var', dendrogram=True, n_genes=5, save='marker_genes_other_fine.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a224552c-c812-41a6-a573-6a19ff162b19",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "redo_data = fibroblast[fibroblast.obs['predicted_celltype']!=fibroblast.obs['Tier3']].copy()\n",
    "sc.pl.rank_genes_groups_dotplot(redo_data, standard_scale='var', dendrogram=True, n_genes=5, save='marker_genes_disagreement_original.pdf')\n",
    "redo_data.obs['Tier3'] = redo_data.obs['predicted_celltype']\n",
    "sc.pl.rank_genes_groups_dotplot(redo_data, standard_scale='var', dendrogram=True, n_genes=5, save='marker_genes_disagreement_resolvi.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e665b7d-7da0-42f4-8bde-ce6e07e3cad7",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.write_h5ad('/external_data/other/resolvi_final_other_files/figure4/query_trained_data_d0_35.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a164d6a3-b906-4630-8228-5fcf5ff5e82b",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data = sc.read_h5ad('/external_data/other/resolvi_final_other_files/figure4/query_trained_data_d0_35.h5ad')\n",
    "resolvae_both = RESOLVI.load_query_data(all_data, reference_model='mouse_colitis/original_q25_semisupervised_redo3/resolvae_query')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1fcc2fb-e5c9-4777-9f7f-6eacdac493b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_corr = resolvae_both.sample_posterior_predictive(\n",
    "    model=resolvae_both.module.model_corrected,\n",
    "    return_sites=['px_rate', 'obs'],\n",
    "    summary_fun={\"post_sample_means\": np.mean, \"post_sample_q50\": np.median},\n",
    "    num_samples=10, return_samples=False, batch_size=2000, batch_steps=10)\n",
    "samples_corr = pd.DataFrame(samples_corr).T\n",
    "\n",
    "samples = resolvae_both.sample_posterior_predictive(\n",
    "    model=resolvae_both.module.model_residuals,\n",
    "    return_sites=[\n",
    "        'mixture_proportions', 'mean_poisson', 'per_gene_background', \n",
    "        'diffusion_mixture_proportion', 'per_neighbor_diffusion', 'px_r_inv'\n",
    "        ],\n",
    "    summary_fun={\"post_sample_means\": np.mean, \"post_sample_q50\": np.median},\n",
    "    num_samples=10, return_samples=False, batch_size=2000, batch_steps=10)\n",
    "samples = pd.DataFrame(samples).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22cfadd1-4486-4f52-913b-3b93627e0b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse import csr_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a64e3e32-1df4-48e6-aec2-2671da246a0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.layers['generated_expression'] = csr_matrix(samples_corr.loc['post_sample_post_sample_q50', 'obs'])\n",
    "all_data.layers['imputed_expression'] = csr_matrix(samples_corr.loc['post_sample_post_sample_q50', 'px_rate'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78391995-0953-4799-86eb-8bae1b956f95",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.obs['true_proportions'] = samples.loc['post_sample_post_sample_q50', 'mixture_proportions'].T[0, :]\n",
    "all_data.obs['diffusion_proportions'] = samples.loc['post_sample_post_sample_q50', 'mixture_proportions'].T[1, :]\n",
    "all_data.obs['background_proportions'] = samples.loc['post_sample_post_sample_q50', 'mixture_proportions'].T[2, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "957e512c-42d7-4844-b931-7d26fc034995",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.obs['total_counts'] = np.array(all_data.layers['raw_counts'].sum(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "152c3c95-d780-4fd1-9571-e08a8aa18fae",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.obs['true_counts'] = all_data.obs['true_proportions'] * all_data.obs['total_counts']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4943f0f0-5798-4e2a-af57-f06e7ba4f2aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(all_data[all_data.obs['predicted_celltype_prob']>0.6], color=[\"Tier3\", \"predicted_celltype\", \"predicted_celltype_prob\", \"true_counts\"])\n",
    "sc.pl.umap(all_data[all_data.obs['predicted_celltype_prob']<0.6], color=[\"Tier3\", \"predicted_celltype\", \"predicted_celltype_prob\", \"true_counts\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88fb144e-20c9-438f-b14f-4f0d3e25ba05",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.write_h5ad('/external_data/other/resolvi_final_other_files/figure4/query_trained_data_d0_35.h5ad')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdcad9ac-16ca-4495-ad15-20ca676bcf54",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Downstream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3774799b-a543-43ac-a302-e2bba7dacb10",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data = sc.read_h5ad('/external_data/other/resolvi_final_other_files/figure4/query_trained_data_d0_35.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebedd3ec-5955-42f7-b808-b47d60a2e6db",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data.X = all_data.layers['generated_expression'].copy()\n",
    "sc.pp.normalize_total(all_data)\n",
    "sc.pp.log1p(all_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cdcf86f-91dc-4e2e-868d-d099abe94748",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub = all_data[all_data.obs['true_counts']>20].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61681259-c21f-4a5d-b35c-18bf34f454f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.layers['counts'] = sub.layers['raw_counts'].copy()\n",
    "sc.pp.normalize_total(sub, layer='counts')\n",
    "sc.pp.log1p(sub, layer='counts')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e65af0-e1b9-4a2f-b36d-9a823e41e9b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvae_both = RESOLVI.load_query_data(sub, reference_model='mouse_colitis/original_q25_semisupervised_redo3/resolvae_query')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae7d1622-c62e-492e-9ad9-d2633c39a8a9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sub.obs['timepoint'] = [i.split('_')[1] for i in sub.obs['Mouse_ID']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55941aaa-1e50-4cd2-87c6-11ea108e225e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in ['D0', 'D35', 'D3', 'D9', 'D21']:\n",
    "    sc.pl.umap(sub[\n",
    "        (sub.obs['timepoint']==i) & (sub.obs['predicted_celltype_prob']>0.4)\n",
    "    ], color=['predicted_celltype_prob', 'Tier3', 'predicted_celltype'], vmax='p99', ncols=3, size=2, title=i)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff992965-72cc-47ab-8683-51858a3c5a82",
   "metadata": {},
   "outputs": [],
   "source": [
    "fibroblast = sub[(sub.obs['Tier1']=='Fibroblast') & (sub.obs['predicted_celltype_coarse']=='Fibroblast')].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f251bdff-75e8-4d05-8eb9-68cba67f4f83",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(fibroblast, groupby='predicted_celltype')\n",
    "sc.pl.rank_genes_groups_dotplot(fibroblast, standard_scale='var', dendrogram=True, n_genes=5)\n",
    "sc.pl.rank_genes_groups_dotplot(fibroblast, standard_scale='var', dendrogram=True, n_genes=5, layer='generated_expression')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1da423f2-d18e-44d4-b959-3d6d1147c80b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(fibroblast, groupby='timepoint', layer='generated_expression')\n",
    "sc.pl.rank_genes_groups_dotplot(fibroblast, standard_scale='var', dendrogram=True, n_genes=10, layer='generated_expression')\n",
    "sc.pl.rank_genes_groups_dotplot(fibroblast, standard_scale='var', dendrogram=True, n_genes=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80da1412-2faf-4ab5-b80c-297ca71529e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "epithelial = sub[(sub.obs['Tier1']=='Epithelial') & (sub.obs['predicted_celltype_coarse']=='Epithelial')].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27b34795-6fdc-4272-9865-d844e957884b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(epithelial, groupby='predicted_celltype')\n",
    "sc.pl.rank_genes_groups_dotplot(epithelial, standard_scale='var', dendrogram=True, n_genes=5)\n",
    "sc.pl.rank_genes_groups_dotplot(epithelial, standard_scale='var', dendrogram=True, n_genes=5, layer='generated_expression')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d2f8b9-4be8-49dd-8ebd-02798fef56b7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for slice in epithelial.obs['Slice_ID'].unique():\n",
    "    sc.pl.spatial(epithelial[epithelial.obs['Slice_ID']==slice], spot_size=15, color=['predicted_celltype', 'Tier3'], title=slice, ncols=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19b243a1-6c40-4003-b49d-f5db70a0c3ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(fibroblast[fibroblast.obs['Slice_ID']=='062921_D0_m3a_1_slice_2'], spot_size=15, color=['predicted_celltype', 'Tier3'], title='062921_D0_m3a_1_slice_2')\n",
    "sc.pl.spatial(fibroblast[fibroblast.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=15, color=['predicted_celltype', 'Tier3'], title='072523_D35_m6_1_slice_3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59142bcb-dea7-4d10-a4d8-27a745fdab57",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(epithelial[epithelial.obs['Slice_ID']=='062921_D0_m3a_1_slice_2'], spot_size=7, color=['predicted_celltype', 'Tier3'], title='062921_D0_m3a_1_slice_2')\n",
    "sc.pl.spatial(epithelial[epithelial.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=7, color=['predicted_celltype', 'Tier3'], title='072523_D35_m6_1_slice_3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66de40a3-9095-44aa-ad81-5254556961bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(sub[sub.obs['Slice_ID']=='062921_D0_m3a_1_slice_2'], spot_size=10, color=['Il22ra1', 'Rorc', 'Il10rb', 'Edn1'], title='062921_D0_m3a_1_slice_2', vmax=[3, 4, 3, 2])\n",
    "sc.pl.spatial(sub[sub.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=10, color=['Il22ra1', 'Rorc', 'Il10rb', 'Edn1'], title='072523_D35_m6_1_slice_3', vmax=[3, 4, 3, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "992385b5-72a8-48bf-8cd3-40ddf5e94c35",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(sub[sub.obs['Slice_ID']=='062921_D0_m3a_1_slice_2'], spot_size=7, layer='generated_expression',\n",
    "              color=['Il18', 'Tlr1', 'Il18', 'Acvr1b', 'L1cam'], title='062921_D0_m3a_1_slice_2', vmax=[3, 2, 3, 4, 2])\n",
    "sc.pl.spatial(sub[sub.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=7, layer='generated_expression',\n",
    "              color=['Il18', 'Tlr1', 'Il18', 'Acvr1b', 'L1cam'], title='072523_D35_m6_1_slice_3', vmax=[3, 2, 3, 4, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "254423e9-35b9-42d1-acdd-2c852fc89ae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.write_h5ad('/external_data/other/resolvi_final_other_files/figure4/query_trained_data_d0_35_filtered.h5ad')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24e0c88d-da66-4dca-bb70-a03ef82e14ae",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a113b30e-a9de-4f21-bd1d-d1271d3aeb80",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sub = sc.read_h5ad('/external_data/other/resolvi_final_other_files/figure4/query_trained_data_d0_35_filtered.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef621227-547a-4707-8cd4-f968223ab46a",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(sub, color='timepoint', save='timepoint_d0-35.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04ffeb4d-4a16-4229-b72d-760f0b63866d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "resolvae_both = RESOLVI.load_query_data(sub, reference_model='mouse_colitis/original_q25_semisupervised_redo3/resolvae_query')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d49723b7-cafd-4b68-a879-373dacda17e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import squidpy as sq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a4a2a30-dbe6-4962-b0f5-6c740beb8251",
   "metadata": {},
   "outputs": [],
   "source": [
    "sq.gr.spatial_neighbors(sub, coord_type=\"generic\", library_key='Slice_ID', n_neighs=30, delaunay=False, transform=None, set_diag=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e01dd7b-2964-41e4-8c3c-e5998d8ca2f5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cell_types_prediction = sub.obsm['celltype_predicted'].values\n",
    "most_common_type = []\n",
    "connectivities = sub.obsp['spatial_connectivities']\n",
    "weighted_ratios = np.zeros(cell_types_prediction.shape)\n",
    "\n",
    "for i in tqdm(range(connectivities.shape[0])):\n",
    "    neighbors = connectivities[i].nonzero()[1]\n",
    "    neighbor_types = cell_types_prediction[neighbors, :]\n",
    "    weighted_ratios[i, :] = np.sum(neighbor_types, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a196d199-9216-4ea0-91fc-a2ebcbf756dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obsm['celltypes_neighborhood'] = pd.DataFrame(weighted_ratios, columns=sub.obsm['celltype_predicted'].columns, index=sub.obs_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79ac8318-60f8-4cb8-8c91-1d0eebacb84e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_types_prediction = sub.obsm['celltype_predicted'].values\n",
    "most_common_type = []\n",
    "connectivities = sub.obsp['spatial_connectivities']\n",
    "weighted_ratios = np.zeros(cell_types_prediction.shape)\n",
    "\n",
    "for i in tqdm(range(connectivities.shape[0])):\n",
    "    neighbors = connectivities[i].nonzero()[1]\n",
    "    neighbor_types = cell_types_prediction[neighbors, :]\n",
    "    neighbor_weights = connectivities[i, neighbors].toarray().flatten()\n",
    "    weighted_counts = neighbor_types @ neighbor_weights\n",
    "\n",
    "    total_weight = neighbor_weights.sum()\n",
    "    weighted_ratios[i, :] = weighted_counts / total_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ebad9c-8c8f-40ba-8c8e-e1b1c617314f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obsm['celltypes_neighborhood_weighted'] = pd.DataFrame(weighted_ratios, columns=sub.obsm['celltype_predicted'].columns, index=sub.obs_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f4e859f-e61f-489b-9d50-e255ea910e3f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pp.neighbors(sub, use_rep='celltypes_neighborhood', n_neighbors=50, method='rapids')\n",
    "sub.obsm['X_tsne'] = sc.tl.umap(sub, min_dist=0.5, init_pos='spectral', copy=True).obsm['X_umap']\n",
    "sc.pl.tsne(sub, color=['Tier1', 'Tier3'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aee1c5a8-2a0c-4ecd-b179-ad5fa5d20cec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.tl.louvain(sub, resolution=0.8, key_added='leiden_spatial', flavor='rapids')\n",
    "sc.pl.tsne(sub, color=['leiden_spatial'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2453da68-49bf-4c00-874b-6d6f0526bc5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "palette = {\n",
    "    'FOL': '#7a4900',\n",
    "    'FOL1': '#1ce6ff',\n",
    "    'FOL2': '#ff34ff',\n",
    "    'FOL3': '#ff4a46',\n",
    "    'LUM': '#008941',\n",
    "    'ME': '#997d87',\n",
    "    'ME1': '#a30059',\n",
    "    'ME2': '#ffdbe5',\n",
    "    'ME3': '#ffff00',\n",
    "    'ME4': '#0000a6',\n",
    "    'MES1': '#00c2a0',\n",
    "    'MES2': '#b79762',\n",
    "    'MES3': '#004d43',\n",
    "    'MU': '#8fb0ff',\n",
    "    'MU1': '#006fa6',\n",
    "    'MU2': '#5a0007',\n",
    "    'MU3': '#d16100',\n",
    "    'MU4': '#6a3a4c',\n",
    "    'MU5': '#63ffac',\n",
    "    'MU6': '#b903aa',\n",
    "    'MU7': '#3b5dff',\n",
    "    'MU8': '#4a3b53',\n",
    "    'MU9': '#ff2f80',\n",
    "    'MU10': '#61615a',\n",
    "    'MU11': '#ba0900',\n",
    "    'SM': '#6b7900',\n",
    "    'SM1': '#1b4400',\n",
    "    'SM2': '#ffaa92',\n",
    "    'SM3': '#ff90c9',\n",
    "    'other': '#4fc601',\n",
    "    'others': '#d16100'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6c23e97-39d7-49f3-b327-d228995ba838",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obs['Leiden_neigh_re'] = sub.obs['Leiden_neigh']\n",
    "sub.obs.loc[sub.obs['timepoint']=='D35', 'Leiden_neigh_re'] = None\n",
    "sub.obs.loc[sub.obs['Leiden_neigh_re']=='others', 'Leiden_neigh_re'] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8408fee7-ed99-4e3c-b4d6-193d0f76e78c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.tsne(sub, color=['timepoint', 'predicted_celltype_coarse', 'leiden_neigh_resolvi'], save='niches_embedding.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27de691a-2f25-4430-b84a-d90f953a3b8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.tsne(sub[sub.obs['predicted_celltype_coarse']=='Fibroblast'], color=['timepoint', 'predicted_celltype', 'leiden_neigh_resolvi'], save='niches_embedding_fibroblast.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e13e1dc-5e93-416c-92fd-4d5aaa75f9bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.tsne(sub[sub.obs['predicted_celltype']=='Fibro 2'], color=['timepoint', 'predicted_celltype', 'leiden_neigh_resolvi'], save='niches_embedding_fibro2.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aab6763-f9c6-4930-84d3-384cd8c9ec2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_types = sub.obs['Leiden_neigh_re'].cat.codes.values\n",
    "most_common_type = []\n",
    "connectivities = sub.obsp['connectivities']\n",
    "weighted_ratios = []\n",
    "\n",
    "for i in tqdm(range(connectivities.shape[0])):\n",
    "    neighbors = connectivities[i].nonzero()[1]\n",
    "    neighbor_types = cell_types[neighbors]\n",
    "    neighbor_weights = connectivities[i, neighbors].toarray().flatten()\n",
    "    unique_types, counts = np.unique(neighbor_types, return_counts=True)\n",
    "    weighted_counts = np.zeros(len(sub.obs['Leiden_neigh_re'].cat.categories))\n",
    "\n",
    "    for j, type_index in enumerate(unique_types):\n",
    "        weighted_counts[type_index] = neighbor_weights[neighbor_types == type_index].sum()\n",
    "\n",
    "    most_common = weighted_counts[:-1].argmax()\n",
    "    total_weight = neighbor_weights.sum()\n",
    "    if weighted_counts[:-1].sum() > 0:\n",
    "        most_common_type.append(most_common)\n",
    "        weighted_ratios.append(weighted_counts[most_common] / total_weight)\n",
    "    else:\n",
    "        most_common_type.append(-1)\n",
    "        weighted_ratios.append(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a45a0aad-ec6d-4798-860b-7ddcdde44b0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obs['leiden_neigh_resolvi'] = sub.obs['Leiden_neigh_re'].cat.categories[most_common_type]\n",
    "sub.obs['uncertainty_leiden_neigh'] = 1 - np.array(weighted_ratios)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43ace60d-c6cb-4c47-b63e-61911809c9ff",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pl.tsne(sub, color=['uncertainty_leiden_neigh', 'Leiden_neigh', 'leiden_neigh_resolvi'], palette=palette, ncols=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1170a0ad-2ddf-4b2a-92d1-f9164e1201a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.tsne(sub, color=['Leiden_neigh_re'], legend_loc='on data', legend_fontsize='small')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52792485-ece0-4995-8051-f22f81c885aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.tsne(sub[sub.obs['timepoint']=='D35'], color=['Leiden_neigh_re', 'leiden_neigh_resolvi', 'leiden_spatial', 'uncertainty_leiden_neigh'], legend_loc='on data', legend_fontsize='small', title='D35')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66366e32-a58f-4d36-8a45-073be1a42e47",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in ['D35', 'D0', 'D3', 'D9', 'D21']:\n",
    "    sc.pl.tsne(sub[sub.obs['timepoint']==i], color=['leiden_neigh_resolvi', 'leiden_spatial'], legend_loc='on data', legend_fontsize='small', title=i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fee98d6-3163-4ea7-9e61-6875452597ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in ['D35', 'D0', 'D3', 'D9', 'D21']:\n",
    "    sc.pl.tsne(sub[sub.obs['timepoint']==i], color='predicted_celltype', groups=['Fibro 2', 'Stem cells', 'Fibro 6', 'Fibro 4', 'Fibro 7', 'Colonocytes', 'Goblet 1', 'Goblet 2', 'TA'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8138e480-f745-486c-80ce-11fdbb1fcae8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sub.write_h5ad(f'figure4_new/processed_adata_all_final_niche_final.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6425b852-a059-48b6-b23d-f368c97a3145",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sub = sc.read_h5ad('figure4_new/processed_adata_all_final_niche_final.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167ae6b6-4b61-4502-89dd-d158a5d34cf0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "resolvae_both = RESOLVI.load_query_data(sub, reference_model='mouse_colitis/original_q25_semisupervised_redo3/resolvae_query')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eadd1a74-2a9f-4082-8372-ba4908ed68fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "tier2_3_ordered = ['Stem cells',\n",
    "       'TA',\n",
    "       'Colonocytes',\n",
    "       'Goblet 1',\n",
    "       'Goblet 2',\n",
    "       'EEC',\n",
    "       'M cells',\n",
    "       'Epithelial (Clu+)',\n",
    "       'IAE 1',\n",
    "       'IAE 2',\n",
    "       'IAE 3',\n",
    "       'Repair associated  (Arg1+)',\n",
    "       'Fibro 6',\n",
    "       'Fibro 2',\n",
    "       'Fibro 13',\n",
    "       'Fibro 7',\n",
    "       'Fibro 4',\n",
    "       'Fibro 1',\n",
    "       'Fibro 15',\n",
    "       'Fibro 5',\n",
    "       'Fibro 12',\n",
    "       'FRC',\n",
    "       'Pericyte 1',\n",
    "       'Pericyte 2',\n",
    "       'IAF 1',\n",
    "       'IAF 2',\n",
    "       'IAF 3',\n",
    "       'IAF 4',\n",
    "       'IAF 5',\n",
    "       \n",
    "       'Monocyte', 'Macrophage (Itgax+)','Macrophage (Mrc1+)', 'Macrophage (Lyve1+)','Macrophage (Cxcl10+)',\n",
    "       'cDC1', 'DC (Fscn1+)','DC (Ccl22+)','Neutrophil 1','Neutrophil 2', 'Mast cell',\n",
    "       'T (Cd4+ Ccr7+)','Treg','T (Mki67+)','Tfh','T (Cd8+)','NK','ILC2','ILC3-LTi-like',\n",
    "       'B cell 1','B cell 2','B cell (Aicda+)','Plasma cell',\n",
    "        \n",
    "       'SMC 1','SMC 2', 'SMC Mesentery',\n",
    "       'IASMC 1', 'IASMC 2', 'IASMC 3', \n",
    "       \n",
    "       'Lymphatic EC',\n",
    "       'Lymphatic EC (Ccl21a+)', 'Arterial EC', 'Venous EC',\n",
    "       'Capillarial EC', \n",
    "       \n",
    "       'Glia (Apod+)', 'Glia (Gfap+)', 'Glia (Gfra3+)',\n",
    "       'Glia (Mpz+)', 'Neuron (Chat+)', 'Neuron (Nos1+)', \n",
    "       'ICC 1','ICC 2', 'Adipose']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c72142-1dc3-4522-8d49-9138f3c645af",
   "metadata": {},
   "outputs": [],
   "source": [
    "columns_probabilities = []\n",
    "for ct in tier2_3_ordered:\n",
    "    try:\n",
    "        values = sub.obsm['celltype_predicted'][ct]\n",
    "        values[values<1e-3] = 0\n",
    "        sub.obs[f'prediction_probability_{ct}'] = values\n",
    "        columns_probabilities.append(f'prediction_probability_{ct}')\n",
    "    except:\n",
    "        print(ct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "760cafc8-d27b-4075-a712-930381b5a4a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "list_neighborhoods = ['MU1', 'MU2', 'MU3', 'MU4', 'MU5', 'MU6', 'MU7', 'MU8', 'MU9', 'MU10', 'MU11','LUM',\n",
    "       'SM1', 'SM2', 'SM3', 'ME1', 'ME2', 'ME3', 'ME4',\n",
    "       'FOL1', 'FOL2', 'FOL3','MES1', 'MES2', 'MES3'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f946aa17-c807-492b-9980-2db768e4e6db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pl.dotplot(sub[(~sub.obs['Leiden_neigh'].isin(['ME', 'FOL', 'MU', 'SM', 'other', 'others'])) & (sub.obs['leiden_neigh_resolvi'].astype(str)!=sub.obs['Leiden_neigh'].astype(str))],\n",
    "              var_names=columns_probabilities, groupby='Leiden_neigh', standard_scale='var', dendrogram=False,\n",
    "             categories_order=list_neighborhoods, smallest_dot=6., save='original_niches.pdf')\n",
    "sc.pl.dotplot(sub[(~sub.obs['Leiden_neigh'].isin(['ME', 'FOL', 'MU', 'SM', 'other', 'others'])) & (sub.obs['leiden_neigh_resolvi'].astype(str)!=sub.obs['Leiden_neigh'].astype(str))],\n",
    "              var_names=columns_probabilities, groupby='leiden_neigh_resolvi', standard_scale='var', dendrogram=False,\n",
    "             categories_order=list_neighborhoods, smallest_dot=6., save='resolvi_niches.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90733725-11b0-44cb-aebb-bdfc0560f0ac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pl.dotplot(sub[sub.obs['leiden_neigh_resolvi']!='others'], var_names=columns_probabilities, groupby='leiden_neigh_resolvi', standard_scale='var', dendrogram=False,\n",
    "             categories_order=list_neighborhoods, smallest_dot=6., save='resolvi_niches_all.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffb31292-e80e-4548-8653-50ab1767fbf6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.pl.dotplot(sub[~sub.obs['Leiden_neigh'].isin(['ME', 'FOL', 'MU', 'SM', 'other', 'others'])], var_names=columns_probabilities, groupby='Leiden_neigh', standard_scale='var', dendrogram=False,\n",
    "             categories_order=list_neighborhoods, smallest_dot=6., save='original_niches_all.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8339846a-db97-472a-b3b1-45c9098e964c",
   "metadata": {},
   "outputs": [],
   "source": [
    "da = resolvae_both.differential_niche_abundance(\n",
    "    groupby='timepoint', group1='D35', group2='D0', subset_idx=np.where(sub.obs['predicted_celltype']=='Fibro 2')[0], neighbor_key='index_neighbor', test_mode='three',\n",
    "    delta=0.05, pseudocounts=3e-2)\n",
    "da.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afce1829-b798-40ef-b8f5-4c7188d615b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "da.head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3592a0ea-b88f-4a3f-8499-8a54bcb3b7f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import decoupler as dc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39329a4f-ebbf-4f1a-b91f-33289701e63d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dc.plot_volcano_df(\n",
    "    da,\n",
    "    x='lfc_mean',\n",
    "    y='proba_not_de',\n",
    "    sign_thr=0.5,\n",
    "    lFCs_thr=0.5,\n",
    "    top=30,\n",
    "    figsize=(10, 10),\n",
    "    save='figure4_new/fibro2_da.pdf'\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaf26cef-bb40-48f8-a9ae-1cd30128a77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obs['sub_fibro'] = [i if j=='Fibroblast' else None for i,j in zip(sub.obs['predicted_celltype'], sub.obs['predicted_celltype_coarse'])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ccdda74-c8f8-476c-9065-bf196fccbf66",
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (15, 15), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(\n",
    "        sub[sub.obs['Slice_ID']=='082421_D0_m6_1_slice_1'], spot_size=15, layer='generated_expression', color=['sub_fibro'], title='082421_D0_m6_1_slice_1', ncols=1, save='distribution_fibro2_d0.pdf'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b84ffc5-6851-4439-b48e-e7cd150058b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (15, 15), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(\n",
    "        sub[sub.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=15, layer='generated_expression', color=['sub_fibro'], title='072523_D35_m6_1_slice_3', ncols=1, save='distribution_fibro2_d35.pdf'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ee5e5c6-b349-4b6b-b750-830f1da208ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "de_result_importance = resolvae_both.differential_expression(\n",
    "    adata=sub[(sub.obs['predicted_celltype_coarse']=='Fibroblast') & (sub.obs['timepoint']=='D0')], groupby='leiden_neigh_resolvi', group1='MU3', group2='MU1', weights='importance',\n",
    "    pseudocounts=1e-2, delta=0.05, filter_outlier_cells=True, mode='change', test_mode='three'#, batch_correction=True, batchid1=batch_index, batchid2=batch_index,\n",
    ")\n",
    "de_result_importance.head(50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed2a40a0-14d5-4b36-87eb-d41d77aebb1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dc.plot_volcano_df(\n",
    "    de_result_importance,\n",
    "    x='lfc_mean',\n",
    "    y='proba_not_de',\n",
    "    sign_thr=0.3,\n",
    "    lFCs_thr=0.4,\n",
    "    top=30,\n",
    "    figsize=(10, 10),\n",
    "    save='figure4_new/fibroblast_resolvi_de.pdf'\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79038654-72ea-45f2-9935-41841da016f4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "de_result_uniform = resolvae_both.differential_expression(\n",
    "    adata=sub[sub.obs['predicted_celltype']=='Colonocytes'], groupby='timepoint', group1='D0', group2='D35', weights='importance',\n",
    "    pseudocounts=1e-2, delta=0.05, filter_outlier_cells=True, mode='change', test_mode='three', #batch_correction=True, batchid1=batch_index, batchid2=batch_index,\n",
    ")\n",
    "de_result_uniform.head(30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b67afd3-4f0e-4752-9758-2fd62fe4ea9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dc.plot_volcano_df(\n",
    "    de_result_uniform,\n",
    "    x='lfc_mean',\n",
    "    y='proba_not_de',\n",
    "    sign_thr=0.3,\n",
    "    lFCs_thr=0.4,\n",
    "    top=60,\n",
    "    figsize=(10, 10),\n",
    "    save='figure4_new/colonocyte_resolvi_de.pdf'\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a6efc83-fb8f-4521-8542-af600f982a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obs['leiden_spatial_fibroblast'] = [i if j=='Fibroblast' else None for i, j in zip(sub.obs['leiden_spatial'], sub.obs['Tier1'])]\n",
    "sub.obs['leiden_neigh_resolvi_fibroblast'] = [i if j=='Fibroblast' else None for i, j in zip(sub.obs['leiden_neigh_resolvi'], sub.obs['Tier1'])]\n",
    "sub.obs['Tier3_fibroblast'] = [i if j=='Fibroblast' else None for i, j in zip(sub.obs['Tier3'], sub.obs['Tier1'])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba1fcf68-c9c7-4f51-adeb-2e500239ed84",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for slice in sub.obs['Slice_ID'].unique():\n",
    "    if ('D0' in slice) or ('D35' in slice):\n",
    "        sc.pl.spatial(\n",
    "            sub[(sub.obs['Slice_ID']==slice)],\n",
    "            spot_size=7, color=['leiden_neigh_resolvi_fibroblast', 'leiden_spatial_fibroblast', 'Tier3_fibroblast'], title=slice) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a15bb4fe-452a-44d0-b15c-0ed77c616e7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obs['neighborhood_colonocyte'] = sub.obsm['celltypes_neighborhood']['Colonocytes'] / sub.obsm['celltypes_neighborhood'].sum(1)\n",
    "sub.obs['neighborhood_stem'] = sub.obsm['celltypes_neighborhood']['Stem cells'] / sub.obsm['celltypes_neighborhood'].sum(1)\n",
    "sub.obs['neighborhood_ta'] = sub.obsm['celltypes_neighborhood']['TA'] / sub.obsm['celltypes_neighborhood'].sum(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25ade921-3107-4272-8913-b1b23f0a6d2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.tsne(sub, color=['leiden_spatial', 'neighborhood_colonocyte', 'neighborhood_stem', 'neighborhood_ta'], cmap='Reds', vmax=0.7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea80a801-3a9e-4a82-98eb-f0bfbd1ff769",
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (15, 15), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(\n",
    "        sub[sub.obs['Slice_ID']=='082421_D0_m6_1_slice_1'], spot_size=8, layer='generated_expression',\n",
    "        color=['neighborhood_colonocyte', 'neighborhood_stem', 'neighborhood_ta'], title='082421_D0_m6_1_slice_1', ncols=3, save='zonation_epithelial_d0.pdf'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "949341d2-a224-4296-94e1-fb1c18a02aff",
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (15, 15), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(\n",
    "        sub[sub.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=8, layer='generated_expression',\n",
    "        color=['neighborhood_colonocyte', 'neighborhood_stem', 'neighborhood_ta'], ncols=3, save='zonation_epithelial_d0.pdf'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e06a916e-d72a-4410-b527-8711fe69ebd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obs['neighborhood_stem_colonocyte'] = sub.obs['neighborhood_colonocyte'] - sub.obs['neighborhood_stem']\n",
    "with plt.rc_context({\"figure.figsize\": (15, 15), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(\n",
    "        sub[sub.obs['Slice_ID']=='082421_D0_m6_1_slice_1'], spot_size=8, layer='generated_expression', cmap='bwr', vmax=1, vmin=-1,\n",
    "        color=['neighborhood_stem_colonocyte'], ncols=3, save='gradient_epithelial_d0.pdf'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c59e7df3-c864-4796-aeb3-0e94d5716cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (15, 15), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(\n",
    "        sub[sub.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=8, layer='generated_expression', cmap='bwr', vmax=1, vmin=-1,\n",
    "        color=['neighborhood_stem_colonocyte'], ncols=3, save='gradient_epithelial_d35.pdf'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df895c87-304d-4526-bc76-f9b656685ac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.obs['neighborhood_stem_colonocyte_threshold'] = (sub.obs['neighborhood_stem_colonocyte'] > 0.5).astype(str)\n",
    "sc.pl.spatial(\n",
    "    sub[sub.obs['Slice_ID']=='072523_D35_m6_1_slice_3'], spot_size=8, layer='generated_expression', cmap='bwr', vmax=1, vmin=-1,\n",
    "    color=['neighborhood_stem_colonocyte_threshold'], ncols=3,\n",
    ")\n",
    "sc.pl.spatial(\n",
    "    sub[sub.obs['Slice_ID']=='082421_D0_m6_1_slice_1'], spot_size=8, layer='generated_expression', cmap='bwr', vmax=1, vmin=-1,\n",
    "    color=['neighborhood_stem_colonocyte_threshold'], ncols=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628e7dfb-def1-4f4a-b89f-7baeecfed522",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "de_result_importance = resolvae_both.differential_expression(\n",
    "    adata=sub[(sub.obs['timepoint']=='D35') & (sub.obs['predicted_celltype']=='Colonocytes')], groupby='neighborhood_stem_colonocyte_threshold', group1='True', group2='False', weights='importance',\n",
    "    pseudocounts=1e-2, delta=0.05, filter_outlier_cells=True, mode='change', test_mode='three', #batch_correction=True, batchid1=batch_index, batchid2=batch_index,\n",
    ")\n",
    "de_result_importance.head(50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09dbec3e-4529-4ba3-92f7-73590adb8b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "dc.plot_volcano_df(\n",
    "    de_result_importance,\n",
    "    x='lfc_mean',\n",
    "    y='proba_not_de',\n",
    "    sign_thr=0.5,\n",
    "    lFCs_thr=0.1,\n",
    "    top=20,\n",
    "    figsize=(10, 10),\n",
    "    save='figure4_new/colonocyte_niches_resolvi_de.pdf'\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5a45485-a0e6-4dad-a8c4-f7231e20efcc",
   "metadata": {},
   "source": [
    "# Benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47b443b4-3d7f-4d87-bfaa-5c45fb3730c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from harmony import harmonize\n",
    "sc.tl.pca(sub)\n",
    "sub.obsm['X_pca_harmony'] = harmonize(sub.obsm['X_pca'], sub.obs, batch_key = 'Slice_ID')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06b96c39-6b7a-47ed-9d1f-c09f19e744e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvae = RESOLVI.load_query_data(sub, 'mouse_colitis/original_q25_semisupervised_tier3_d0-d21/resolvae')\n",
    "sub.obsm['X_resolvi_unsupervised'] = resolvae.get_latent_representation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dadaccd1-51f4-4f38-9890-abb5209d3e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub = sc.read_h5ad(f'figure4_new/processed_adata_all_final_niche_final.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "214df07b-aad8-4329-950e-53fe3aac5005",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_d0 = sub[sub.obs['timepoint']=='D0']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d796720d-dd26-48ea-9c54-90fed590ea88",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(sub_d0, color='predicted_celltype_coarse', save='celltypes_d0.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e017da7d-647d-43a9-8a53-b04183c85f5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_d0 = sub_d0[sub_d0.obs['low_quality'] == 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc48044-466c-44e2-810b-873c7c883322",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scib_metrics.benchmark._core import BatchCorrection\n",
    "\n",
    "batch_correction = BatchCorrection(\n",
    "    silhouette_batch=True,\n",
    "    ilisi_knn=True,\n",
    "    kbet_per_label=True,\n",
    "    graph_connectivity=True,\n",
    "    pcr_comparison=False,\n",
    ")\n",
    "\n",
    "from contextlib import contextmanager\n",
    "from scib_metrics.benchmark import Benchmarker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29b0cd65-726a-42c3-b530-9ddd3403b16d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bm = Benchmarker(\n",
    "    sub_d0,\n",
    "    batch_key=\"Slice_ID\",\n",
    "    label_key=\"Tier3\",\n",
    "    batch_correction_metrics=batch_correction,\n",
    "    embedding_obsm_keys=['X_pca', 'X_pca_harmony', 'X_resolVI', 'X_resolvi_unsupervised'],\n",
    "    pre_integrated_embedding_obsm_key='X_pca',\n",
    "    n_jobs=12,\n",
    ")\n",
    "bm.benchmark()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27346401-59a8-42b9-9967-e7a40d48d238",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['font.weight'] = 'normal'\n",
    "plt.rcParams['axes.titleweight'] = 'normal'\n",
    "plt.rcParams['axes.labelweight'] = 'normal'\n",
    "\n",
    "bm.plot_results_table(min_max_scale=False, save_dir=f'figure4_new/')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "resolvi_new",
   "language": "python",
   "name": "resolvi_new"
  },
  "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}