full_training_recreation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "343fc490",
   "metadata": {},
   "source": [
    "# Full Training Recreation Enformer Celltyping\n",
    "\n",
    "This workbook steps through training Enformer Celltyping using the full EpiMap dataset\n",
    "as outlined in our manuscript. Note that the steps involved including downloading data, \n",
    "preprocessing and training are both memory and computationally intensive. \n",
    "\n",
    "These steps are to be run on the command line. We advise running steps in parallel where \n",
    "possible and the use of at least a 24GB GPU for model training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b074e8d",
   "metadata": {},
   "source": [
    "## Step 1. Download Data\n",
    "\n",
    "Multiple environments are used to download and preprocess data\n",
    "to deal with different package dependencies\n",
    "Use the `Makefile` to create the necessary conda environments:\n",
    "\n",
    "```\n",
    "make pyanlyenv &&\\\n",
    "make renv\n",
    "```\n",
    "\n",
    "Next download the data for the analysis by first activating the\n",
    "py_analysis conda environment:\n",
    "\n",
    "```\n",
    "conda activate py_analysis\n",
    "```\n",
    "\n",
    "If this is your first time using py_analysis you will need to \n",
    "install the local package:\n",
    "\n",
    "```\n",
    "pip install -e .\n",
    "```\n",
    "\n",
    "And then run:\n",
    "\n",
    "```\n",
    "python bin/download_data_parallel.py\n",
    "```\n",
    "\n",
    "The files are downloaded to the `data/` folder. There are 103 cell types that have been \n",
    "downloaded forn EpiMap and will be used for model training (see metadata folder for more details \n",
    "on these). **Note** this download will run in parallel but will take some time. It will also need \n",
    "a decent amount of disk space (188 GB).\n",
    "\n",
    "**Note** the downloaded files are in bigWig format taken from [EpiMap](http://compbio.mit.edu/epimap/).\n",
    "BigWig files for chromatin accessibility (ATAC-Seq) or histone marks, contain the\n",
    "-log10 p-value of a peak for each base-pair position of the genome. This differs from BAM\n",
    "files which are just the peak regions listed at a certain threshold. Note that\n",
    "although the bigWig files contain a value for every base-pair in the genome, these\n",
    "have been averaged at 25 base-pairs by the peak caller (e.g. MACS2).\n",
    "\n",
    "We use bigWig files so we can predict continuous values with our model (regression problem)\n",
    "rather than 0/1 peak in a region (classification problem). This approach has been proven to\n",
    "give [better performance](https://www.nature.com/articles/s42256-022-00570-9)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "210cbfe2",
   "metadata": {},
   "source": [
    "## Step 2. Preprocess Data\n",
    "\n",
    "### 2.1 Liftover hg38 bigwigs to hg19\n",
    "\n",
    "The analysis is done in hg19 but certain files are in hg38. Lift these over using the\n",
    "following command, again with the py_analysis environment:\n",
    "\n",
    "```\n",
    "bash bin/liftover_to_hg19.sh\n",
    "```\n",
    "\n",
    "\n",
    "### 2.2 Convert all bigwigs to 128 base-pair (bp) resolution files\n",
    "\n",
    "The model's predicitons are in 128bp resolution. However, the downloaded files are all 25bp \n",
    "resolution. So the data tracks need to be adjusted to this resolution. Use the `r_bioc` \n",
    "environment to do this and run the following script. This will take some time and will\n",
    "need roughly the same amount of space as currently being used in the data folder since all\n",
    "bigWigs need to be converted to 128 bp. Also, this will take a long time to run, feel\n",
    "free to change the `--jobs X` amount in `bin/avg_bigwig_tracks.sh` if you have free RAM\n",
    "to run more in parallel and speed it up (running on a 64 CPU, 128GB RAM machine took ~1.5 \n",
    "hours):\n",
    "\n",
    "```\n",
    "bash bin/avg_bigwig_tracks.sh\n",
    "```\n",
    "\n",
    "Note you might need to change permissions on all bigwigs before running this:\n",
    "\n",
    "```\n",
    "chmod 777 ./data/*\n",
    "```\n",
    "\n",
    "**Note** I would advise to now delete all the bigWig files other than the `*_128.bigWigs`\n",
    "files to save on space. \n",
    "\n",
    "\n",
    "### Normalise assay signals\n",
    "\n",
    "The differing data for the assays correspond to the -log10(adjusted p-value) from MACS2 peak calling.\n",
    "This indicates the statistical significance of a genomic position i.e. is the site likely to be a real\n",
    "binding site when compared against a background. However, if we were to directly predict these scores, \n",
    "a model would perform poorly since our training samples have differing read depths.\n",
    "\n",
    "To reduce the effect of outliers, we use arcsinh-transformed signals:\n",
    "\n",
    "sinh^−1 𝑥=ln( 𝑥+sqrt(1+𝑥^2) )\n",
    "\n",
    "This has been used to train [Avocado](https://genomebiology.biomedcentral.com/articles/10.1186/s13059-020-01977-6) \n",
    "model and for all evaluations. Other models, such as [PREDICTD](https://www.nature.com/articles/s41467-018-03635-9)\n",
    "and [Segway](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3340533/), also use this transformation to\n",
    "sharpen the effect of the shape of the signal while diminishing the effect of large values.\n",
    "\n",
    "Nothing needs to be done for this step since it is taken care of in the data generator (see \n",
    "`EnformerCelltyping/utils.py` and the `generate_data()` function)\n",
    "\n",
    "\n",
    "### 2.3 Create average epigenetic mark signals\n",
    "\n",
    "Create average epigenetic mark signals from the **training cells**. This step creates an average\n",
    "histone mark signal, quantile of the distribution and delta friom the 103 training cell types. The\n",
    "delta here is the difference between the average and the cell type speicific histone mark signal at\n",
    "a given genomic location. These three data streams are used for model training. Moreover, the \n",
    "average chromatin accessibilty (ATAC-seq) signal is used benchamrk performance in our manuscript.\n",
    "Use the `py_analysis` environment and run:\n",
    "\n",
    "```\n",
    "python3 bin/calculate_avg_track.py -m epi_mark\n",
    "```\n",
    "\n",
    "replacing epi_mark with your chosen mark to average. Or run:\n",
    "\n",
    "```\n",
    "bash bin/calculate_avg_track.sh\n",
    "```\n",
    "\n",
    "to calculate an average for all marks. Again, if you have free RAM to run in parallel it might be \n",
    "worthwhile (running on a 64 CPU, 128GB RAM machine took ~6 hours).\n",
    "\n",
    "### 2.4 Install Enformer\n",
    "\n",
    "If you have not done so already install Enformer. All use-cases for Enformer Celltyping involve \n",
    "transfer learning from the enformer model (with frozen weights) so download this model from \n",
    "tensorflow hub:\n",
    "\n",
    "```\n",
    "mkdir data/enformer_model/ &&\\\n",
    "cd data/enformer_model &&\\\n",
    "wget -O enformer_model.tar.gz https://tfhub.dev/deepmind/enformer/1?tf-hub-format=compressed &&\\\n",
    "tar xvzf enformer_model.tar.gz &&\\\n",
    "rm enformer_model.tar.gz &&\\\n",
    "cd ../../\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a1af2de",
   "metadata": {},
   "source": [
    "## Step 3. Train Enformer Celltyping model\n",
    "\n",
    "### 3.1 Precompute training data\n",
    "\n",
    "You now have all the data downloaded and preprocessed to train Enformer Celltyping to predict histone marks at\n",
    "different genomic regions based on DNA, local chromatin accessibility and global chromatin accessibility signals.\n",
    "\n",
    "There is a custom built, general purpose data loader `generate_data()` (in `./EnformerCelltyping/utils.py`) which\n",
    "can be used to load the input (X) and output (Y) for any genomic position. This can be used with a tensorflow\n",
    "data loader to load positions on-the-fly while training. However, this is not every efficient since the genomic information needs to be loaded from bigWigs containing DNA, chromatin accessibility or histone mark signals for the whole genome. Moreover, this includes passing the DNA data through the pre-trained enformer model for transfer learning each time the data is loaded.\n",
    "\n",
    "A faster approach is to precompute the data for every training position and then save them in an efficient access file\n",
    "type (compressed numpy files - `.npz`). When training Enformer Celltyping, we used the following approach to \n",
    "identify training positions:\n",
    "\n",
    "  1. Bin genome based on predictive window\n",
    "  2. Filter bins to select training set based on DNA and cell type filters. \n",
    "      * DNA filters:\n",
    "          1. Leave buffer at start/end of chromosome large enough for DNA and local chromatin accessibility windows\n",
    "          2. Not in blacklist regions\n",
    "      * Cell type filters:\n",
    "          1. Coverage for the histone mark > 12.5% of the returned window to prioritise training on regions with peaks.\n",
    "  3. Down sample resulting regions to equal the lowest count of regions for any histone mark so each hist mark has equal representation. This avoids the model biasing training on one mark. \n",
    "  \n",
    "\n",
    "This results in 67,007 training & validation positions (cell type and genomic region combinations) and 14,188 unique genomic positions which is similar to number of positions basenji & enformer trained on (14,533). The approach ensures model sees peaks for all histone marks. The validation set positons are randomly shifted by up to a quarter of the predictive window so the model's performance doesn't overfit to the initial genomic bins.\n",
    "\n",
    "These chosen training regions can be downloaded (running within the downloaded repo - `EnformerCelltyping/`):\n",
    "\n",
    "```\n",
    "mkdir checkpoint &&\\\n",
    "wget -O ./checkpoint/reg_cov.csv.gz https://figshare.com/ndownloader/files/39116387\n",
    "```\n",
    "\n",
    "Alternatively, the script below will create the file if it is not found. To save the data for the training regions \n",
    "run the following with the `EnformerCelltyping` environment:\n",
    "\n",
    "```\n",
    "python ./bin/precomute_train_data.py  -s 0 -e 1000\n",
    "```\n",
    "\n",
    "Note this is quite time intensive so should be run in parallel and on GPU(s) if possible. It is set up to run\n",
    "in parallel already as you input the row numbers of the csv of the 67,007 training regions to create."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b078882",
   "metadata": {},
   "source": [
    "### 3.2 Train the model\n",
    "\n",
    "Finally, we can train the model with the following.\n",
    "\n",
    "Training Enformer Celltyping was done in **two** steps:\n",
    "\n",
    "#### 1. Pretrain a split model\n",
    "\n",
    "Since the enformer model is being used to embed the DNA information in Enformer Celltyping, the weights for the\n",
    "DNA information are already sensibly initiated. To avoid the DNA information dominating the final predictions in\n",
    "Enformer Celltyping (and the model just ignoring the chromatin accessibility information), we first trained \n",
    "Enformer Celltyping with a split architecture as some 'warm-up' steps. For a given genomic position, the model predicts the (1) average signal at this position and the (2) quantile distribution of the signal for the 104 trianing cell types from the DNA information and the (3) delta i.e. difference between the average signal and the cell type-specific signal for the cell type of interest from the chromatin accessibility information. We show how this two-stage training approach improves performance in our manuscript.\n",
    "\n",
    "This warm-up stage is run for 1,000 steps where each step contains the full batch size of 128 positions.\n",
    "\n",
    "```\n",
    "python pretrain.py\n",
    "```\n",
    "\n",
    "**Note**: If you want to skip this warm-up step, you can download the h5 file of the resulting weights to use with the second step of training with:\n",
    "\n",
    "```\n",
    "wget https://figshare.com/ndownloader/files/42284757\n",
    "```\n",
    "\n",
    "Just make sure to update the path to these in `train.py` \n",
    "\n",
    "For the manuscript, we pretrained the model for 1,000 steps of 128 batch size (128,000 positions) taking 4 days on a 80GB RAM GPU.\n",
    "\n",
    "#### 2. Train a combined model\n",
    "\n",
    "Now we have sensibly initiated the wieghts for both the DNA and chromatin accessibility streams, we can combine them\n",
    "to give a cell type-specific prediciton as an output. To run this training use:\n",
    "\n",
    "```\n",
    "python train.py\n",
    "```\n",
    "\n",
    "For the manuscript, we trained the model for 4 full epochs (6940 steps of 128 batch size - 888,320 positions) stopping when it started to overfit, taking 1.5 days on a 80GB RAM GPU.\n",
    "\n",
    "Also see [training_demo.ipynb](https://github.com/neurogenomics/EnformerCelltyping/blob/master/training_demo.ipynb) for an interactive training script on a small demo dataset. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EnformerCelltyping",
   "language": "python",
   "name": "enformercelltyping"
  },
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}