# Copyright (c) 2016. Mount Sinai School of Medicine
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import print_function, division, absolute_import
from collections import defaultdict, OrderedDict
import logging
from six import string_types
import pandas as pd
import numpy as np
from typechecks import require_iterable_of
from sklearn.cross_validation import KFold
from .common import geometric_mean
from .dataset_helpers import (
prepare_pMHC_affinity_arrays,
load_dataframe
)
from .peptide_encoding import fixed_length_index_encoding
from .imputation_helpers import (
check_dense_pMHC_array,
prune_dense_matrix_and_labels,
dense_pMHC_matrix_to_nested_dict,
)
class Dataset(object):
"""
Peptide-MHC binding dataset with helper methods for constructing
different representations (arrays, DataFrames, dictionaries, &c).
Design considerations:
- want to allow multiple measurements for each pMHC pair (which can
be dynamically combined)
- optional sample weights associated with each pMHC measurement
"""
def __init__(self, df):
"""
Constructs a DataSet from a pandas DataFrame with the following
columns:
- allele
- peptide
- affinity
Also, there is an optional column:
- sample_weight
If `sample_weight` is missing then it is filled with a default value
of 1.0
Parameters
----------
df : pandas.DataFrame
"""
columns = set(df.columns)
for expected_column_name in {"allele", "peptide", "affinity"}:
if expected_column_name not in columns:
raise ValueError("Missing column '%s' from DataFrame")
# make allele and peptide columns the index, and copy it
# so we can add a column without any observable side-effect in
# the calling code
df = df.set_index(["allele", "peptide"], drop=False)
if "sample_weight" not in columns:
df["sample_weight"] = np.ones(len(df), dtype=float)
self._df = df
self._alleles = np.asarray(df["allele"])
self._peptides = np.asarray(df["peptide"])
self._affinities = np.asarray(df["affinity"])
self._sample_weights = np.asarray(df["sample_weight"])
def to_dataframe(self):
"""
Returns DataFrame representation of data contained in Dataset
"""
return self._df
@property
def peptides(self):
"""
Array of peptides from pMHC measurements.
"""
return self._df["peptide"].values
@property
def alleles(self):
"""
Array of MHC allele names from pMHC measurements.
"""
return self.to_dataframe()["allele"].values
@property
def affinities(self):
"""
Array of affinities from pMHC measurements.
"""
return self.to_dataframe()["affinity"].values
@property
def sample_weights(self):
"""
Array of sample weights for each pMHC measurement.
"""
return self.to_dataframe()["sample_weight"].values
def __len__(self):
return len(self.to_dataframe())
def __str__(self):
return "Dataset(n=%d, alleles=%s)" % (
len(self), list(sorted(self.unique_alleles())))
def __repr__(self):
return str(self)
def __eq__(self, other):
"""
Two datasets are equal if they contain the same number of samples
with the same properties and values.
"""
if type(other) is not Dataset:
return False
elif len(self) != len(other):
return False
columns = self.columns
if len(columns) != len(other.columns):
return False
elif set(columns) != set(other.columns):
return False
# test for equality of the rows of the two DataFrames regardless
# of order
my_dict = self.allele_and_peptide_pair_to_row_dictionary()
other_dict = other.allele_and_peptide_pair_to_row_dictionary()
if set(my_dict.keys()) != set(other_dict.keys()):
return False
for key, my_row in my_dict.items():
for column in columns:
if my_row[column] != other_dict[key][column]:
return False
return True
def iterrows(self):
"""
Iterate over tuples containing: (allele, peptide), other_fields
for each pMHC measurement.
"""
return self.to_dataframe().iterrows()
def allele_and_peptide_pair_to_row_dictionary(self):
"""
Returns a dictionary mapping (allele, peptide) pairs to rows.
"""
return {key: row for (key, row) in self.iterrows()}
@property
def columns(self):
return self.to_dataframe().columns
def unique_alleles(self):
"""
Returns the set of allele names contained in this Dataset.
"""
return set(self.alleles)
def unique_peptides(self):
"""
Returns the set of peptide sequences contained in this Dataset.
"""
return set(self.peptides)
def unique_allele_peptide_pairs(self):
"""
Returns set of every unique pMHC pairing in the dataset.
"""
return set(zip(self.alleles, self.peptides))
def groupby_allele(self):
"""
Yields a sequence of tuples of allele names with Datasets containing
entries just for that allele.
"""
for (allele_name, group_df) in self.to_dataframe().groupby("allele"):
yield (allele_name, Dataset(group_df))
def groupby_allele_dictionary(self):
"""
Returns dictionary mapping each allele name to a Dataset containing
only entries from that allele.
"""
return dict(self.groupby_allele())
def allele_counts_dictionary(self):
"""
Returns a dictionary mapping each allele name to the number of entries
associated with it.
"""
return {
allele_name: len(allele_dataset)
for allele_name, allele_dataset
in self.groupby_allele()
}
def filter_alleles_by_count(self, min_peptides_per_allele=0):
return self.concat([
allele_dataset
for (_, allele_dataset)
in self.groupby_allele()
if len(allele_dataset) >= min_peptides_per_allele])
def to_nested_dictionary(self, combine_fn=geometric_mean):
"""
Returns a dictionary mapping from allele name to a dictionary which
maps from peptide to measured value. Caution, this eliminates sample
weights!
Parameters
----------
combine_fn : function
How to combine multiple measurements for the same pMHC complex.
Takes affinities and optional `weights` argument.
"""
allele_to_peptide_to_affinity_dict = {}
for allele, allele_dataset in self.groupby_allele():
# dictionary mapping each peptide to a list of affinities
peptide_to_affinity_dict = defaultdict(list)
peptide_to_weight_dict = defaultdict(list)
for (allele, peptide), row in allele_dataset.iterrows():
affinity = row["affinity"]
sample_weight = row["sample_weight"]
peptide_to_affinity_dict[peptide].append(affinity)
peptide_to_weight_dict[peptide].append(sample_weight)
allele_to_peptide_to_affinity_dict[allele] = {
peptide: combine_fn(
peptide_to_affinity_dict[peptide],
peptide_to_weight_dict[peptide])
for peptide in peptide_to_affinity_dict.keys()
}
return allele_to_peptide_to_affinity_dict
@classmethod
def from_sequences(
cls,
alleles,
peptides,
affinities,
sample_weights=None,
extra_columns={}):
"""
Parameters
----------
alleles : numpy.ndarray, pandas.Series, or list
Name of allele for that pMHC measurement
peptides : numpy.ndarray, pandas.Series, or list
Sequence of peptide in that pMHC measurement.
affinities : numpy.ndarray, pandas.Series, or list
Affinity value (typically IC50 concentration) for that pMHC
sample_weights : numpy.ndarray of float, optional
extra_columns : dict
Dictionary of any extra properties associated with a
pMHC measurement
"""
alleles, peptides, affinities, sample_weights = \
prepare_pMHC_affinity_arrays(
alleles=alleles,
peptides=peptides,
affinities=affinities,
sample_weights=sample_weights)
df = pd.DataFrame()
df["allele"] = alleles
df["peptide"] = peptides
df["affinity"] = affinities
df["sample_weight"] = sample_weights
for column_name, column in extra_columns.items():
if len(column) != len(alleles):
raise ValueError(
"Wrong length for column '%s', expected %d but got %d" % (
column_name,
len(alleles),
len(column)))
df[column_name] = np.asarray(column)
return cls(df)
@classmethod
def from_single_allele_dataframe(cls, allele_name, single_allele_df):
"""
Construct a Dataset from a single MHC allele's DataFrame
"""
df = single_allele_df.copy()
df["allele"] = allele_name
return cls(df)
@classmethod
def from_nested_dictionary(
cls,
allele_to_peptide_to_affinity_dict):
"""
Given nested dictionaries mapping allele -> peptide -> affinity,
construct a Dataset with uniform sample weights.
"""
alleles = []
peptides = []
affinities = []
for allele, allele_dict in allele_to_peptide_to_affinity_dict.items():
for peptide, affinity in allele_dict.items():
alleles.append(allele)
peptides.append(peptide)
affinities.append(affinity)
return cls.from_sequences(
alleles=alleles,
peptides=peptides,
affinities=affinities)
@classmethod
def create_empty(cls):
"""
Returns an empty Dataset containing no pMHC entries.
"""
return cls.from_nested_dictionary({})
@classmethod
def from_single_allele_dictionary(
cls,
allele_name,
peptide_to_affinity_dict):
"""
Given a peptide->affinity dictionary for a single allele,
create a Dataset.
"""
return cls.from_nested_dictionary({allele_name: peptide_to_affinity_dict})
@classmethod
def from_csv(
cls,
filename,
sep=None,
allele_column_name=None,
peptide_column_name=None,
affinity_column_name=None):
df, allele_column_name, peptide_column_name, affinity_column_name = \
load_dataframe(
filename=filename,
sep=sep,
allele_column_name=allele_column_name,
peptide_column_name=peptide_column_name,
affinity_column_name=affinity_column_name)
df = df.rename(columns={
allele_column_name: "allele",
peptide_column_name: "peptide",
affinity_column_name: "affinity"})
return cls(df)
def get_allele(self, allele_name):
"""
Get Dataset for a single allele
"""
if allele_name not in self.unique_alleles():
raise KeyError("Allele '%s' not found" % (allele_name,))
df = self.to_dataframe()
df_allele = df[df.allele == allele_name]
return self.__class__(df_allele)
def get_alleles(self, allele_names):
"""
Restrict Dataset to several allele names.
"""
datasets = []
for allele_name in allele_names:
datasets.append(self.get_allele(allele_name))
return self.concat(datasets)
@classmethod
def concat(cls, datasets):
"""
Concatenate several datasets into a single object.
"""
dataframes = [dataset.to_dataframe() for dataset in datasets]
return cls(pd.concat(dataframes))
def replace_allele(self, allele_name, new_dataset):
"""
Replace data for given allele with new entries.
"""
if allele_name not in self.unique_alleles():
raise ValueError("Allele '%s' not found" % (allele_name,))
df = self.to_dataframe()
df_without = df[df.allele != allele_name]
new_df = new_dataset.to_dataframe()
combined_df = pd.concat([df_without, new_df])
return self.__class__(combined_df)
def flatmap_peptides(self, peptide_fn):
"""
Create zero or more peptides from each pMHC entry. The affinity of all
new peptides is identical to the original, but sample weights are
divided across the number of new peptides.
Parameters
----------
peptide_fn : function
Maps each peptide to a list of peptides.
"""
columns = self.to_dataframe().columns
new_data_dict = OrderedDict(
(column_name, [])
for column_name in columns
)
if "original_peptide" not in new_data_dict:
create_original_peptide_column = True
new_data_dict["original_peptide"] = []
for (allele, peptide), row in self.iterrows():
new_peptides = peptide_fn(peptide)
n = len(new_peptides)
weight = row["sample_weight"]
# we're either going to create a fresh original peptide column
# or extend the existing original peptide tuple that tracks
# the provenance of entries in the new Dataset
original_peptide = row.get("original_peptide")
if original_peptide is None:
original_peptide = ()
elif isinstance(original_peptide, string_types):
original_peptide = (original_peptide,)
else:
original_peptide = tuple(original_peptide)
for new_peptide in new_peptides:
for column_name in columns:
if column_name == "peptide":
new_data_dict["peptide"].append(new_peptide)
elif column_name == "sample_weight":
new_data_dict["sample_weight"].append(weight / n)
elif column_name == "original_peptide":
new_data_dict["original_peptide"] = original_peptide + (peptide,)
else:
new_data_dict[column_name].append(row[column_name])
if create_original_peptide_column:
new_data_dict["original_peptide"].append((peptide,))
df = pd.DataFrame(new_data_dict)
return self.__class__(df)
def kmer_index_encoding(
self,
kmer_size=9,
allow_unknown_amino_acids=True):
"""
Encode peptides in this dataset using a fixed-length vector
representation.
Parameters
----------
kmer_size : int
Length of encoding for each peptide
allow_unknown_amino_acids : bool
If True, then extend shorter amino acids using "X" character,
otherwise fill in all possible combinations of real amino acids.
Returns:
- 2d array of encoded kmers
- 1d array of affinity value corresponding to the source
peptide for each kmer
- sample_weights (1 / kmer count per peptide)
- indices of original peptides from which kmers were extracted
"""
if len(self.peptides) == 0:
return (
np.empty((0, kmer_size), dtype=int),
np.empty((0,), dtype=float),
np.empty((0,), dtype=float),
np.empty((0,), dtype=int)
)
X_index, _, original_peptide_indices, counts = \
fixed_length_index_encoding(
peptides=self.peptides,
desired_length=kmer_size,
start_offset_shorten=0,
end_offset_shorten=0,
start_offset_extend=0,
end_offset_extend=0,
allow_unknown_amino_acids=allow_unknown_amino_acids)
original_peptide_indices = np.asarray(original_peptide_indices)
counts = np.asarray(counts)
kmer_affinities = self.affinities[original_peptide_indices]
kmer_sample_weights = self.sample_weights[original_peptide_indices]
assert len(original_peptide_indices) == len(kmer_affinities)
assert len(counts) == len(kmer_affinities)
assert len(kmer_sample_weights) == len(kmer_affinities)
# combine the original sample weights of varying length peptides
# with a 1/n_kmers factor for the number of kmers pulled out of each
# original peptide
combined_sample_weights = kmer_sample_weights * (1.0 / counts)
return X_index, kmer_affinities, combined_sample_weights, original_peptide_indices
def to_dense_pMHC_affinity_matrix(
self,
min_observations_per_peptide=1,
min_observations_per_allele=1):
"""
Returns a tuple with a dense matrix of affinities, a dense matrix of
sample weights, a list of peptide labels for each row and a list of
allele labels for each column.
Parameters
----------
min_observations_per_peptide : int
Drop peptide rows with fewer than this number of observed values.
min_observations_per_allele : int
Drop allele columns with fewer than this number of observed values.
"""
allele_to_peptide_to_affinity_dict = self.to_nested_dictionary()
peptides_list = list(sorted(self.unique_peptides()))
peptide_order = {p: i for (i, p) in enumerate(peptides_list)}
n_peptides = len(peptides_list)
alleles_list = list(sorted(self.unique_alleles()))
allele_order = {a: i for (i, a) in enumerate(alleles_list)}
n_alleles = len(alleles_list)
shape = (n_peptides, n_alleles)
X = np.ones(shape, dtype=float) * np.nan
for (allele, allele_dict) in allele_to_peptide_to_affinity_dict.items():
column_index = allele_order[allele]
for (peptide, affinity) in allele_dict.items():
row_index = peptide_order[peptide]
X[row_index, column_index] = affinity
check_dense_pMHC_array(X, peptides_list, alleles_list)
# drop alleles and peptides with small amounts of data
return prune_dense_matrix_and_labels(
X, peptides_list, alleles_list,
min_observations_per_peptide=min_observations_per_peptide,
min_observations_per_allele=min_observations_per_allele)
def slice(self, indices):
"""
Create a new Dataset by slicing through all columns of this dataset
with the given indices.
"""
max_index = indices.max()
n_total = len(self)
if max_index >= len(self):
raise ValueError("Invalid index %d for Dataset of size %d" % (
max_index, n_total))
df = self.to_dataframe()
df_subset = pd.DataFrame()
for column_name in df.columns:
df_subset[column_name] = np.asarray(df[column_name].values)[indices]
return self.__class__(df_subset)
def random_split(self, n=None):
"""
Randomly split the Dataset into smaller Dataset objects.
Parameters
----------
n : int, optional
Size of the left split, half of the dataset if omitted.
Returns a pair of Dataset objects.
"""
n_total = len(self)
if n is None:
n = n_total // 2
elif n >= n_total:
raise ValueError(
"Training subset can't have more than %d samples (given n=%d)" % (
n_total - 1,
n))
all_indices = np.arange(n_total)
np.random.shuffle(all_indices)
left = self.slice(all_indices[:n])
right = self.slice(all_indices[n:])
return left, right
def cross_validation_iterator(
self, test_allele=None, n_folds=3, shuffle=True):
"""
Yields a sequence of training/test splits of this dataset.
If test_allele is None then split across all pMHC entries, otherwise
only split the measurements of the specified allele (other alleles
will then always be included in the training datasets).
"""
if test_allele is None:
candidate_test_indices = np.arange(len(self))
elif test_allele not in self.unique_alleles():
raise ValueError("Allele '%s' not in Dataset" % test_allele)
else:
candidate_test_indices = np.where(self.alleles == test_allele)[0]
n_candidate_test_samples = len(candidate_test_indices)
n_total = len(self)
for _, subindices in KFold(
n=n_candidate_test_samples,
n_folds=n_folds,
shuffle=shuffle):
test_indices = candidate_test_indices[subindices]
train_mask = np.ones(n_total, dtype=bool)
train_mask[test_indices] = False
train_data = self.slice(train_mask)
test_data = self.slice(test_indices)
yield train_data, test_data
def split_allele_randomly_and_impute_training_set(
self, allele, n_training_samples=None, **kwargs):
"""
Split an allele into training and test sets, and then impute values
for peptides missing from the training set using data from other alleles
in this Dataset.
(apologies for the wordy name, this turns out to be a common operation)
Parameters
----------
allele : str
Name of allele
n_training_samples : int, optional
Size of the training set to return for this allele.
**kwargs : dict
Extra keyword arguments passed to Dataset.impute_missing_values
Returns three Dataset objects:
- training set with original pMHC affinities for given allele
- larger imputed training set for given allele
- test set
"""
dataset_allele = self.get_allele(allele)
dataset_allele_train, dataset_allele_test = dataset_allele.random_split(
n=n_training_samples)
full_dataset_without_test_samples = self.difference(dataset_allele_test)
imputed_dataset = full_dataset_without_test_samples.impute_missing_values(**kwargs)
imputed_dataset_allele = imputed_dataset.get_allele(allele)
return dataset_allele_train, imputed_dataset_allele, dataset_allele_test
def drop_allele_peptide_lists(self, alleles, peptides):
"""
Drop all allele-peptide pairs in the given lists.
Parameters
----------
alleles : list of str
peptides : list of str
The two arguments are assumed to be the same length.
Returns Dataset of equal or smaller size.
"""
if len(alleles) != len(peptides):
raise ValueError(
"Expected alleles to be same length (%d) as peptides (%d)" % (
len(alleles), len(peptides)))
return self.drop_allele_peptide_pairs(list(zip(alleles, peptides)))
def drop_allele_peptide_pairs(self, allele_peptide_pairs):
"""
Drop all allele-peptide tuple pairs in the given list.
Parameters
----------
allele_peptide_pairs : list of (str, str) tuples
The two arguments are assumed to be the same length.
Returns Dataset of equal or smaller size.
"""
require_iterable_of(allele_peptide_pairs, tuple)
keys_to_remove_set = set(allele_peptide_pairs)
remove_mask = np.array([
(k in keys_to_remove_set)
for k in zip(self.alleles, self.peptides)
])
keep_mask = ~remove_mask
return self.slice(keep_mask)
def difference(self, other_dataset):
"""
Remove all pMHC pairs in the other dataset from this one.
Parameters
----------
other_dataset : Dataset
Returns a new Dataset object of equal or lesser size.
"""
return self.drop_allele_peptide_lists(
alleles=other_dataset.alleles,
peptides=other_dataset.peptides)
def intersection(self, other_dataset):
not_in_other = self.difference(other_dataset)
return self.difference(not_in_other)
def impute_missing_values(
self,
imputation_method,
log_transform=True,
min_observations_per_peptide=1,
min_observations_per_allele=1):
"""
Synthesize new measurements for missing pMHC pairs using the given
imputation_method.
Parameters
----------
imputation_method : object
Expected to have a method called `complete` which takes a 2d array
of floats and replaces some or all NaN values with synthetic
affinities.
log_transform : function
Transform affinities with to log10 values before imputation
(and then transform back afterward).
min_observations_per_peptide : int
Drop peptide rows with fewer than this number of observed values.
min_observations_per_allele : int
Drop allele columns with fewer than this number of observed values.
Returns Dataset with original pMHC affinities and additional
synthetic samples.
"""
X_incomplete, peptide_list, allele_list = self.to_dense_pMHC_affinity_matrix(
min_observations_per_peptide=min_observations_per_peptide,
min_observations_per_allele=min_observations_per_allele)
if log_transform:
X_incomplete = np.log(X_incomplete)
if np.isnan(X_incomplete).sum() == 0:
# if all entries in the matrix are already filled in then don't
# try using an imputation algorithm since it might raise an
# exception.
logging.warn("No missing values, using original data instead of imputation")
X_complete = X_incomplete
else:
X_complete = imputation_method.complete(X_incomplete)
if log_transform:
X_complete = np.exp(X_complete)
allele_to_peptide_to_affinity_dict = dense_pMHC_matrix_to_nested_dict(
X=X_complete,
peptide_list=peptide_list,
allele_list=allele_list)
return self.from_nested_dictionary(allele_to_peptide_to_affinity_dict)