recipes.py

#!/usr/bin/env python
# recipes.py
#
# Entry Point for Variant Filtration with Neural Nets
# The design philosophy is to have small functions in here.
# High-level descriptions to create training sets, architectures, and evaluations.
# Each function can be considered a recipe for cooking up a neural net.
# The gory details of finding data should be in training_data.py
# Model architectures and optimizations should be in models.py
#
# December 2016
# Sam Friedman 
# sam@broadinstitute.org

# Python 2/3 friendly
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

# Imports
import os
import sys
import vcf
import h5py
import time
import plots
import pysam
import models
import defines
import inspect
import operator
import arguments
import numpy as np
import training_data as td
from collections import Counter


def run(file_fxns):
	'''Dispatch on args.mode command-line supplied recipe
	
	Any function defined in this file is considered a recipe.
	They must take the args namespace as input.
	Recipes are run via:

	python recipes.py <NAME_OF_RECIPE_FUNCTION> [Additional arguments]

	Arguments:
		file_fxns: a dict mapping function names as strings to function objects
	'''	
	args = arguments.parse_args()

	if args.mode in file_fxns:
		ret = file_fxns[args.mode](args)
		print(args.mode, 'Recipe completed, returned:', ret)
	else:
		raise ValueError('Unknown recipe mode:', args.mode)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ~~~~~~~ Variant Calling Recipes ~~~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def train_calling_model(args):
	'''Trains the variant calling as 1D segmentation CNN architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train. 

	This architecture looks at read_tensors and predicts site_labels.
	Tensors must be generated by calling td.write_calling_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.
	'''
	args.labels = defines.calling_labels
	model = models.build_2d_cnn_calling_segmentation_1d(args)
	
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths_all(args)
	generate_train = td.calling_tensors_generator(args, train_paths)
	generate_valid = td.calling_tensors_generator(args, valid_paths)
	generate_test = td.calling_tensors_generator(args, test_paths)

	weight_path = arguments.weight_path_from_args(args)
	if args.inspect_model:
		models.inspect_model(args, model, generate_train, generate_valid, weight_path.replace('.hd5', '.png'))
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test_tensors = np.zeros((args.iterations*args.batch_size,) + defines.tensor_shape_from_args(args))
	test_labels = np.zeros((args.iterations*args.batch_size, args.window_size, len(args.labels)))

	for i in range(args.iterations):
		next_batch = next(generate_test)
		test_tensors[i*args.batch_size:(i+1)*args.batch_size,:,:,:] = next_batch[0][args.tensor_map]
		test_labels[i*args.batch_size:(i+1)*args.batch_size,:] = next_batch[1]

	predictions = model.predict(test_tensors)
	print('prediction shape:', predictions.shape)

	melt_shape = (predictions.shape[0]*predictions.shape[1], predictions.shape[2])
	predictions = predictions.reshape(melt_shape)
	test_truth = test_labels.reshape(melt_shape)
	print('predictions reshaped:', predictions.shape, 'np sum truth:\n', np.sum(test_truth, axis=0), '\nnp sum pred:\n', np.sum(predictions, axis=0))
		
	plots.plot_precision_recall_per_class_predictions(predictions, test_truth, args.labels, args.id)


def train_calling_model_full(args):
	'''Trains the variant calling as 1D segmentation CNN architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train. 

	This architecture looks at read_tensors and predicts site_labels.
	Tensors must be generated by calling td.write_calling_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.
	'''
	args.labels = defines.calling_labels
	model = models.build_2d_cnn_calling_segmentation_full_2d(args)

	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths_all(args)

	generate_train = td.calling_tensors_generator(args, train_paths)
	generate_valid = td.calling_tensors_generator(args, valid_paths)
	generate_test = td.calling_tensors_generator(args, test_paths)

	weight_path = arguments.weight_path_from_args(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test_tensors = np.zeros((args.iterations*args.batch_size,) + defines.tensor_shape_from_args(args))
	test_labels = np.zeros((args.iterations*args.batch_size, args.window_size, len(args.labels)))

	for i in range(args.iterations):
		next_batch = next(generate_test)
		test_tensors[i*args.batch_size:(i+1)*args.batch_size,:,:,:] = next_batch[0][args.tensor_map]
		test_labels[i*args.batch_size:(i+1)*args.batch_size,:] = next_batch[1]

	predictions = model.predict(test_tensors)
	print('prediction shape:', predictions.shape)

	melt_shape = (predictions.shape[0]*predictions.shape[1], predictions.shape[2])
	predictions = predictions.reshape(melt_shape)
	test_truth = test_labels.reshape(melt_shape)
	print('predictions reshaped:', predictions.shape, 'np sum truth:', np.sum(test_truth, axis=0), '\nnp sum pred :', np.sum(predictions, axis=0))
		
	plots.plot_precision_recall_per_class_predictions(predictions, test_truth, args.labels, args.id)


def train_calling_model_1d(args):
	'''Trains the variant calling as 1D segmentation CNN architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train. 

	This architecture looks at read_tensors and predicts site_labels.
	Tensors must be generated by calling td.write_calling_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.
	'''
	args.labels = defines.calling_labels

	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths_all(args)

	generate_train = td.calling_pileup_tensors_generator(args, train_paths)
	generate_valid = td.calling_pileup_tensors_generator(args, valid_paths)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_1d_cnn_calling_segmentation_1d(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	args.batch_size = args.samples # hack
	generate_test = td.calling_pileup_tensors_generator(args, test_paths)
	test_batch = next(generate_test)
	predictions = model.predict(test_batch[0])
	for i in range(30):
		print('\n\n\npredictions ', i,' is:\n', np.argmax(predictions[i,:,:], axis =-1), '\n for truth labels:\n', np.argmax(test_batch[1][i,:,:], axis=-1))
	
	print('prediction shape:', predictions.shape)

	melt_shape = (predictions.shape[0]*predictions.shape[1], predictions.shape[2])
	predictions = predictions.reshape(melt_shape)
	test_truth = test_batch[1].reshape(melt_shape)
	print('prediction shape:', predictions.shape)
	print('np sum', np.sum(test_truth, axis=0), 'np sum pred:', np.sum(predictions, axis=0))
	
	pred_subset = []
	truth_subset = []
	for i in range(predictions.shape[0]):
		if test_truth[i][0] == 1 and np.argmax(predictions[i]) == 0:
			continue
		pred_subset.append(predictions[i])
		truth_subset.append(test_truth[i])

	print('np sum', np.sum(np.array(truth_subset), axis=0), 'np sum pred:', np.sum(np.array(pred_subset), axis=0))
	
	#plots.plot_roc_per_class(model, test_batch[0], test_batch[1], args.labels, title_suffix, melt=True)
	plots.plot_roc_per_class_predictions(np.array(pred_subset), np.array(truth_subset), args.labels, args.id+'_no_reference_tp')
	plots.plot_roc_per_class_predictions(predictions, test_truth, args.labels, args.id)


def train_pileup_filter(args):
	'''Trains a variant filtration CNN architecture on pileup tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train. 

	This architecture looks at pileup_tensors centered at a variant and predicts if the variant is real.
	Tensors must be generated by calling td.write_pileup_tensors() before this function is used.
	After training with early stopping ROC curves are plotted on the test dataset.
	'''

	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	generate_train = td.pileup_tensor_generator(args, train_paths)
	generate_valid = td.pileup_tensor_generator(args, valid_paths)
	generate_test = td.pileup_tensor_generator(args, test_paths)

	test_pileups = np.zeros((args.iterations*args.batch_size, args.window_size, defines.get_reference_and_read_channels(args)))
	test_labels = np.zeros((args.iterations*args.batch_size, len(args.labels)))

	for i in range(args.iterations):
		next_batch = next(generate_test)
		test_pileups[i*args.batch_size:(i+1)*args.batch_size,:,:] = next_batch[0]['pileup_tensor']
		test_labels[i*args.batch_size:(i+1)*args.batch_size,:] = next_batch[1]

	weight_path = arguments.weight_path_from_args(args)
	
	model = models.build_pileup_filter(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
	
	plots.plot_roc_per_class(model, test_pileups, test_labels, args.labels, args.id)


def test_caller_pileup(args):
	'''Tests the variant calling as 1D segmentation CNN architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train. 

	This architecture looks at pileup_tensors and predicts site_labels.
	Tensors must be generated by calling td.write_calling_tensors() before this function is used.
	Performance curves are plotted on the test dataset.
	'''
	args.labels = defines.calling_labels

	_, _, test_paths = td.get_train_valid_test_paths_all(args)
	generate_test = td.calling_pileup_tensors_generator(args, test_paths)

	model = models.build_1d_cnn_calling_segmentation_1d(args)

	test_pileups = np.zeros((args.iterations*args.batch_size, args.window_size, defines.get_reference_and_read_channels(args)))
	test_labels = np.zeros((args.iterations*args.batch_size, args.window_size, len(args.labels)))

	for i in range(args.iterations):
		next_batch = next(generate_test)
		test_pileups[i*args.batch_size:(i+1)*args.batch_size,:,:] = next_batch[0]['pileup_tensor']
		test_labels[i*args.batch_size:(i+1)*args.batch_size,:] = next_batch[1]

	predictions = model.predict(test_pileups)
	print('prediction shape:', predictions.shape)

	melt_shape = (predictions.shape[0]*predictions.shape[1], predictions.shape[2])
	predictions = predictions.reshape(melt_shape)
	test_truth = test_labels.reshape(melt_shape)
	print('prediction shape:', predictions.shape)
	print('np sum truth:', np.sum(test_truth, axis=0), '\nnp sum pred :', np.sum(predictions, axis=0))
	
	pred_subset = []
	truth_subset = []
	for i in range(predictions.shape[0]):
		if test_truth[i][0] == 1 and np.argmax(predictions[i]) == 0:
			continue
		pred_subset.append(predictions[i])
		truth_subset.append(test_truth[i])

	print('np sum truth:', np.sum(np.array(truth_subset), axis=0), '\nnp sum pred :', np.sum(np.array(pred_subset), axis=0))
	
	plots.plot_roc_per_class_predictions(np.array(pred_subset), np.array(truth_subset), args.labels, args.id+'_no_reference_tp')
	plots.plot_roc_per_class_predictions(predictions, test_truth, args.labels, args.id)
	plots.plot_precision_recall_per_class_predictions(predictions, test_truth, args.labels, args.id)


def test_caller_2d(args):
	'''Tests the variant calling as 1D segmentation CNN architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train. 

	This architecture looks at read_tensors and predicts site_labels.
	Tensors are generated by calling td.write_calling_tensors() before this function is used.
	Performance curves are plotted on the test dataset.
	'''
	args.labels = defines.calling_labels

	_, _, test_paths = td.get_train_valid_test_paths_all(args)
	generate_test = td.calling_tensors_generator(args, test_paths)

	#model = models.build_2d_cnn_calling_segmentation_1d(args)
	model = models.load_model(args.weights_hd5, custom_objects=models.get_all_custom_objects(args.labels))
	model.summary()

	test_tensors = np.zeros((args.iterations*args.batch_size,) + defines.tensor_shape_from_args(args))
	test_labels = np.zeros((args.iterations*args.batch_size, args.window_size, len(args.labels)))

	for i in range(args.iterations):
		next_batch = next(generate_test)
		test_tensors[i*args.batch_size:(i+1)*args.batch_size,:,:,:] = next_batch[0][args.tensor_map]
		test_labels[i*args.batch_size:(i+1)*args.batch_size,:] = next_batch[1]

	predictions = model.predict(test_tensors)
	print('prediction shape:', predictions.shape)

	melt_shape = (predictions.shape[0]*predictions.shape[1], predictions.shape[2])
	predictions = predictions.reshape(melt_shape)
	test_truth = test_labels.reshape(melt_shape)
	print('prediction shape:', predictions.shape)
	print('np sum truth:', np.sum(test_truth, axis=0), '\nnp sum pred :', np.sum(predictions, axis=0))
	
	pred_subset = []
	truth_subset = []
	for i in range(predictions.shape[0]):
		if test_truth[i][0] == 1 and np.argmax(predictions[i]) == 0:
			continue
		pred_subset.append(predictions[i])
		truth_subset.append(test_truth[i])

	print('np sum truth:', np.sum(np.array(truth_subset), axis=0), '\nnp sum pred :', np.sum(np.array(pred_subset), axis=0))
	
	plots.plot_roc_per_class_predictions(np.array(pred_subset), np.array(truth_subset), args.labels, args.id+'_no_reference_tp')
	plots.plot_roc_per_class_predictions(predictions, test_truth, args.labels, args.id)
	plots.plot_precision_recall_per_class_predictions(predictions, test_truth, args.labels, args.id)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ~~~~~~~ Variant Filtering Recipes ~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def train_annotation_multilayer_perceptron(args):
	'''Train a Multilayer Perceptron (MLP) Annotation architecture (no convolutions).

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		
	Annotation tensors must be generated by calling 
	td.write_dna_tensors(include_dna=False) before this function is used.
	Performance curves for MLP are plotted on the test dataset.
	'''		
	args.tensor_map = None
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.annotation_multilayer_perceptron_from_args(args,
											fc_layers = [128, 128, 256, 256, 128, 128],
											dropout = 0.3,
											skip_connection = True,
											batch_normalize_input = True)
	
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
	
	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	plots.plot_roc_per_class(model, [test[0][args.annotation_set]], test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, [test[0][args.annotation_set]], test[1], args.labels)


def train_mlp_selu(args):
	'''Train a Multilayer Perceptron (MLP) Annotation architecture (no convolutions).

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		
	Annotation tensors must be generated by calling 
	td.write_dna_tensors(include_dna=False) before this function is used.
	Performance curves for MLP are plotted on the test dataset.
	'''		
	args.tensor_map = None
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.annotation_multilayer_perceptron_from_args(args,
											fc_layers = [32, 48, 64, 96, 128, 96, 64, 48, 32],
											dropout = 0,
											activation = 'selu',
											#batch_normalization = True,
											batch_normalize_input = True,
											skip_connection = False)

	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
	
	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	plots.plot_roc_per_class(model, [test[0][args.annotation_set]], test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, [test[0][args.annotation_set]], test[1], args.labels)


def train_reference_only(args):
	'''Train a 1D Convolution architecture on reference tracks.

	Arguments
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		
	Tensors must be generated by calling 
	td.write_dna_multisource_annotations() 
	before this function is used.
	Performance curves for CNN are plotted on the test dataset.
	'''	
	args.tensor_map = 'reference'
	args.annotation_set = '_'
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_reference_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	plots.plot_roc_per_class(model, [test[0][args.tensor_map]], test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, [test[0][args.tensor_map]], test[1], args.labels)


def train_reference_annotation(args):
	'''Train a 1D Convolution plus reference tracks and MLP Annotation architecture.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		
	Annotation tensors must be generated by calling 
	td.write_dna_multisource_annotations() before this function is used.
	Performance curves for CNN are plotted on the test dataset.
	'''
	args.tensor_map = 'reference'
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	#model = models.build_reference_plus_model(args)
	model = models.build_reference_annotation_skip_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


def train_reference_annotation_b(args):
	'''Train a 1D Convolution plus MLP Annotation with skip connection architecture.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		
	Annotation tensors must be generated by calling 
	td.write_dna_multisource_annotations() before this function is used.
	Performance curves for CNN are plotted on the test dataset.
	'''
	args.tensor_map = 'reference'

	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_reference_annotation_1d_model_from_args(args, 
											conv_width = 7, 
											conv_layers = [256, 216, 128, 64, 32],
											conv_dropout = 0.1, # was .3
											conv_batch_normalize = True,
											spatial_dropout = True,
											max_pools = [],
											padding = 'same',
											annotation_units = 64,
											annotation_shortcut = True,
											fc_layers = [64, 64],
											fc_dropout = 0.2,
											annotation_batch_normalize = not args.normalize_annotations,
											fc_batch_normalize = False)
	
	model.save(weight_path)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


def train_reference_annotation_c(args):
	'''Train a 1D Convolution plus MLP Annotation with skip connection architecture.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		
	Annotation tensors must be generated by calling 
	td.write_dna_multisource_annotations() before this function is used.
	Performance curves for CNN are plotted on the test dataset.
	'''
	args.tensor_map = 'reference'

	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_reference_annotation_1d_model_from_args(args, 
											conv_width = 11, 
											conv_layers = [216, 180, 128, 64, 32],
											conv_dropout = 0.0, # was .3
											conv_batch_normalize = False,
											max_pools = [],
											padding = 'valid',
											activation = 'relu',
											annotation_units = 64,
											annotation_shortcut = True,
											annotation_batch_normalize = True,
											fc_layers = [64, 64],
											fc_dropout = 0.2,
											fc_batch_normalize = False)
	
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)




def train_reference_annotation_1layer(args):
	'''Train a 1D Convolution plus reference tracks and MLP Annotation architecture.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		
	Annotation tensors must be generated by calling 
	td.write_dna_multisource_annotations() before this function is used.
	Performance curves for CNN are plotted on the test dataset.
	'''	
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_reference_1d_1layer_model(args)
	for i in range(args.epochs):
		model.save(args.output_dir + args.id + '_epoch_' + str(i) + '.hd5')
		model.fit_generator(generate_train, 
			steps_per_epoch=args.training_steps, epochs=1, verbose=1, 
			validation_steps=args.validation_steps, validation_data=generate_valid)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


def train_ref_read_model(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	in_channels = defines.total_input_channels_from_args(args)
	if args.channels_last:
		tensor_shape = (args.read_limit, args.window_size, in_channels)
	else:
		tensor_shape = (in_channels, args.read_limit, args.window_size) 

	generate_train = td.tensor_generator(args, train_paths, tensor_shape)
	generate_valid = td.tensor_generator(args, valid_paths, tensor_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_read_tensor_2d_model(args)
	models.serialize_model_semantics(args, weight_path)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.load_tensors_from_class_dirs(args, test_paths, per_class_max=800)
	plots.plot_roc_per_class(model, [test[0]], test[1], args.labels, args.id)


def train_ref_read_b(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	in_channels = defines.total_input_channels_from_args(args)
	if args.channels_last:
		tensor_shape = (args.read_limit, args.window_size, in_channels)
	else:
		tensor_shape = (in_channels, args.read_limit, args.window_size) 

	generate_train = td.tensor_generator(args, train_paths, tensor_shape)
	generate_valid = td.tensor_generator(args, valid_paths, tensor_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_model_from_args(args, 
									conv_width = 5, 
									conv_height = 5,
									conv_layers = [256, 192, 128, 96, 64],
									conv_dropout = 0.1,
									conv_batch_normalize = False,
									spatial_dropout = True,
									max_pools = [(2,1), (2,1), (2,1)],
									padding='same',
									fc_layers = [32],
									fc_dropout = 0.2,
									fc_batch_normalize = False)
	
	models.serialize_model_semantics(args, weight_path)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.load_tensors_from_class_dirs(args, test_paths, per_class_max=800)
	plots.plot_roc_per_class(model, [test[0]], test[1], args.labels, args.id)


def train_ref_read_c(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	in_channels = defines.total_input_channels_from_args(args)
	if args.channels_last:
		tensor_shape = (args.read_limit, args.window_size, in_channels)
	else:
		tensor_shape = (in_channels, args.read_limit, args.window_size) 

	generate_train = td.tensor_generator(args, train_paths, tensor_shape)
	generate_valid = td.tensor_generator(args, valid_paths, tensor_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_model_from_args(args, 
									conv_width = 7, 
									conv_height = 5,
									conv_layers =  [216, 192, 160, 128, 108, 96, 64, 48],
									conv_dropout = 0.1,
									conv_batch_normalize = False,
									spatial_dropout = True,
									max_pools = [(3,1), (3,1), (3,1)],
									padding='valid',
									fc_layers = [32],
									fc_dropout = 0.2,
									fc_batch_normalize = False)
	
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.load_tensors_from_class_dirs(args, test_paths, per_class_max=800,dataset_id='read_tensor',tensor_shape=tensor_shape)
	plots.plot_roc_per_class(model,[test[0]], test[1], args.labels, args.id)



def train_ref_read_resnet(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	args.annotation_set = '_'
	
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)
	generate_train = td.tensor_generator_from_label_dirs_and_args(args, train_paths)
	generate_valid = td.tensor_generator_from_label_dirs_and_args(args, valid_paths)
	generate_test = td.tensor_generator_from_label_dirs_and_args(args, test_paths, with_positions=True)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_read_tensor_keras_resnet(args)

	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.big_batch_from_minibatch_generator(args, generate_test)
	plots.plot_roc_per_class(model, [test[0][args.tensor_map]], test[1], args.labels, args.id, batch_size=args.batch_size)


def train_ref_read_anno_resnet(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''	
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)
	generate_train = td.tensor_generator_from_label_dirs_and_args(args, train_paths)
	generate_valid = td.tensor_generator_from_label_dirs_and_args(args, valid_paths)
	generate_test = td.tensor_generator_from_label_dirs_and_args(args, test_paths, with_positions=True)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_ref_read_anno_keras_resnet(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id, batch_size=args.batch_size)


def train_ref_read_inception_model(args):
	'''Trains a reference and read based inception architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	in_channels = defines.total_input_channels_from_args(args)
	if args.channels_last:
		tensor_shape = (args.read_limit, args.window_size, in_channels)
	else:
		tensor_shape = (in_channels, args.read_limit, args.window_size) 

	generate_train = td.tensor_generator(args, train_paths, tensor_shape)
	generate_valid = td.tensor_generator(args, valid_paths, tensor_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_read_tensor_2d_inception_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.load_tensors_from_class_dirs(args, test_paths, per_class_max=800)
	plots.plot_roc_per_class(model, test[0], test[1], args.labels, args.id, batch_size=args.batch_size)


def train_ref_read_dilated_model(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	in_channels = defines.total_input_channels_from_args(args)
	if args.channels_last:
		tensor_shape = (args.read_limit, args.window_size, in_channels)
	else:
		tensor_shape = (in_channels, args.read_limit, args.window_size) 

	generate_train = td.tensor_generator(args, train_paths, tensor_shape)
	generate_valid = td.tensor_generator(args, valid_paths, tensor_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_read_tensor_2d_dilated_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.load_tensors_from_class_dirs(args, test_paths, per_class_max=800)
	plots.plot_roc_per_class(model, test[0], test[1], args.labels, args.id, batch_size=args.batch_size)


def train_ref_read_annotation_model(args):
	'''Trains a reference, read, and annotation CNN architecture on tensors at the supplied data directory.

	This architecture looks at reads, read flags, reference sequence, and variant annotations.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_read_tensor_2d_and_annotations_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


def train_ref_read_annotation_exome_model(args):
	'''Trains a reference, read, and annotation CNN architecture on tensors at the supplied data directory.

	This architecture looks at reads, read flags, reference sequence, and variant annotations.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''	
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	in_channels = defines.total_input_channels_from_args(args)
	if args.channels_last:
		tensor_shape = (args.read_limit, args.window_size, in_channels)
	else:
		tensor_shape = (in_channels, args.read_limit, args.window_size) 

	generate_train = td.tensor_annotation_generator(args, train_paths, tensor_shape)
	generate_valid = td.tensor_annotation_generator(args, valid_paths, tensor_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_read_tensor_2d_annotations_exome_model(args)

	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	test = td.load_tensors_and_annotations_from_class_dirs(args, test_paths, per_class_max=args.samples)
	plots.plot_roc_per_class(model, [test[0], test[1]], test[2], args.labels, args.id, batch_size=args.batch_size)


def train_ref_read_anno_b(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_annotation_model_from_args(args, 
									conv_width = 3,
									conv_height = 11,
									conv_layers = [128, 96, 64, 48],
									conv_dropout = 0.2,
									conv_batch_normalize = False,
									spatial_dropout = True,
									kernel_single_channel = False,
									max_pools = [(3,1),(3,1),(3,1)],
									padding='same',
									annotation_units = 16,
									annotation_shortcut = True,
									fc_layers = [24],
									fc_dropout = 0.3,
									fc_batch_normalize = False)

	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
	
	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


def train_ref_read_anno_c(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_annotation_model_from_args(args, 
									conv_width = 5,
									conv_height = 5,
									conv_layers = [256, 192, 128, 128, 108, 108, 64, 64],
									conv_dropout = 0.1,
									conv_batch_normalize = False,
									kernel_single_channel = True,
									spatial_dropout = True,
									max_pools = [(3,1),(3,1)],
									padding='valid',
									annotation_units = 16,
									annotation_shortcut = True,
									fc_layers = [28],
									fc_dropout = 0.2,
									fc_batch_normalize = False)

	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test[0], test[1], args.labels)

def train_ref_read_anno_residual(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_annotation_residual_from_args(args, 
									conv_width = 5,
									conv_height = 5,
									conv_layers = [128, 128, 128, 128, 128, 128, 128, 128],
									conv_dropout = 0.0,
									conv_batch_normalize = False,
									kernel_single_channel = True,
									spatial_dropout = True,
									max_pools = [(2,2),(2,2),(2,2),(2,2)],
									padding='same',
									annotation_units = 32,
									annotation_shortcut = False,
									fc_layers = [32, 32],
									fc_dropout = 0.0,
									fc_batch_normalize = False)

	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)

def train_ref_read_residual(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	args.annotation_set = '_'
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)
	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_residual_from_args(args, 
									conv_width = 5,
									conv_height = 5,
									conv_layers = [128, 128, 128, 128, 128, 128, 128, 128],
									conv_dropout = 0.0,
									conv_batch_normalize = False,
									kernel_single_channel = True,
									spatial_dropout = True,
									max_pools = [(2,2),(2,2),(2,2),(2,2)],
									padding='same',
									fc_layers = [32, 32],
									fc_dropout = 0.0,
									fc_batch_normalize = False)

	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	plots.plot_roc_per_class(model, test[0], test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test[0], test[1], args.labels)



def train_ref_read_anno_d(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_annotation_model_from_args(args, 
									conv_width = 5, 
									conv_height = 5,
									conv_layers = [256, 256, 128, 128, 96, 96, 64, 64],
									conv_dropout = 0.15,
									conv_batch_normalize = False,
									kernel_single_channel = True,
									spatial_dropout = True,
									max_pools = [(3,1),(3,1),(3,1)],
									padding='same',
									annotation_units = 16,
									annotation_shortcut = True,
									fc_layers = [32],
									fc_dropout = 0.3,
									fc_batch_normalize = False)
	
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


def train_ref_read_anno_x(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.separable_2d_annotation_model_from_args(args,
									annotation_units = 32,
									annotation_shortcut = True,
									fc_layers = [32],
									fc_dropout = 0.0,
									fc_batch_normalize = False)
	
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)

def train_ref_read_anno_small(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_annotation_model_from_args(args, 
									conv_width = 25,
									conv_height = 25,
									conv_layers = [64, 48, 32, 24],
									conv_dropout = 0.1,
									conv_batch_normalize = False,
									spatial_dropout = True,
									kernel_single_channel = True,
									max_pools = [(3,1),(3,1)],
									padding='valid',
									annotation_units = 64,
									annotation_shortcut = False,
									fc_layers = [24],
									fc_dropout = 0.3,
									fc_batch_normalize = False)
	model.save(weight_path)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


def train_ref_read_anno_tiny(args):
	'''Trains a reference and read based architecture on tensors at the supplied data directory.

	This architecture looks at reads, and read flags.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	
	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	'''
	generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)

	weight_path = arguments.weight_path_from_args(args)
	model = models.read_tensor_2d_annotation_model_from_args(args, 
									conv_width = 5,
									conv_height = 5,
									conv_layers = [64, 48, 32],
									conv_dropout = 0.1,
									conv_batch_normalize = False,
									spatial_dropout = True,
									kernel_single_channel = True,
									max_pools = [(3,1),(3,1)],
									padding='valid',
									annotation_units = 64,
									annotation_shortcut = False,
									fc_layers = [24],
									fc_dropout = 0.3,
									fc_batch_normalize = False)
	model.save(weight_path)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)

	_, _, generate_test = td.train_valid_test_generators_from_args(args, with_positions=True)
	test = td.big_batch_from_minibatch_generator(args, generate_test)
	test_data = [test[0][args.tensor_map], test[0][args.annotation_set]]
	plots.plot_roc_per_class(model, test_data, test[1], args.labels, args.id)
	return plots.get_per_class_auc(model, test_data, test[1], args.labels)


# def load_and_save_model(args):
# 	model = models.load_model(args.weights_hd5, custom_objects=models.get_metric_dict(args.labels))
# 	weight_path = arguments.weight_path_from_args(args)
# 	model.save(weight_path)
# 	if not os.path.exists(os.path.dirname(weight_path)):
# 		os.makedirs(os.path.dirname(weight_path))
# 	model.summary()
# 	models.serialize_model_semantics(args, weight_path)


def bqsr_train_tensor(args):
	'''Trains the bqsr tensor architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	This architecture looks at reads, flags and annotations.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.
	'''
	args.labels = defines.base_labels_binary
	args.input_symbols = defines.bqsr_tensor_channel_map()

	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	generate_train = td.bqsr_tensor_generator(args, train_paths)
	generate_valid = td.bqsr_tensor_generator(args, valid_paths)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_bqsr_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
		
	test = td.load_bqsr_tensors_from_class_dirs(args, test_paths, per_class_max=1800)
	plots.plot_roc_per_class(model, test[0], test[1], args.labels, args.id)


def bqsr_train_annotation_tensor(args):
	'''Trains the bqsr tensor architecture on read tensors and annotations at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	This architecture looks at reads, flags and annotations.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.
	'''
	args.labels = defines.base_labels_binary
	args.input_symbols = defines.bqsr_tensor_channel_map()
	args.annotations = defines.bqsr_annotations

	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	generate_train = td.bqsr_tensor_annotation_generator(args, train_paths)
	generate_valid = td.bqsr_tensor_annotation_generator(args, valid_paths)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_bqsr_annotation_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
		
	test = td.load_bqsr_tensors_annotations_from_class_dirs(args, test_paths, per_class_max=1800)
	plots.plot_roc_per_class(model, [test[0], test[1]], test[2], args.labels, args.id)


def bqsr_lstm_train_tensor(args):
	'''Trains the bqsr tensor architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 

	This architecture looks at reads, flags and annotations.
	Tensors must be generated by calling td.write_tensors() before this function is used.
	After training with early stopping performance curves are plotted on the test dataset.
	'''
	args.labels = defines.base_labels_binary
	args.input_symbols = defines.bqsr_tensor_channel_map()

	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	generate_train = td.bqsr_tensor_generator(args, train_paths)
	generate_valid = td.bqsr_tensor_generator(args, valid_paths)

	model = models.build_bqsr_lstm_model(args)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
		
	test = td.load_bqsr_tensors_from_class_dirs(args, test_paths, per_class_max=1800)
	plots.plot_roc_per_class(model, test[0], test[1], args.labels, args.id)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ~~~~~~~ Plotting Recipes ~~~~~~~~~~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def plot_vcf_roc(args):
	'''Plot ROC curves from a vcf annotated with scores.'''
	snp_scores, snp_truth, indel_scores, indel_truth = td.scores_from_vcf(args, args.score_keys)

	snp_suffix = args.id+'_true_'+str(np.sum(snp_truth))+'_false_'+str(len(snp_truth)-np.sum(snp_truth))
	plots.plot_rocs_from_scores(snp_truth, snp_scores, 'SNP_ROC_'+snp_suffix)
	plots.plot_precision_recall_from_scores(snp_truth, snp_scores, 'SNP_Precision_Recall_' + snp_suffix)

	indel_suffix = args.id+'_true_'+str(np.sum(indel_truth))+'_false_'+str(len(indel_truth)-np.sum(indel_truth))
	plots.plot_rocs_from_scores(indel_truth, indel_scores, 'INDEL_ROC_'+indel_suffix)
	plots.plot_precision_recall_from_scores(indel_truth, indel_scores, 'INDEL_Precision_Recall_' + indel_suffix)


def plot_vcf_roc_gnomad_scores(args):
	'''Plot ROC and Precision-Recall curves from gnomAD vcfs annotated with scores.'''
	if args.single_sample_vqsr:
		args.score_keys.append('VQSR Single Sample')

	snp_scores, snp_truth, indel_scores, indel_truth = td.scores_from_gnomad_vcf(args, args.score_keys)
	
	snp_suffix = args.id+'_true_'+str(np.sum(snp_truth))+'_false_'+str(len(snp_truth)-np.sum(snp_truth))
	plots.plot_rocs_from_scores(snp_truth, snp_scores, 'SNP_ROC_'+snp_suffix)
	plots.plot_precision_recall_from_scores(snp_truth, snp_scores, 'SNP_Precision_Recall_'+snp_suffix)
	
	indel_suffix = args.id+'_true_'+str(np.sum(indel_truth))+'_false_'+str(len(indel_truth)-np.sum(indel_truth))
	plots.plot_rocs_from_scores(indel_truth, indel_scores, 'INDEL_ROC_'+indel_suffix)
	plots.plot_precision_recall_from_scores(indel_truth, indel_scores, 'INDEL_Precision_Recall_'+indel_suffix)


def roc_curve_through_learning(args):
	"""Plot ROC curves during optimization.

	Arguments:
		args.data_dir tensors to train and evaluate model
		args.samples number of ROC curves to plot during training.

	Tensors must be generated by calling 
	td.nist_samples_to_tensors_flags_and_annotations(args) 
	before this function is used.
	"""		
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	in_channels = defines.total_input_channels_from_args(args)
	if args.channels_last:
		tensor_shape = (args.read_limit, args.window_size, in_channels)
	else:
		tensor_shape = (in_channels, args.read_limit, args.window_size) 

	generate_train = td.tensor_annotation_generator(args, train_paths, tensor_shape)
	generate_valid = td.tensor_annotation_generator(args, valid_paths, tensor_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.build_read_tensor_2d_and_annotations_model(args)

	test = td.load_tensors_and_annotations_from_class_dirs(args, test_paths, per_class_max=args.samples)
	for i in range(args.iterations):
		model.fit_generator(generate_train, 
			samples_per_epoch=args.batch_size*2, nb_epoch=1, verbose=1, 
			nb_val_samples=args.batch_size, validation_data=generate_valid,
			callbacks=models.get_callbacks(weight_path, patience=4))
		plots.plot_roc_per_class(model, [test[0], test[1], test[2]], test[3], args.labels, args.id+str(i), prefix='./figures/animations/')


def animate_learning_segmentation(args):
	"""Plot ROC curves during optimization of segmenting architecture.

	Arguments:
		args.data_dir tensors to train and evaluate model
		args.samples number of ROC curves to plot during training.
	"""		
	args.labels = defines.calling_labels
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths_all(args)

	generate_train = td.calling_tensors_generator(args, train_paths)
	generate_valid = td.calling_tensors_generator(args, valid_paths)
	generate_test = td.calling_tensors_generator(args, test_paths)

	#model = models.build_2d_cnn_calling_segmentation_1d(args)
	model = models.build_2d_cnn_calling_segmentation_full_2d(args)

	test_tensors = np.zeros((args.iterations*args.batch_size,) + defines.tensor_shape_from_args(args))
	test_labels = np.zeros((args.iterations*args.batch_size, args.window_size, len(args.labels)))

	for i in range(args.iterations):
		next_batch = next(generate_test)
		test_tensors[i*args.batch_size:(i+1)*args.batch_size,:,:,:] = next_batch[0][args.tensor_map]
		test_labels[i*args.batch_size:(i+1)*args.batch_size,:] = next_batch[1]

	melt_shape = (test_labels.shape[0]*test_labels.shape[1], test_labels.shape[2])
	predictions = model.predict(test_tensors)
	predictions = predictions.reshape(melt_shape)
	print('random predictions are:\n', predictions)
	for i in range(args.samples):
		predictions = model.predict(test_tensors)
		predictions = predictions.reshape(melt_shape)
		test_truth = test_labels.reshape(melt_shape)
		plots.plot_precision_recall_per_class_predictions(predictions, test_truth, args.labels, args.id+str(i), 
			prefix='./figures/animations/seg_prauc2/pr_')
		plots.plot_roc_per_class_predictions(predictions, test_truth, args.labels, args.id+str(i), 
			prefix='./figures/animations/seg_roc2/roc_')

		model.fit_generator(generate_train, 
			steps_per_epoch=args.batch_size, epochs=1, verbose=1, 
			validation_steps=args.batch_size*8, validation_data=generate_valid)


def depristo_inception(args):
	"""Trains the deep variant architecture on tensors at the supplied data directory.

	Arguments:
		args.data_dir: must be set to an appropriate directory with
			subdirectories of test, valid and train, each containing
			subdirectories for each label with tensors stored as hd5 files. 
		args.read_limit: Max number of reads to use (height of the png).
		args.window_size: sequence window around variant (width of the png).
	This architecture looks at pngs of read level data.
	(See: http://biorxiv.org/content/early/2016/12/14/092890).
	Tensors must be generated by calling: td.nist_samples_to_png(args)
	before this function is used.
	After training with early stopping a performance curves are plotted on the test dataset.
	"""	
	train_paths, valid_paths, test_paths = td.get_train_valid_test_paths(args)

	image_shape = (args.read_limit, args.window_size)
	generate_train = td.image_generator(args, train_paths, shape=image_shape)
	generate_valid = td.image_generator(args, valid_paths, shape=image_shape)

	weight_path = arguments.weight_path_from_args(args)
	model = models.inception_v3_max(args, architecture=args.weights_hd5)
	model = models.train_model_from_generators(args, model, generate_train, generate_valid, weight_path)
		
	test = td.load_images_from_class_dirs(args, test_paths, shape=image_shape, per_class_max=args.samples)
	plots.plot_roc_per_class(model, test[0], test[1], args.labels, args.id)


def test_score_consistency():
	model = models.load_model(args.weights_hd5, custom_objects=models.get_all_custom_objects(args.labels))
	model.summary()	



# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# ~~~~~~~ Testing Recipes ~~~~~~~~~~~~~~~~~~
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def find_architecture_discordance(args):
	'''Evaluate architectures defined by architecture configuration files

	Compares against VQSR, gnomAD VQSR and Random Forest, and deep variant.

	Arguments:
			args.architectures: list of architecture semantics config files
	'''
	cnn_snp_dicts = {}
	cnn_indel_dicts = {}
	number_of_sites = 10
	for a in args.architectures:	
		print('Processing architecture:', a)
		model = models.set_args_and_get_model_from_semantics(args, a)
		_, _, test_generator = td.train_valid_test_generators_from_args(args, with_positions=True)
		test, labels, positions = td.big_batch_from_minibatch_generator(args, test_generator)
		cnn_predictions = model.predict(test, batch_size=args.batch_size)
		cnn_snp_dicts[a], cnn_indel_dicts[a] = models.predictions_to_snp_indel_scores(args, cnn_predictions, positions)

	for a in args.architectures:
		for b in args.architectures:
			if a == b:
				continue
			print('Comparing architectures:', a, 'to:', b)
			diff_list =	[(p, abs(cnn_snp_dicts[a][p]-cnn_snp_dicts[b][p]), l) for l,p in zip(labels, positions) if l[0] or l[2]]
			sorted_cnn = sorted(diff_list, key=lambda s: -s[1])
			for i in range(min(number_of_sites, len(sorted_cnn))):	
				p = sorted_cnn[i][0]
				print('Discordant CNN SNP:', sorted_cnn[i], 'a, then b:', cnn_snp_dicts[a][p], cnn_snp_dicts[b][p])
			
			diff_list =	[(p, abs(cnn_indel_dicts[a][p]-cnn_indel_dicts[b][p]), l) for l,p in zip(labels, positions) if l[1] or l[3]]
			sorted_cnn = sorted(diff_list, key=lambda s: -s[1])
			for i in range(min(number_of_sites, len(sorted_cnn))):	
				p = sorted_cnn[i][0]
				print('Discordant CNN Indel:', sorted_cnn[i], 'a, then b:', cnn_indel_dicts[a][p], cnn_indel_dicts[b][p])


def test_architectures(args):
	'''Evaluate architectures defined by architecture configuration files

	Compares against VQSR, gnomAD VQSR and Random Forest, and deep variant.

	Arguments:
			args.architectures: list of architecture semantics config files
	'''
	cnn_snp_dicts = {}
	cnn_indel_dicts = {}
	for a in args.architectures:	
		print('Processing architecture:', a)
		model = models.set_args_and_get_model_from_semantics(args, a)
		
		if args.inspect_model:
			generate_train, generate_valid, _ = td.train_valid_test_generators_from_args(args, with_positions=False)
			image_path = a.replace('.json','.png') if args.image_dir is None else args.image_dir + a.replace('.json','.png')
			models.inspect_model(args, model, generate_train, generate_valid, image_path=image_path)

		_, _, test_generator = td.train_valid_test_generators_from_args(args, with_positions=True)
		test, _, positions = td.big_batch_from_minibatch_generator(args, test_generator)
		cnn_predictions = model.predict(test, batch_size=args.batch_size)
		cnn_snp_dicts[a], cnn_indel_dicts[a] = models.predictions_to_snp_indel_scores(args, cnn_predictions, positions)

	compare_snp, compare_indel = score_dicts_from_positions(args, positions)

	if 'SNP' in args.labels:
		snp_key_sets = [set(compare_snp[k].keys()) for k in compare_snp]
		shared_snp_keys = set.intersection(*snp_key_sets)
		snp_truth = [compare_snp[list(compare_snp)[0]][p][1] for p in shared_snp_keys]
		snp_scores = score_dict_from_shared_positions(args, 'SNP', shared_snp_keys, cnn_snp_dicts, compare_snp)
		title_suffix = args.id+'_true_'+str(np.sum(snp_truth))+'_false_'+str(len(snp_truth)-np.sum(snp_truth))
		plots.plot_rocs_from_scores(snp_truth, snp_scores, 'SNP_ROC_'+title_suffix)
		plots.plot_precision_recall_from_scores(snp_truth, snp_scores, 'SNP_Precision_Recall_'+title_suffix, args.baseline_key)		
	
	if 'INDEL' in args.labels:
		indel_key_sets = [set(compare_indel[k].keys()) for k in compare_indel]
		shared_indel_keys = set.intersection(*indel_key_sets)
		indel_truth = [compare_indel[list(compare_indel)[0]][p][1] for p in shared_indel_keys]
		indel_scores = score_dict_from_shared_positions(args, 'INDEL', shared_indel_keys, cnn_indel_dicts, compare_indel)
		title_suffix = args.id+'_true_'+str(np.sum(indel_truth))+'_false_'+str(len(indel_truth)-np.sum(indel_truth))
		plots.plot_rocs_from_scores(indel_truth, indel_scores, 'INDEL_ROC_'+title_suffix)
		plots.plot_precision_recall_from_scores(indel_truth, indel_scores, 'INDEL_Precision_Recall_'+title_suffix, args.baseline_key)


def score_dicts_from_positions(args, positions):
	compare_snp = {}
	compare_indel = {}
	
	if args.gnomad_compare:
		compare_snp['VQSR gnomAD'], compare_indel['VQSR gnomAD'] = td.gnomad_scores_from_positions(args, positions)
		compare_snp['RF gnomAD'], compare_indel['RF gnomAD']  = td.gnomad_scores_from_positions(args, positions, score_key='AS_RF')

	if args.single_sample_vqsr:
		compare_snp['VQSR Single Sample'], compare_indel['VQSR Single Sample'] = td.scores_from_positions(args, positions)

	if args.deep_variant_vcf:
		compare_snp['Deep Variant'], compare_indel['Deep Variant'] = td.scores_from_positions(args, positions, 'QUAL', args.deep_variant_vcf)

	if args.hard_filter_compare:
		compare_snp['GATK Hard Filter Score'], compare_indel['GATK Hard Filter Score'] = td.scores_from_gatk_hard_filters(args, positions, distance_score=True)
		compare_snp['GATK Hard Filter'], compare_indel['GATK Hard Filter'] = td.scores_from_gatk_hard_filters(args, positions)
		compare_snp['Heng Li Hard Filter'], compare_indel['Heng Li Hard Filter'] = td.scores_from_heng_li_filters(args, positions)
	
	return compare_snp, compare_indel


def score_dict_from_shared_positions(args, kind, shared_positions, cnn_scores, compare_scores):
	scores = {}

	for a in args.architectures:
		scores['GATK4_'+td.plain_name(a)] = [cnn_scores[a][p] for p in shared_positions]

	for k in compare_scores:
		scores[k] = [compare_scores[k][p][0] for p in shared_positions]

	if args.emit_interesting_sites:
		emit_interesting_sites(args, kind, shared_positions, cnn_scores, compare_scores)

	return scores


def emit_interesting_sites(args, kind, shared_positions, cnn_scores, compare_scores, number_of_sites=5):
	for a in args.architectures:
		sorted_cnn = sorted([(p, cnn_scores[a][p]) for p in shared_positions], key=lambda s: s[1])
		for i in range(len(sorted_cnn)):
			p = sorted_cnn[i][0]
			if p in shared_positions:
				for k in compare_scores:
					over = (sorted_cnn[i][1] > 0 and compare_scores[k][p][1] == 0) 
					under = (sorted_cnn[i][1] < 0 and compare_scores[k][p][1] == 1) 
					if over or under:
						print('CNN Wrong :', sorted_cnn[i], 'but truth is: ', compare_scores[k][p][1], k, 'scored:', compare_scores[k][p][0])

		for i in range(min(number_of_sites, len(sorted_cnn))):	
			p = sorted_cnn[i][0]
			print(kind, 'Bad CNN score:', sorted_cnn[i])
			if p in shared_positions:
				for k in compare_scores:
					print(k, 'score, truth:', compare_scores[k][p])

		for i in range(len(sorted_cnn)-1, max(0,len(sorted_cnn)-number_of_sites), -1):	
			p = sorted_cnn[i][0]
			print(kind, 'Good CNN score:', sorted_cnn[i])

			if p in shared_positions:
				for k in compare_scores:
					print(k, 'score, truth:', compare_scores[k][p])	


# Back to the top!
if "__main__" == __name__:
	file_fxns = { name:obj for name,obj in inspect.getmembers(sys.modules[__name__]) if inspect.isfunction(obj) }
	run(file_fxns)