class1_presentation_predictor.py

from __future__ import print_function

from os.path import join, exists
from os import mkdir
from socket import gethostname
from getpass import getuser

import time
import collections
import logging
from six import string_types

import numpy
import pandas
import sklearn
import sklearn.linear_model


try:
    import tqdm
except ImportError:
    tdqm = None

from .version import __version__
from .class1_affinity_predictor import Class1AffinityPredictor
from .class1_processing_predictor import Class1ProcessingPredictor
from .class1_neural_network import DEFAULT_PREDICT_BATCH_SIZE
from .encodable_sequences import EncodableSequences
from .regression_target import from_ic50, to_ic50
from .downloads import get_default_class1_presentation_models_dir


MAX_ALLELES_PER_SAMPLE = 6
PREDICT_BATCH_SIZE = DEFAULT_PREDICT_BATCH_SIZE
PREDICT_CHUNK_SIZE = 100000  # currently used only for cleavage prediction


class Class1PresentationPredictor(object):
    """
    A logistic regression model over predicted binding affinity (BA) and antigen
    processing (AP) score.

    Instances of this class delegate to Class1AffinityPredictor and
    Class1ProcessingPredictor instances to generate BA and AP predictions.
    These predictions are combined using a logistic regression model to give
    a "presentation score" prediction.

    See load() and predict() methods for basic usage.
    """
    model_inputs = ["affinity_score", "processing_score"]

    def __init__(
            self,
            affinity_predictor=None,
            processing_predictor_with_flanks=None,
            processing_predictor_without_flanks=None,
            weights_dataframe=None,
            metadata_dataframes=None):

        self.affinity_predictor = affinity_predictor
        self.processing_predictor_with_flanks = processing_predictor_with_flanks
        self.processing_predictor_without_flanks = processing_predictor_without_flanks
        self.weights_dataframe = weights_dataframe
        self.metadata_dataframes = (
            dict(metadata_dataframes) if metadata_dataframes else {})
        self._models_cache = {}

    @property
    def supported_alleles(self):
        """
        List of alleles supported by the underlying Class1AffinityPredictor
        """
        return self.affinity_predictor.supported_alleles

    @property
    def supported_peptide_lengths(self):
        """
        (min, max) of supported peptide lengths, inclusive.
        """
        return self.affinity_predictor.supported_peptide_lengths

    def predict_affinity(
            self,
            peptides,
            alleles,
            experiment_names=None,
            include_affinity_percentile=False,
            verbose=1,
            throw=True):
        """
        Predict binding affinities.

        Parameters
        ----------
        peptides : list of string
        experiment_names : list of string [same length as peptides]
            Sample names corresponding to each peptide. These are used to
            lookup the alleles for each peptide in the alleles dict.
        alleles : dict of string -> list of string
            Keys are experiment names, values are the alleles for that sample
        include_affinity_percentile : bool
            Whether to include affinity percentile ranks
        verbose : int
            Set to 0 for quiet.
        throw : verbose
            Whether to throw exception (vs. just log a warning) on invalid
            peptides, etc.

        Returns
        -------
        pandas.DataFrame : predictions
        """
        df = pandas.DataFrame({
            "peptide": numpy.array(peptides, copy=False),
        })
        if experiment_names is None:
            peptides = EncodableSequences.create(peptides)
            all_alleles = set()
            for lst in alleles.values():
                all_alleles.update(lst)

            iterator = sorted(all_alleles)

            if verbose > 0:
                print("Predicting affinities.")
                if tqdm is not None:
                    iterator = tqdm.tqdm(iterator, total=len(all_alleles))

            predictions_df = pandas.DataFrame(index=df.index)
            for allele in iterator:
                predictions_df[allele] = self.affinity_predictor.predict(
                    peptides=peptides,
                    allele=allele,
                    model_kwargs={'batch_size': PREDICT_BATCH_SIZE},
                    throw=throw)

            dfs = []
            for (experiment_name, experiment_alleles) in alleles.items():
                new_df = df.copy()
                new_df["experiment_name"] = experiment_name
                new_df["affinity"] = predictions_df[
                    experiment_alleles
                ].min(1).values
                new_df["best_allele"] = predictions_df[
                    experiment_alleles
                ].idxmin(1).values

            result_df = pandas.concat(dfs, ignore_index=True)
        else:
            df["experiment_name"] = numpy.array(experiment_names, copy=False)

            iterator = df.groupby("experiment_name")
            if verbose > 0:
                print("Predicting affinities.")
                if tqdm is not None:
                    iterator = tqdm.tqdm(
                        iterator, total=df.experiment_name.nunique())

            for (experiment, sub_df) in iterator:
                predictions_df = pandas.DataFrame(index=sub_df.index)
                experiment_peptides = EncodableSequences.create(sub_df.peptide.values)
                for allele in alleles[experiment]:
                    predictions_df[allele] = self.affinity_predictor.predict(
                        peptides=experiment_peptides,
                        allele=allele,
                        model_kwargs={'batch_size': PREDICT_BATCH_SIZE},
                        throw=throw)
                df.loc[
                    sub_df.index, "affinity"
                ] = predictions_df.min(1).values
                df.loc[
                    sub_df.index, "best_allele"
                ] = predictions_df.idxmin(1).values

            result_df = df

        if include_affinity_percentile:
            result_df["affinity_percentile"] = (
                self.affinity_predictor.percentile_ranks(
                    df.affinity.values,
                    alleles=df.best_alleles.values,
                    throw=False))

        return result_df

    def predict_processing(
            self, peptides, n_flanks=None, c_flanks=None, verbose=1):
        """
        Predict antigen processing scores for individual peptides, optionally
        including flanking sequences for better cleavage prediction.

        Parameters
        ----------
        peptides : list of string
        n_flanks : list of string [same length as peptides]
        c_flanks : list of string [same length as peptides]
        verbose  : int

        Returns
        -------
        numpy.array : Antigen processing scores for each peptide
        """

        if (n_flanks is None) != (c_flanks is None):
            raise ValueError("Specify both or neither of n_flanks, c_flanks")

        if n_flanks is None:
            if self.processing_predictor_without_flanks is None:
                raise ValueError("No processing predictor without flanks")
            predictor = self.processing_predictor_without_flanks
            n_flanks = [""] * len(peptides)
            c_flanks = n_flanks
        else:
            if self.processing_predictor_with_flanks is None:
                raise ValueError("No processing predictor with flanks")
            predictor = self.processing_predictor_with_flanks

        num_chunks = int(numpy.ceil(float(len(peptides)) / PREDICT_CHUNK_SIZE))
        peptide_chunks = numpy.array_split(peptides, num_chunks)
        n_flank_chunks = numpy.array_split(n_flanks, num_chunks)
        c_flank_chunks = numpy.array_split(c_flanks, num_chunks)

        iterator = zip(peptide_chunks, n_flank_chunks, c_flank_chunks)
        if verbose > 0:
            print("Predicting processing.")
            if tqdm is not None:
                iterator = tqdm.tqdm(iterator, total=len(peptide_chunks))

        result_chunks = []
        for (peptide_chunk, n_flank_chunk, c_flank_chunk) in iterator:
            result_chunk = predictor.predict(
                peptides=peptide_chunk,
                n_flanks=n_flank_chunk,
                c_flanks=c_flank_chunk,
                batch_size=PREDICT_BATCH_SIZE)
            result_chunks.append(result_chunk)
        return numpy.concatenate(result_chunks)

    def fit(
            self,
            targets,
            peptides,
            experiment_names,
            alleles,
            n_flanks=None,
            c_flanks=None,
            verbose=1):
        """
        Fit the presentation score logistic regression model.

        Parameters
        ----------
        targets : list of int/float
            1 indicates hit, 0 indicates decoy
        peptides : list of string [same length as targets]
        experiment_names : list of string [same length as targets]
        alleles : dict of string -> list of string
            Keys are experiment names, values are the alleles for that sample
        n_flanks : list of string [same length as targets]
        c_flanks : list of string [same length as targets]
        verbose : int
        """

        df = self.predict_affinity(
            peptides=peptides,
            alleles=alleles,
            experiment_names=experiment_names,
            verbose=verbose)
        df["affinity_score"] = from_ic50(df.affinity)
        df["target"] = numpy.array(targets, copy=False)

        if (n_flanks is None) != (c_flanks is None):
            raise ValueError("Specify both or neither of n_flanks, c_flanks")

        with_flanks_list = []
        if self.processing_predictor_without_flanks is not None:
            with_flanks_list.append(False)

        if n_flanks is not None and self.processing_predictor_with_flanks is not None:
            with_flanks_list.append(True)

        if not with_flanks_list:
            raise RuntimeError("Can't fit any models")

        if self.weights_dataframe is None:
            self.weights_dataframe = pandas.DataFrame()

        for with_flanks in with_flanks_list:
            model_name = 'with_flanks' if with_flanks else "without_flanks"
            if verbose > 0:
                print("Training variant", model_name)

            df["processing_score"] = self.predict_processing(
                peptides=df.peptide.values,
                n_flanks=n_flanks if with_flanks else None,
                c_flanks=c_flanks if with_flanks else None,
                verbose=verbose)

            model = self.get_model()
            if verbose > 0:
                print("Fitting LR model.")
                print(df)

            model.fit(
                X=df[self.model_inputs].values,
                y=df.target.astype(float))

            self.weights_dataframe.loc[model_name, "intercept"] = model.intercept_
            for (name, value) in zip(self.model_inputs, numpy.squeeze(model.coef_)):
                self.weights_dataframe.loc[model_name, name] = value
            self._models_cache[model_name] = model

    def get_model(self, name=None):
        """
        Load or instantiate a new logistic regression model. Private helper
        method.

        Parameters
        ----------
        name : string
            If None (the default), an un-fit LR model is returned. Otherwise the
            weights are loaded for the specified model.

        Returns
        -------
        sklearn.linear_model.LogisticRegression
        """
        if name is None or name not in self._models_cache:
            model = sklearn.linear_model.LogisticRegression(solver="lbfgs")
            if name is not None:
                row = self.weights_dataframe.loc[name]
                model.intercept_ = row.intercept
                model.coef_ = numpy.expand_dims(
                    row[self.model_inputs].values, axis=0)
                model.classes_ = numpy.array([0, 1])
        else:
            model = self._models_cache[name]
        return model

    def predict(
            self,
            peptides,
            alleles,
            experiment_names=None,
            n_flanks=None,
            c_flanks=None,
            verbose=1):
        """
        Predict presentation scores across a set of peptides.

        Presentation scores combine predictions for MHC I binding affinity
        and antigen processing.

        For intermediate results, see the `predict_to_dataframe` method.

        Parameters
        ----------
        peptides : list of string
            Peptide sequences
        alleles : list of string or string -> string dict
            If you are predicting for a single sample, pass a list of strings
            (up to 6) indicating the genotype. If you are predicting across
            multiple samples, pass a dict where the keys are (arbitrary)
            experiment names and the values are the alleles to predict for that
            sample.
        experiment_names : list of string [same length as peptides]
            If you are passing a dict for 'alleles', use this argument to
            specify which peptides go with which experiments.
        n_flanks : list of string [same length as peptides]
            Upstream sequences before the peptide. Sequences of any length can
            be given and a suffix of the size supported by the model will be
            used.
        c_flanks : list of string [same length as peptides]
            Downstream sequences after the peptide. Sequences of any length can
            be given and a prefix of the size supported by the model will be
            used.
        verbose : int
            Set to 0 for quiet mode.

        Returns
        -------
        numpy.array

        Presentation scores for each peptide. Scores range from 0 to 1, with
        higher values indicating more favorable presentation likelihood.
        """
        return self.predict_to_dataframe(
            peptides=peptides,
            alleles=alleles,
            experiment_names=experiment_names,
            n_flanks=n_flanks,
            c_flanks=c_flanks,
            verbose=verbose).presentation_score.values

    def predict_to_dataframe(
            self,
            peptides,
            alleles,
            experiment_names=None,
            n_flanks=None,
            c_flanks=None,
            include_affinity_percentile=False,
            verbose=1,
            throw=True):
        """
        Predict presentation scores across a set of peptides.

        Presentation scores combine predictions for MHC I binding affinity
        and antigen processing.

        This method returns a pandas.DataFrame giving presentation scores plus
        the binding affinity and processing predictions and other intermediate
        results.

        Parameters
        ----------
        peptides : list of string
            Peptide sequences
        alleles : list of string or string -> string dict
            If you are predicting for a single sample, pass a list of strings
            (up to 6) indicating the genotype. If you are predicting across
            multiple samples, pass a dict where the keys are (arbitrary)
            experiment names and the values are the alleles to predict for that
            sample.
        experiment_names : list of string [same length as peptides]
            If you are passing a dict for 'alleles', use this argument to
            specify which peptides go with which experiments.
        n_flanks : list of string [same length as peptides]
            Upstream sequences before the peptide. Sequences of any length can
            be given and a suffix of the size supported by the model will be
            used.
        c_flanks : list of string [same length as peptides]
            Downstream sequences after the peptide. Sequences of any length can
            be given and a prefix of the size supported by the model will be
            used.
        include_affinity_percentile : bool
            Whether to include affinity percentile ranks
        verbose : int
            Set to 0 for quiet.
        throw : verbose
            Whether to throw exception (vs. just log a warning) on invalid
            peptides, etc.

        Returns
        -------
        pandas.DataFrame

        Presentation scores and intermediate results.
        """

        if isinstance(peptides, string_types):
            raise TypeError("peptides must be a list not a string")
        if isinstance(alleles, string_types):
            raise TypeError("alleles must be a list or dict")

        if isinstance(alleles, dict):
            if experiment_names is None:
                raise ValueError(
                    "experiment_names must be supplied when alleles is a dict")
        else:
            if experiment_names is not None:
                raise ValueError(
                    "alleles must be a dict when experiment_names is specified")
            alleles = numpy.array(alleles, copy=False)
            if len(alleles) > MAX_ALLELES_PER_SAMPLE:
                raise ValueError(
                    "When alleles is a list, it must have at most %d elements. "
                    "These alleles are taken to be a genotype for an "
                    "individual, and the strongest prediction across alleles "
                    "will be taken for each peptide. Note that this differs "
                    "from Class1AffinityPredictor.predict(), where alleles "
                    "is expected to be the same length as peptides."
                    % MAX_ALLELES_PER_SAMPLE)

            experiment_names = ["experiment1"] * len(peptides)
            alleles = {
                "experiment1": alleles,
            }

        if (n_flanks is None) != (c_flanks is None):
            raise ValueError("Specify both or neither of n_flanks, c_flanks")

        processing_scores = self.predict_processing(
            peptides=peptides,
            n_flanks=n_flanks,
            c_flanks=c_flanks,
            verbose=verbose)

        df = self.predict_affinity(
            peptides=peptides,
            experiment_names=experiment_names,
            alleles=alleles,
            include_affinity_percentile=include_affinity_percentile,
            verbose=verbose,
            throw=throw)
        df["affinity_score"] = from_ic50(df.affinity)
        df["processing_score"] = processing_scores
        if c_flanks is not None:
            df.insert(1, "c_flank", c_flanks)
        if n_flanks is not None:
            df.insert(1, "n_flank", n_flanks)

        model_name = 'with_flanks' if n_flanks is not None else "without_flanks"
        model = self.get_model(model_name)
        df["presentation_score"] = model.predict_proba(
            df[self.model_inputs].values)[:,1]
        del df["affinity_score"]
        return df

    def predict_sequences(
            self,
            sequences,
            alleles,
            result="best",
            comparison_quantity="presentation_score",
            filter_value=None,
            peptide_lengths=[8, 9, 10, 11],
            use_flanks=True,
            include_affinity_percentile=False,
            verbose=1,
            throw=True):
        """
        Predict presentation across protein sequences.

        Parameters
        ----------
        sequences : str, list of string, or string -> string dict
            Protein sequences. If a dict is given, the keys are arbitrary (
            e.g. protein names), and the values are the amino acid sequences.
        alleles : str, list of string, list of list of string, or string -> string dict
            MHC I alleles. Can be: (1) a string (a single allele), (2) a list of
            strings (a single genotype), (3) a list of list of strings
            (multiple genotypes, where the total number of genotypes must equal
            the number of sequences), or (4) a dict (in which case the keys must
            match the sequences dict keys).
        result : string
            One of:
             - "best": return the strongest peptide for each sequence
             - "all": return predictions for all peptides
             - "filtered": return predictions where comparison_quantity is
             stronger (i.e (<) for affinity, (>) for scores) than filter_value.
        comparison_quantity : string
            One of "presentation_score", "processing_score", or "affinity".
            Quantity to use to rank (if result is "best") or filter (if result
            is "filtered") results.
        filter_value : float
            Threshold value to use, only relevant when result is "filtered".
            If comparison_quantity is "affinity", then all results less than
            (i.e. tighter than) the specified nM affinity are retained. If it's
            "presentation_score" or "processing_score" then results greater than
            the indicated filter_value are retained.
        peptide_lengths : list of int
            Peptide lengths to predict for.
        use_flanks : bool
            Whether to include flanking sequences when running the AP predictor
            (for better cleavage prediction).
        include_affinity_percentile : bool
            Whether to include affinity percentile ranks in output.
        verbose : int
            Set to 0 for quiet mode.
        throw : boolean
            Whether to throw exceptions (vs. log warnings) on invalid inputs.

        Returns
        -------
        pandas.DataFrame with columns:
            peptide, n_flank, c_flank, sequence_name, affinity, best_allele,
            processing_score, presentation_score
        """

        processing_predictor = self.processing_predictor_with_flanks
        if not use_flanks or processing_predictor is None:
            processing_predictor = self.processing_predictor_without_flanks

        supported_sequence_lengths = processing_predictor.sequence_lengths
        n_flank_length = supported_sequence_lengths["n_flank"]
        c_flank_length = supported_sequence_lengths["c_flank"]

        sequence_names = []
        n_flanks = [] if use_flanks else None
        c_flanks = [] if use_flanks else None
        peptides = []

        if isinstance(sequences, string_types):
            sequences = [sequences]

        if not isinstance(sequences, dict):
            sequences = collections.OrderedDict(
                ("sequence_%04d" % (i + 1), sequence)
                for (i, sequence) in enumerate(sequences))

        if isinstance(alleles, string_types):
            alleles = [alleles]

        if not isinstance(alleles, dict):
            if all([isinstance(item, string_types) for item in alleles]):
                alleles = dict((name, alleles) for name in sequences.keys())
            elif len(alleles) != len(sequences):
                raise ValueError(
                    "alleles must be (1) a string (a single allele), (2) a list of "
                    "strings (a single genotype), (3) a list of list of strings ("
                    "(multiple genotypes, where the total number of genotypes "
                    "must equal the number of sequences), or (4) a dict (in which "
                    "case the keys must match the sequences dict keys). Here "
                    "it seemed like option (3) was being used, but the length "
                    "of alleles (%d) did not match the length of sequences (%d)."
                    % (len(alleles), len(sequences)))
            else:
                alleles = dict(zip(sequences.keys(), alleles))

        missing = [key for key in sequences if key not in alleles]
        if missing:
            raise ValueError(
                "Sequence names missing from alleles dict: ", missing)

        for (name, sequence) in sequences.items():
            if not isinstance(sequence, string_types):
                raise ValueError("Expected string, not %s (%s)" % (
                    sequence, type(sequence)))
            for peptide_start in range(len(sequence) - min(peptide_lengths)):
                n_flank_start = max(0, peptide_start - n_flank_length)
                for peptide_length in peptide_lengths:
                    c_flank_end = (
                        peptide_start + peptide_length + c_flank_length)
                    sequence_names.append(name)
                    peptides.append(
                        sequence[peptide_start: peptide_start + peptide_length])
                    if use_flanks:
                        n_flanks.append(
                            sequence[n_flank_start : peptide_start])
                        c_flanks.append(
                            sequence[peptide_start + peptide_length : c_flank_end])

        result_df = self.predict_to_dataframe(
            peptides=peptides,
            alleles=alleles,
            n_flanks=n_flanks,
            c_flanks=c_flanks,
            experiment_names=sequence_names,
            include_affinity_percentile=include_affinity_percentile,
            verbose=verbose,
            throw=throw)

        result_df = result_df.rename(
            columns={"experiment_name": "sequence_name"})

        comparison_is_score = comparison_quantity.endswith("score")

        result_df = result_df.sort_values(
            comparison_quantity,
            ascending=not comparison_is_score)

        if result == "best":
            result_df = result_df.drop_duplicates(
                "sequence_name", keep="first").sort_values("sequence_name")
        elif result == "filtered":
            if comparison_is_score:
                result_df = result_df.loc[
                    result_df[comparison_quantity] >= filter_value
                ]
            else:
                result_df = result_df.loc[
                    result_df[comparison_quantity] <= filter_value
                ]
        elif result == "all":
            pass
        else:
            raise ValueError(
                "Unknown result: %s. Valid choices are: best, filtered, all"
                % result)

        result_df = result_df.reset_index(drop=True)
        result_df = result_df.copy()

        return result_df

    def save(self, models_dir):
        """
        Save the predictor to a directory on disk. If the directory does
        not exist it will be created.

        The wrapped Class1AffinityPredictor and Class1ProcessingPredictor
        instances are included in the saved data.

        Parameters
        ----------
        models_dir : string
            Path to directory. It will be created if it doesn't exist.
        """

        if self.weights_dataframe is None:
            raise RuntimeError("Can't save before fitting")

        if not exists(models_dir):
            mkdir(models_dir)

        # Save underlying predictors
        self.affinity_predictor.save(join(models_dir, "affinity_predictor"))
        if self.processing_predictor_with_flanks is not None:
            self.processing_predictor_with_flanks.save(
                join(models_dir, "processing_predictor_with_flanks"))
        if self.processing_predictor_without_flanks is not None:
            self.processing_predictor_without_flanks.save(
                join(models_dir, "processing_predictor_without_flanks"))

        # Save model coefficients.
        self.weights_dataframe.to_csv(join(models_dir, "weights.csv"))

        # Write "info.txt"
        info_path = join(models_dir, "info.txt")
        rows = [
            ("trained on", time.asctime()),
            ("package   ", "mhcflurry %s" % __version__),
            ("hostname  ", gethostname()),
            ("user      ", getuser()),
        ]
        pandas.DataFrame(rows).to_csv(
            info_path, sep="\t", header=False, index=False)

        if self.metadata_dataframes:
            for (name, df) in self.metadata_dataframes.items():
                metadata_df_path = join(models_dir, "%s.csv.bz2" % name)
                df.to_csv(metadata_df_path, index=False, compression="bz2")


    @classmethod
    def load(cls, models_dir=None, max_models=None):
        """
        Deserialize a predictor from a directory on disk.

        This will also load the wrapped Class1AffinityPredictor and
        Class1ProcessingPredictor instances.

        Parameters
        ----------
        models_dir : string
            Path to directory. If unspecified the default downloaded models are
            used.

        max_models : int, optional
            Maximum number of affinity and processing (counted separately)
            models to load

        Returns
        -------
        `Class1PresentationPredictor` instance
        """
        if models_dir is None:
            models_dir = get_default_class1_presentation_models_dir()

        affinity_predictor = Class1AffinityPredictor.load(
            join(models_dir, "affinity_predictor"), max_models=max_models)

        processing_predictor_with_flanks = None
        if exists(join(models_dir, "processing_predictor_with_flanks")):
            processing_predictor_with_flanks = Class1ProcessingPredictor.load(
                join(models_dir, "processing_predictor_with_flanks"),
                max_models=max_models)
        else:
            logging.warning(
                "Presentation predictor is missing processing predictor: %s",
                join(models_dir, "processing_predictor_with_flanks"))

        processing_predictor_without_flanks = None
        if exists(join(models_dir, "processing_predictor_without_flanks")):
            processing_predictor_without_flanks = Class1ProcessingPredictor.load(
                join(models_dir, "processing_predictor_without_flanks"),
                max_models=max_models)
        else:
            logging.warning(
                "Presentation predictor is missing processing predictor: %s",
                join(models_dir, "processing_predictor_without_flanks"))

        weights_dataframe = pandas.read_csv(
            join(models_dir, "weights.csv"),
            index_col=0)

        result = cls(
            affinity_predictor=affinity_predictor,
            processing_predictor_with_flanks=processing_predictor_with_flanks,
            processing_predictor_without_flanks=processing_predictor_without_flanks,
            weights_dataframe=weights_dataframe)
        return result