"""
Generate certain RST files used in documentation.
"""
from __future__ import print_function
import sys
import argparse
from collections import OrderedDict, defaultdict
import os
from os.path import join, exists
from os import mkdir
import pandas
import logomaker
import tqdm
from matplotlib import pyplot
from mhcflurry.downloads import get_path
from mhcflurry.amino_acid import COMMON_AMINO_ACIDS
from mhcflurry.class1_affinity_predictor import Class1AffinityPredictor
AMINO_ACIDS = sorted(COMMON_AMINO_ACIDS)
parser = argparse.ArgumentParser(usage=__doc__)
parser.add_argument(
"--class1-models-dir",
metavar="DIR",
default=get_path(
"models_class1_pan", "models.combined", test_exists=False),
help="Class1 models. Default: %(default)s",
)
parser.add_argument(
"--logo-cutoff",
default=0.01,
type=float,
help="Fraction of top to use for motifs",
)
parser.add_argument(
"--length-cutoff",
default=0.01,
type=float,
help="Fraction of top to use for length distribution",
)
parser.add_argument(
"--length-distribution-lengths",
nargs="+",
default=[8, 9, 10, 11, 12, 13, 14, 15],
type=int,
help="Peptide lengths for length distribution plots",
)
parser.add_argument(
"--motif-lengths",
nargs="+",
default=[8, 9, 10, 11],
type=int,
help="Peptide lengths for motif plots",
)
parser.add_argument(
"--out-dir",
metavar="DIR",
required=True,
help="Directory to write RSTs and images to",
)
parser.add_argument(
"--max-alleles",
default=None,
type=int,
metavar="N",
help="Only use N alleles (for testing)",
)
def model_info(models_dir):
allele_to_sequence = Class1AffinityPredictor.load(
models_dir).allele_to_sequence
length_distributions_df = pandas.read_csv(
join(models_dir, "length_distributions.csv.bz2"))
frequency_matrices_df = pandas.read_csv(
join(models_dir, "frequency_matrices.csv.bz2"))
try:
train_data_df = pandas.read_csv(
join(models_dir, "train_data.csv.bz2"))
observations_per_allele = (
train_data_df.groupby("allele").peptide.nunique().to_dict())
except IOError:
observations_per_allele = None
distribution = frequency_matrices_df.loc[
(frequency_matrices_df.cutoff_fraction == 1.0), AMINO_ACIDS
].mean(0)
normalized_frequency_matrices = frequency_matrices_df.copy()
normalized_frequency_matrices.loc[:, AMINO_ACIDS] = (
normalized_frequency_matrices[AMINO_ACIDS] / distribution)
sequence_to_alleles = defaultdict(list)
for allele in normalized_frequency_matrices.allele.unique():
sequence = allele_to_sequence[allele]
sequence_to_alleles[sequence].append(allele)
allele_equivalance_classes = sorted([
sorted(equivalence_group)
for equivalence_group in sequence_to_alleles.values()
], key=lambda equivalence_group: equivalence_group[0])
return {
'length_distributions': length_distributions_df,
'normalized_frequency_matrices': normalized_frequency_matrices,
'observations_per_allele': observations_per_allele,
'allele_equivalance_classes': allele_equivalance_classes,
}
def write_logo(
normalized_frequency_matrices,
allele,
lengths,
cutoff,
models_label,
out_dir):
fig = pyplot.figure(figsize=(8,10))
for (i, length) in enumerate(lengths):
ax = pyplot.subplot(len(lengths), 1, i + 1)
matrix = normalized_frequency_matrices.loc[
(normalized_frequency_matrices.allele == allele) &
(normalized_frequency_matrices.length == length) &
(normalized_frequency_matrices.cutoff_fraction == cutoff)
].set_index("position")[AMINO_ACIDS]
if matrix.shape[0] == 0:
return None
matrix = (matrix.T / matrix.sum(1)).T # row normalize
ss_logo = logomaker.Logo(
matrix,
width=.8,
vpad=.05,
fade_probabilities=True,
stack_order='small_on_top',
ax=ax,
)
pyplot.title(
"%s %d-mer (%s)" % (allele, length, models_label), y=0.85)
pyplot.xticks(matrix.index.values)
pyplot.tight_layout()
name = "%s.motifs.%s.png" % (
allele.replace("*", "-").replace(":", "-"), models_label)
filename = os.path.abspath(join(out_dir, name))
pyplot.savefig(filename)
print("Wrote: ", filename)
fig.clear()
pyplot.close(fig)
return name
def write_length_distribution(
length_distributions_df, allele, lengths, cutoff, models_label, out_dir):
length_distribution = length_distributions_df.loc[
(length_distributions_df.allele == allele) &
(length_distributions_df.cutoff_fraction == cutoff)
]
if length_distribution.shape[0] == 0:
return None
length_distribution = length_distribution.set_index(
"length").reindex(lengths).fillna(0.0).reset_index()
fig = pyplot.figure(figsize=(8, 2))
length_distribution.plot(x="length", y="fraction", kind="bar", color="black")
pyplot.title("%s (%s)" % (allele, models_label))
pyplot.xlabel("")
pyplot.xticks(rotation=0)
pyplot.gca().get_legend().remove()
name = "%s.lengths.%s.png" % (
allele.replace("*", "-").replace(":", "-"), models_label)
filename = os.path.abspath(join(out_dir, name))
pyplot.savefig(filename)
print("Wrote: ", filename)
fig.clear()
pyplot.close(fig)
return name
def go(argv):
args = parser.parse_args(argv)
if not exists(args.out_dir):
mkdir(args.out_dir)
predictors = [
("combined", args.class1_models_dir),
]
info_per_predictor = OrderedDict()
alleles = set()
for (label, models_dir) in predictors:
if not models_dir:
continue
info_per_predictor[label] = model_info(models_dir)
alleles.update(
info_per_predictor[label]["normalized_frequency_matrices"].allele.unique())
lines = []
def w(*pieces):
lines.extend(pieces)
w('Motifs and length distributions from the pan-allele predictor')
w('=' * 80, "")
w(
"Length distributions and binding motifs were calculated by ranking a "
"large set of random peptides (an equal number of peptides for each "
"length 8-15) by predicted affinity for each allele. "
"For length distribution, the top %g%% of peptides were collected and "
"their length distributions plotted. For sequence motifs, sequence "
"logos for the top %g%% "
"peptides for each length are shown.\n" % (
args.length_cutoff * 100.0,
args.logo_cutoff * 100.0,
))
w(".. contents:: :local:", "")
def image(name):
if name is None:
return ""
return '.. image:: %s\n' % name
alleles = sorted(alleles, key=lambda a: ("HLA" not in a, a))
if args.max_alleles:
alleles = alleles[:args.max_alleles]
for allele in tqdm.tqdm(alleles):
w(allele, "-" * 80, "")
for (label, info) in info_per_predictor.items():
length_distribution = info["length_distributions"]
normalized_frequency_matrices = info["normalized_frequency_matrices"]
length_distribution_image_path = write_length_distribution(
length_distributions_df=length_distribution,
allele=allele,
lengths=args.length_distribution_lengths,
cutoff=args.length_cutoff,
out_dir=args.out_dir,
models_label=label)
if not length_distribution_image_path:
continue
w("*%s*\n")
if info['observations_per_allele'] is not None:
w("Training observations (unique peptides): %d" % (
info['observations_per_allele'].get(allele, 0)))
w("\n")
w(image(length_distribution_image_path))
w(image(write_logo(
normalized_frequency_matrices=normalized_frequency_matrices,
allele=allele,
lengths=args.motif_lengths,
cutoff=args.logo_cutoff,
out_dir=args.out_dir,
models_label=label,
)))
w("")
document_path = join(args.out_dir, "allele_motifs.rst")
with open(document_path, "w") as fd:
for line in lines:
fd.write(line)
fd.write("\n")
print("Wrote", document_path)
if __name__ == "__main__":
go(sys.argv[1:])