Fine-tune Enformer to perform binary classification on human snATAC-seq data#
In this tutorial, we learn how to fine-tune the Enformer model (https://www.nature.com/articles/s41592-021-01252-x) on the CATLAS single-nucleus ATAC-seq dataset from human tissues (http://catlas.org/humanenhancer/).
We will perform binary classification, in which the model learns to predict the probability that a given sequence is accessible in the different cell types of the dataset. For an example of regression modeling, see the next tutorial.
import anndata
import os
import importlib
import pandas as pd
import numpy as np
%matplotlib inline
Set experiment parameters#
experiment='tutorial_2'
if not os.path.exists(experiment):
os.makedirs(experiment)
Load data#
We download the CATlas ATAC-seq binary cell type x peak matrix from http://catlas.org/catlas_downloads/humantissues/cCRE_by_cell_type/ and load it as an AnnData object.
!wget http://catlas.org/catlas_downloads/humantissues/cCRE_by_cell_type/matrix.tsv.gz
ad = anndata.read_mtx('matrix.tsv.gz').T
# Prepare ad.obs
ad.obs = pd.read_table('http://catlas.org/catlas_downloads/humantissues/cCRE_by_cell_type/celltypes.txt.gz', header=None, names=['cell type'])
ad.obs_names = ad.obs['cell type']
# Prepare ad.var
var = pd.read_table('http://catlas.org/catlas_downloads/humantissues/cCRE_hg38.tsv.gz')
var.columns = ['chrom', 'start', 'end', 'cre_class', 'in_fetal', 'in_adult', 'cre_module']
var.index = var.index.astype(str)
ad.var = var
print(ad.shape)
--2024-11-20 00:43:05-- http://catlas.org/catlas_downloads/humantissues/cCRE_by_cell_type/matrix.tsv.gz
Resolving catlas.org (catlas.org)... 132.239.162.129
Connecting to catlas.org (catlas.org)|132.239.162.129|:80... connected.
HTTP request sent, awaiting response... 200 OK
Length: 38772708 (37M) [application/x-gzip]
Saving to: ‘matrix.tsv.gz.1’
matrix.tsv.gz.1 100%[===================>] 36.98M 44.5MB/s in 0.8s
2024-11-20 00:43:06 (44.5 MB/s) - ‘matrix.tsv.gz.1’ saved [38772708/38772708]
(222, 1154611)
This contains a binary matrix representing the accessibility of 1154611 CREs measured in 222 cell types. Let us look at the components of this object:
ad.var.head()
| chrom | start | end | cre_class | in_fetal | in_adult | cre_module | |
|---|---|---|---|---|---|---|---|
| 0 | chr1 | 9955 | 10355 | Promoter Proximal | yes | yes | 146 |
| 1 | chr1 | 29163 | 29563 | Promoter | yes | yes | 37 |
| 2 | chr1 | 79215 | 79615 | Distal | no | yes | 75 |
| 3 | chr1 | 102755 | 103155 | Distal | no | yes | 51 |
| 4 | chr1 | 115530 | 115930 | Distal | yes | no | 36 |
ad.obs.head()
| cell type | |
|---|---|
| cell type | |
| Follicular | Follicular |
| Fibro General | Fibro General |
| Acinar | Acinar |
| T Lymphocyte 1 (CD8+) | T Lymphocyte 1 (CD8+) |
| T lymphocyte 2 (CD4+) | T lymphocyte 2 (CD4+) |
The contents of this anndata object are binary values (0 or 1). 1 indicates accessibility of the peak in the cell type.
ad.X[:5, :5].todense()
matrix([[1., 0., 0., 0., 0.],
[1., 0., 0., 0., 0.],
[1., 0., 0., 0., 0.],
[1., 0., 0., 0., 0.],
[1., 0., 0., 0., 0.]], dtype=float32)
Filter peaks#
We will perform filtering of this dataset using the grelu.data.preprocess module.
First, we filter peaks within autosomes (chromosomes 1-22) or chromsomes X/Y. You can also supply autosomes or autosomesX to further restrict the chromosomes.
import grelu.data.preprocess
ad = grelu.data.preprocess.filter_chromosomes(ad, 'autosomesXY')
Keeping 1154464 intervals
Next, we drop peaks overlapping with the ENCODE blacklist regions for the hg38 genome.
ad = grelu.data.preprocess.filter_blacklist(ad, genome='hg38')
Keeping 1154464 intervals
Visualize data#
Next, we can plot the distribution of the data in various ways.
import grelu.visualize
%matplotlib inline
In how many accessible cell types is each peak accessible?
cell_types_per_peak = np.array(np.sum(ad.X > 0, axis=0))
grelu.visualize.plot_distribution(
cell_types_per_peak,
title='number of cell types per peak',
method='density', # Alternative: histogram
figsize=(3.6, 2), # width, height
)
How many peaks are accessible in each cell type?
fraction_accessible_per_cell_type = np.array(np.mean(ad.X > 0, axis=1))
grelu.visualize.plot_distribution(
fraction_accessible_per_cell_type,
title='Fraction of peaks accessible per cell type',
method='histogram', # alternative: density
binwidth=0.005,
figsize=(3.6, 2), # width, height
)
It seems that some cell types have very few accessible peaks. We will drop these cell types from the dataset.
print(ad.shape)
ad = ad[ad.X.mean(axis=1) > .03, :]
print(ad.shape)
(222, 1154464)
(203, 1154464)
We can plot the distribution once again to see the effect of this filtering:
fraction_accessible_per_cell_type = np.array(np.mean(ad.X > 0, axis=1))
grelu.visualize.plot_distribution(
fraction_accessible_per_cell_type,
title='Fraction of peaks accessible per cell type',
method='histogram', # alternative: density
binwidth=0.005,
figsize=(3.6, 2), # width, height
)
Resize peaks#
Finally, since the ATAC-seq peaks can have different lengths, we resize all of them to a constant length in order to train the model. Here, we take 200 bp. We can use the resize function in grelu.sequence.utils, which contains functions to manipulate DNA sequences.
import grelu.sequence.utils
seq_len = 200
ad.var = grelu.sequence.utils.resize(ad.var, seq_len)
ad.var.head(3)
| chrom | start | end | cre_class | in_fetal | in_adult | cre_module | |
|---|---|---|---|---|---|---|---|
| 0 | chr1 | 10055 | 10255 | Promoter Proximal | yes | yes | 146 |
| 1 | chr1 | 29263 | 29463 | Promoter | yes | yes | 37 |
| 2 | chr1 | 79315 | 79515 | Distal | no | yes | 75 |
Split data#
We will split the peaks by chromosome to create separate sets for training, validation and testing.
train_chroms='autosomes'
val_chroms=['chr10']
test_chroms=['chr11']
ad_train, ad_val, ad_test = grelu.data.preprocess.split(
ad,
train_chroms=train_chroms, val_chroms=val_chroms, test_chroms=test_chroms,
)
Selecting training samples
Keeping 1007912 intervals
Selecting validation samples
Keeping 56974 intervals
Selecting test samples
Keeping 56433 intervals
Final sizes: train: (203, 1007912), val: (203, 56974), test: (203, 56433)
Make labeled sequence datasets#
grelu.data.dataset contains PyTorch Dataset classes that can load and process genomic data for training a deep learning model. Here, we use the AnnDataSeqDataset class which loads data from an AnnData object. Other available dataset classes include DFSeqDataset and BigWigSeqDataset.
We first make the training dataset. To increase model robustness we use several forms of data augmentation here: rc=True (reverse complementing the input sequence), and max_seq_shift=1 (shifting the coordinates of the input sequence by upto 1 bp in either direction; also known as jitter). We use augmentation_mode="random" which means that at each iteration, the model sees a randomly selected augmented version for each sequence.
import grelu.data.dataset
train_dataset = grelu.data.dataset.AnnDataSeqDataset(
ad_train.copy(),
genome='hg38',
rc=True, # reverse complement
max_seq_shift=1, # Shift the sequence
augment_mode="random", # Randomly select which augmentations to apply
)
We do not apply any augmentations to the validation and test datasets (although it is possible to do so).
val_dataset = grelu.data.dataset.AnnDataSeqDataset(ad_val.copy(), genome='hg38')
test_dataset = grelu.data.dataset.AnnDataSeqDataset(ad_test.copy(), genome='hg38')
Build the enformer model#
gReLU contains many model architectures. One of these is a class called EnformerPretrainedModel. This class creates a model identical to the published Enformer model and initialized with the trained weights, but where you can change the number of transformer layers and the output head.
Models are created using the grelu.lightning module. In order to instantiate a model, we need to specify model_params (parameters for the model architecture) and train_params (parameters for training).
model_params = {
'model_type':'EnformerPretrainedModel', # Type of model
'n_tasks': ad.shape[0], # Number of cell types to predict
'crop_len':0, # No cropping of the model output
'n_transformers': 1, # Number of transformer layers; the published Enformer model has 11
}
train_params = {
'task':'binary', # binary classification
'lr':1e-4, # learning rate
'logger': 'csv', # Logs will be written to a CSV file
'batch_size': 1024,
'num_workers': 8,
'devices': 0, # GPU index
'save_dir': experiment,
'optimizer': 'adam',
'max_epochs': 10,
'checkpoint': True, # Save checkpoints
}
import grelu.lightning
model = grelu.lightning.LightningModel(model_params=model_params, train_params=train_params)
Train the model#
We train the model on the new training dataset using the train_on_dataset method. Note that here, we update the weights of the entire model during training. If you want to hold the enformer weights fixed and only learn a linear layer from the enformer embedding to the outputs, see the tune_on_dataset method.
# See the 'tutorial_2' folder for logs
trainer = model.train_on_dataset(
train_dataset=train_dataset,
val_dataset=val_dataset,
)
Load best model from checkpoint#
During training, the performance of the model on the validation set is checked after each epoch. We will load the version of the model that had the best validation set performance.
best_checkpoint = trainer.checkpoint_callback.best_model_path
print(best_checkpoint)
model = grelu.lightning.LightningModel.load_from_checkpoint(best_checkpoint)
tutorial_2/2024_06_11_05_13/version_0/checkpoints/epoch=6-step=6895.ckpt
Calculate performance metrics on the test set#
We calculate global performance metrics using the test_on_dataset method.
test_metrics = model.test_on_dataset(
test_dataset,
devices=0,
num_workers=8,
batch_size=1024,
)
Testing DataLoader 0: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 56/56 [00:06<00:00, 8.23it/s]
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ DataLoader 0 ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ test_accuracy │ 0.9459772706031799 │ │ test_auroc │ 0.9041877388954163 │ │ test_avgprec │ 0.6062116622924805 │ │ test_best_f1 │ 0.5677600502967834 │ │ test_loss │ 0.1546824723482132 │ └───────────────────────────┴───────────────────────────┘
test_metrics is a dataframe containing metrics for each model task.
test_metrics.head()
| test_accuracy | test_auroc | test_avgprec | test_best_f1 | |
|---|---|---|---|---|
| Follicular | 0.909113 | 0.879024 | 0.627018 | 0.577536 |
| Fibro General | 0.916698 | 0.882387 | 0.624694 | 0.575148 |
| Acinar | 0.944483 | 0.913085 | 0.649766 | 0.599485 |
| T Lymphocyte 1 (CD8+) | 0.961742 | 0.930787 | 0.635401 | 0.593375 |
| T lymphocyte 2 (CD4+) | 0.967306 | 0.939841 | 0.622047 | 0.577923 |
Visualize performance metrics#
We can plot the distribution of each metric across all cell types:
grelu.visualize.plot_distribution(
test_metrics.test_best_f1,
method='histogram',
title='Best F1 Score per task',
binwidth=0.005,
figsize=(3,2),
)
grelu.visualize.plot_distribution(
test_metrics.test_avgprec,
method='histogram',
title='Average Precision per task',
binwidth=0.005,
figsize=(3,2),
)
grelu.visualize.plot_distribution(
test_metrics.test_auroc,
method='histogram',
title='AUROC per task',
binwidth=0.005,
figsize=(3,2),
)
Testing DataLoader 0: 0%| | 0/56 [00:20<?, ?it/s]
Run inference on the test set#
Instead of overall metrics, we can also get the individual predictions for each test set example.
probs = model.predict_on_dataset(
test_dataset,
devices=0,
num_workers=8,
batch_size=1024,
return_df=True # Return the output as a pandas dataframe
)
probs.head()
Predicting DataLoader 0: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 56/56 [00:03<00:00, 14.34it/s]
| Follicular | Fibro General | Acinar | T Lymphocyte 1 (CD8+) | T lymphocyte 2 (CD4+) | Natural Killer T | Naive T | Fibro Epithelial | Cardiac Pericyte 1 | Pericyte General 1 | ... | Fetal Cardiac Fibroblast | Fetal Fibro General 2 | Fetal Fibro Muscle 1 | Fetal Fibro General 3 | Fetal Mesangial 2 | Fetal Stellate | Fetal Alveolar Epithelial 1 | Fetal Cilliated | Fetal Excitatory Neuron 1 | Fetal Excitatory Neuron 2 | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | 0.420895 | 0.674017 | 0.506232 | 0.685084 | 0.514758 | 0.443587 | 0.465559 | 0.661670 | 0.419363 | 0.730238 | ... | 0.111079 | 0.121903 | 0.171487 | 0.156433 | 0.180907 | 0.060901 | 0.169595 | 0.100527 | 0.205600 | 0.062775 |
| 1 | 0.008851 | 0.021444 | 0.001757 | 0.069185 | 0.060921 | 0.048642 | 0.016652 | 0.005472 | 0.051451 | 0.004445 | ... | 0.070789 | 0.017949 | 0.016753 | 0.014286 | 0.008884 | 0.019222 | 0.017777 | 0.007793 | 0.020279 | 0.007875 |
| 2 | 0.022793 | 0.037817 | 0.010379 | 0.037814 | 0.043917 | 0.056750 | 0.018970 | 0.037695 | 0.112500 | 0.018786 | ... | 0.054257 | 0.017340 | 0.026617 | 0.041380 | 0.023223 | 0.019772 | 0.022399 | 0.014391 | 0.248927 | 0.029103 |
| 3 | 0.764947 | 0.803988 | 0.857578 | 0.784091 | 0.808821 | 0.677050 | 0.721701 | 0.707757 | 0.749479 | 0.606714 | ... | 0.755405 | 0.768155 | 0.714758 | 0.770326 | 0.734850 | 0.619608 | 0.853031 | 0.790041 | 0.696608 | 0.416579 |
| 4 | 0.174569 | 0.210124 | 0.301300 | 0.363085 | 0.266278 | 0.162586 | 0.191470 | 0.195322 | 0.103195 | 0.108766 | ... | 0.104694 | 0.099994 | 0.098529 | 0.112873 | 0.094192 | 0.091924 | 0.176032 | 0.184840 | 0.101290 | 0.029382 |
5 rows × 203 columns
Since this is a binary classification model, the output takes the form of probabilities ranging from 0 to 1. We interpret these as the predicted probabilities of the element being accessible in each cell type.
Plot additional visualizations of the test set predictions#
We can plot a calibration curve for all the tasks. This shows us the fraction of true positive examples, for different levels of model-predicted probability.
grelu.visualize.plot_calibration_curve(
probs, labels=test_dataset.labels, aggregate=False, show_legend=False
)
We can also pick any cell type and compare the predictions on accessible and non-accessible elements.
grelu.visualize.plot_binary_preds(
probs,
labels=test_dataset.labels,
tasks='Microglia',
figure_size=(3, 2) # Width, height
)
Interpret model predictions (for microglia) using TF-modisco#
Suppose we want to focus specifically on Microglia. We can create a transform - a class that takes in the model’s prediction and returns a function of the prediction, e.g. the prediction for only a subset of cell types that we are interested in. Then, all subsquent analyses that we do will be based only on this subset of the full prediction matrix.
Here, we use the Aggregate class to create a transform that will take in the full set of predictions produced by the model and subset only the predictions in Microglia.
from grelu.transforms.prediction_transforms import Aggregate
microglia_score = Aggregate(
tasks = ["Microglia"],
model = model,
)
Let us now identify all peaks in the test set that are accessible in microglia.
is_accessible_in_mcg = np.array(ad_test["Microglia", :].X.todense()==1).squeeze()
mcg_peaks = ad_test.var[is_accessible_in_mcg]
len(mcg_peaks)
2273
We now run TF-Modisco on these peaks. TF-modisco identifies motifs that consistently contribute to the model’s output. Since we are using the microglia_score filter to limit the model’s prediction to microglia, we will only get motifs relevant to this cell type. We also use TOMTOM to match the TF-Modisco motifs to a set of reference motifs. Here, we use a reference set of non-redundant motifs (https://www.vierstra.org/resources/motif_clustering) that are provided with gReLU as consensus.
%%time
import grelu.interpret.modisco
grelu.interpret.modisco.run_modisco(
model,
seqs=mcg_peaks,
genome="hg38",
prediction_transform=microglia_score, # Base importance scores will be calculated with respect to this output
meme_file="hocomoco_v12", # We will compare the Modisco CWMs to HOCOMOCO motifs
method="ism", # Base-level attribution scores will be calculated using ISM. You can also use "saliency".
out_dir=experiment,
batch_size=1024,
devices=0,
num_workers=16,
window=100, # ISM scores will be calculated over the central 100 bp of each peak
seed=0,
)
Performing ISM
Predicting DataLoader 0: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 16.16it/s]
Testing DataLoader 0: 0%| | 0/56 [01:56<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:56<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:56<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:56<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:56<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [01:57<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [02:10<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [02:11<?, ?it/s]
Testing DataLoader 0: 0%| | 0/56 [02:11<?, ?it/s]
Predicting DataLoader 0: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 666/666 [00:47<00:00, 14.13it/s]
Running modisco
Testing DataLoader 0: 0%| | 0/56 [02:24<?, ?it/s]
Writing modisco output
Creating sequence logos
Creating html report
Running TOMTOM
CPU times: user 3h 25min 17s, sys: 1min 21s, total: 3h 26min 39s
Wall time: 9min 9s
Load TOMTOM output for modisco motifs#
The full output of TF-Modisco and TOMTOM can be found in the experiment folder. Here, we read the output of TOMTOM and list the significant TOMTOM matches, i.e. known TF motifs that are similar to those found by TF-MoDISco.
tomtom_file = os.path.join(experiment, 'tomtom.csv')
tomtom = pd.read_csv(tomtom_file, index_col=0)
tomtom[tomtom['q-value'] < 1e-4] # Display most significant matches
| Query_ID | Target_ID | Optimal_offset | p-value | E-value | q-value | Overlap | Query_consensus | Target_consensus | Orientation | |
|---|---|---|---|---|---|---|---|---|---|---|
| 950 | pos_pattern_0 | SPI1.H12CORE.0.P.B | 2.0 | 1.938174e-08 | 2.796785e-05 | 3.729047e-05 | 7.0 | ACTTCCT | AAAAGAGGAAGTGA | - |
| 1861 | pos_pattern_1 | IRF4.H12CORE.0.P.B | 2.0 | 8.912515e-11 | 1.286076e-07 | 5.144304e-07 | 14.0 | TTTCACTTCCTCTT | AAAGAGGAACTGAAACT | - |
| 1870 | pos_pattern_1 | IRF8.H12CORE.0.P.B | 5.0 | 3.314362e-10 | 4.782625e-07 | 9.565249e-07 | 14.0 | TTTCACTTCCTCTT | AAGAGGAAGTGAAAGTAAA | - |
| 2395 | pos_pattern_1 | SPIB.H12CORE.0.P.B | 1.0 | 3.247408e-08 | 4.686009e-05 | 3.748807e-05 | 14.0 | TTTCACTTCCTCTT | AAAGAGGAAGTGAAAG | - |
| 2995 | pos_pattern_2 | CTCFL.H12CORE.0.P.B | -1.0 | 5.496209e-08 | 7.931030e-05 | 5.287353e-05 | 15.0 | CGCCGCCCCCTCCTGG | GCCGCCAGGGGGCGCC | - |
| 3822 | pos_pattern_2 | SP2.H12CORE.1.P.B | 5.0 | 2.985629e-08 | 4.308262e-05 | 3.748807e-05 | 16.0 | CGCCGCCCCCTCCTGG | CCGGCGGGGGGCGGGGCCGGG | - |