rename Dataset to AffinityMeasurementDataset, fix bug in tests/test_ensemble.py

downloads-generation/data_combined_iedb_kim2014/create-combined-class1-dataset.py

@@ -17,7 +17,7 @@
 """
 Combine 2013 Kim/Peters NetMHCpan dataset[*] with more recent IEDB entries
 
-* = "Dataset size and composition impact the reliability..."
+* = "AffinityMeasurementDataset size and composition impact the reliability..."
 """
  
 from __future__ import (

mhcflurry/dataset.py → mhcflurry/affinity_measurement_dataset.py

@@ -37,7 +37,7 @@ from .imputation_helpers import (
 )
 
 
-class Dataset(object):
+class AffinityMeasurementDataset(object):
     """
     Peptide-MHC binding dataset with helper methods for constructing
     different representations (arrays, DataFrames, dictionaries, &c).
@@ -53,7 +53,7 @@ class Dataset(object):
     """
     def __init__(self, df):
         """
-        Constructs a Dataset from a pandas DataFrame with the following
+        Constructs a AffinityMeasurementDataset from a pandas DataFrame with the following
         columns:
             - allele
             - peptide
@@ -92,7 +92,7 @@ class Dataset(object):
 
     def to_dataframe(self):
         """
-        Returns DataFrame representation of data contained in Dataset
+        Returns DataFrame representation of data contained in AffinityMeasurementDataset
         """
         return self._df
 
@@ -128,7 +128,7 @@ class Dataset(object):
         return len(self.to_dataframe())
 
     def __str__(self):
-        return "Dataset(n=%d, alleles=%s)" % (
+        return "AffinityMeasurementDataset(n=%d, alleles=%s)" % (
             len(self), list(sorted(self.unique_alleles())))
 
     def __repr__(self):
@@ -139,7 +139,7 @@ class Dataset(object):
         Two datasets are equal if they contain the same number of samples
         with the same properties and values.
         """
-        if type(other) is not Dataset:
+        if type(other) is not AffinityMeasurementDataset:
             return False
         elif len(self) != len(other):
             return False
@@ -183,13 +183,13 @@ class Dataset(object):
 
     def unique_alleles(self):
         """
-        Returns the set of allele names contained in this Dataset.
+        Returns the set of allele names contained in this AffinityMeasurementDataset.
         """
         return set(self.alleles)
 
     def unique_peptides(self):
         """
-        Returns the set of peptide sequences contained in this Dataset.
+        Returns the set of peptide sequences contained in this AffinityMeasurementDataset.
         """
         return set(self.peptides)
 
@@ -205,11 +205,11 @@ class Dataset(object):
         entries just for that allele.
         """
         for (allele_name, group_df) in self.to_dataframe().groupby("allele"):
-            yield (allele_name, Dataset(group_df))
+            yield (allele_name, AffinityMeasurementDataset(group_df))
 
     def groupby_allele_dictionary(self):
         """
-        Returns dictionary mapping each allele name to a Dataset containing
+        Returns dictionary mapping each allele name to a AffinityMeasurementDataset containing
         only entries from that allele.
         """
         return dict(self.groupby_allele())
@@ -317,7 +317,7 @@ class Dataset(object):
     @classmethod
     def from_single_allele_dataframe(cls, allele_name, single_allele_df):
         """
-        Construct a Dataset from a single MHC allele's DataFrame
+        Construct a AffinityMeasurementDataset from a single MHC allele's DataFrame
         """
         df = single_allele_df.copy()
         df["allele"] = allele_name
@@ -329,7 +329,7 @@ class Dataset(object):
             allele_to_peptide_to_affinity_dict):
         """
         Given nested dictionaries mapping allele -> peptide -> affinity,
-        construct a Dataset with uniform sample weights.
+        construct a AffinityMeasurementDataset with uniform sample weights.
         """
         alleles = []
         peptides = []
@@ -347,7 +347,7 @@ class Dataset(object):
     @classmethod
     def create_empty(cls):
         """
-        Returns an empty Dataset containing no pMHC entries.
+        Returns an empty AffinityMeasurementDataset containing no pMHC entries.
         """
         return cls.from_nested_dictionary({})
 
@@ -358,7 +358,7 @@ class Dataset(object):
             peptide_to_affinity_dict):
         """
         Given a peptide->affinity dictionary for a single allele,
-        create a Dataset.
+        create a AffinityMeasurementDataset.
         """
         return cls.from_nested_dictionary({allele_name: peptide_to_affinity_dict})
 
@@ -385,7 +385,7 @@ class Dataset(object):
 
     def get_allele(self, allele_name):
         """
-        Get Dataset for a single allele
+        Get AffinityMeasurementDataset for a single allele
         """
         if allele_name not in self.unique_alleles():
             raise KeyError("Allele '%s' not found, available alleles: %s" % (
@@ -396,7 +396,7 @@ class Dataset(object):
 
     def get_alleles(self, allele_names):
         """
-        Restrict Dataset to several allele names.
+        Restrict AffinityMeasurementDataset to several allele names.
         """
         datasets = []
         for allele_name in allele_names:
@@ -449,7 +449,7 @@ class Dataset(object):
             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
+            # the provenance of entries in the new AffinityMeasurementDataset
             original_peptide = row.get("original_peptide")
             if original_peptide is None:
                 original_peptide = ()
@@ -572,14 +572,14 @@ class Dataset(object):
 
     def slice(self, indices):
         """
-        Create a new Dataset by slicing through all columns of this dataset
+        Create a new AffinityMeasurementDataset by slicing through all columns of this dataset
         with the given indices.
         """
         indices = np.asarray(indices)
         max_index = indices.max()
         n_total = len(self)
         if max_index >= len(self):
-            raise ValueError("Invalid index %d for Dataset of size %d" % (
+            raise ValueError("Invalid index %d for AffinityMeasurementDataset of size %d" % (
                 max_index, n_total))
 
         df = self.to_dataframe()
@@ -590,7 +590,7 @@ class Dataset(object):
 
     def random_split(self, n=None, stratify_fn=None):
         """
-        Randomly split the Dataset into smaller Dataset objects.
+        Randomly split the AffinityMeasurementDataset into smaller AffinityMeasurementDataset objects.
 
         Parameters
         ----------
@@ -601,7 +601,7 @@ class Dataset(object):
             Function that takes a row and returns bool, stratifying sampling
             into two groups.
 
-        Returns a pair of Dataset objects.
+        Returns a pair of AffinityMeasurementDataset objects.
         """
         n_total = len(self)
         if n is None:
@@ -649,7 +649,7 @@ class Dataset(object):
             test_samples = self
             test_sample_indices = np.arange(n_total)
         elif test_allele not in self.unique_alleles():
-            raise ValueError("Allele '%s' not in Dataset" % test_allele)
+            raise ValueError("Allele '%s' not in AffinityMeasurementDataset" % test_allele)
 
         else:
             test_sample_indices = np.where(self.alleles == test_allele)[0]
@@ -688,7 +688,7 @@ class Dataset(object):
         """
         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.
+        in this AffinityMeasurementDataset.
 
         (apologies for the wordy name, this turns out to be a common operation)
 
@@ -701,13 +701,13 @@ class Dataset(object):
             Size of the training set to return for this allele.
 
         stratify_fn : function
-            Function mapping from rows of the Dataset to booleans for stratifying
+            Function mapping from rows of the AffinityMeasurementDataset to booleans for stratifying
             by two groups.
 
         **kwargs : dict
-            Extra keyword arguments passed to Dataset.impute_missing_values
+            Extra keyword arguments passed to AffinityMeasurementDataset.impute_missing_values
 
-        Returns three Dataset objects:
+        Returns three AffinityMeasurementDataset objects:
             - training set with original pMHC affinities for given allele
             - larger imputed training set for given allele
             - test set
@@ -732,7 +732,7 @@ class Dataset(object):
 
         The two arguments are assumed to be the same length.
 
-        Returns Dataset of equal or smaller size.
+        Returns AffinityMeasurementDataset of equal or smaller size.
         """
         if len(alleles) != len(peptides):
             raise ValueError(
@@ -749,7 +749,7 @@ class Dataset(object):
         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.
+        Returns AffinityMeasurementDataset of equal or smaller size.
         """
         require_iterable_of(allele_peptide_pairs, tuple)
         keys_to_remove_set = set(allele_peptide_pairs)
@@ -766,7 +766,7 @@ class Dataset(object):
 
         Parameters
         ----------
-        other_dataset : Dataset
+        other_dataset : AffinityMeasurementDataset
 
         Returns a new Dataset object of equal or lesser size.
         """
@@ -805,7 +805,7 @@ class Dataset(object):
         min_observations_per_allele : int
             Drop allele columns with fewer than this number of observed values.
 
-        Returns Dataset with original pMHC affinities and additional
+        Returns AffinityMeasurementDataset with original pMHC affinities and additional
         synthetic samples.
         """
         if isinstance(imputation_method, string_types):

mhcflurry/antigen_presentation/presentation_component_models/mhc_binding_component_model_base.py

@@ -47,7 +47,7 @@ class MHCBindingComponentModelBase(PresentationComponentModel):
     fallback_predictor : function: (allele, peptides) -> predictions
         Used when missing an allele.
 
-    iedb_dataset : mhcflurry.Dataset
+    iedb_dataset : mhcflurry.AffinityMeasurementDataset
         IEDB data for this allele. If not specified no iedb data is used.
 
     decoys_per_hit : int

mhcflurry/antigen_presentation/presentation_component_models/mhcflurry_trained_on_hits.py

@@ -3,7 +3,7 @@ from copy import copy
 import pandas
 from numpy import log, exp, nanmean
 
-from ...dataset import Dataset
+from ...affinity_measurement_dataset import AffinityMeasurementDataset
 from ...class1_allele_specific import Class1BindingPredictor
 from ...common import normalize_allele_name
 
@@ -22,7 +22,7 @@ class MHCflurryTrainedOnHits(MHCBindingComponentModelBase):
 
     Parameters
     ------------
-    iedb_dataset : mhcflurry.Dataset
+    iedb_dataset : mhcflurry.AffinityMeasurementDataset
         IEDB data for this allele. If not specified no iedb data is used.
 
     mhcflurry_hyperparameters : dict
@@ -96,7 +96,7 @@ class MHCflurryTrainedOnHits(MHCBindingComponentModelBase):
         if self.iedb_dataset is not None:
             df = df.append(iedb_dataset_df, ignore_index=True)
 
-        dataset = Dataset(
+        dataset = AffinityMeasurementDataset(
             df.sample(frac=1))  # shuffle dataframe
         print("Train data: ", dataset)
         model = Class1BindingPredictor(

mhcflurry/class1_allele_specific/class1_allele_specific_kmer_ic50_predictor_base.py

@@ -168,12 +168,12 @@ class Class1AlleleSpecificKmerIC50PredictorBase(IC50PredictorBase):
 
         Parameters
         ----------
-        dataset : Dataset
+        dataset : AffinityMeasurementDataset
 
-        pretraining_dataset : Dataset
+        pretraining_dataset : AffinityMeasurementDataset
 
         sample_censored_affinities : bool
-            If a column named 'inequality' is in the Dataset then every
+            If a column named 'inequality' is in the AffinityMeasurementDataset then every
             peptide with a value of '>' on each training epoch, gets a
             randomly sampled IC50 between its observed value and the
             max_ic50 of the predictor. Default is False.

mhcflurry/class1_allele_specific/cross_validation.py

@@ -102,12 +102,12 @@ def cross_validation_folds(
         },
         parallel_backend=None):
     '''
-    Split a Dataset into n_folds cross validation folds for each allele,
+    Split a AffinityMeasurementDataset into n_folds cross validation folds for each allele,
     optionally performing imputation.
 
     Parameters
     -----------
-    train_data : mhcflurry.Dataset
+    train_data : mhcflurry.AffinityMeasurementDataset
 
     alleles : string list, optional
         Alleles to run cross validation on. Default: all alleles in
@@ -125,7 +125,7 @@ def cross_validation_folds(
         Imputer to use. If not specified, no imputation is done.
 
     impute_kwargs : dict, optional
-        Additional kwargs to pass to mhcflurry.Dataset.impute_missing_values.
+        Additional kwargs to pass to mhcflurry.AffinityMeasurementDataset.impute_missing_values.
 
     parallel_backend : mhcflurry.parallelism.ParallelBackend, optional
         Futures implementation to use for running on multiple threads,

mhcflurry/class1_allele_specific/cv_and_train_command.py

@@ -49,7 +49,7 @@ import pickle
 import numpy
 
 from .. import parallelism
-from ..dataset import Dataset
+from ..affinity_measurement_dataset import AffinityMeasurementDataset
 from ..imputation_helpers import imputer_from_name
 from .cross_validation import cross_validation_folds
 from .train import (
@@ -216,12 +216,12 @@ def go(args):
         logging.info(
             "Subselected to %d model architectures" % len(model_architectures))
 
-    train_data = Dataset.from_csv(args.train_data)
+    train_data = AffinityMeasurementDataset.from_csv(args.train_data)
     logging.info("Loaded training dataset: %s" % train_data)
 
     test_data = None
     if args.test_data:
-        test_data = Dataset.from_csv(args.test_data)
+        test_data = AffinityMeasurementDataset.from_csv(args.test_data)
         logging.info("Loaded testing dataset: %s" % test_data)
 
     if args.min_samples_per_allele:

mhcflurry/class1_allele_specific/train.py

@@ -52,10 +52,10 @@ def impute_and_select_allele(dataset, imputer, allele=None, **kwargs):
 
     Parameters
     -----------
-    dataset : mhcflurry.Dataset
+    dataset : mhcflurry.AffinityMeasurementDataset
 
     imputer : object or string
-        See Dataset.impute_missing_values
+        See AffinityMeasurementDataset.impute_missing_values
 
     allele : string [optional]
         Allele name to subselect to after imputation
@@ -121,14 +121,14 @@ def train_and_test_one_model_one_fold(
     -----------
     model_description : dict of model parameters
 
-    train_dataset : mhcflurry.Dataset
-        Dataset to train on. Must include only one allele.
+    train_dataset : mhcflurry.AffinityMeasurementDataset
+        AffinityMeasurementDataset to train on. Must include only one allele.
 
-    test_dataset : mhcflurry.Dataset, optional
-        Dataset to test on. Must include only one allele. If not specified
+    test_dataset : mhcflurry.AffinityMeasurementDataset, optional
+        AffinityMeasurementDataset to test on. Must include only one allele. If not specified
         no testing is performed.
 
-    imputed_train_dataset : mhcflurry.Dataset, optional
+    imputed_train_dataset : mhcflurry.AffinityMeasurementDataset, optional
         Required only if model_description["impute"] == True
 
     return_train_scores : boolean

mhcflurry/class1_allele_specific_ensemble/train_command.py

@@ -44,7 +44,7 @@ import traceback
 import signal
 
 from .. import parallelism
-from ..dataset import Dataset
+from ..affinity_measurement_dataset import AffinityMeasurementDataset
 
 from .class1_ensemble_multi_allele_predictor import (
     Class1EnsembleMultiAllelePredictor)
@@ -204,7 +204,7 @@ def go(args):
         logging.info(
             "Subselected to %d model architectures" % len(model_architectures))
 
-    train_dataset = Dataset.from_csv(args.train_data)
+    train_dataset = AffinityMeasurementDataset.from_csv(args.train_data)
     logging.info("Loaded training data: %s" % train_dataset)
 
     if args.alleles:

mhcflurry/ic50_predictor_base.py

@@ -21,7 +21,7 @@ from __future__ import (
 import numpy as np
 
 from .regression_target import regression_target_to_ic50, MAX_IC50
-from .dataset import Dataset
+from .affinity_measurement_dataset import AffinityMeasurementDataset
 from .hyperparameters import HyperparameterDefaults
 
 
@@ -71,7 +71,7 @@ class IC50PredictorBase(object):
         peptide_to_ic50_dict : dict
             Dictionary that maps peptides to IC50 values.
         """
-        dataset = Dataset.from_peptide_to_affinity_dictionary(
+        dataset = AffinityMeasurementDataset.from_peptide_to_affinity_dictionary(
             allele_name=self.name,
             peptide_to_affinity_dict=peptide_to_ic50_dict)
         return self.fit_dataset(dataset, **kwargs)
@@ -84,7 +84,7 @@ class IC50PredictorBase(object):
             alleles=None, **kwargs):
         if alleles is None:
             alleles = [self.name] * len(peptides)
-        dataset = Dataset.from_sequences(
+        dataset = AffinityMeasurementDataset.from_sequences(
             alleles=alleles,
             peptides=peptides,
             affinities=affinities,

mhcflurry/measurement_collection.py

 from sklearn.model_selection import StratifiedKFold
 import pandas
 
-from .dataset import Dataset
+from .affinity_measurement_dataset import AffinityMeasurementDataset
 from .imputation_helpers import imputer_from_name
 
 COLUMNS = [
@@ -32,9 +32,9 @@ class MeasurementCollection(object):
     A single measurement collection may have both MS hits and affinity
     measurements.
 
-    This is more general than a Dataset since it supports MS hits. It is also
+    This is more general than a AffinityMeasurementDataset since it supports MS hits. It is also
     simpler, as the user is expected to manipulate the underlying dataframe.
-    Later we may want to retire Dataset or combine it with this class.
+    Later we may want to retire AffinityMeasurementDataset or combine it with this class.
     """
 
     def __init__(self, df, check=True):
@@ -50,7 +50,7 @@ class MeasurementCollection(object):
     @staticmethod
     def from_dataset(dataset):
         """
-        Given a Dataset, return a MeasurementCollection
+        Given a AffinityMeasurementDataset, return a MeasurementCollection
         """
         dataset_df = dataset.to_dataframe()
         df = dataset_df.reset_index(drop=True)[["allele", "peptide"]].copy()
@@ -141,7 +141,7 @@ class MeasurementCollection(object):
             ms_hit_affinity=1.0,
             ms_decoy_affinity=20000):
         """
-        Return a Dataset containing the observations in the collection.
+        Return a AffinityMeasurementDataset containing the observations in the collection.
         Mass-spec data are converted to affinities according to
         ms_hit_affinity and ms_decoy_affinity.
 
@@ -160,10 +160,10 @@ class MeasurementCollection(object):
 
         Returns
         -------------
-        Dataset instance
+        AffinityMeasurementDataset instance
         """
         if include_ms:
-            dataset = Dataset(pandas.DataFrame({
+            dataset = AffinityMeasurementDataset(pandas.DataFrame({
                 "allele": self.df.allele,
                 "peptide": self.df.peptide,
                 "affinity": [
@@ -180,7 +180,7 @@ class MeasurementCollection(object):
                 (self.df.measurement_type == "affinity") &
                 (self.df.measurement_source == "in_vitro_affinity_assay")
             ]
-            dataset = Dataset(pandas.DataFrame({
+            dataset = AffinityMeasurementDataset(pandas.DataFrame({
                 "allele": df.allele,
                 "peptide": df.peptide,
                 "affinity": df.measurement_value,

test/test_class1_binding_predictor_A0205.py

 import numpy as np
 np.random.seed(0)
 
-from mhcflurry.dataset import Dataset
+from mhcflurry.affinity_measurement_dataset import AffinityMeasurementDataset
 from mhcflurry import Class1BindingPredictor
 
 from nose.tools import eq_
@@ -11,10 +11,10 @@ from mhcflurry.downloads import get_path
 
 
 def test_class1_binding_predictor_A0205_training_accuracy():
-    dataset = Dataset.from_csv(get_path(
+    dataset = AffinityMeasurementDataset.from_csv(get_path(
         "data_combined_iedb_kim2014", "combined_human_class1_dataset.csv"))
     dataset_a0205_all_lengths = dataset.get_allele("HLA-A0205")
-    dataset_a0205 = Dataset(
+    dataset_a0205 = AffinityMeasurementDataset(
         dataset_a0205_all_lengths._df.ix[
             dataset_a0205_all_lengths._df.peptide.str.len() == 9])
 

test/test_dataset.py

 from nose.tools import eq_
-from mhcflurry.dataset import Dataset
+from mhcflurry.affinity_measurement_dataset import AffinityMeasurementDataset
 
 def test_create_allele_data_from_single_allele_dict():
     peptide_to_ic50_dict = {
         ("A" * 10): 1.2,
         ("C" * 9): 1000,
     }
-    dataset = Dataset.from_single_allele_dictionary(
+    dataset = AffinityMeasurementDataset.from_single_allele_dictionary(
         allele_name="A0201",
         peptide_to_affinity_dict=peptide_to_ic50_dict)
-    assert isinstance(dataset, Dataset)
+    assert isinstance(dataset, AffinityMeasurementDataset)
 
     eq_(len(peptide_to_ic50_dict), len(dataset))
     expected_peptides = set([
@@ -22,7 +22,7 @@ def test_create_allele_data_from_single_allele_dict():
         eq_(pi, pj)
 
 def test_dataset_random_split():
-    dataset = Dataset.from_nested_dictionary({
+    dataset = AffinityMeasurementDataset.from_nested_dictionary({
         "H-2-Kb": {
             "SIINFEKL": 10.0,
             "FEKLSIIN": 20000.0,
@@ -33,15 +33,15 @@ def test_dataset_random_split():
     assert len(right) == 1
 
 def test_dataset_difference():
-    dataset1 = Dataset.from_nested_dictionary({
+    dataset1 = AffinityMeasurementDataset.from_nested_dictionary({
         "H-2-Kb": {
             "SIINFEKL": 10.0,
             "FEKLSIIN": 20000.0,
             "SIFEKLIN": 50000.0,
         }})
-    dataset2 = Dataset.from_nested_dictionary({"H-2-Kb": {"SIINFEKL": 10.0}})
+    dataset2 = AffinityMeasurementDataset.from_nested_dictionary({"H-2-Kb": {"SIINFEKL": 10.0}})
     dataset_diff = dataset1.difference(dataset2)
-    expected_result = Dataset.from_nested_dictionary({
+    expected_result = AffinityMeasurementDataset.from_nested_dictionary({
         "H-2-Kb": {
             "FEKLSIIN": 20000.0,
             "SIFEKLIN": 50000.0,
@@ -50,20 +50,20 @@ def test_dataset_difference():
 
 
 def test_dataset_intersection():
-    dataset1 = Dataset.from_nested_dictionary({
+    dataset1 = AffinityMeasurementDataset.from_nested_dictionary({
         "H-2-Kb": {
             "SIINFEKL": 10.0,
             "FEKLSIIN": 20000.0,
             "SIFEKLIN": 50000.0,
         }})
-    dataset2 = Dataset.from_nested_dictionary({"H-2-Kb": {"SIINFEKL": 30.0}})
+    dataset2 = AffinityMeasurementDataset.from_nested_dictionary({"H-2-Kb": {"SIINFEKL": 30.0}})
     dataset_intersection = dataset1.intersection(dataset2)
-    expected_result = Dataset.from_nested_dictionary({
+    expected_result = AffinityMeasurementDataset.from_nested_dictionary({
         "H-2-Kb": {"SIINFEKL": 10.0}})
     eq_(dataset_intersection, expected_result)
 
 def test_dataset_cross_validation():
-    dataset = Dataset.from_nested_dictionary({
+    dataset = AffinityMeasurementDataset.from_nested_dictionary({
         "H-2-Kb": {
             "SIINFEKL": 10.0,
             "FEKLSIIN": 20000.0,

test/test_ensemble.py

@@ -12,7 +12,7 @@ from nose.tools import eq_
 from mhcflurry.class1_allele_specific import scoring
 from mhcflurry.measurement_collection import MeasurementCollection
 from mhcflurry.class1_allele_specific_ensemble import train_command
-from mhcflurry.dataset import Dataset
+from mhcflurry.affinity_measurement_dataset import AffinityMeasurementDataset
 from mhcflurry.downloads import get_path
 from mhcflurry \
     .class1_allele_specific_ensemble \
@@ -34,9 +34,9 @@ def test_basic():
         hyperparameters_to_search=model_hyperparameters)
     print(model)
 
-    dataset = Dataset.from_csv(get_path(
+    dataset = AffinityMeasurementDataset.from_csv(get_path(
         "data_combined_iedb_kim2014", "combined_human_class1_dataset.csv"))
-    sub_dataset = Dataset(
+    sub_dataset = AffinityMeasurementDataset(
         dataset._df.ix[
             (dataset._df.allele.isin(["HLA-A0101", "HLA-A0205"])) &
             (dataset._df.peptide.str.len() == 9)
@@ -83,7 +83,7 @@ def test_basic():
         shutil.rmtree(tmpdir)
 
     eq_(model.ensemble_size, model2.ensemble_size)
-    eq_(model.supported_alleles(), model2.supported_alleles())
+    eq_(model.supported_alleles, model2.supported_alleles)
     eq_(model.hyperparameters_to_search, model2.hyperparameters_to_search)
     ic50_pred2 = model.predict(mc)
     assert_allclose(ic50_pred, ic50_pred2, rtol=1e-06)

test/test_imputation.py

@@ -2,7 +2,7 @@ import numpy as np
 np.random.seed(0)
 
 from mhcflurry.imputation_helpers import imputer_from_name
-from mhcflurry.dataset import Dataset
+from mhcflurry.affinity_measurement_dataset import AffinityMeasurementDataset
 from mhcflurry import Class1BindingPredictor
 
 from fancyimpute import MICE, KNN, SoftImpute, IterativeSVD
@@ -12,13 +12,13 @@ from mhcflurry.downloads import get_path
 
 
 def test_create_imputed_datasets_empty():
-    empty_dataset = Dataset.create_empty()
+    empty_dataset = AffinityMeasurementDataset.create_empty()
     result = empty_dataset.impute_missing_values(MICE(n_imputations=25))
     eq_(result, empty_dataset)
 
 
 def test_create_imputed_datasets_two_alleles():
-    dataset = Dataset.from_nested_dictionary({
+    dataset = AffinityMeasurementDataset.from_nested_dictionary({
         "HLA-A*02:01": {
             "A" * 9: 20.0,
             "C" * 9: 40000.0,
@@ -38,7 +38,7 @@ def test_create_imputed_datasets_two_alleles():
 def test_performance_improves_for_A0205_with_pretraining():
     # test to make sure that imputation improves predictive accuracy after a
     # small number of training iterations (5 epochs)
-    dataset = Dataset.from_csv(
+    dataset = AffinityMeasurementDataset.from_csv(
         get_path("data_combined_iedb_kim2014", "combined_human_class1_dataset.csv"))
     print("Full dataset: %d pMHC entries" % len(dataset))
 
@@ -51,7 +51,7 @@ def test_performance_improves_for_A0205_with_pretraining():
 
     a0205_data_without_imputation = dataset.get_allele("HLA-A0205")
 
-    print("Dataset with only A0205, # entries: %d" % (
+    print("AffinityMeasurementDataset with only A0205, # entries: %d" % (
         len(a0205_data_without_imputation)))
 
     predictor_without_imputation = Class1BindingPredictor(

test/test_load_allele_datasets.py

-from mhcflurry.dataset import Dataset
+from mhcflurry.affinity_measurement_dataset import AffinityMeasurementDataset
 from nose.tools import eq_
 from . import data_path
 
 
 def load_csv(filename):
-    return Dataset.from_csv(data_path(filename))
+    return AffinityMeasurementDataset.from_csv(data_path(filename))
 
 
 def test_load_allele_datasets_8mer():