Set upper limit on tensorflow version. Fixes #234

mhcflurry/version.py

@@ -1 +1 @@
-__version__ = "2.1.0"
+__version__ = "2.1.1"

requirements.txt

@@ -1,6 +1,6 @@
 six
 pandas>=0.20.3
-tensorflow>=2.12.0
+tensorflow>=2.12.0,<2.16.0
 appdirs
 scikit-learn
 mhcgnomes

setup.py

@@ -41,7 +41,7 @@
         "mhcgnomes>=0.8.4",
         "pyyaml",
         "tqdm",
-        "tensorflow>=2.12.0",
+        "tensorflow>=2.12.0,<2.16.0",
     ]
 
     setup(

test/test_calibrate_percentile_ranks_command.py

@@ -8,6 +8,7 @@
 import shutil
 import tempfile
 import subprocess
+import pytest
 
 from numpy.testing import assert_equal
 
@@ -17,8 +18,12 @@
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
 from mhcflurry.testing_utils import cleanup, startup
-teardown = cleanup
-setup = startup
+
+pytest.fixture(autouse=True, scope="module")
+def setup_module():
+    startup()
+    yield
+    cleanup()
 
 
 def run_and_check(n_jobs=0, delete=True, additional_args=[]):

test/test_changing_allele_representations.py

@@ -2,14 +2,19 @@
 initialize()
 
 import pandas
+import pytest
 
 from mhcflurry.class1_affinity_predictor import Class1AffinityPredictor
 from mhcflurry.downloads import get_path
 
 
 from mhcflurry.testing_utils import cleanup, startup
-teardown = cleanup
-setup = startup
+
+pytest.fixture(autouse=True, scope="module")
+def setup_module():
+    startup()
+    yield
+    cleanup()
 
 ALLELE_TO_SEQUENCE = pandas.read_csv(
     get_path(

test/test_class1_neural_network.py

@@ -4,6 +4,8 @@
 import numpy
 from numpy import testing
 
+import pytest
+
 from nose.tools import eq_, assert_less, assert_greater, assert_almost_equal
 
 import pandas
@@ -13,11 +15,15 @@
 from mhcflurry.common import random_peptides
 
 from mhcflurry.testing_utils import cleanup, startup
-teardown = cleanup
-setup = startup
+
+@pytest.fixture(scope="module")
+def setup_module():
+    startup()
+    yield
+    cleanup()
 
 
-def test_class1_neural_network_a0205_training_accuracy():
+def test_class1_neural_network_a0205_training_accuracy(setup_module):
     # Memorize the dataset.
     hyperparameters = dict(
         activation="tanh",
@@ -72,7 +78,7 @@ def test_class1_neural_network_a0205_training_accuracy():
     eq_(predictor.network().to_json(), predictor2.network().to_json())
 
 
-def test_inequalities():
+def test_inequalities(setup_module):
     # Memorize the dataset.
     hyperparameters = dict(
         peptide_amino_acid_encoding="one-hot",

test/test_class1_pan.py

@@ -6,6 +6,7 @@
 
 from sklearn.metrics import roc_auc_score
 import pandas
+import pytest
 
 from numpy.testing import assert_, assert_equal
 
@@ -14,8 +15,12 @@
 from mhcflurry.downloads import get_path
 
 from mhcflurry.testing_utils import cleanup, startup
-teardown = cleanup
-setup = startup
+
+pytest.fixture(autouse=True, scope="module")
+def setup_module():
+    startup()
+    yield
+    cleanup()
 
 
 HYPERPARAMETERS = {

test/test_class1_presentation_predictor.py

@@ -7,6 +7,7 @@
 
 from numpy.testing import assert_, assert_equal, assert_allclose, assert_array_equal
 from nose.tools import assert_greater, assert_less
+import pytest
 import numpy
 
 from sklearn.metrics import roc_auc_score
@@ -20,44 +21,28 @@
 
 from . import data_path
 
-AFFINITY_PREDICTOR = None
-CLEAVAGE_PREDICTOR = None
-CLEAVAGE_PREDICTOR_NO_FLANKING = None
-PRESENTATION_PREDICTOR = None
 
-
-def setup():
-    global AFFINITY_PREDICTOR
-    global CLEAVAGE_PREDICTOR
-    global CLEAVAGE_PREDICTOR_NO_FLANKING
-    global PRESENTATION_PREDICTOR
+@pytest.fixture(scope="module")
+def predictors():
     startup()
-    AFFINITY_PREDICTOR = Class1AffinityPredictor.load(
-        get_path("models_class1_pan", "models.combined"),
-        optimization_level=0,
-        max_models=1)
-    CLEAVAGE_PREDICTOR = Class1ProcessingPredictor.load(
-        get_path("models_class1_processing", "models.selected.with_flanks"),
-        max_models=1)
-    CLEAVAGE_PREDICTOR_NO_FLANKING = Class1ProcessingPredictor.load(
-        get_path("models_class1_processing", "models.selected.no_flank"),
-        max_models=1)
-    PRESENTATION_PREDICTOR = Class1PresentationPredictor.load()
-
-
-def teardown():
-    global AFFINITY_PREDICTOR
-    global CLEAVAGE_PREDICTOR
-    global CLEAVAGE_PREDICTOR_NO_FLANKING
-    global PRESENTATION_PREDICTOR
-    AFFINITY_PREDICTOR = None
-    CLEAVAGE_PREDICTOR = None
-    CLEAVAGE_PREDICTOR_NO_FLANKING = None
-    PRESENTATION_PREDICTOR = None
+    predictors = {
+        'affinity_predictor': Class1AffinityPredictor.load(
+            get_path("models_class1_pan", "models.combined"),
+            optimization_level=0,
+            max_models=1),
+        'cleavage_predictor': Class1ProcessingPredictor.load(
+            get_path("models_class1_processing", "models.selected.with_flanks"),
+            max_models=1),
+        'cleavage_predictor_no_flanking': Class1ProcessingPredictor.load(
+            get_path("models_class1_processing", "models.selected.no_flank"),
+            max_models=1),
+        'presentation_predictor': Class1PresentationPredictor.load()
+    }
+    yield predictors
     cleanup()
 
 
-def test_basic():
+def test_basic(predictors):
     df = pandas.read_csv(data_path("multiallelic.benchmark.small.csv.bz2"))
     train_df = df.loc[
         df.sample_id.isin(sorted(df.sample_id.unique())[:3])
@@ -71,9 +56,9 @@ def test_basic():
         df.drop_duplicates("sample_id").set_index("sample_id").hla.str.split().to_dict())
 
     predictor = Class1PresentationPredictor(
-        affinity_predictor=AFFINITY_PREDICTOR,
-        processing_predictor_without_flanks=CLEAVAGE_PREDICTOR_NO_FLANKING,
-        processing_predictor_with_flanks=CLEAVAGE_PREDICTOR)
+        affinity_predictor=predictors["affinity_predictor"],
+        processing_predictor_without_flanks=predictors['cleavage_predictor_no_flanking'],
+        processing_predictor_with_flanks=predictors['cleavage_predictor'])
 
     predictor.fit(
         targets=train_df.hit.values,
@@ -159,11 +144,11 @@ def add_prediction_cols(test_df, predictor):
             test_df["prediction2"], other_test_df["prediction2"], decimal=6)
 
 
-def test_downloaded_predictor_small():
-    global PRESENTATION_PREDICTOR
+def test_downloaded_predictor_small(predictors):
+    presentation_predictor = predictors['presentation_predictor']
 
     # Test sequence scanning
-    scan_results = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results = presentation_predictor.predict_sequences(
         sequences=[
             "MESLVPGFN",
             "QPYVFIKRS",
@@ -178,7 +163,7 @@ def test_downloaded_predictor_small():
     print(scan_results)
     assert_equal(len(scan_results), 6)
 
-    scan_results = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results = presentation_predictor.predict_sequences(
         sequences=[
             "MESLVPGFN",
             "QPYVFIKRS",
@@ -193,7 +178,7 @@ def test_downloaded_predictor_small():
     print(scan_results)
     assert_equal(len(scan_results), 6)
 
-    scan_results = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results = presentation_predictor.predict_sequences(
         sequences=[
             "MESLVPGFN",
             "QPYVFIKRS",
@@ -208,7 +193,7 @@ def test_downloaded_predictor_small():
     print(scan_results)
     assert_equal(len(scan_results), 6)
 
-    scan_results = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results = presentation_predictor.predict_sequences(
         sequences=[
             "MESLVPGFN",
             "QPYVFIKRS",
@@ -223,7 +208,7 @@ def test_downloaded_predictor_small():
     print(scan_results)
     assert_equal(len(scan_results), 18)
 
-    scan_results = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results = presentation_predictor.predict_sequences(
         sequences=[
             "MESLVPGFN",
             "QPYVFIKRS",
@@ -239,11 +224,11 @@ def test_downloaded_predictor_small():
     assert_equal(len(scan_results), 0)
 
 
-def test_downloaded_predictor():
-    global PRESENTATION_PREDICTOR
+def test_downloaded_predictor(predictors):
+    presentation_predictor = predictors['presentation_predictor']
 
     # Test sequence scanning
-    scan_results1 = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results1 = presentation_predictor.predict_sequences(
         sequences=[
             "MESLVPGFNEKTHVQLSLPVLQVRDVLVRGFGDSVEEVLSEARQHLKDGTCGLVEVEKGVLPQLE",
             "QPYVFIKRSDARTAPHGHVMVELVAELEGIQYGRSGETLGVLVPHVGEIPVAYRKVLLRKNGNKG",
@@ -263,7 +248,7 @@ def test_downloaded_predictor():
     assert (scan_results1.affinity < 200).all(), str(scan_results1)
     assert (scan_results1.presentation_score > 0.7).all(), str(scan_results1)
 
-    scan_results2 = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results2 = presentation_predictor.predict_sequences(
         result="filtered",
         filter_value=500,
         comparison_quantity="affinity",
@@ -285,7 +270,7 @@ def test_downloaded_predictor():
     assert len(scan_results2) > 10
     assert (scan_results2.affinity <= 500).all()
 
-    scan_results3 = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results3 = presentation_predictor.predict_sequences(
         result="filtered",
         filter_value=0.9,
         comparison_quantity="presentation_score",
@@ -307,7 +292,7 @@ def test_downloaded_predictor():
     assert len(scan_results3) >= 5, len(scan_results3)
     assert (scan_results3.presentation_score >= 0.9).all()
 
-    scan_results4 = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results4 = presentation_predictor.predict_sequences(
         result="all",
         comparison_quantity="affinity",
         sequences={
@@ -337,7 +322,7 @@ def test_downloaded_predictor():
             "AGGHSYGADLKSFDLGDELGTDPYEDFQENWNTKHSSGVTRELMRELNGGAYTRYVDNNFCGPDG",
     }
 
-    scan_results5 = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results5 = presentation_predictor.predict_sequences(
         result="all",
         comparison_quantity="affinity",
         sequences=sequences,
@@ -363,7 +348,7 @@ def test_downloaded_predictor():
     assert_equal(len(scan_results5), len(scan_results4) * 2)
 
     # Test case-insensitive.
-    scan_results6 = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results6 = presentation_predictor.predict_sequences(
         result="all",
         comparison_quantity="affinity",
         sequences=dict((k, v.lower()) for (k, v) in sequences.items()),
@@ -399,7 +384,7 @@ def test_downloaded_predictor():
         scan_results6.presentation_score.values,
         scan_results5.presentation_score.values)
 
-    scan_results7 = PRESENTATION_PREDICTOR.predict_sequences(
+    scan_results7 = presentation_predictor.predict_sequences(
         result="all",
         comparison_quantity="affinity",
         sequences={
@@ -422,8 +407,8 @@ def test_downloaded_predictor():
     assert "DNNFCGPdg" in scan_results7.peptide.values, scan_results7.peptide
 
 
-def test_downloaded_predictor_invalid_peptides():
-    global PRESENTATION_PREDICTOR
+def test_downloaded_predictor_invalid_peptides(predictors):
+    presentation_predictor = predictors['presentation_predictor']
 
     peptides = [
         "SIINFEKL",
@@ -441,11 +426,11 @@ def test_downloaded_predictor_invalid_peptides():
 
     numpy.testing.assert_raises(
         ValueError,
-        PRESENTATION_PREDICTOR.predict,
+        presentation_predictor.predict,
         peptides=peptides,
         alleles=alleles)
 
-    results1 = PRESENTATION_PREDICTOR.predict(
+    results1 = presentation_predictor.predict(
         peptides=peptides,
         alleles=alleles,
         throw=False).presentation_score.values

test/test_class1_processing_neural_network.py

@@ -10,15 +10,20 @@
 from nose.tools import eq_, assert_less, assert_greater, assert_almost_equal
 
 import pandas
+import pytest
 
 from mhcflurry.class1_processing_neural_network import Class1ProcessingNeuralNetwork
 from mhcflurry.common import random_peptides
 from mhcflurry.amino_acid import BLOSUM62_MATRIX
 from mhcflurry.flanking_encoding import FlankingEncoding
 
 from mhcflurry.testing_utils import cleanup, startup
-teardown = cleanup
-setup = startup
+
+@pytest.fixture(autouse=True, scope="module")
+def setup_module():
+    startup()
+    yield
+    cleanup()
 
 table = dict([
     (tuple(encoding), amino_acid)

test/test_custom_loss.py

@@ -5,12 +5,16 @@
 
 import numpy
 import tensorflow as tf
+import pytest
 from mhcflurry.custom_loss import CUSTOM_LOSSES, MultiallelicMassSpecLoss
 
 from mhcflurry.testing_utils import cleanup, startup
-teardown = cleanup
-setup = startup
 
+pytest.fixture(autouse=True, scope="module")
+def setup_module():
+    startup()
+    yield
+    cleanup()
 
 def evaluate_loss(loss, y_true, y_pred):
     y_true = tf.convert_to_tensor(y_true, dtype='float32', name='y_true')