{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scanpy==1.6.0 anndata==0.7.4 umap==0.4.6 numpy==1.18.5 scipy==1.4.1 pandas==1.1.3 scikit-learn==0.23.2 statsmodels==0.12.0 python-igraph==0.7.1 leidenalg==0.8.2\n"
     ]
    }
   ],
   "source": [
    "# import the transcriptome data\n",
    "import os\n",
    "import pandas as pd\n",
    "import scanpy as sc\n",
    "sc.logging.print_header()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/lustre/scratch117/cellgen/team297/kt16/newcastle_covid/notebooks'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir('/lustre/scratch117/cellgen/team297/kt16/newcastle_covid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 74437 × 1051\n",
       "    obs: 'Site', 'doublet', 'patient_id', 'sample_id', 'batch', 'Resample', 'Collection_Day', 'Sex', 'Age', 'Swab_result', 'Status', 'Smoker', 'Status_on_day_collection', 'Status_on_day_collection_summary', 'Days_from_onset', 'time_after_LPS', 'Worst_Clinical_Status', 'Outcome', 'leiden', 'consensus', 'initial_clustering', 'study_id', 'initial_clustering_B', 'n_genes_by_counts', 'total_counts', 'total_counts_mt', 'pct_counts_mt', 'leiden_B', 'celltype_B', 'celltype_B_v2', 'Status_on_day_collection_summary_v2', 'filter_rna', 'has_bcr', 'filter_bcr_quality', 'filter_bcr_heavy', 'filter_bcr_light', 'bcr_QC_pass', 'filter_bcr', 'leiden_B2'\n",
       "    var: 'feature_types', 'mt', 'n_cells_by_counts', 'mean_counts', 'pct_dropout_by_counts', 'total_counts', 'highly_variable', 'means', 'dispersions', 'dispersions_norm', 'mean', 'std'\n",
       "    uns: 'Site_colors', 'Status_on_day_collection_summary_colors', 'Status_on_day_collection_summary_v2_colors', 'bcr_QC_pass_colors', 'celltype_B_colors', 'celltype_B_v2_colors', 'consensus_colors', \"dendrogram_['leiden_B']\", 'filter_bcr_colors', 'hvg', 'initial_clustering_colors', 'leiden', 'leiden_B2_colors', 'leiden_B_colors', 'leiden_colors', 'neighbors', 'pca', 'rank_genes_groups', 'umap'\n",
       "    obsm: 'X_pca', 'X_pca_harmony', 'X_umap'\n",
       "    varm: 'PCs'\n",
       "    obsp: 'connectivities', 'distances'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Read in B cell data\n",
    "bdata = sc.read_h5ad('h5ad/covid_jan_2021_bcells.h5ad')\n",
    "bdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/nfs/users/nfs_k/kt16/miniconda2/envs/dandelion/lib/python3.7/site-packages/anndata/_core/anndata.py:1094: FutureWarning: is_categorical is deprecated and will be removed in a future version.  Use is_categorical_dtype instead\n",
      "  if not is_categorical(df_full[k]):\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 64265 × 1051\n",
       "    obs: 'Site', 'doublet', 'patient_id', 'sample_id', 'batch', 'Resample', 'Collection_Day', 'Sex', 'Age', 'Swab_result', 'Status', 'Smoker', 'Status_on_day_collection', 'Status_on_day_collection_summary', 'Days_from_onset', 'time_after_LPS', 'Worst_Clinical_Status', 'Outcome', 'leiden', 'consensus', 'initial_clustering', 'study_id', 'initial_clustering_B', 'n_genes_by_counts', 'total_counts', 'total_counts_mt', 'pct_counts_mt', 'leiden_B', 'celltype_B', 'celltype_B_v2', 'Status_on_day_collection_summary_v2', 'filter_rna', 'has_bcr', 'filter_bcr_quality', 'filter_bcr_heavy', 'filter_bcr_light', 'bcr_QC_pass', 'filter_bcr', 'leiden_B2'\n",
       "    var: 'feature_types', 'mt', 'n_cells_by_counts', 'mean_counts', 'pct_dropout_by_counts', 'total_counts', 'highly_variable', 'means', 'dispersions', 'dispersions_norm', 'mean', 'std'\n",
       "    uns: 'Site_colors', 'Status_on_day_collection_summary_colors', 'Status_on_day_collection_summary_v2_colors', 'celltype_B_colors', 'celltype_B_v2_colors', 'consensus_colors', \"dendrogram_['leiden_B']\", 'filter_bcr_colors', 'hvg', 'initial_clustering_colors', 'leiden', 'leiden_B_colors', 'leiden_colors', 'neighbors', 'pca', 'rank_genes_groups', 'umap'\n",
       "    obsm: 'X_pca', 'X_pca_harmony', 'X_umap'\n",
       "    varm: 'PCs'\n",
       "    obsp: 'connectivities', 'distances'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# subset to just D0 samples\n",
    "bdata = bdata[bdata.obs['Collection_Day'] == 'D0'].copy()\n",
    "# also subset to just healthy and COVID\n",
    "bdata = bdata[~(bdata.obs['Status_on_day_collection_summary_v2'].isin(['LPS', 'Non_covid']))].copy()\n",
    "bdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Moderate        19098\n",
       "Mild            16192\n",
       "Severe          11577\n",
       "Healthy          7680\n",
       "Critical         7051\n",
       "Asymptomatic     2667\n",
       "Name: Status_on_day_collection_summary_v2, dtype: int64"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bdata.obs['Status_on_day_collection_summary_v2'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "def exportDEres(adata, key = 'rank_genes_groups', column = None, filename = None):\n",
    "    scores = pd.DataFrame(data = adata.uns[key]['scores'][column], index = adata.uns[key]['names'][column])\n",
    "    lfc = pd.DataFrame(data = adata.uns[key]['logfoldchanges'][column], index = adata.uns[key]['names'][column])\n",
    "    pvals = pd.DataFrame(data = adata.uns[key]['pvals'][column], index = adata.uns[key]['names'][column])\n",
    "    padj = pd.DataFrame(data = adata.uns[key]['pvals_adj'][column], index = adata.uns[key]['names'][column])\n",
    "    try:\n",
    "        pts = pd.DataFrame(data = adata.uns[key]['pts'][column], index = adata.uns[key]['names'][column])\n",
    "    except:\n",
    "        pass\n",
    "    scores = scores.loc[scores.index.dropna()]\n",
    "    lfc = lfc.loc[lfc.index.dropna()]\n",
    "    pvals = pvals.loc[pvals.index.dropna()]\n",
    "    padj = padj.loc[padj.index.dropna()]\n",
    "    try:\n",
    "        pts = pts.loc[pts.index.dropna()]\n",
    "    except:\n",
    "        pass\n",
    "    try:\n",
    "        dfs = [scores, lfc, pvals, padj, pts]\n",
    "    except:\n",
    "        dfs = [scores, lfc, pvals, padj]\n",
    "    df_final = functools.reduce(lambda left,right: pd.merge(left,right,left_index = True, right_index = True), dfs)\n",
    "    try:\n",
    "        df_final.columns = ['scores', 'logfoldchanges', 'pvals', 'pvals_adj', 'pts']\n",
    "    except:\n",
    "        df_final.columns = ['scores', 'logfoldchanges', 'pvals', 'pvals_adj']\n",
    "\n",
    "    df_final.to_csv(filename, sep = '\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plasma_cell\n",
      "B_exhausted\n",
      "B_naive\n",
      "B_immature\n",
      "B_non-switched_memory\n",
      "B_switched_memory\n",
      "Plasmablast\n"
     ]
    }
   ],
   "source": [
    "for x in list(set(bdata.obs['celltype_B_v2'])):\n",
    "    print(x)\n",
    "    bdatax = bdata[bdata.obs['celltype_B_v2'] == x].copy()\n",
    "    sc.tl.rank_genes_groups(bdatax, groupby = 'Status_on_day_collection_summary_v2', method = 'wilcoxon', reference = 'Healthy', n_genes = 40000)\n",
    "    for y in ['Asymptomatic', 'Mild', 'Moderate', 'Severe', 'Critical']:\n",
    "        filepath = 'DEG/wilcoxon'\n",
    "        if not os.path.exists(filepath):\n",
    "            os.makedirs(filepath)\n",
    "        exportDEres(bdatax, column = y, filename = filepath +'/'+y+'_'+x+'_wilcoxon_v2.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
