parse variable number of keys

faroit · Aug 11, 2016 · 18ebc6b · 18ebc6b
1 parent f4aba9e
commit 18ebc6b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/dsdtools/evaluate.py b/dsdtools/evaluate.py
@@ -25,10 +25,10 @@ def to_pickle(self, filename):
     def import_mat(self, filename):
         mat = scipy.io.loadmat(filename)
         mdata = mat['result']
-        ndata = {'Dev': mdata['dev'][0, 0], 'Test': mdata['test'][0, 0]}
+        ndata = {n.title(): mdata[n][0, 0] for n in mdata.dtype.names}
         s = []
-        for subset in ['Dev', 'Test']:
-            data = ndata[subset]['results']
+        for subset, subset_data in ndata.items():
+            data = subset_data['results']
             for track in range(data.shape[1]):
                 tdata = data[0, track][0, 0]
                 for target in [