<jupyter_start><jupyter_text>Large-scale softmax classification for the drug screen<jupyter_code>%load_ext blackcellmagic
%load_ext autoreload
%autoreload 2
import numpy as np 
import pandas as pd
import tqdm
import anndata as ad
import seaborn as sns
from matplotlib import pyplot as plt

import torch
import torch.nn as nn 
import torch.optim as optim 
import torchvision.transforms as transforms 
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import IterableDataset, DataLoader

import holoviews as hv 
from holoviews.operation.datashader import datashade, rasterize
from holoviews.operation import gridmatrix
import hvplot.pandas
import colorcet as cc
hv.extension('bokeh')

import warnings 
warnings.filterwarnings('ignore')

seed = 6398629
torch.manual_seed(seed)
np.random.seed(seed)
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
from magma import models as mm
from magma import utils as mu
mu.set_plotting_style_plt()<jupyter_output><empty_output><jupyter_text>In brief, in this notebook we will build a cell classifier model.<jupyter_code>from IPython.display import Image

url = 'https://github.com/manuflores/sandbox/blob/master/figs/cell_clf.png?raw=true'

Image(
    url = url, format = 'png', width = 500, height = 500
)<jupyter_output><empty_output><jupyter_text>Digression on using other models As we talked during the meeting, we are using this classifier model as a pre-training step in order to have a good vectorized representation of the cells amenable to use the contrastive learning procedure. We could in theory use other variables of interest. For example, we can train a model to predict the drug target or the pathway of the molecule that perturbed the cell. We have information about metadata about the drugs. We could also use other models like Variational Autoencoder or any other type of encoder model.  About the dataset. As you can see from the diagram above, the input for this model will be a cell (transcriptome) and the output will be the label of the drug that perturbed it. We're going to use a subset of our drug screen dataset that contains data of only 100 drugs, from **myeloid cells subjected to the CD3 treatment**. We could in fact use the data from all cell types, but this is just a proof-of-concept. . The idea is that we're would later use the data from the other samples for validation for the joint embedding. For context, the data is in log-normalized, i.e. normalized by the total number of counts per cell and then scaled by a factor of $10^4$, to be finally set in log scale. In other words we used the following equation on the counts of each gene:  $$\tilde{g_{i}} = \mathrm{ln} \left( \frac{g_i \times 10^4}{\sum_{i}^n g_i} + 1 \right)$$ where $\tilde{g_i}$ is the log-normalized counts for gene $i$ for a given cell. The single-cell data count matrix is in [`anndata`](https://anndata.readthedocs.io/) format. This is the python analog of the Seurat object from R. [I highly recommend this tutorial](https://falexwolf.de/blog/171223_AnnData_indexing_views_HDF5-backing/) from <NAME>, one of the main developers of AnnData, that introduces how to use numpy-like indexing (numerical indexing) and pandas-like (indexing with boolean Series meeting categorical criteria). I have another tutorial working with single-cell RNA seq data [here](https://github.com/manuflores/sandbox/blob/master/notebooks/qc_cv_filering.ipynb) if you're interested, but the basic idea is that you can use indexing both in the rows and columns of the `anndata`.   Load cell dataset: CD3 (+) myeloid cells Let's load the data and take a look.<jupyter_code># Import dataset
a = ad.read_h5ad('mult_cd3_100_train.h5ad')
# Take a look
a
type(a)
type(a.obs)
type(a.var)
# Contains metadata from the cells 
a.obs.head(2)
# Contains metadata from the genes. 
a.var.head()
print('We have data from cells subjected to %d perturbations.'%a.obs.drug_name.unique().shape[0])<jupyter_output><empty_output><jupyter_text>Very briefly, I'll assume you have no experience with anndata and filter both rows and columns. 1. Filter to get only nilotinib cells (rows).2. Filter to get only data from the first 50 cells and genes with [coefficient of variation (CV)](https://en.wikipedia.org/wiki/Coefficient_of_variation) higher than 5.<jupyter_code># 1. notice how we get only 700 cells
a[a.obs.drug_name == 'nilotinib']
# 2 - notice how we get only 50 cells and 1045 genes
a[:50, a.var.cv > 5]<jupyter_output><empty_output><jupyter_text>All right let's proceed.<jupyter_code># Check that we only have myeloid cells
a.obs.cell_type.unique()
# Check that we only have CD3+ data
a.obs.CD3.unique()<jupyter_output><empty_output><jupyter_text>Similarly to the molecule classification case, we need to make sure that the data is balanced. Let's check the samples with the most and the lowest number of cells.<jupyter_code># Samples with most cells
a.obs.sample_id.value_counts().head(3)
# Samples with least number of cells
a.obs.sample_id.value_counts().tail(3)
# Remove Nilotinib, Ketoprofen and Eupatilin for out-of-sample predictions
#a = a[~a.obs.drug_name.isin(['Nilotinib', 'Ketoprofen', 'Eupatilin'])]<jupyter_output><empty_output><jupyter_text>Let's put a cap on using at most 500 cells and sample !<jupyter_code>%%time
n_samples = 500

sampling_ix = (
    a.obs.groupby(["sample_id"])
    .apply(
        lambda group_df: group_df.sample(
            group_df.shape[0] if group_df.shape[0] < n_samples else n_samples,
            replace = False)
    )
    .index.get_level_values(1) # Get the numerical index :) 
)<jupyter_output><empty_output><jupyter_text>Now we can use this index to select the data numpy-style.<jupyter_code># Select cells
ada = a[sampling_ix].copy()<jupyter_output><empty_output><jupyter_text>Let's confirm that we indeed selected the data correctly.<jupyter_code>ada.obs.sample_id.value_counts().head()
ada.obs.sample_id.value_counts().tail()<jupyter_output><empty_output><jupyter_text>Excellent let's proceed to make categorical indices of the samples.<jupyter_code>codes, uniques = pd.factorize(ada.obs["sample_id"].values.to_list())
ada.obs['sample_codes'] = codes<jupyter_output><empty_output><jupyter_text>Now we can make our train-test split using stratified sampling, keeping the same proportion of the drug perturbations and the drug target in each fold.<jupyter_code>from sklearn.model_selection import StratifiedShuffleSplit
# Initialize stratified splitter
splitter = StratifiedShuffleSplit(n_splits = 1, test_size = 0.4, random_state = seed)

# Get indices 
ixs = list(splitter.split(ada.X, ada.obs[['sample_codes', 'target']]))

# Extract train and test indices
train_ixs, val_ixs = ixs[0][0], ixs[0][1]
train_adata = ada[train_ixs].copy()
test_adata = ada[val_ixs].copy()
# train_adata = a[train_ixs].copy()
# test_adata = a[val_ixs].copy()<jupyter_output><empty_output><jupyter_text>Let's check that we have the same proportion of<jupyter_code>train_adata.obs.sample_id.value_counts().tail()/ train_adata.n_obs
test_adata.obs.sample_id.value_counts().tail()/ test_adata.n_obs<jupyter_output><empty_output><jupyter_text>Making torch datasets and dataloaders As in the graph case, we also need to construct torch datasets and dataloaders. The most important I think is the torch dataset, because it is highly dependent on the type of data you're using (images, text, speech, or numerical vectors) and is almost always in the need of customization. I have written code to make torch datasets using anndata so we're good. The basic idea though is that we need to override the `torch.Dataset` base class with a  `__getitem__()` method, that allows to extract data by indexing (we'll show an example below). You can imagine that if you're using a supervised method one has to also write down a way to extract the label $y_i$ for each input data point $x_i$. This can be done in a highly efficient manner as to read data from disk and just feed it by minibatches through the `DataLoader` generator. So that's the reason that they exist, they allow you to feed data into a model using minibatches from disk, that could in theory not be handled if loaded in memory. I actually wrote a function to do this using numpy arrays instead of `anndataset`s for extremely large count matrices but I hope that this dataset would fit fine in memory as it is 0.1 GB in size.<jupyter_code># Initialize torch dataset 
train_dataset = mu.adata_torch_dataset(
    train_adata,
    transform = transforms.ToTensor(),
    supervised = True,
    target_col = 'sample_codes'
)

test_dataset = mu.adata_torch_dataset(
    test_adata,
    transform = transforms.ToTensor(),
    supervised = True,
    target_col = 'sample_codes',
    #multilabel = True
)
# Extract the first datapoint from the training set:
train_dataset[0]<jupyter_output><empty_output><jupyter_text>We can see that we can index the train dataset object and it returns a tuple of a) a tensor containing the log-normed counts and b) the label for the drug perturbation. With this torch dataset in place, we can go ahead and initialize the dataloaders.<jupyter_code># check your processors, in case you don't know yet !
import multiprocessing as mp
n_cores = mp.cpu_count()

print('We have %d available cores'%n_cores)
batch_size = 64  # increase batch size because of large dataset

# Initialize DataLoader for minibatching
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    drop_last=True,
    shuffle=False,
    num_workers=n_cores,
)

val_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    drop_last=True,
    shuffle=False,
    num_workers=n_cores,
)<jupyter_output><empty_output><jupyter_text>All right, we're ready to initialize our classifier. We will be using a simple deep MLP stored in the `mm.supervised_model()` class. We only need to specify the dimensionalities of the input, output and intermediate layers of the MLP. We will start with a dimensionality equal to the number of genes, and the final output dimension is the number of drugs to predict, in this case 100. Finally we will add three intermediate layers to enable the capacity of learning optimized representations of size 512, 256, and 64.<jupyter_code>#help(mm.supervised_model)
number_of_drugs = len(uniques)
number_of_genes = ada.n_vars

# Initialize dimensionalities list
dims = [number_of_genes, 512, 256, 64, number_of_drugs]

dims<jupyter_output><empty_output><jupyter_text>Now we can initialize the model, we will set the model type to `multiclass`, and we won't use dropout for the model.<jupyter_code>model = mm.supervised_model(dims, model = 'multiclass', dropout = False)<jupyter_output><empty_output><jupyter_text>Finally, let's initialize the weights of our model (if you want a nice explanation of why this is important [check this lecture notes](https://www.deeplearning.ai/ai-notes/initialization/index.html)), and initialize our optimizer and error function.<jupyter_code># Initialize weights according to Xavier 
model = mu.initialize_network_weights(model, method = 'xavier_normal')
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3, weight_decay = 0)
criterion = nn.NLLLoss()
# Take a look at the method. 
model<jupyter_output><empty_output><jupyter_text>It will take around 30mins - 1 hour to train the model. The trainer for MLPs currently has much more functionality than for GCNs. We can save models every epoch by supplying a valid directory in the `model_dir` kwarg and a corresponding `model_name`. We can also add a ratio for early stopping if we don't want to monitor the validation loss for convergence. The early stopping is activated when the validation loss increases a fraction of `early_stopping_tol` (by default 0.2) with respect to a previous epoch. If you want to let the model run substantially, set the `n_epochs` parameter to 200 and let the early stopping handle the rest. If for some reason you don't want early stopping, set the `early_stopping_tol` parameter to 50.   Extra note for running on a GPU We havent covered this but we could also run the model in a GPU. We do this under the hood with the trainer function but you can explicity do this by runnning the following cell.<jupyter_code># device = nm.try_gpu()
# model = model.to(device)
# You can set a directory to save the weights per each epoch.
#model_dir = path + '../models/droog_100/'
n_epochs = 4
%%time
train_loss, val_loss = mu.supervised_trainer(
    n_epochs,
    train_loader,
    val_loader,
    model,
    criterion,
    optimizer,
    multiclass = True,
    n_classes = number_of_drugs, 
    model_dir = None, # for example './' to save in current directory
    model_name = None, # for example 'classiffier_100_drugs'
    early_stopping_tol = 0.2
)<jupyter_output><empty_output><jupyter_text>On a previous run, the model at 2 epochs was the best.<jupyter_code>#model = nm.supervised_model(dims, model = 'multiclass', dropout = False)
#model
# trained_weights = torch.load(path + 'droog_clf_100.pt')
# model.load_state_dict(trained_weights)<jupyter_output><empty_output><jupyter_text>Make predictions Let's check our classification accuracy.<jupyter_code>%%time
# Make predictions in test set 
with torch.no_grad():
    model.eval()
    y_pred = mu.supervised_model_predict(
        model,
        val_loader,
        criterion,
        multiclass = True, 
        n_outputs = n_cats,
        score = True
    )
acc = (y_pred.argmax(axis = 1) == test_adata.obs.sample_codes.values).sum() / test_adata.n_obs
print('Accuracy of the model is %.3f'%acc)<jupyter_output><empty_output><jupyter_text>As in the GCN example, we can compute the confusion matrix of the predictions.<jupyter_code>conf_mat = mu.confusion_matrix(df_proj.y_pred, df_proj.sample_codes)
df_conf_mat = pd.DataFrame(
    conf_mat, 
    columns = uniques,
    index = uniques
)
df_conf_mat.head()<jupyter_output><empty_output><jupyter_text>Project to latent space Let's initialize a new data loader to get the cell low dimensional embeddings.<jupyter_code>projection_dataset = mu.adata_torch_dataset(
    data = a,
    transform = transforms.ToTensor(),
    supervised = True,
    target_col = 'sample_codes'
)

# Initialize DataLoader for projection
projection_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    drop_last=False,
    shuffle=False,
    num_workers=8,
)<jupyter_output><empty_output><jupyter_text>I didn't go through this in the last tutorial, but the signature call of the projection function returns a [generator](https://realpython.com/introduction-to-python-generators/) to make it a bit more efficient, that's why the model is wrapped under a list call, and then under a `np.array()` function to return the embeddings as a matrix of size `n_cells` by `n_dimensions`. In our case, the dimension of the last hidden layer was 64.<jupyter_code># Returns a generator
x = model.project_to_latent_space(
    data_loader = projection_loader,
    n_feats = n_genes, 
    latent_dim = dims[-2] # last hidden dimension (64)
)
type(x)
# Get the first embedding
embedding = next(iter(x))
print('Size of the cell embeddings is %d.'%embedding.shape[0])<jupyter_output><empty_output><jupyter_text>Let's get the actual embeddings this time and save them into a dataframe for further visualization.<jupyter_code>%%time
with torch.no_grad():
    model.eval()

    projection_arr = np.array(
        list(model.project_to_latent_space(projection_loader, dims[0], dims[-2]))
    )
df = pd.DataFrame(
    projection_arr, columns = ['latent_' + str(i) for i in range(1, dims[-2] + 1)]
)
df_proj = pd.concat(
    [test_adata.obs, df.set_index(test_adata.obs.index)], axis=1
)

# Add model predictions to df
df_proj['y_pred'] = y_pred.argmax(axis = 1)
df_proj.head()<jupyter_output><empty_output><jupyter_text>Reproducibility<jupyter_code>%load_ext watermark

%watermark -m -v -p numpy,torch,pandas,anndata,sklearn,scipy<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08302999874482239, 'fraction_numerical': 0.009162796535709804, 'mean_word_length': 5.02913356034809, 'pattern_counts': {'":': 0, '<': 103, '<?xml version=': 0, '>': 103, 'https://': 7, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '25', 'dir': 'jupyter-structured-clean-dedup', 'id': '395532', 'license': 'MIT', 'n_tokens_mistral': 5095, 'n_tokens_neox': 4838, 'n_words': 1808, 'path': 'notebooks/cell_softmax_clf_100.ipynb', 'repo_name': 'manuflores/sandbox', 'size': '15937'}