Merge pull request #99 from hammerlab/add-speed-script

Add speed test script

.travis.yml

@@ -45,6 +45,6 @@ script:
   # download data and models, then run tests
   - mhcflurry-downloads fetch
   - mhcflurry-downloads info  # just to test this command works
-  - nosetests test --with-coverage --cover-package=mhcflurry  && ./lint.sh
+  - nosetests test -sv --with-coverage --cover-package=mhcflurry  && ./lint.sh
 after_success:
   coveralls

test/test_speed.py

+import numpy
+numpy.random.seed(0)
+import time
+import cProfile
+import pstats
+
+import pandas
+
+from mhcflurry import Class1AffinityPredictor
+from mhcflurry.common import random_peptides
+
+NUM = 10000
+
+DOWNLOADED_PREDICTOR = Class1AffinityPredictor.load()
+
+
+def test_speed(profile=False):
+    starts = {}
+    timings = {}
+    profilers = {}
+
+    def start(name):
+        starts[name] = time.time()
+        if profile:
+            profilers[name] = cProfile.Profile()
+            profilers[name].enable()
+
+    def end(name):
+        timings[name] = time.time() - starts[name]
+        if profile:
+            profilers[name].disable()
+
+    start("first")
+    DOWNLOADED_PREDICTOR.predict(["SIINFEKL"], allele="HLA-A*02:01")
+    end("first")
+
+    peptides = random_peptides(NUM)
+    start("pred_%d" % NUM)
+    DOWNLOADED_PREDICTOR.predict(peptides, allele="HLA-A*02:01")
+    end("pred_%d" % NUM)
+
+    print("SPEED BENCHMARK")
+    print("Results:\n%s" % str(pandas.Series(timings)))
+
+    return dict(
+        (key, pstats.Stats(value)) for (key, value) in profilers.items())
+
+
+if __name__ == '__main__':
+    # If run directly from python, do profiling and leave the user in a shell
+    # to explore results.
+
+    result = test_speed(profile=True)
+    result["pred_%d" % NUM].sort_stats("cumtime").reverse_order().print_stats()
+
+    # Leave in ipython
+    locals().update(result)
+    import ipdb ; ipdb.set_trace()