small updates for unit testing

Author: janvanrijn

openml/runs/__init__.py

@@ -1,6 +1,6 @@
 from .run import OpenMLRun
 from .trace import OpenMLRunTrace, OpenMLTraceIteration
-from .functions import (run_task, get_run, list_runs, get_runs,
+from .functions import (run_task, get_run, list_runs, get_runs, get_run_trace,
                         initialize_model_from_run, initialize_model_from_trace)
 
 __all__ = ['OpenMLRun', 'run_task', 'get_run', 'list_runs', 'get_runs']

openml/runs/functions.py

@@ -97,6 +97,13 @@ def run_task(task, model, avoid_duplicate_runs=True, flow_tags=None, seed=None):
 
     return run
 
+
+def get_run_trace(run_id):
+    trace_xml = _perform_api_call('run/trace/%d' % run_id)
+    run_trace = _create_trace_from_description(trace_xml)
+    return run_trace
+
+
 def initialize_model_from_run(run_id):
     '''
     Initialized a model based on a run_id (i.e., using the exact
@@ -133,18 +140,18 @@ def initialize_model_from_trace(run_id, repeat, fold, iteration=None):
         The fold nr (column in trace file)
 
     iteration: int
-        The iteration nr (column in trace file)
+        The iteration nr (column in trace file). If None, the
+        best (selected) iteration will be searched (slow)
 
     Returns
     -------
     model : sklearn model
         the scikitlearn model with all parameters initailized
     '''
-    run = get_run(run_id)
-    if 'trace' not in run.output_files:
-        raise PyOpenMLError('Run does not contain trace file')
-    trace_xml = _perform_api_call('run/trace/%d' %run_id)
-    run_trace = _create_trace_from_description(trace_xml)
+    run_trace = get_run_trace(run_id)
+
+    if iteration is None:
+        iteration = run_trace.get_selected_iteration(repeat, fold)
 
     request = (repeat, fold, iteration)
     if request not in run_trace.trace_iterations:

openml/runs/trace.py

@@ -13,6 +13,12 @@ def __init__(self, run_id, trace_iterations):
         self.run_id = run_id
         self.trace_iterations = trace_iterations
 
+    def get_selected_iteration(self, fold, repeat):
+        for (r, f, i) in self.trace_iterations:
+            if r == repeat and f == fold and self.trace_iterations[(r, f, i)].selected == True:
+                return i
+        raise ValueError('Could not find the selected iteration for rep/fold %d/%d' %(repeat,fold))
+
     def __str__(self):
         return '[Run id: %d, %d trace iterations]' %(self.run_id, len(self.trace_iterations))