"""
Code to train a model
"""
import os
import sys
import logging
import argparse
import copy
import functools
import itertools
import numpy as np
import pandas as pd
import scipy.spatial
import scanpy as sc
import matplotlib.pyplot as plt
from skorch.helper import predefined_split
import torch
import torch.nn as nn
import torch.nn.functional as F
import skorch
import skorch.helper
torch.backends.cudnn.deterministic = True # For reproducibility
torch.backends.cudnn.benchmark = False
SRC_DIR = os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "babel"
)
assert os.path.isdir(SRC_DIR)
sys.path.append(SRC_DIR)
MODELS_DIR = os.path.join(SRC_DIR, "models")
assert os.path.isdir(MODELS_DIR)
sys.path.append(MODELS_DIR)
import sc_data_loaders
import adata_utils
import model_utils
import autoencoders
import loss_functions
import layers
import activations
import plot_utils
import utils
import metrics
import interpretation
logging.basicConfig(level=logging.INFO)
OPTIMIZER_DICT = {
"adam": torch.optim.Adam,
"rmsprop": torch.optim.RMSprop,
}
def build_parser():
"""Build argument parser"""
parser = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
input_group = parser.add_mutually_exclusive_group(required=True)
input_group.add_argument(
"--data", "-d", type=str, nargs="*", help="Data files to train on",
)
input_group.add_argument(
"--snareseq",
action="store_true",
help="Data in SNAREseq format, use custom data loading logic for separated RNA ATC files",
)
input_group.add_argument(
"--shareseq",
nargs="+",
type=str,
choices=["lung", "skin", "brain"],
help="Load in the given SHAREseq datasets",
)
parser.add_argument(
"--nofilter",
action="store_true",
help="Whether or not to perform filtering (only applies with --data argument)",
)
parser.add_argument(
"--linear",
action="store_true",
help="Do clustering data splitting in linear instead of log space",
)
parser.add_argument(
"--clustermethod",
type=str,
choices=["leiden", "louvain"],
default="leiden",
help="Clustering method to determine data splits",
)
parser.add_argument(
"--validcluster", type=int, default=0, help="Cluster ID to use as valid cluster"
)
parser.add_argument(
"--testcluster", type=int, default=1, help="Cluster ID to use as test cluster"
)
parser.add_argument(
"--outdir", "-o", required=True, type=str, help="Directory to output to"
)
parser.add_argument(
"--naive",
"-n",
action="store_true",
help="Use a naive model instead of lego model",
)
parser.add_argument(
"--hidden", type=int, nargs="*", default=[16], help="Hidden dimensions"
)
parser.add_argument(
"--pretrain",
type=str,
default="",
help="params.pt file to use to warm initialize the model (instead of starting from scratch)",
)
parser.add_argument(
"--lossweight",
type=float,
nargs="*",
default=[1.33],
help="Relative loss weight",
)
parser.add_argument(
"--optim",
type=str,
default="adam",
choices=OPTIMIZER_DICT.keys(),
help="Optimizer to use",
)
parser.add_argument(
"--lr", "-l", type=float, default=[0.01], nargs="*", help="Learning rate"
)
parser.add_argument(
"--batchsize", "-b", type=int, nargs="*", default=[512], help="Batch size"
)
parser.add_argument(
"--earlystop", type=int, default=25, help="Early stopping after N epochs"
)
parser.add_argument(
"--seed", type=int, nargs="*", default=[182822], help="Random seed to use"
)
parser.add_argument("--device", default=0, type=int, help="Device to train on")
parser.add_argument(
"--ext",
type=str,
choices=["png", "pdf", "jpg"],
default="pdf",
help="Output format for plots",
)
return parser
def plot_loss_history(history, fname: str):
"""Create a plot of train valid loss"""
fig, ax = plt.subplots(dpi=300)
ax.plot(
np.arange(len(history)), history[:, "train_loss"], label="Train",
)
ax.plot(
np.arange(len(history)), history[:, "valid_loss"], label="Valid",
)
ax.legend()
ax.set(
xlabel="Epoch", ylabel="Loss",
)
fig.savefig(fname)
return fig
def main():
"""Run the script"""
parser = build_parser()
args = parser.parse_args()
args.outdir = os.path.abspath(args.outdir)
if not os.path.isdir(os.path.dirname(args.outdir)):
os.makedirs(os.path.dirname(args.outdir))
# Specify output log file
logger = logging.getLogger()
fh = logging.FileHandler(f"{args.outdir}_training.log", "w")
fh.setLevel(logging.INFO)
logger.addHandler(fh)
# Log parameters and pytorch version
if torch.cuda.is_available():
logging.info(f"PyTorch CUDA version: {torch.version.cuda}")
for arg in vars(args):
logging.info(f"Parameter {arg}: {getattr(args, arg)}")
# Borrow parameters
logging.info("Reading RNA data")
if args.snareseq:
rna_data_kwargs = copy.copy(sc_data_loaders.SNARESEQ_RNA_DATA_KWARGS)
elif args.shareseq:
logging.info(f"Loading in SHAREseq RNA data for: {args.shareseq}")
rna_data_kwargs = copy.copy(sc_data_loaders.SNARESEQ_RNA_DATA_KWARGS)
rna_data_kwargs["fname"] = None
rna_data_kwargs["reader"] = None
rna_data_kwargs["cell_info"] = None
rna_data_kwargs["gene_info"] = None
rna_data_kwargs["transpose"] = False
# Load in the datasets
shareseq_rna_adatas = []
for tissuetype in args.shareseq:
shareseq_rna_adatas.append(
adata_utils.load_shareseq_data(
tissuetype,
dirname="/data/wukevin/commonspace_data/GSE140203_SHAREseq",
mode="RNA",
)
)
shareseq_rna_adata = shareseq_rna_adatas[0]
if len(shareseq_rna_adatas) > 1:
shareseq_rna_adata = shareseq_rna_adata.concatenate(
*shareseq_rna_adatas[1:],
join="inner",
batch_key="tissue",
batch_categories=args.shareseq,
)
rna_data_kwargs["raw_adata"] = shareseq_rna_adata
else:
rna_data_kwargs = copy.copy(sc_data_loaders.TENX_PBMC_RNA_DATA_KWARGS)
rna_data_kwargs["fname"] = args.data
if args.nofilter:
rna_data_kwargs = {
k: v for k, v in rna_data_kwargs.items() if not k.startswith("filt_")
}
rna_data_kwargs["data_split_by_cluster_log"] = not args.linear
rna_data_kwargs["data_split_by_cluster"] = args.clustermethod
sc_rna_dataset = sc_data_loaders.SingleCellDataset(
valid_cluster_id=args.validcluster,
test_cluster_id=args.testcluster,
**rna_data_kwargs,
)
sc_rna_train_dataset = sc_data_loaders.SingleCellDatasetSplit(
sc_rna_dataset, split="train",
)
sc_rna_valid_dataset = sc_data_loaders.SingleCellDatasetSplit(
sc_rna_dataset, split="valid",
)
sc_rna_test_dataset = sc_data_loaders.SingleCellDatasetSplit(
sc_rna_dataset, split="test",
)
# ATAC
logging.info("Aggregating ATAC clusters")
if args.snareseq:
atac_data_kwargs = copy.copy(sc_data_loaders.SNARESEQ_ATAC_DATA_KWARGS)
elif args.shareseq:
logging.info(f"Loading in SHAREseq ATAC data for {args.shareseq}")
atac_data_kwargs = copy.copy(sc_data_loaders.SNARESEQ_ATAC_DATA_KWARGS)
atac_data_kwargs["reader"] = None
atac_data_kwargs["fname"] = None
atac_data_kwargs["cell_info"] = None
atac_data_kwargs["gene_info"] = None
atac_data_kwargs["transpose"] = False
atac_adatas = []
for tissuetype in args.shareseq:
atac_adatas.append(
adata_utils.load_shareseq_data(
tissuetype,
dirname="/data/wukevin/commonspace_data/GSE140203_SHAREseq",
mode="ATAC",
)
)
atac_bins = [a.var_names for a in atac_adatas]
if len(atac_adatas) > 1:
atac_bins_harmonized = sc_data_loaders.harmonize_atac_intervals(*atac_bins)
atac_adatas = [
sc_data_loaders.repool_atac_bins(a, atac_bins_harmonized)
for a in atac_adatas
]
shareseq_atac_adata = atac_adatas[0]
if len(atac_adatas) > 1:
shareseq_atac_adata = shareseq_atac_adata.concatenate(
*atac_adatas[1:],
join="inner",
batch_key="tissue",
batch_categories=args.shareseq,
)
atac_data_kwargs["raw_adata"] = shareseq_atac_adata
else:
atac_parsed = [
utils.sc_read_10x_h5_ft_type(fname, "Peaks") for fname in args.data
]
if len(atac_parsed) > 1:
atac_bins = sc_data_loaders.harmonize_atac_intervals(
atac_parsed[0].var_names, atac_parsed[1].var_names
)
for bins in atac_parsed[2:]:
atac_bins = sc_data_loaders.harmonize_atac_intervals(
atac_bins, bins.var_names
)
logging.info(f"Aggregated {len(atac_bins)} bins")
else:
atac_bins = list(atac_parsed[0].var_names)
atac_data_kwargs = copy.copy(sc_data_loaders.TENX_PBMC_ATAC_DATA_KWARGS)
atac_data_kwargs["fname"] = rna_data_kwargs["fname"]
atac_data_kwargs["pool_genomic_interval"] = 0 # Do not pool
atac_data_kwargs["reader"] = functools.partial(
utils.sc_read_multi_files,
reader=lambda x: sc_data_loaders.repool_atac_bins(
utils.sc_read_10x_h5_ft_type(x, "Peaks"), atac_bins,
),
)
atac_data_kwargs["cluster_res"] = 0 # Do not bother clustering ATAC data
sc_atac_dataset = sc_data_loaders.SingleCellDataset(
predefined_split=sc_rna_dataset, **atac_data_kwargs
)
sc_atac_train_dataset = sc_data_loaders.SingleCellDatasetSplit(
sc_atac_dataset, split="train",
)
sc_atac_valid_dataset = sc_data_loaders.SingleCellDatasetSplit(
sc_atac_dataset, split="valid",
)
sc_atac_test_dataset = sc_data_loaders.SingleCellDatasetSplit(
sc_atac_dataset, split="test",
)
sc_dual_train_dataset = sc_data_loaders.PairedDataset(
sc_rna_train_dataset, sc_atac_train_dataset, flat_mode=True,
)
sc_dual_valid_dataset = sc_data_loaders.PairedDataset(
sc_rna_valid_dataset, sc_atac_valid_dataset, flat_mode=True,
)
sc_dual_test_dataset = sc_data_loaders.PairedDataset(
sc_rna_test_dataset, sc_atac_test_dataset, flat_mode=True,
)
sc_dual_full_dataset = sc_data_loaders.PairedDataset(
sc_rna_dataset, sc_atac_dataset, flat_mode=True,
)
# Model
param_combos = list(
itertools.product(
args.hidden, args.lossweight, args.lr, args.batchsize, args.seed
)
)
for h_dim, lw, lr, bs, rand_seed in param_combos:
outdir_name = (
f"{args.outdir}_hidden_{h_dim}_lossweight_{lw}_lr_{lr}_batchsize_{bs}_seed_{rand_seed}"
if len(param_combos) > 1
else args.outdir
)
if not os.path.isdir(outdir_name):
assert not os.path.exists(outdir_name)
os.makedirs(outdir_name)
assert os.path.isdir(outdir_name)
with open(os.path.join(outdir_name, "rna_genes.txt"), "w") as sink:
for gene in sc_rna_dataset.data_raw.var_names:
sink.write(gene + "\n")
with open(os.path.join(outdir_name, "atac_bins.txt"), "w") as sink:
for atac_bin in sc_atac_dataset.data_raw.var_names:
sink.write(atac_bin + "\n")
# Write dataset
### Full
sc_rna_dataset.size_norm_counts.write_h5ad(
os.path.join(outdir_name, "full_rna.h5ad")
)
sc_rna_dataset.size_norm_log_counts.write_h5ad(
os.path.join(outdir_name, "full_rna_log.h5ad")
)
sc_atac_dataset.data_raw.write_h5ad(os.path.join(outdir_name, "full_atac.h5ad"))
### Train
sc_rna_train_dataset.size_norm_counts.write_h5ad(
os.path.join(outdir_name, "train_rna.h5ad")
)
sc_atac_train_dataset.data_raw.write_h5ad(
os.path.join(outdir_name, "train_atac.h5ad")
)
### Valid
sc_rna_valid_dataset.size_norm_counts.write_h5ad(
os.path.join(outdir_name, "valid_rna.h5ad")
)
sc_atac_valid_dataset.data_raw.write_h5ad(
os.path.join(outdir_name, "valid_atac.h5ad")
)
### Test
sc_rna_test_dataset.size_norm_counts.write_h5ad(
os.path.join(outdir_name, "truth_rna.h5ad")
)
sc_atac_dataset.data_raw.write_h5ad(os.path.join(outdir_name, "full_atac.h5ad"))
sc_atac_test_dataset.data_raw.write_h5ad(
os.path.join(outdir_name, "truth_atac.h5ad")
)
# Instantiate and train model
model_class = (
autoencoders.NaiveSplicedAutoEncoder
if args.naive
else autoencoders.AssymSplicedAutoEncoder
)
spliced_net = autoencoders.SplicedAutoEncoderSkorchNet(
module=model_class,
module__hidden_dim=h_dim, # Based on hyperparam tuning
module__input_dim1=sc_rna_dataset.data_raw.shape[1],
module__input_dim2=sc_atac_dataset.get_per_chrom_feature_count(),
module__final_activations1=[
activations.Exp(),
activations.ClippedSoftplus(),
],
module__final_activations2=nn.Sigmoid(),
module__flat_mode=True,
module__seed=rand_seed,
lr=lr, # Based on hyperparam tuning
criterion=loss_functions.QuadLoss,
criterion__loss2=loss_functions.BCELoss, # handle output of encoded layer
criterion__loss2_weight=lw, # numerically balance the two losses with different magnitudes
criterion__record_history=True,
optimizer=OPTIMIZER_DICT[args.optim],
iterator_train__shuffle=True,
device=utils.get_device(args.device),
batch_size=bs, # Based on hyperparam tuning
max_epochs=500,
callbacks=[
skorch.callbacks.EarlyStopping(patience=args.earlystop),
skorch.callbacks.LRScheduler(
policy=torch.optim.lr_scheduler.ReduceLROnPlateau,
**model_utils.REDUCE_LR_ON_PLATEAU_PARAMS,
),
skorch.callbacks.GradientNormClipping(gradient_clip_value=5),
skorch.callbacks.Checkpoint(
dirname=outdir_name, fn_prefix="net_", monitor="valid_loss_best",
),
],
train_split=skorch.helper.predefined_split(sc_dual_valid_dataset),
iterator_train__num_workers=8,
iterator_valid__num_workers=8,
)
if args.pretrain:
# Load in the warm start parameters
spliced_net.load_params(f_params=args.pretrain)
spliced_net.partial_fit(sc_dual_train_dataset, y=None)
else:
spliced_net.fit(sc_dual_train_dataset, y=None)
fig = plot_loss_history(
spliced_net.history, os.path.join(outdir_name, f"loss.{args.ext}")
)
plt.close(fig)
logging.info("Evaluating on test set")
logging.info("Evaluating RNA > RNA")
sc_rna_test_preds = spliced_net.translate_1_to_1(sc_dual_test_dataset)
sc_rna_test_preds_anndata = sc.AnnData(
sc_rna_test_preds,
var=sc_rna_test_dataset.data_raw.var,
obs=sc_rna_test_dataset.data_raw.obs,
)
sc_rna_test_preds_anndata.write_h5ad(
os.path.join(outdir_name, "rna_rna_test_preds.h5ad")
)
fig = plot_utils.plot_scatter_with_r(
sc_rna_test_dataset.size_norm_counts.X,
sc_rna_test_preds,
one_to_one=True,
logscale=True,
density_heatmap=True,
title="RNA > RNA (test set)",
fname=os.path.join(outdir_name, f"rna_rna_scatter_log.{args.ext}"),
)
plt.close(fig)
logging.info("Evaluating ATAC > ATAC")
sc_atac_test_preds = spliced_net.translate_2_to_2(sc_dual_test_dataset)
sc_atac_test_preds_anndata = sc.AnnData(
sc_atac_test_preds,
var=sc_atac_test_dataset.data_raw.var,
obs=sc_atac_test_dataset.data_raw.obs,
)
sc_atac_test_preds_anndata.write_h5ad(
os.path.join(outdir_name, "atac_atac_test_preds.h5ad")
)
fig = plot_utils.plot_auroc(
sc_atac_test_dataset.data_raw.X,
sc_atac_test_preds,
title_prefix="ATAC > ATAC",
fname=os.path.join(outdir_name, f"atac_atac_auroc.{args.ext}"),
)
plt.close(fig)
logging.info("Evaluating ATAC > RNA")
sc_atac_rna_test_preds = spliced_net.translate_2_to_1(sc_dual_test_dataset)
sc_atac_rna_test_preds_anndata = sc.AnnData(
sc_atac_rna_test_preds,
var=sc_rna_test_dataset.data_raw.var,
obs=sc_rna_test_dataset.data_raw.obs,
)
sc_atac_rna_test_preds_anndata.write_h5ad(
os.path.join(outdir_name, "atac_rna_test_preds.h5ad")
)
fig = plot_utils.plot_scatter_with_r(
sc_rna_test_dataset.size_norm_counts.X,
sc_atac_rna_test_preds,
one_to_one=True,
logscale=True,
density_heatmap=True,
title="ATAC > RNA (test set)",
fname=os.path.join(outdir_name, f"atac_rna_scatter_log.{args.ext}"),
)
plt.close(fig)
logging.info("Evaluating RNA > ATAC")
sc_rna_atac_test_preds = spliced_net.translate_1_to_2(sc_dual_test_dataset)
sc_rna_atac_test_preds_anndata = sc.AnnData(
sc_rna_atac_test_preds,
var=sc_atac_test_dataset.data_raw.var,
obs=sc_atac_test_dataset.data_raw.obs,
)
sc_rna_atac_test_preds_anndata.write_h5ad(
os.path.join(outdir_name, "rna_atac_test_preds.h5ad")
)
fig = plot_utils.plot_auroc(
sc_atac_test_dataset.data_raw.X,
sc_rna_atac_test_preds,
title_prefix="RNA > ATAC",
fname=os.path.join(outdir_name, f"rna_atac_auroc.{args.ext}"),
)
plt.close(fig)
del spliced_net
if __name__ == "__main__":
main()
Datasets
Models
