StereoSeq_Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50281cd8-a393-48bb-bba5-f3456ee5173f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "\n",
    "\n",
    "import datetime\n",
    "import os\n",
    "import re\n",
    "import random\n",
    "\n",
    "# Disable INFO and WARNING errors from tensorflow\n",
    "# Most nodes do not have GPU's and this prevents warnings about that being printed.\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
    "\n",
    "import anndata as ad\n",
    "import collections as c\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scanpy as sc\n",
    "import seaborn as sns\n",
    "\n",
    "from pathlib import Path\n",
    "from matplotlib import rcParams\n",
    "\n",
    "import stereo as st\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "random.seed('cbd-bioinf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ce8d44-1f33-4b87-9fc6-116768c7a38e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sc.settings.cachedir = '/staging/leuven/stg_00003/cbd-bioinf/cache/'\n",
    "sc.settings.verbosity = 3\n",
    "sc.logging.print_header()\n",
    "sc.settings.set_figure_params(dpi=120, facecolor='white')\n",
    "rcParams[\"figure.dpi\"] = 120\n",
    "\n",
    "projects_data_dir = Path('./')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "468e13cd-6fb1-4504-8028-c5ae2c859812",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the wrapper class for single cell samples\n",
    "\n",
    "def load_stereoseq_data(project, sample, species, bin_data, bin_sizes=[50]):\n",
    "    if bin_data.endswith('gef') or bin_data.endswith('gef.gz'):\n",
    "        print(f'Detected GEF file')\n",
    "        load_func = st.io.read_gef\n",
    "    elif bin_data.endswith('gem') or bin_data.endswith('gem.gz'):\n",
    "        print(f'Detected GEM file')\n",
    "        load_func = st.io.read_gem\n",
    "    else:\n",
    "        raise IOError('File type not recognised.')\n",
    "    \n",
    "    return_datas = {}\n",
    "    for bin_size in bin_sizes:\n",
    "        stereoseq_data = load_func(bin_data, bin_size=bin_size, bin_type='bins')\n",
    "        return_datas[f'{sample}_bin{bin_size}'] = SingleCellSample(\n",
    "            project,\n",
    "            f'{sample}_bin{bin_size}',\n",
    "            species,\n",
    "            st.io.stereo_to_anndata(stereoseq_data, sample_id=sample))\n",
    "        return_datas[f'{sample}_bin{bin_size}'].spatial = True\n",
    "        \n",
    "    return return_datas\n",
    "    \n",
    "\n",
    "class SingleCellSample:\n",
    "    '''\n",
    "    This class contains the main functions used in the rest of the notebook. \n",
    "    Each function is typically a wrapper around one or multiple scanpy functions.\n",
    "    \n",
    "    '''\n",
    "    \n",
    "    SCS_VER = '2.0.5'\n",
    "    \n",
    "    def __init__(self, project, sample, species, adata, data_path=None):\n",
    "        self.SCS_VER = SingleCellSample.SCS_VER\n",
    "        self.sample = sample\n",
    "        self.adata = adata\n",
    "        self.project = project\n",
    "        self.species = species\n",
    "        self.data_path = data_path\n",
    "        self.multiplexed = False\n",
    "        self.spatial = False\n",
    "        self.pca_n = 0\n",
    "        self.filters = {}\n",
    "\n",
    "    @classmethod\n",
    "    def load(cls, filename):\n",
    "        with open(filename, 'rb') as fh:\n",
    "            data = pickle.load(fh)\n",
    "            if data.SCS_VER != SingleCellSample.SCS_VER:\n",
    "                print(f'WARNING: Loaded SCS_VER ({data.SCS_VER} does not equal notebook version ({SingleCellSample.SCS_VER})')\n",
    "            return data     \n",
    "        \n",
    "    @classmethod\n",
    "    def read_10x(cls, project, sample, species, data_path, **adata_kwargs):\n",
    "        adata = sc.read_10x_mtx(TenXDir, var_names='gene_symbols', cache=True, **adata_kwargs)\n",
    "        scs = SingleCellSample(project, sample, species, adata, data_path=data_path)\n",
    "        return scs\n",
    "        \n",
    "        \n",
    "    def save(self, filename):\n",
    "        if not filename.endswith('pickle'):\n",
    "            filename = f'{filename}.pickle'\n",
    "        with open(filename, 'wb') as fh:\n",
    "            pickle.dump(self, fh)\n",
    "\n",
    "    def save_h5ad(self, filename_prefix, save_real_raw=False):\n",
    "        if filename_prefix.endswith('.h5ad'):\n",
    "            filename_prefix = filename_prefix.replace('.h5ad', '')\n",
    "        self.adata.write(f'{filename_prefix}.h5ad')\n",
    "        if save_real_raw:\n",
    "            self.raw.write(f'{filename_prefix}_raw.h5ad')\n",
    "            \n",
    "\n",
    "    def preprocess_adata(self):\n",
    "        \"\"\" Preprocess data on basic level.\n",
    "\n",
    "        Fix any known issues for gene names (i.e. nan)\n",
    "\n",
    "        \"\"\"\n",
    "\n",
    "        self.adata.var_names_make_unique()\n",
    "\n",
    "        if 'gene_ids' not in self.adata.var.keys():\n",
    "            self.adata.var['gene_ids'] = self.adata.var_names\n",
    "\n",
    "        # This is an issue in Drosophila genomes. The symbol for 'nanchung' is 'nan' and gets cast to np.NaN\n",
    "        if np.NaN in self.adata.var['gene_ids']:\n",
    "            nan_loc = self.adata.var.index.get_loc(np.NaN)\n",
    "            self.adata.var['gene_ids'][nan_loc] = 'nan'\n",
    "            newIx = list(self.adata.var.index)\n",
    "            newIx[nan_loc] = 'nan'\n",
    "            self.adata.var.index = newIx\n",
    "\n",
    "\n",
    "    def append_sample_names(self):\n",
    "        \"\"\" Append sample name to cells for easier identification later \"\"\"\n",
    "        self.adata.obs.index = ['-'.join([x.split('-')[0], '__'.join(self.sample.split('__')[-2:])]) for x in self.adata.obs.index]\n",
    "                  \n",
    "\n",
    "    def calc_qc(self, mito_identifier):\n",
    "                \n",
    "        self.adata.var['mt'] = self.adata.var_names.str.startswith(mito_identifier)\n",
    "        sc.pp.calculate_qc_metrics(self.adata, qc_vars=['mt'], percent_top=None, log1p=False, inplace=True)\n",
    "    \n",
    "    def plot_qc(self):\n",
    "        \n",
    "        sc.pl.violin(self.adata, ['n_genes_by_counts', 'total_counts', 'pct_counts_mt'],\n",
    "                     jitter=0.4, multi_panel=True)\n",
    "        \n",
    "        fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(16, 4))\n",
    "                \n",
    "        sc.pl.scatter(self.adata, x='total_counts', y='pct_counts_mt', ax=ax1, show=False)\n",
    "        sc.pl.scatter(self.adata, x='total_counts', y='n_genes_by_counts', ax=ax2, show=False)\n",
    "        plt.tight_layout()\n",
    "\n",
    "    def filter_cells(self, min_genes = 1000, min_cells = 3, n_genes = 7000, percent_mito = 7.5, n_counts = 0, mito_identifier = 'MT-'):\n",
    "        \n",
    "        if 'mt' not in self.adata.obs.columns:\n",
    "            self.calc_qc(mito_identifier=mito_identifier)\n",
    "        \n",
    "        if percent_mito < 1:\n",
    "            print(f'WARNING: Automatically adjusting percent_mito... {percent_mito} -> {percent_mito * 100}')\n",
    "            percent_mito = percent_mito * 100\n",
    "                    \n",
    "        \"\"\" Perform basic filtering steps \"\"\"\n",
    "        self.filters['min_genes'] = min_genes\n",
    "        self.filters['min_cells'] = min_cells\n",
    "        self.filters['n_genes'] = n_genes\n",
    "        self.filters['percent_mito'] = percent_mito\n",
    "        \n",
    "        sc.pp.filter_cells(self.adata, min_genes=min_genes)\n",
    "        sc.pp.filter_genes(self.adata, min_cells=min_cells)\n",
    "\n",
    "        self.adata = self.adata[self.adata.obs['n_genes_by_counts'] < n_genes, :]\n",
    "        self.adata = self.adata[self.adata.obs['pct_counts_mt'] < percent_mito, :]\n",
    "        if n_counts > 0:\n",
    "            self.adata = self.adata[self.adata.obs['total_counts'] < n_counts, :]\n",
    "            self.filters['n_counts'] = n_counts\n",
    "\n",
    "\n",
    "        print(f\"{self.sample} post-filter stats\")\n",
    "        self.plot_qc()\n",
    "\n",
    "        sc.pl.highest_expr_genes(self.adata)\n",
    "\n",
    "                  \n",
    "    def normalize_and_log_transform(self):\n",
    "        print('Storing raw counts in self.raw')\n",
    "        self.raw = self.adata.copy()\n",
    "        sc.pp.normalize_total(self.adata, target_sum=1e4)\n",
    "        sc.pp.log1p(self.adata)\n",
    "        self.adata.raw = self.adata\n",
    "                  \n",
    "    \n",
    "    def detect_high_variable_genes(self):\n",
    "        sc.pp.highly_variable_genes(self.adata, min_mean=0.0125, max_mean=3, min_disp=0.5)\n",
    "        sc.pl.highly_variable_genes(self.adata, )\n",
    "        self.hvgs = self.adata.var.index[self.adata.var['highly_variable']]\n",
    "        self.adata = self.adata[:, self.hvgs]\n",
    "        self.filters['highly_variable_genes'] = True\n",
    "\n",
    "                  \n",
    "    def regress_and_scale(self, features = ['total_counts', 'pct_counts_mt'], n_jobs=None):\n",
    "        for feature in features:\n",
    "            if feature not in self.adata.obs.columns:\n",
    "                print(f'Removing feature {feature} as it is not in observations')\n",
    "                features.remove(feature)\n",
    "        sc.pp.regress_out(self.adata, features, n_jobs=n_jobs)\n",
    "        sc.pp.scale(self.adata, max_value=10)\n",
    "        \n",
    "    \n",
    "    def perform_pca(self, max_pcs = 150, colour = None):\n",
    "        sc.tl.pca(self.adata, svd_solver='arpack', n_comps=max_pcs)\n",
    "\n",
    "        # Adjust the figure size to show 150 components correctly\n",
    "        _old_figsize = rcParams['figure.figsize']\n",
    "        rcParams['figure.figsize'] = (12.0, 4.0)\n",
    "        sc.pl.pca_variance_ratio(self.adata, log=True, n_pcs=max_pcs)\n",
    "        rcParams['figure.figsize'] = _old_figsize\n",
    "\n",
    "        sc.pl.pca(self.adata, color = colour)\n",
    "        \n",
    "        \n",
    "    def reduce_dimensions(self, n_jobs=None, use_paga_init=False, use_harmony=True):\n",
    "        random.seed('cbd-bioinf')\n",
    "        if 'X_pca_harmony' in self.adata.obsm and use_harmony:\n",
    "            print('Using Harmony corrected PCA results. Pass \"use_harmony=False\" to disable.')\n",
    "            if self.pca_n > 0:\n",
    "                print(f\"Creating self.adata.obsm['X_pca_harmony_trimmed'] with {self.pca_n} components\")\n",
    "                self.adata.obsm['X_pca_harmony_trimmed'] = self.adata.obsm['X_pca_harmony'][:, :self.pca_n]\n",
    "                rep_to_use = 'X_pca_harmony_trimmed'\n",
    "            else:\n",
    "                rep_to_use = 'X_pca_harmony'\n",
    "        else:\n",
    "            rep_to_use = None\n",
    "        sc.pp.neighbors(self.adata, n_neighbors=15, n_pcs=self.pca_n, use_rep=rep_to_use)\n",
    "        if use_paga_init:\n",
    "            sc.tl.paga(self.adata)\n",
    "            sc.pl.paga(self.adata, plot=False)\n",
    "            sc.tl.umap(self.adata, init_pos='paga')\n",
    "        else:\n",
    "            sc.tl.umap(self.adata)\n",
    "        sc.tl.tsne(self.adata, n_pcs=self.pca_n, n_jobs=n_jobs, use_rep=rep_to_use)\n",
    "        \n",
    "\n",
    "    def detect_clusters(self, res=None, method='leiden', calc_markers=True, diff_method='wilcoxon', n_genes=None):\n",
    "        key_added = f'{method}_res{res}' if res != None else None\n",
    "\n",
    "        if method == 'leiden':\n",
    "            sc.tl.leiden(self.adata, resolution = res, key_added=key_added)\n",
    "        elif method == 'louvain':\n",
    "            sc.tl.louvain(self.adata, resolution = res, key_added=key_added)\n",
    "        if calc_markers:\n",
    "            sc.tl.rank_genes_groups(self.adata, key_added, method=diff_method, \n",
    "                                    n_genes=self.adata.raw.shape[1] if n_genes == None else n_genes,\n",
    "                                    key_added=f'rank_genes_groups_{key_added}')\n",
    "\n",
    "        \n",
    "    def plot_qc(self):\n",
    "        \n",
    "        sc.pl.violin(self.adata, ['n_genes_by_counts', 'total_counts', 'pct_counts_mt'],\n",
    "                     jitter=0.4, multi_panel=True)\n",
    "        \n",
    "        fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(16, 4))\n",
    "                \n",
    "        sc.pl.scatter(self.adata, x='total_counts', y='pct_counts_mt', ax=ax1, show=False)\n",
    "        sc.pl.scatter(self.adata, x='total_counts', y='n_genes_by_counts', ax=ax2, show=False)\n",
    "        plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01d15a26-3002-4c00-818b-b5209cd19971",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# This should match a pre-determined project code from https://docs.google.com/spreadsheets/d/1ctn8ULH3T6YHzGdxXDVAO6YG6E2qQN-25tJ7z4zyVgk\n",
    "project = 'NovaST'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3abd00a-3a0f-43b4-b4cc-2bc89f0589da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "species = \"Mus musculus\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e4731ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "adatas = {}\n",
    "\n",
    "\n",
    "nst_datas = load_stereoseq_data(\n",
    "    project,\n",
    "    'StereoSeq_MouseBrain',\n",
    "    species,\n",
    "    f\"StereoSeq_20231121/00.RAW/SS200000366BL_E4/Downsampled/SS200000366BL_E4_Saw_Output_v7/02.count/SS200000366BL_E4.raw.gef\",\n",
    "    bin_sizes=[\n",
    "        50,\n",
    "        100,\n",
    "        200\n",
    "    ],  # Bin 200\n",
    ")\n",
    "for k, v in nst_datas.items():\n",
    "    adatas[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62f6bad6-9193-4fd1-bf81-a4db84bf4c40",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    data.preprocess_adata()\n",
    "    data.adata.var_names_make_unique()\n",
    "    data.append_sample_names()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87492032-3460-4b37-b557-941d648b803a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Human: MT- | Mouse: mt- | Drosophila: mt:\n",
    "mito_identifier = 'mt-'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8544c3cb-a3a8-4266-a531-ad2120a4fbee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    data.calc_qc(mito_identifier=mito_identifier)\n",
    "    data.plot_qc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8196244e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    sc.pl.embedding(data.adata, color='total_counts', basis='spatial')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c61f2ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 4))\n",
    "    sp_1_counts = sorted(zip(data.adata.obsm['spatial'][:, 0], data.adata.obs['total_counts']), key=lambda x: x[0])\n",
    "    sp_2_counts = sorted(zip(data.adata.obsm['spatial'][:, 1], data.adata.obs['total_counts']), key=lambda x: x[0])\n",
    "    # average counts per coord\n",
    "    sp_1_counts = [(coord, np.sum([x[1] for x in sp_1_counts if x[0] == coord])) for coord in set([x[0] for x in sp_1_counts])]\n",
    "    sp_2_counts = [(coord, np.sum([x[1] for x in sp_2_counts if x[0] == coord])) for coord in set([x[0] for x in sp_2_counts])]\n",
    "    \n",
    "    ax1.scatter([x[0] for x in sp_1_counts], [x[1] for x in sp_1_counts])\n",
    "    ax2.scatter([x[0] for x in sp_2_counts], [x[1] for x in sp_2_counts])\n",
    "\n",
    "    # title\n",
    "    plt.suptitle(sample)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10431ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    data.adata = data.adata[np.logical_and(\n",
    "        np.logical_and(data.adata.obsm['spatial'][:, 0] > 7000, data.adata.obsm['spatial'][:, 0] < 21000),\n",
    "        np.logical_and(data.adata.obsm['spatial'][:, 1] > 2000, data.adata.obsm['spatial'][:, 1] < 23000)\n",
    "    )]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac5ee621",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    sc.pl.embedding(data.adata, color=['total_counts', 'n_genes_by_counts'], basis='spatial')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d03cc670-6c56-4360-acd5-ad080b2d505f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    print(sample)\n",
    "    print(data.adata.obs['total_counts'].median(), data.adata.obs['n_genes_by_counts'].median())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d0c167-aec3-4fbb-98b1-5565bf58e553",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    sc.pl.violin(data.adata, ['n_genes_by_counts', 'total_counts', 'pct_counts_mt'], jitter=False, multi_panel=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aa3d807-8a76-47a7-9fb7-b97875098ced",
   "metadata": {},
   "outputs": [],
   "source": [
    "gene_thresholds = {\n",
    "    'bin50': 200,\n",
    "    'bin100': 750,\n",
    "    'bin200': 2250,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef72770-b049-4a21-9c7c-689d0d9fce93",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    bin_size = sample.split('_')[-1]\n",
    "    data.filter_cells(n_counts = 1000000, n_genes=25000, percent_mito=100, min_genes=gene_thresholds[bin_size], mito_identifier=mito_identifier)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20edce98",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    print(data.adata.obs['total_counts'].median(), data.adata.obs['n_genes_by_counts'].median())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "533904e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "umi_dfs = {}\n",
    "gene_dfs = {}\n",
    "\n",
    "for sample in ['StereoSeq_MouseBrain']:\n",
    "    umi_df = pd.DataFrame(columns=['bin', 'count'])\n",
    "    gene_df = pd.DataFrame(columns=['bin', 'count'])\n",
    "    for bin_size in [50, 100, 200]:\n",
    "        umi_df = umi_df.append(pd.DataFrame({\n",
    "            'bin': [f'Bin {bin_size}'] * adatas[f'{sample}_bin{bin_size}'].adata.shape[0],\n",
    "            'count': adatas[f'{sample}_bin{bin_size}'].adata.obs['total_counts'].values\n",
    "        }), ignore_index=True)\n",
    "        gene_df = gene_df.append(pd.DataFrame({\n",
    "            'bin': [f'Bin {bin_size}'] * adatas[f'{sample}_bin{bin_size}'].adata.shape[0],\n",
    "            'count': adatas[f'{sample}_bin{bin_size}'].adata.obs['n_genes_by_counts'].values\n",
    "        }), ignore_index=True)\n",
    "        umi_df['bin'] = umi_df['bin'].astype('category')\n",
    "        gene_df['bin'] = gene_df['bin'].astype('category')\n",
    "        umi_df['count'] = umi_df['count'].astype('int')\n",
    "        gene_df['count'] = gene_df['count'].astype('int')\n",
    "    umi_dfs[sample] = umi_df\n",
    "    gene_dfs[sample] = gene_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9ef4ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample in ['StereoSeq_MouseBrain']:\n",
    "    sns.set_style('white')\n",
    "    df1 = gene_dfs[sample]\n",
    "    df2 = umi_dfs[sample]\n",
    "    df1['metric'] = 1\n",
    "    df2['metric'] = 2\n",
    "    fig, axs = plt.subplots(1, 3, figsize=(15, 5))\n",
    "\n",
    "    df1['bin'] = df1['bin'].astype('category')\n",
    "    df2['bin'] = df2['bin'].astype('category')\n",
    "\n",
    "    for i, bin_size in enumerate([50, 100, 200]):\n",
    "\n",
    "        df_g = df1[df1['bin'] == f'Bin {bin_size}']\n",
    "        df_u = df2[df2['bin'] == f'Bin {bin_size}']\n",
    "\n",
    "        df_g_dummy = pd.DataFrame({'bin': pd.Categorical.from_codes([0], [f'Bin {bin_size}']), 'count': [np.nan]})\n",
    "        df_u_dummy = pd.DataFrame({'bin': pd.Categorical.from_codes([0], [f'Bin {bin_size}']), 'count': [np.nan]})\n",
    "        df_g_dummy['metric'] = 2\n",
    "        df_u_dummy['metric'] = 1\n",
    "\n",
    "        ax = axs[i]\n",
    "        sns.violinplot(x='bin', \n",
    "                       y='count', \n",
    "                       hue='metric', \n",
    "                       split=True, \n",
    "                       data=pd.concat([df_g, df_g_dummy]), \n",
    "                       ax=ax\n",
    "                       )\n",
    "        ax.legend_.remove()\n",
    "        ax.set_ylabel('Gene count')\n",
    "        \n",
    "        ax2 = ax.twinx()\n",
    "        sns.violinplot(x='bin', \n",
    "                       y='count', \n",
    "                       hue='metric', \n",
    "                       split=True, \n",
    "                       data=pd.concat([df_u, df_u_dummy]), \n",
    "                       ax=ax2\n",
    "                       )\n",
    "\n",
    "        ax2.set_ylabel('UMI count')\n",
    "        sns.despine(top=True, right=False, left=False, bottom=True, ax=ax)\n",
    "        sns.despine(top=True, right=False, left=False, bottom=True, ax=ax2)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(f'{sample}_StdScale.pdf')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23ae57cd-d133-4d99-8ded-94ed531973ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize=(36, 12))\n",
    "for n, (sample, data) in enumerate(adatas.items()):\n",
    "    sc.pl.embedding(data.adata, color='n_genes_by_counts', basis='spatial', ax=axs[n], s=13 * ((n+1)*3), show=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8f7f151-2844-442b-aa0a-3647fd8f51a1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    data.normalize_and_log_transform()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaa94ded-9503-4e5f-8cc4-8f955c10fe78",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    data.detect_high_variable_genes()\n",
    "    data.regress_and_scale(n_jobs=24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2782050-2c65-49bc-96c6-7f3de4504eb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    data.perform_pca()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dae364d",
   "metadata": {},
   "outputs": [],
   "source": [
    "adatas['StereoSeq_MouseBrain_bin50'].pca_n = 22\n",
    "adatas['StereoSeq_MouseBrain_bin100'].pca_n = 27\n",
    "adatas['StereoSeq_MouseBrain_bin200'].pca_n = 35"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67fb9452-3dc3-4bc4-8980-bd46fd3afa5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    data.reduce_dimensions(n_jobs=24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fde7a45-bef7-4067-87a6-8f3c42a3fec2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    for res in [0.8, 1.0, 2.0, 4.0, 8.0]:\n",
    "        data.detect_clusters(res=res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "474545aa-f471-4b6f-b76d-b2dfea5cb45e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    print(sample)\n",
    "    sc.pl.umap(data.adata, color=['total_counts', 'n_genes_by_counts', 'pct_counts_mt', 'leiden_res1.0'], wspace=0.35, ncols=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d580afdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    print(sample)\n",
    "    sc.pl.embedding(data.adata, basis='spatial', color=['total_counts', 'n_genes_by_counts', 'pct_counts_mt', 'leiden_res1.0'], wspace=0.35, ncols=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9da187ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    # rotate spatial 90 degrees clockwise, x, y -> y, -x\n",
    "    data.adata.obsm['spatial'][:, 0] = -data.adata.obsm['spatial'][:, 0]\n",
    "    data.adata.obsm['spatial'] = np.flip(data.adata.obsm['spatial'], axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "039b9626-eb61-4b5c-bffe-6a1d48f01da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    print(sample)\n",
    "    sc.pl.embedding(data.adata, basis='spatial', color=['total_counts', 'n_genes_by_counts', 'pct_counts_mt', 'leiden_res1.0'], wspace=0.35, ncols=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6cfbf4d-999a-46ff-b70f-bb77693a59aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    print(sample)\n",
    "    sc.pl.embedding(data.adata, basis='spatial', color=['leiden_res1.0', 'leiden_res2.0', 'leiden_res4.0'], wspace=0.35, ncols=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64008f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96afd4a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(columns=['Count', 'Type', 'Bin Size'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb827aa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n, (sample, data) in enumerate(adatas.items()):\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(12, 6))\n",
    "    sc.pl.violin(data.adata, keys=['n_genes_by_counts'], jitter=False, show=False, ax=axs[0])\n",
    "    sc.pl.violin(data.adata, keys=['total_counts'], jitter=False, show=False, ax=axs[1])\n",
    "    fig.suptitle(sample)\n",
    "    axs[0].set_ylim(0, 4000 * (n+1))\n",
    "    axs[1].set_ylim(0, 40000 * (n+1))\n",
    "    plt.tight_layout()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "803d8549",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    bin_size = sample.split('_')[-1]\n",
    "    df = df.append(pd.DataFrame(data={'Count': data.adata.obs['n_genes_by_counts'], 'Type': 'Genes', 'Bin Size': bin_size}))\n",
    "    df = df.append(pd.DataFrame(data={'Count': data.adata.obs['total_counts'], 'Type': 'UMIs', 'Bin Size': bin_size}))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb1277d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['Count'] = df['Count'].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae1e410",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.sort_values(by=['Type'], inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d3b239",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.violinplot(data=df[df['Type'] == 'Genes'], x=\"Bin Size\", y=\"Count\", inner=\"quart\")\n",
    "plt.title('Genes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e41a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.violinplot(data=df[df['Type'] == 'UMIs'], x=\"Bin Size\", y=\"Count\", inner=\"quart\")\n",
    "plt.title('UMIs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "210276ea-88aa-4e2c-8533-034147136c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import loompy as lp\n",
    "import datetime\n",
    "\n",
    "\n",
    "def dfToNamedMatrix(df):\n",
    "    arr_ip = [tuple(i) for i in df.values]\n",
    "    dtyp = np.dtype(list(zip(df.dtypes.index, df.dtypes)))\n",
    "    arr = np.array(arr_ip, dtype=dtyp)\n",
    "    return arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6cebc6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs_dir = './outputs/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce30460a-8b1d-4281-b39b-09474d8338fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, data in adatas.items():\n",
    "    \n",
    "    if data.spatial:\n",
    "        mainEmbedding = pd.DataFrame(data.adata.obsm['spatial'], columns=['_X', '_Y'])\n",
    "\n",
    "        Embeddings_X = pd.DataFrame()\n",
    "        Embeddings_Y = pd.DataFrame()\n",
    "        \n",
    "        Embeddings_X[\"1\"] = pd.DataFrame(data.adata.obsm['X_umap'])[0]\n",
    "        Embeddings_Y[\"1\"] = pd.DataFrame(data.adata.obsm['X_umap'])[1]\n",
    "        \n",
    "        Embeddings_X[\"2\"] = pd.DataFrame(data.adata.obsm['X_tsne'])[0]\n",
    "        Embeddings_Y[\"2\"] = pd.DataFrame(data.adata.obsm['X_tsne'])[1]\n",
    "\n",
    "        Embeddings_X[\"3\"] = pd.DataFrame(data.adata.obsm['X_pca'])[0]\n",
    "        Embeddings_Y[\"3\"] = pd.DataFrame(data.adata.obsm['X_pca'])[1]\n",
    "\n",
    "        metaJson = {}\n",
    "\n",
    "        metaJson['embeddings'] = [\n",
    "            {\n",
    "                \"id\": -1,\n",
    "                \"name\": \"Spatial\"\n",
    "            },\n",
    "            {\n",
    "                \"id\": 1,\n",
    "                \"name\": f\"Scanpy UMAP {data.pca_n}PC\"\n",
    "            },\n",
    "            {\n",
    "                \"id\": 2,\n",
    "                \"name\": f\"Scanpy t-SNE {data.pca_n}PC\"\n",
    "            },\n",
    "            {\n",
    "                \"id\": 3,\n",
    "                \"name\": \"Scanpy PC1/PC2\"\n",
    "            }\n",
    "        ]\n",
    "    else:\n",
    "        mainEmbedding = pd.DataFrame(data.adata.obsm['X_umap'], columns=['_X', '_Y'])\n",
    "\n",
    "        Embeddings_X = pd.DataFrame()\n",
    "        Embeddings_Y = pd.DataFrame()\n",
    "\n",
    "        Embeddings_X[\"1\"] = pd.DataFrame(data.adata.obsm['X_tsne'])[0]\n",
    "        Embeddings_Y[\"1\"] = pd.DataFrame(data.adata.obsm['X_tsne'])[1]\n",
    "\n",
    "        Embeddings_X[\"2\"] = pd.DataFrame(data.adata.obsm['X_pca'])[0]\n",
    "        Embeddings_Y[\"2\"] = pd.DataFrame(data.adata.obsm['X_pca'])[1]\n",
    "\n",
    "        metaJson = {}\n",
    "\n",
    "        metaJson['embeddings'] = [\n",
    "            {\n",
    "                \"id\": -1,\n",
    "                \"name\": f\"Scanpy UMAP {data.pca_n}PC\"\n",
    "            },\n",
    "            {\n",
    "                \"id\": 1,\n",
    "                \"name\": f\"Scanpy t-SNE {data.pca_n}PC\"\n",
    "            },\n",
    "            {\n",
    "                \"id\": 2,\n",
    "                \"name\": \"Scanpy PC1/PC2\"\n",
    "            }\n",
    "        ]\n",
    "    \n",
    "    \n",
    "    col_attrs = {\"CellID\": np.array(data.adata.obs.index),\n",
    "                 \"nUMI\": np.array(data.adata.obs['total_counts'].values),\n",
    "                 \"nGene\": np.array(data.adata.obs['n_genes_by_counts'].values),\n",
    "                 \"Percent_mito\": np.array(data.adata.obs['pct_counts_mt'].values),\n",
    "                 \"Embedding\": dfToNamedMatrix(mainEmbedding),\n",
    "                 \"Embeddings_X\": dfToNamedMatrix(Embeddings_X),\n",
    "                 \"Embeddings_Y\": dfToNamedMatrix(Embeddings_Y)\n",
    "                }\n",
    "    \n",
    "    if not data.spatial:\n",
    "        col_attrs[\"Doublet_score\"] = np.array(data.adata.obs['doublet_score'].values)\n",
    "    \n",
    "    row_attrs = {\"Gene\": np.array(data.raw.var.index)}\n",
    "    \n",
    "    metaJson[\"clusterings\"] = []\n",
    "    clusterings = pd.DataFrame()\n",
    "    \n",
    "    clusterings_to_process = []    \n",
    "    for key in data.adata.obs.columns:\n",
    "        if 'leiden' in key or 'louvain' in key:\n",
    "            clusterings_to_process.append(key)\n",
    "            \n",
    "        \n",
    "\n",
    "#     for n, res in enumerate([1.0, 2.0]):\n",
    "    for n, key in enumerate(clusterings_to_process):\n",
    "        if n == 1:\n",
    "            col_attrs[\"ClusterID\"] = np.array(data.adata.obs[key].values)\n",
    "\n",
    "        ClusterMarkers = pd.DataFrame(index=data.raw.var.index, columns=[str(x) for x in range(max(set([int(x) for x in data.adata.obs[key]])) + 1)])\n",
    "        ClusterMarkers_avg_logFC = pd.DataFrame(index=data.raw.var.index, columns=[str(x) for x in range(max(set([int(x) for x in data.adata.obs[key]])) + 1)])\n",
    "        ClusterMarkers_pval = pd.DataFrame(index=data.raw.var.index, columns=[str(x) for x in range(max(set([int(x) for x in data.adata.obs[key]])) + 1)])\n",
    "\n",
    "        ClusterMarkers.fillna(0, inplace=True)\n",
    "        ClusterMarkers_avg_logFC.fillna(0, inplace=True)\n",
    "        ClusterMarkers_pval.fillna(0, inplace=True)\n",
    "\n",
    "        for i in range(max(set([int(x) for x in data.adata.obs[key]])) + 1):\n",
    "            i = str(i)\n",
    "            tot_genes = len(data.adata.uns[f'rank_genes_groups_{key}']['pvals_adj'][i])\n",
    "            sigGenes = data.adata.uns[f'rank_genes_groups_{key}']['pvals_adj'][i] < 0.05\n",
    "            deGenes = np.logical_and(np.logical_or(data.adata.uns[f'rank_genes_groups_{key}']['logfoldchanges'][i] >= 1.5, data.adata.uns[f'rank_genes_groups_{key}']['logfoldchanges'][i] <= -1.5), np.isfinite(data.adata.uns[f'rank_genes_groups_{key}']['logfoldchanges'][i]))\n",
    "            sigAndDE = np.logical_and(sigGenes, deGenes)\n",
    "\n",
    "            names = data.adata.uns[f'rank_genes_groups_{key}']['names'][i][sigAndDE]\n",
    "            ClusterMarkers.loc[names, i] = 1\n",
    "            ClusterMarkers_avg_logFC.loc[names, i] = np.around(data.adata.uns[f'rank_genes_groups_{key}']['logfoldchanges'][i][sigAndDE], decimals=6)\n",
    "            ClusterMarkers_pval.loc[names, i] = np.around(data.adata.uns[f'rank_genes_groups_{key}']['pvals_adj'][i][sigAndDE], decimals=6)\n",
    "            \n",
    "        row_attrs[f\"ClusterMarkers_{n}\"] = dfToNamedMatrix(ClusterMarkers)\n",
    "        row_attrs[f\"ClusterMarkers_{n}_avg_logFC\"] = dfToNamedMatrix(ClusterMarkers_avg_logFC)\n",
    "        row_attrs[f\"ClusterMarkers_{n}_pval\"] = dfToNamedMatrix(ClusterMarkers_pval)\n",
    "        cluster_meta = {\n",
    "                \"id\": n,\n",
    "                \"group\": \"Scanpy\",\n",
    "                \"name\": f\"Scanpy {key}\",\n",
    "                \"clusters\": [],\n",
    "                \"clusterMarkerMetrics\": [{\n",
    "                    \"accessor\": \"avg_logFC\", \n",
    "                    \"name\": \"Avg. logFC\",\n",
    "                    \"description\": \"Average log fold change from Wilcox test\"\n",
    "                }, {\n",
    "                    \"accessor\": \"pval\", \n",
    "                    \"name\": \"Adj. P-Value\",\n",
    "                    \"description\": \"Adj. P-Value from Wilcox test\"\n",
    "                }\n",
    "                ]\n",
    "            }\n",
    "        metaJson[\"clusterings\"].append(cluster_meta)\n",
    "        \n",
    "        for i in range(max(set([int(x) for x in data.adata.obs[key]])) + 1):\n",
    "            clustDict = {}\n",
    "            clustDict['id'] = i\n",
    "            clustDict['description'] = f\"Unannotated Cluster {str(i+1)}\"\n",
    "            metaJson['clusterings'][n]['clusters'].append(clustDict)\n",
    "\n",
    "        clusterings[str(n)] = data.adata.obs[key].values.astype(np.int64)\n",
    "    col_attrs['Clusterings'] = dfToNamedMatrix(clusterings)\n",
    "\n",
    "    metaJson[\"metrics\"] = [\n",
    "            {\n",
    "                \"name\": \"nUMI\"\n",
    "            }, {\n",
    "                \"name\": \"nGene\"\n",
    "            }, {\n",
    "                \"name\": \"Percent_mito\"\n",
    "            }\n",
    "    ]\n",
    "    \n",
    "    if not data.spatial:\n",
    "        metaJson[\"metrics\"].append(\n",
    "            {\n",
    "                \"name\": \"Doublet_score\"\n",
    "            })\n",
    "\n",
    "    metaJson[\"annotations\"] = []\n",
    "    \n",
    "    for key in []:\n",
    "        if key in data.adata.obs.keys():\n",
    "            col_attrs[key] = np.array([str(x) for x in data.adata.obs[key].values])\n",
    "            metaJson[\"annotations\"].append(\n",
    "            {\n",
    "                \"name\": key,\n",
    "                \"values\": [str(x) for x in sorted(list(set(data.adata.obs[key].values)))]\n",
    "            }\n",
    "            )\n",
    "        else:\n",
    "            print(f'{key} not found in {sample}!')\n",
    "    \n",
    "    \n",
    "    attrs = {\"title\": sample,\n",
    "             \"MetaData\": json.dumps(metaJson),\n",
    "             \"Genome\": '',\n",
    "             \"SCopeTreeL1\": \"Deep_Seq\",\n",
    "             \"SCopeTreeL2\": \"\",\n",
    "             \"SCopeTreeL3\": \"\"\n",
    "            }\n",
    "\n",
    "        \n",
    "    attrs['MetaData'] = json.dumps(metaJson)\n",
    "    lp.create(filename=f'{outputs_dir}/{datetime.date.today().strftime(\"%Y%m%d\")}_{sample}.loom', layers=data.raw.X.T.todense(), row_attrs=row_attrs, col_attrs=col_attrs, file_attrs=attrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43cf0ec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 3, figsize=(15, 10))\n",
    "ax_n = 0\n",
    "for sample, data in adatas.items():\n",
    "    if not sample.endswith(\"bin2912\"):\n",
    "        continue\n",
    "    ax = axs[ax_n // 3, ax_n % 3]\n",
    "    ax.set_title(f'{sample}\\n{int(data.adata.obs[\"total_counts\"].sum()) / 1000000:.2f}M UMIs')\n",
    "    sc.pl.violin(data.adata, \"n_genes_by_counts\", ax=ax, show=False)\n",
    "    ax_n += 1\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3.8.10/R4.2.1 - CBD Bioinf Analysis v4",
   "language": "Python",
   "name": "cbd_bioinf_v4"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}