figure2_final.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyro"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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,
   "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 = 'figure2'\n",
    "plt.rcParams['pdf.fonttype'] = 'truetype'\n",
    "plt.rcParams['svg.fonttype'] = 'none'\n",
    "plt.rcParams['pdf.use14corefonts'] = True "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append('.')\n",
    "import _utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Load datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sc.read_h5ad('xenium_brain/original_high_lr/complete_adata.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdata = sc.read_h5ad('xenium_brain/original_semisupervised/complete_adata.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "resolvi = RESOLVI.load('xenium_brain/original_high_lr/resolvae')\n",
    "resolvi_semisupervised = RESOLVI.load('xenium_brain/original_semisupervised/resolvae')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = resolvi.history['elbo_train'].plot()\n",
    "_ = resolvi_semisupervised.history['elbo_train'].plot()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bdata = adata.copy()\n",
    "bdata.X = bdata.layers['raw_counts']\n",
    "sc.pp.downsample_counts(bdata, counts_per_cell=20, random_state=0)\n",
    "bdata.layers['raw_counts'] = bdata.X\n",
    "\n",
    "bdata.obsm[\"x_resolVI_downsampled_30\"] = resolvi.get_latent_representation(adata=bdata)\n",
    "_utils.compute_umap_embedding(bdata, representation_key=\"x_resolVI_downsampled_20\", n_comps=None, show=True, key='resolvi_latent_downsampled20', n_neighbors=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_utils.compute_umap_embedding(adata, representation_key=\"X_resolVI\", n_comps=None, show=True, key='resolvi_latent', n_neighbors=20)\n",
    "_utils.compute_umap_embedding(adata, representation_key=\"counts\", show=True, key='raw_counts', n_neighbors=20)\n",
    "_utils.compute_umap_embedding(adata, representation_key=\"generated_expression\", show=True, key='resolvi_generated', n_neighbors=20)\n",
    "_utils.compute_umap_embedding(adata, representation_key=\"corrected_counts\", show=True, key='resolvi_corrected', n_neighbors=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_utils.compute_umap_embedding(cdata, representation_key=\"corrected_counts\", show=True, key='resolvi_corrected', n_neighbors=20)\n",
    "_utils.compute_umap_embedding(cdata, representation_key=\"generated_expression\", show=True, key='resolvi_generated', n_neighbors=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## SCIB metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.options.display.max_columns = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scvi\n",
    "\n",
    "scvi.model.SCVI.setup_anndata(adata, layer=\"raw_counts\")\n",
    "vae = scvi.model.SCVI(adata, gene_likelihood=\"nb\", n_layers=2, n_latent=10)\n",
    "vae.train()\n",
    "adata.obsm[\"scVI\"] = vae.get_latent_representation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lvae = scvi.model.SCANVI.from_scvi_model(\n",
    "    vae,\n",
    "    adata=adata,\n",
    "    labels_key=\"predicted_celltype\",\n",
    "    unlabeled_category=\"Unknown\",\n",
    ")\n",
    "lvae.train(max_epochs=20, n_samples_per_label=100)\n",
    "adata.obsm[\"scANVI\"] = lvae.get_latent_representation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scib_metrics\n",
    "from scib_metrics.benchmark import Benchmarker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obsm['Unintegrated'] = adata.obsm['X_pca_raw_counts']\n",
    "adata.obsm['resolVI'] = adata.obsm['X_resolVI']\n",
    "adata.obsm['resolVI Supervised'] = cdata.obsm['X_resolVI']\n",
    "adata.obsm['Generated Expression'] = adata.obsm['X_pca_resolvi_generated']\n",
    "adata.obsm['Generated Expression Supervised'] = cdata.obsm['X_pca_resolvi_generated']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sc.read('xenium_brain/full_data_concatenated.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scib_metrics\n",
    "batch_correction = scib_metrics.benchmark._core.BatchCorrection(pcr_comparison=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['diffusion_cells'] = [str(i<5120) for i in adata.obs['x_centroid']]\n",
    "adata.obs['diffusion_cells'] = [str(i>0.8) for i in adata.obs['true_proportion']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(adata, spot_size=30, color='diffusion_cells')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bm = Benchmarker(\n",
    "    adata,\n",
    "    batch_key=\"diffusion_cells\",\n",
    "    label_key=\"predicted_celltype\",\n",
    "    embedding_obsm_keys=[\"Unintegrated\", \"Generated Expression\", \"resolVI\", \"scVI\", \"scANVI\", \"resolVI Supervised\", \"Generated Expression Supervised\"],\n",
    "    pre_integrated_embedding_obsm_key='Unintegrated',\n",
    "    n_jobs=12,\n",
    ")\n",
    "bm.benchmark()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bm._results = bm._results.drop('pcr_comparison', axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm.get_results(min_max_scale=False).to_csv('xenium_brain/scib_results_all.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "\n",
    "@contextmanager\n",
    "def default_rcparams():\n",
    "    default_params = plt.rcParams.copy()  # Store current rcParams\n",
    "    plt.rcdefaults()   # Reset all rcParams to their defaults\n",
    "    yield\n",
    "    plt.rcParams.update(default_params)   # Restore rcParams to their original values\n",
    "    plt.rcParams['svg.fonttype'] = 'none'\n",
    "\n",
    "# Example usage\n",
    "with default_rcparams():\n",
    "    bm.plot_results_table(min_max_scale=False, save_dir='figure2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm.get_results().to_csv('scib_results_filtered.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.umap(adata, color=['diffusion_proportion', 'background_proportion'], ncols=1, size=1, save='semisupervised_proportions.pdf')\n",
    "sc.pl.umap(adata[adata.obs['true_proportion']>0.8], color=['predicted_celltype'], ncols=1, size=1, save='semisupervised_low_diffusion.pdf')\n",
    "sc.pl.umap(adata[adata.obs['true_proportion']<0.8], color=['predicted_celltype'], ncols=1, size=1, save='semisupervised_high_diffusion.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pynndescent import PyNNDescentTransformer\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.pipeline import make_pipeline\n",
    "\n",
    "complete_data = adata # [adata.obs['diffusion_proportion']<0.2]\n",
    "train_X = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obsm[\"scVI\"]\n",
    "train_Y = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obs[\"predicted_celltype\"].to_numpy()\n",
    "knn = make_pipeline(\n",
    "    PyNNDescentTransformer(\n",
    "        n_neighbors=15,\n",
    "        parallel_batch_queries=True,\n",
    "    ),\n",
    "    KNeighborsClassifier(metric=\"precomputed\", weights=\"uniform\"),\n",
    ")\n",
    "knn.fit(train_X, train_Y)\n",
    "adata.obs['predicted_celltypes_scvi'] = knn.predict(adata.obsm[\"scVI\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pynndescent import PyNNDescentTransformer\n",
    "from sklearn.neighbors import KNeighborsClassifiersdc\n",
    "from sklearn.pipeline import make_pipeline\n",
    "\n",
    "complete_data = adata # [adata.obs['diffusion_proportion']<0.2]\n",
    "train_X = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obsm[\"resolVI\"]\n",
    "train_Y = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obs[\"predicted_celltype\"].to_numpy()\n",
    "knn = make_pipeline(\n",
    "    PyNNDescentTransformer(\n",
    "        n_neighbors=15,\n",
    "        parallel_batch_queries=True,\n",
    "    ),\n",
    "    KNeighborsClassifier(metric=\"precomputed\", weights=\"uniform\"),\n",
    ")\n",
    "knn.fit(train_X, train_Y)\n",
    "adata.obs['predicted_celltypes'] = knn.predict(adata.obsm[\"resolVI\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pynndescent import PyNNDescentTransformer\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.pipeline import make_pipeline\n",
    "\n",
    "complete_data = adata # [adata.obs['diffusion_proportion']<0.2]\n",
    "train_X = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obsm[\"resolVI Supervised\"]\n",
    "train_Y = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obs[\"predicted_celltype\"].to_numpy()\n",
    "knn = make_pipeline(\n",
    "    PyNNDescentTransformer(\n",
    "        n_neighbors=15,\n",
    "        parallel_batch_queries=True,\n",
    "    ),\n",
    "    KNeighborsClassifier(metric=\"precomputed\", weights=\"uniform\"),\n",
    ")\n",
    "knn.fit(train_X, train_Y)\n",
    "adata.obs['predicted_celltypes_supervised'] = knn.predict(adata.obsm[\"resolVI Supervised\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pynndescent import PyNNDescentTransformer\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.pipeline import make_pipeline\n",
    "\n",
    "complete_data = adata # [adata.obs['diffusion_proportion']<0.2]\n",
    "train_X = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obsm[\"scANVI\"]\n",
    "train_Y = complete_data[~pd.isna(complete_data.obs[\"predicted_celltype\"])].obs[\"predicted_celltype\"].to_numpy()\n",
    "knn = make_pipeline(\n",
    "    PyNNDescentTransformer(\n",
    "        n_neighbors=15,\n",
    "        parallel_batch_queries=True,\n",
    "    ),\n",
    "    KNeighborsClassifier(metric=\"precomputed\", weights=\"uniform\"),\n",
    ")\n",
    "knn.fit(train_X, train_Y)\n",
    "adata.obs['predicted_celltypes_scanvi'] = knn.predict(adata.obsm[\"scANVI\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvi_semisupervised = RESOLVI.load('xenium_brain/original_semisupervised/resolvae')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(\n",
    "    model,\n",
    "    adata = None,\n",
    "    indices = None,\n",
    "    soft: bool = False,\n",
    "    batch_size: int | None = 500,\n",
    "    num_samples: int | None = 30\n",
    ") -> np.ndarray | pd.DataFrame:\n",
    "    adata = model._validate_anndata(adata)\n",
    "\n",
    "    if indices is None:\n",
    "        indices = np.arange(adata.n_obs)\n",
    "\n",
    "    sampled_prediction = model.sample_posterior_predictive(\n",
    "        adata=adata,\n",
    "        indices=indices,\n",
    "        model=model.module.model_corrected,\n",
    "        return_sites=['probs_prediction'],\n",
    "        num_samples=num_samples,\n",
    "        return_samples=False,\n",
    "        batch_size=batch_size,\n",
    "        batch_steps=10\n",
    "    )\n",
    "    y_pred = sampled_prediction['post_sample_means']['probs_prediction']\n",
    "\n",
    "    if not soft:\n",
    "        y_pred = y_pred.argmax(axis=1)\n",
    "        predictions = [model._code_to_label[p] for p in y_pred]\n",
    "        return np.array(predictions)\n",
    "    else:\n",
    "        n_labels = len(y_pred[0])\n",
    "        predictions = pd.DataFrame(\n",
    "            y_pred,\n",
    "            columns=model._label_mapping[:n_labels],\n",
    "            index=adata.obs_names[indices],\n",
    "        )\n",
    "        return predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "adata.obs['celltype_predicted'] = predict(resolvi_semisupervised)\n",
    "subset = adata[~adata.obs['predicted_celltype'].isin(['Junk', 'Endothelial-Astrocyte', 'Thalamus_Glia', 'Thalamus_Oligodendrocyte'])].copy()\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "# Create a DataFrame from 'cluster' and 'celltype_predicted' columns\n",
    "df = pd.DataFrame({\n",
    "    'cluster': subset.obs['cluster'],\n",
    "    'celltype_predicted': subset.obs['celltype_predicted']\n",
    "})\n",
    "\n",
    "# Group by 'cluster' and compute accuracy for each group\n",
    "accuracy = df.groupby('cluster').apply(lambda x: accuracy_score(x['cluster'], x['celltype_predicted']))\n",
    "\n",
    "# Plot accuracy for each cluster\n",
    "plt.figure(figsize=(10, 6))\n",
    "accuracy.plot(kind='bar')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.ylim(0., 1)\n",
    "plt.title('Accuracy for each label in cluster')\n",
    "plt.savefig('figure2/accuracy_resolvi_prediction.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f1_score(df['cluster'], df['celltype_predicted'], average='macro')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset.obs['celltype_predicted_scanvi'] = lvae.predict(adata=subset)\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "# Create a DataFrame from 'cluster' and 'celltype_predicted' columns\n",
    "df2 = pd.DataFrame({\n",
    "    'cluster': subset.obs['cluster'],\n",
    "    'celltype_predicted': subset.obs['celltype_predicted_scanvi']\n",
    "})\n",
    "\n",
    "# Group by 'cluster' and compute accuracy for each group\n",
    "accuracy = df2.groupby('cluster').apply(lambda x: accuracy_score(x['cluster'], x['celltype_predicted']))\n",
    "\n",
    "# Plot accuracy for each cluster\n",
    "plt.figure(figsize=(10, 6))\n",
    "accuracy.plot(kind='bar')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.ylim(0., 1)\n",
    "plt.title('Accuracy for each label in cluster')\n",
    "plt.savefig('figure2/accuracy_scanvi_prediction.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.write_h5ad('xenium_brain/full_data_concatenated.h5ad')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Spatial display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sc.read_h5ad('xenium_brain/full_data_concatenated.h5ad')\n",
    "sc.pl.spatial(adata, color='cluster', spot_size=20, sort_order=False, save='spatial_celltypes.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (8, 8), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(adata, color='celltype_predicted', spot_size=30, sort_order=False, palette=sc.plotting.palettes.default_102, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resolvi = RESOLVI.load('xenium_brain/original_high_lr/resolvae')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_res = resolvi.sample_posterior_predictive(\n",
    "    model=resolvi.module.model_residuals,\n",
    "    return_sites=['px_rate', 'obs'],\n",
    "    num_samples=10, return_samples=False, batch_size=2000, batch_steps=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_corr = resolvi.sample_posterior_predictive(\n",
    "    model=resolvi.module.model_corrected,\n",
    "    return_sites=['px_rate', 'obs'],\n",
    "    num_samples=10, return_samples=False, batch_size=5000, batch_steps=10)\n",
    "samples_corr = pd.DataFrame(samples_corr).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.layers['generated_expression_small'] = samples_corr.loc['post_sample_q50', 'obs']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.normalize_total(adata, layers=['generated_expression_small', 'generated_expression', 'corrected_counts', 'raw_counts'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "mpl.rcParams.update(mpl.rcParamsDefault)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (8, 8), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(adata[adata.obs['celltype_predicted']=='Microglia'],\n",
    "                  color=['Gfap', 'Slc17a6', 'Slc17a7', 'Trem2', 'Aqp4', 'Pecam1', 'diffusion_proportion'], spot_size=100, layer='raw_counts', vmax='p95', sort_order=False)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (8, 8), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(adata[adata.obs['celltype_predicted']=='Microglia'],\n",
    "                  color=['Gfap', 'Slc17a6', 'Slc17a7', 'Trem2', 'Aqp4', 'Pecam1', 'diffusion_proportion'], spot_size=100, layer='generated_expression', vmax='p95', sort_order=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (8, 8), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(adata,\n",
    "                  color=['Slc17a6', 'Trem2', 'diffusion_proportion'], spot_size=30, layer='raw_counts', vmax=[10, 20], sort_order=False, save='spatial_all_slc17a6_trem2.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (8, 8), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(adata[adata.obs['celltype_predicted']=='Microglia'],\n",
    "                  color=['Slc17a6', 'Trem2', 'diffusion_proportion'], spot_size=100, layer='raw_counts', vmax=[10, 20], sort_order=False, save='spatial_microglia_slc17a6_trem2.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (8, 8), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(adata,\n",
    "                  color=['Slc17a6', 'Trem2'], spot_size=30, layer='generated_expression', vmax=[10, 20], sort_order=False, save='spatial_all_slc17a6_trem2_generated.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with plt.rc_context({\"figure.figsize\": (8, 8), \"figure.dpi\": (300)}):\n",
    "    sc.pl.spatial(adata[adata.obs['celltype_predicted']=='Microglia'],\n",
    "                  color=['Slc17a6', 'Trem2'], spot_size=100, layer='generated_expression', vmax=[10, 20], sort_order=False, save='spatial_microglia_slc17a6_trem2_generated.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = sc.pl.scatter(adata, x='Slc17a6', y='Trem2', layers='raw_counts', size=30, color='cluster', groups=['Microglia', 'Excitatory Neurons Thalamus'], show=False)\n",
    "ax.set_xlim(0, 100)\n",
    "ax.set_ylim(0, 40)\n",
    "plt.savefig('figure2/scatter_Slc17a6_Trem2_raw.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.layers['generated_expression'][adata.layers['generated_expression']>70] = 70"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.scatter(adata, x='Slc17a6', y='Trem2', layers='generated_expression', size=30, color='cluster', groups=['Microglia', 'Excitatory Neurons Thalamus'], save='Slc17a6_Trem2_generated.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.write_h5ad('xenium_brain/full_data.h5ad')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single cell reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sc.read_h5ad('xenium_brain/full_data.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scanpy as sc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "single_cell_reference = sc.read_loom('xenium_brain/single_cell_reference.loom')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "single_cell_reference.var_names_make_unique()\n",
    "single_cell_reference.obs_names_make_unique()\n",
    "single_cell_reference = single_cell_reference[:, np.intersect1d(adata.var_names, single_cell_reference.var_names)].copy()\n",
    "single_cell_reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "single_cell_reference.layers['counts'] = single_cell_reference.X.copy()\n",
    "single_cell_reference.obsm['counts'] = pd.DataFrame(single_cell_reference.layers['counts'].A, columns=single_cell_reference.var_names, index=single_cell_reference.obs_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_utils.double_positive_pmm(single_cell_reference, single_cell_reference.var_names, layer_key=\"counts\", output_dir='figure2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "single_cell_reference.obsm['counts'].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "single_cell_reference.obsm['positive_pmm_counts']['celltype'] = single_cell_reference.obs['Class']\n",
    "per_celltype_positive = single_cell_reference.obsm['positive_pmm_counts'].groupby('celltype').mean()\n",
    "per_celltype_positive.drop('PeripheralGlia', inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "per_celltype_positive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "celltype_gene_dict = {}\n",
    "\n",
    "# Iterate over each column\n",
    "for col in per_celltype_positive.columns:\n",
    "    # Check if only one value is above 0.2 and all other values are below 0.05\n",
    "    if (per_celltype_positive[col] > 0.05).sum() == 1 and (per_celltype_positive[col] < 0.01).sum() == len(per_celltype_positive) - 1:\n",
    "        # Get the celltype for which the value is above 0.2\n",
    "        celltype = per_celltype_positive[per_celltype_positive[col] > 0.05].index[0]\n",
    "        # If the celltype is not in the result dictionary, add it with an empty list\n",
    "        if celltype not in celltype_gene_dict:\n",
    "            celltype_gene_dict[celltype] = []\n",
    "        # Append the column (gene) to the list of genes for this celltype\n",
    "        celltype_gene_dict[celltype].append(col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "celltype_gene_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "celltype_gene_dict = {'Vascular': ['Adgrl4', 'Cldn5', 'Emcn', 'Nostrin', 'Pln', 'Slfn5', 'Sox17'],\n",
    " 'Neurons': ['Bcl11b',\n",
    "  'Cabp7',\n",
    "  'Cbln1',\n",
    "  'Cbln4',\n",
    "  'Chrm2',\n",
    "  'Cntnap4',\n",
    "  'Cpne4',\n",
    "  'Fibcd1',\n",
    "  'Gsg1l',\n",
    "  'Hs3st2',\n",
    "  'Lamp5',\n",
    "  'Ndst4',\n",
    "  'Necab1',\n",
    "  'Nell1',\n",
    "  'Neurod6',\n",
    "  'Nwd2',\n",
    "  'Plcxd3',\n",
    "  'Rxfp1',\n",
    "  'Satb2',\n",
    "  'Slc17a6',\n",
    "  'Sncg',\n",
    "  'Syt2',\n",
    "  'Syt6'],\n",
    " 'Immune': ['Cd53', 'Ikzf1', 'Lyz2', 'Siglech', 'Spi1', 'Trem2'],\n",
    " 'Oligos': ['Sema3d'],\n",
    " 'Ependymal': ['Spag16', 'Trp73']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "marker_list = sum(celltype_gene_dict.values(), [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.normalize_total(adata, layers=['generated_expression', 'raw_counts', 'corrected_counts'])\n",
    "adata.obsm['counts'] = pd.DataFrame(adata[:, marker_list].layers['raw_counts'].A, columns=marker_list, index=adata.obs_names)\n",
    "adata.obsm['generated_expression'] = pd.DataFrame(np.array(adata[:, marker_list].layers['generated_expression'].A), columns=marker_list, index=adata.obs_names)\n",
    "adata.obsm['corrected_counts'] = pd.DataFrame(np.array(adata[:, marker_list].layers['corrected_counts'].A), columns=marker_list, index=adata.obs_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_utils.cosine_distance_celltype(single_cell_reference, celltype_gene_dict, layer_key=\"counts\", output_dir='figure2')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_utils.cosine_distance_celltype(adata, celltype_gene_dict, layer_key=\"generated_expression\", output_dir='figure2', vmax=0.2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_utils.double_positive_pmm(adata, marker_list, marker_dict=celltype_gene_dict, layer_key=\"generated_expression\", output_dir='figure2')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_utils.double_positive_pmm(adata[adata.obs['true_proportion']>0.9], marker_list, marker_dict=celltype_gene_dict, layer_key=\"generated_expression\", output_dir='figure2', file_save='_high_wrong')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_utils.double_positive_pmm(single_cell_reference, marker_list, marker_dict=celltype_gene_dict, layer_key='counts', output_dir='figure2', file_save='_single_cell')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def double_positive_boxplot(adata, gene_pairs, save_key='', show=False):\n",
    "    ranges = [0] + [i/10 for i in np.arange(5, 11)]\n",
    "    index = pd.MultiIndex.from_tuples(gene_pairs)\n",
    "    dp_ct_counts = pd.DataFrame(index=index, columns=ranges[1:])\n",
    "    dp_ct_generated = pd.DataFrame(index=index, columns=ranges[1:])\n",
    "\n",
    "    for index, i in enumerate(ranges[1:]):\n",
    "        for gene_x, gene_y in gene_pairs:\n",
    "            subset = adata[np.logical_and(adata.obs['true_proportion']>ranges[index], adata.obs['true_proportion']<ranges[index+1])] \n",
    "            positives_counts = subset.obsm['positive_pmm_counts'][[gene_x, gene_y]].sum(1)\n",
    "            positives_generated = subset.obsm[f'positive_pmm_generated_expression'][[gene_x, gene_y]].sum(1)\n",
    "            dp_ct_counts.loc[(gene_x, gene_y), i] = (np.sum(positives_counts==2) / np.sum(positives_counts>0) if np.sum(positives_counts)>0 else -0.01)\n",
    "            dp_ct_generated.loc[(gene_x, gene_y), i] = (np.sum(positives_generated==2) / np.sum(positives_generated>0) if np.sum(positives_generated)>0 else -0.01)\n",
    "\n",
    "    dp_ct_counts_df = pd.DataFrame(dp_ct_counts).melt()\n",
    "    dp_ct_generated_df = pd.DataFrame(dp_ct_generated).melt()\n",
    "\n",
    "    dp_ct_counts_df['source'] = 'Measured'\n",
    "    dp_ct_generated_df['source'] = 'Generated'\n",
    "\n",
    "    # Concatenate the dataframes\n",
    "    df = pd.concat([dp_ct_counts_df, dp_ct_generated_df])\n",
    "\n",
    "    # Create a color palette\n",
    "    palette = {'Measured': (1, 0, 0, 0.2), 'Generated': (0, 0, 1, 0.2)}  # red and blue with alpha=0.2\n",
    "    palette2 = {'Measured': (0.5, 0.5, 0.5, 0.2), 'Generated': (0.5, 0.5, 0.5, 0.2)}  # red and blue with alpha=0.2\n",
    "\n",
    "    # Create the dotplot\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    sns.set(style='white')\n",
    "    violin_parts = sns.violinplot(df, y='value', x='variable', hue='source', palette=palette, split=True, inner=None)\n",
    "    for pc in violin_parts.collections:\n",
    "        pc.set_alpha(0.8)\n",
    "\n",
    "    # Create the boxplot with a third of the width and black color\n",
    "    sns.boxplot(df, y='value', x='variable', hue='source', width=0.6, palette=palette2, fliersize=1.5, gap=0.5)\n",
    "\n",
    "    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')\n",
    "    plt.savefig(f'figure2/overlapping_{save_key}.pdf')\n",
    "\n",
    "    if show:\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import combinations, chain\n",
    "\n",
    "# Get all pairs across all lists\n",
    "all_genes = list(chain.from_iterable(celltype_gene_dict.values()))\n",
    "\n",
    "# Get all pairs within each list\n",
    "#celltype_gene_dict.pop('Vascular')\n",
    "within_pairs = {key: list(combinations(value, 2)) for key, value in celltype_gene_dict.items()}\n",
    "within_pairs = sum(within_pairs.values(), [])\n",
    "across_pairs = list(set(combinations(all_genes, 2)) - set(within_pairs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "within_ct_counts_reference = single_cell_reference.uns['double_positive_counts'].reset_index().melt(id_vars=['index'], var_name='Gene', value_name='Value')\n",
    "\n",
    "# Set a MultiIndex with both the row and column labels\n",
    "within_ct_counts_reference.set_index(['index', 'Gene'], inplace=True)\n",
    "within_ct_counts_reference = within_ct_counts_reference[within_ct_counts_reference.index.get_level_values('Gene') != within_ct_counts_reference.index.get_level_values('index')]\n",
    "# Display the resulting DataFrame\n",
    "print(within_ct_counts_reference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "within_ct_counts_reference[within_ct_counts_reference>0.1].dropna(how='all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = within_ct_counts_reference[within_ct_counts_reference>0.2].dropna(how='all')\n",
    "index_list = subset.index.to_list()\n",
    "within_pairs_coexpressed = index_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "within_pairs_coexpressed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "within_pairs_coexpressed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "double_positive_boxplot(adata, within_pairs_coexpressed, save_key='original_true_coexpressed', show=True)\n",
    "double_positive_boxplot(adata, across_pairs, save_key='original_false_coexpressed', show=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ls xenium_brain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other segmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bdata = sc.read_h5ad(f'xenium_brain/original_high_lr/complete_adata.h5ad')\n",
    "sc.tl.rank_genes_groups(bdata, groupby='predicted_celltype')\n",
    "sc.pl.rank_genes_groups_dotplot(bdata, n_genes=3, save='original_segmentation.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdata = sc.read_h5ad(f'xenium_brain/proseg_nucleus/complete_adata.h5ad')\n",
    "cdata.uns['dendrogram_predicted_celltype'] = bdata.uns['dendrogram_predicted_celltype']\n",
    "sc.tl.rank_genes_groups(cdata, groupby='predicted_celltype')\n",
    "sc.pl.rank_genes_groups_dotplot(cdata, n_genes=3, save='proseg_segmentation.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in ['original_high_lr', 'proseg_nucleus', 'original_nucleus', 'baysor_prior']:\n",
    "    print(i)\n",
    "    bdata = sc.read_h5ad(f'xenium_brain/{i}/complete_adata.h5ad')\n",
    "    sc.pp.normalize_total(bdata, layers=['generated_expression', 'raw_counts'])\n",
    "    sc.pl.spatial(bdata[bdata.obs['predicted_celltype']=='Microglia'], save=f'_spatial_{i}_microglia_generated.pdf',\n",
    "                  color=['Gfap', 'Slc17a6', 'Slc17a7', 'Trem2', 'Aqp4', 'Pecam1', 'diffusion_proportion'], spot_size=100, layer='generated_expression', vmax='p95', sort_order=False)\n",
    "    sc.pl.spatial(bdata[bdata.obs['predicted_celltype']=='Microglia'], save=f'_spatial_{i}_microglia_raw.pdf',\n",
    "                  color=['Gfap', 'Slc17a6', 'Slc17a7', 'Trem2', 'Aqp4', 'Pecam1', 'diffusion_proportion'], spot_size=100, layer='raw_counts', vmax='p95', sort_order=False)\n",
    "    bdata.obsm['counts'] = pd.DataFrame(bdata[:, marker_list].layers['raw_counts'].A, columns=marker_list, index=bdata.obs_names)\n",
    "    bdata.obsm['generated_expression'] = pd.DataFrame(np.array(bdata[:, marker_list].layers['generated_expression'].A), columns=marker_list, index=bdata.obs_names)\n",
    "    _utils.double_positive_pmm(bdata, marker_list, marker_dict=celltype_gene_dict, layer_key=\"generated_expression\", output_dir='figure2', file_save=f'_{i}')\n",
    "    double_positive_boxplot(bdata, within_pairs_coexpressed, save_key=f'{i}_true_coexpressed', show=True)\n",
    "    double_positive_boxplot(bdata, across_pairs, save_key=f'{i}_false_coexpressed', show=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "celltype_gene_dict_extended = {}\n",
    "\n",
    "# Iterate over each column\n",
    "for col in per_celltype_positive.columns:\n",
    "    # Check if only one value is above 0.2 and all other values are below 0.05\n",
    "    if (per_celltype_positive[col] > 0.1).sum() == 1 and (per_celltype_positive[col] < 0.05).sum() == len(per_celltype_positive) - 1:\n",
    "        # Get the celltype for which the value is above 0.2\n",
    "        celltype = per_celltype_positive[per_celltype_positive[col] > 0.1].index[0]\n",
    "        # If the celltype is not in the result dictionary, add it with an empty list\n",
    "        if celltype not in celltype_gene_dict_extended:\n",
    "            celltype_gene_dict_extended[celltype] = []\n",
    "        # Append the column (gene) to the list of genes for this celltype\n",
    "        celltype_gene_dict_extended[celltype].append(col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bdata = sc.read_h5ad(f'xenium_brain/proseg_nucleus/complete_adata.h5ad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('figure2/celltype_markers_sc_ref_extended.json', 'w') as fp:\n",
    "    json.dump(celltype_gene_dict_extended, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "celltype_gene_dict = {'Vascular': ['Adgrl4', 'Cldn5', 'Emcn', 'Nostrin', 'Pln', 'Slfn5', 'Sox17'],\n",
    " 'Neurons': ['Bcl11b',\n",
    "  'Cabp7',\n",
    "  'Cbln1',\n",
    "  'Cbln4',\n",
    "  'Chrm2',\n",
    "  'Cntnap4',\n",
    "  'Cpne4',\n",
    "  'Fibcd1',\n",
    "  'Gsg1l',\n",
    "  'Hs3st2',\n",
    "  'Lamp5',\n",
    "  'Ndst4',\n",
    "  'Necab1',\n",
    "  'Nell1',\n",
    "  'Neurod6',\n",
    "  'Nwd2',\n",
    "  'Plcxd3',\n",
    "  'Rxfp1',\n",
    "  'Satb2',\n",
    "  'Slc17a6',\n",
    "  'Sncg',\n",
    "  'Syt2',\n",
    "  'Syt6'],\n",
    " 'Immune': ['Cd53', 'Ikzf1', 'Lyz2', 'Siglech', 'Spi1', 'Trem2'],\n",
    " 'Oligos': ['Sema3d'],\n",
    " 'Ependymal': ['Spag16', 'Trp73']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "marker_list = sum(celltype_gene_dict.values(), [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.normalize_total(bdata, layers=['generated_expression', 'counts', 'estimated'])\n",
    "bdata.obsm['counts'] = pd.DataFrame(bdata[:, marker_list].layers['counts'].A, columns=marker_list, index=bdata.obs_names)\n",
    "bdata.obsm['estimated'] = pd.DataFrame(np.array(bdata[:, marker_list].layers['estimated'].A), columns=marker_list, index=bdata.obs_names)\n",
    "_utils.double_positive_pmm(bdata, marker_list, marker_dict=celltype_gene_dict, layer_key=\"estimated_expression\", output_dir='figure2', file_save=f'proseg_corrected_rates')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bdata.obsm['counts'] = pd.DataFrame(bdata[:, marker_list].layers['counts'].A, columns=marker_list, index=bdata.obs_names)\n",
    "bdata.obsm['generated_expression'] = pd.DataFrame(np.array(bdata[:, marker_list].layers['generated_expression'].A), columns=marker_list, index=bdata.obs_names)\n",
    "_utils.double_positive_pmm(bdata, marker_list, marker_dict=celltype_gene_dict, layer_key=\"generated_expression\", output_dir='figure2', file_save=f'proseg_counts')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 4
}