{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88df5dc6",
   "metadata": {},
   "source": [
    "# Preprocessing Data\n",
    "\n",
    "This notebook outlines the data preprocessing steps required for scDynOmics.\n",
    "Before starting, ensure your data meets the following formatting requirements:\n",
    "* **Gene Identifiers:** The `adata.var` index must use Ensembl IDs.\n",
    "* **ATAC Annotations:** Peak annotations must follow the `chrN:start-end` format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ccf601f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100 85.0M  100 85.0M    0     0   106M      0 --:--:-- --:--:-- --:--:--  106M\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Data loaded successfully.\n",
      "Raw GEX adata: View of AnnData object with n_obs × n_vars = 4878 × 32284\n",
      "Raw PEAKS adata: View of AnnData object with n_obs × n_vars = 4878 × 172193\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import anndata as ad, hdf5plugin\n",
    "import scdynomics.utils.pp as pp\n",
    "from scdynomics.utils.data.reader import scMultiomics_10x_h5\n",
    "\n",
    "temp_data_name = \"temp.h5\"\n",
    "\n",
    "# Download the sample data from 10x Genomics and save it as \"temp.h5\"\n",
    "os.system(\"curl -O https://cf.10xgenomics.com/samples/cell-arc/2.0.0/e18_mouse_brain_fresh_5k/e18_mouse_brain_fresh_5k_filtered_feature_bc_matrix.h5\")\n",
    "os.rename(\"e18_mouse_brain_fresh_5k_filtered_feature_bc_matrix.h5\", temp_data_name)\n",
    "\n",
    "# Process the data and load it into AnnData objects\n",
    "raw_gex, raw_peaks = scMultiomics_10x_h5(temp_data_name)\n",
    "os.remove(temp_data_name)\n",
    "print(\"\\nData loaded successfully.\")\n",
    "print(f\"Raw GEX adata: {raw_gex}\")\n",
    "print(f\"Raw PEAKS adata: {raw_peaks}\")\n",
    "\n",
    "# Set the reference data paths\n",
    "ref_adata = ad.read_h5ad(\"../../data/gene_dict/mm10.pan_promoter.ind.h5ad\")\n",
    "ref_dict_path = \"../../data/gene_dict/mm10.pan_promoter.dict.json.gz\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3fa341",
   "metadata": {},
   "source": [
    "## 1. Pretraining Pipeline\n",
    "\n",
    "The pretraining phase requires aligning your data to a standard reference. \n",
    "\n",
    "**Expected Input:** Multimodal samples containing two modalities that represent distinct time points.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "812e2feb",
   "metadata": {},
   "source": [
    "### 1.1 Mapping Raw Inputs to Reference Feature Space\n",
    "We first align the raw gene expression and peak data to the provided reference genome to ensure consistent feature spaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b8549151",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ag_klughammer/gyu/.conda/envs/pg_0/lib/python3.11/site-packages/scdynomics/utils/pp/align.py:108: UserWarning: Some chromosomes were not found in the reference and were ignored: {'GL456233.1': 1, 'JH584304.1': 1, 'GL456216.1': 1, 'JH584292.1': 1, 'JH584295.1': 1}\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overview of the processed AnnData: AnnData object with n_obs × n_vars = 4878 × 22452\n",
      "    var: 'name', 'type', 'chr', 'start', 'end', 'strand', 'promoters', 'mt'\n",
      "    layers: 'pan_promoter'\n"
     ]
    }
   ],
   "source": [
    "# Align the raw data to the reference data\n",
    "gex_adata = pp.align.align_genes_to_reference(\n",
    "    raw_gex,\n",
    "    ref_adata = ref_adata,\n",
    "    handle_duplicated = 'sum',\n",
    ")\n",
    "\n",
    "# Map the peaks to the reference promoters\n",
    "atac_adata = pp.align.align_peaks_to_reference(\n",
    "    raw_peaks,\n",
    "    ref_adata = ref_adata,\n",
    "    ref_dict_path = ref_dict_path\n",
    ")\n",
    "\n",
    "# Check\n",
    "assert gex_adata.obs.index.equals(atac_adata.obs.index)\n",
    "assert gex_adata.var.index.equals(atac_adata.var.index)\n",
    "\n",
    "# Combining into same AnnData object\n",
    "pt_adata = ad.AnnData(\n",
    "    X = gex_adata.X,\n",
    "    obs = gex_adata.obs,\n",
    "    var = gex_adata.var,\n",
    "    layers = {\n",
    "        \"pan_promoter\": atac_adata.X\n",
    "    }\n",
    ")\n",
    "\n",
    "print(f\"Overview of the processed AnnData: {pt_adata}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c53a2534",
   "metadata": {},
   "source": [
    "### 1.2 Quality Control (QC) and Normalization\n",
    "Next, we filter out low-quality cells, remove mitochondrial genes, and normalize the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffaf18b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Genes are already aligned to the reference.\n",
      "WARNING: adata.X seems to be already log-transformed.\n",
      "\n",
      "Overview of the adata object: AnnData object with n_obs × n_vars = 4850 × 22439\n",
      "    var: 'name', 'type', 'chr', 'start', 'end', 'strand', 'promoters'\n",
      "    uns: 'log1p'\n",
      "    layers: 'pan_promoter' \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# WARNING: adata.X seems to be already log-transformed. -> it is not, so this can be ignored.\n",
    "# https://github.com/scverse/scanpy/issues/1333\n",
    "pped_adata = pp.pipeline(\n",
    "    pt_adata,\n",
    "    ref_adata = ref_adata,\n",
    "    layers = ['x', 'pan_promoter'],\n",
    "    matching_strategy = 'align',\n",
    "    filtering = True,\n",
    "    min_genes = 100, # Only filtering cells\n",
    "    remove_outliers = True,\n",
    "    log1p = True,\n",
    "    normalizing = True,\n",
    "    target_sum = 1e4,\n",
    "    methods = ['tpm', 'cpm_promoter'],\n",
    "    remove_mt = True, # Remove mitochondrial genes \n",
    ")\n",
    "\n",
    "print(f\"\\nOverview of the adata object: {pped_adata}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea02ba44",
   "metadata": {},
   "source": [
    "## 2. Fine-Tuning Pipeline\n",
    "\n",
    "Preparing data for downstream fine-tuning requires a slightly different, more flexible approach compared to strict pretraining:\n",
    "\n",
    "* **Feature Selection (Optional):** To reduce dimensionality, you may optionally subset the dataset to only include Highly Variable Genes (HVGs) or peaks. If you prefer to retain the full feature space, this step can be bypassed.\n",
    "* **Filtering:** Strict cell and gene filtering is highly dependent on your specific downstream task. However, to maintain computational stability, establishing a baseline filter of `min_genes = 1` and `min_cells = 1` is strongly recommended.\n",
    "* **Flexibility:** The fine-tuning preprocessing pipeline is highly modular. It allows you to customize or skip steps based on the state of your input data. For example, if your dataset has already been normalized using a custom external method, you could skip the normalization step.\n",
    "* **Hyperparameters:** Please ensure you adjust the functional pipeline arguments below to best fit the distribution and specific needs of your dataset.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "540d5915",
   "metadata": {},
   "source": [
    "### 2.1: Multimodal Processing\n",
    "\n",
    "This example uses a 10-cell sample from the [Gastrulation dataset](https://scvelo.readthedocs.io/en/stable/scvelo.datasets.gastrulation_e75.html),\n",
    "formatted as processed RNA velocity with scVelo. \n",
    "\n",
    "*Note: For standard scMultiomics data, you must map the feature space first, as demonstrated in the Pretraining section 1.1 .*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "50764265",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overview of the adata object: AnnData object with n_obs × n_vars = 10 × 53801\n",
      "    obs: 'barcode', 'sample', 'stage', 'sequencing.batch', 'theiler', 'doub.density', 'doublet', 'cluster', 'cluster.sub', 'cluster.stage', 'cluster.theiler', 'stripped', 'celltype', 'colour', 'umapX', 'umapY', 'haem_gephiX', 'haem_gephiY', 'haem_subclust', 'endo_gephiX', 'endo_gephiY', 'endo_trajectoryName', 'endo_trajectoryDPT', 'endo_gutX', 'endo_gutY', 'endo_gutDPT', 'endo_gutCluster', 'cell_velocyto_loom'\n",
      "    var: 'Chromosome', 'End', 'Start', 'Strand', 'name'\n",
      "    obsm: 'X_pca', 'X_umap'\n",
      "    layers: 'spliced', 'unspliced' \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Load the raw data and set the index as the Ensembl gene ID (Accession)\n",
    "ft_adata = ad.read_h5ad(\"../../data/sample/raw.gastrulation_e75.h5ad\") # example data included in the scDynOmics github repository relative to this notebook\n",
    "ft_adata.var['name'] = ft_adata.var.index\n",
    "ft_adata.var.index = ft_adata.var['Accession']\n",
    "ft_adata.var = ft_adata.var.drop(columns=['Accession'])\n",
    "\n",
    "print(f\"Overview of the adata object: {ft_adata} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a5978e04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: adata.X seems to be already log-transformed.\n",
      "\n",
      "Overview of the adata object: AnnData object with n_obs × n_vars = 9 × 7870\n",
      "    obs: 'barcode', 'sample', 'stage', 'sequencing.batch', 'theiler', 'doub.density', 'doublet', 'cluster', 'cluster.sub', 'cluster.stage', 'cluster.theiler', 'stripped', 'celltype', 'colour', 'umapX', 'umapY', 'haem_gephiX', 'haem_gephiY', 'haem_subclust', 'endo_gephiX', 'endo_gephiY', 'endo_trajectoryName', 'endo_trajectoryDPT', 'endo_gutX', 'endo_gutY', 'endo_gutDPT', 'endo_gutCluster', 'cell_velocyto_loom'\n",
      "    var: 'Chromosome', 'End', 'Start', 'Strand', 'name.original', 'name', 'type', 'chr', 'start', 'end', 'strand', 'promoters'\n",
      "    uns: 'log1p'\n",
      "    obsm: 'X_pca', 'X_umap'\n",
      "    layers: 'spliced', 'unspliced' \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# WARNING: adata.X seems to be already log-transformed. -> it is not, so this can be ignored.\n",
    "# https://github.com/scverse/scanpy/issues/1333\n",
    "pped_adata = pp.pipeline(\n",
    "    ft_adata,\n",
    "    ref_adata = ref_adata,\n",
    "    layers = ['spliced', 'unspliced'],\n",
    "    matching_strategy = 'stratify',\n",
    "    filtering = True,\n",
    "    min_genes = 1,\n",
    "    min_cells = 1,\n",
    "    remove_outliers = True,\n",
    "    log1p = True,\n",
    "    normalizing = True,\n",
    "    target_sum = 1e4,\n",
    "    methods = ['tpm', 'tpm'],\n",
    "    remove_mt = True,\n",
    ")\n",
    "\n",
    "print(f\"\\nOverview of the adata object: {pped_adata} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bae40fab",
   "metadata": {},
   "source": [
    "### 2.2: Monomodal Processing\n",
    "\n",
    "If you are not using multimodal data, you can process raw monomodal read counts directly.\n",
    "This approach also supports Highly Variable Gene (HVG) selection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a394918f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ag_klughammer/gyu/.conda/envs/pg_0/lib/python3.11/site-packages/scanpy/preprocessing/_simple.py:293: ImplicitModificationWarning: Trying to modify attribute `.var` of view, initializing view as actual.\n",
      "  adata.var[\"n_cells\"] = number\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overview of the adata object: AnnData object with n_obs × n_vars = 8 × 2000\n",
      "    obs: 'barcode', 'sample', 'stage', 'sequencing.batch', 'theiler', 'doub.density', 'doublet', 'cluster', 'cluster.sub', 'cluster.stage', 'cluster.theiler', 'stripped', 'celltype', 'colour', 'umapX', 'umapY', 'haem_gephiX', 'haem_gephiY', 'haem_subclust', 'endo_gephiX', 'endo_gephiY', 'endo_trajectoryName', 'endo_trajectoryDPT', 'endo_gutX', 'endo_gutY', 'endo_gutDPT', 'endo_gutCluster', 'cell_velocyto_loom', 'n_genes', 'n_genes_by_counts', 'log1p_n_genes_by_counts', 'total_counts', 'log1p_total_counts', 'pct_counts_in_top_20_genes', 'outlier'\n",
      "    var: 'Chromosome', 'End', 'Start', 'Strand', 'name.original', 'name', 'type', 'chr', 'start', 'end', 'strand', 'promoters', 'n_cells_by_counts', 'mean_counts', 'log1p_mean_counts', 'pct_dropout_by_counts', 'total_counts', 'log1p_total_counts', 'n_cells', 'highly_variable', 'means', 'dispersions', 'dispersions_norm'\n",
      "    uns: 'log1p', 'hvg'\n",
      "    obsm: 'X_pca', 'X_umap'\n",
      "    layers: 'spliced', 'unspliced' \n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Default to use adata.X. No need to specify if you don't have any other layers.\n",
    "pped_adata = pp.pipeline(\n",
    "    ft_adata,\n",
    "    ref_adata = ref_adata,\n",
    "    matching_strategy = 'stratify',\n",
    "    filtering = True,\n",
    "    min_genes = 1,\n",
    "    min_cells = 1,\n",
    "    remove_outliers = True,\n",
    "    log1p = True,\n",
    "    normalizing = True,\n",
    "    target_sum = 1e4,\n",
    "    methods = ['tpm'],\n",
    "    remove_mt = True,\n",
    "    n_top_genes = 2000,\n",
    "    hvg_flavor = 'seurat',\n",
    ")\n",
    "\n",
    "print(f\"Overview of the adata object: {pped_adata} \\n\")"
   ]
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}