now make_multiindex lets you not to add the header attribs to the mul…

…tiindex

notebooks/user_data_loader.py

@@ -32,27 +32,57 @@
 ind_pos_z = [i for i in index if '_pos_z' in i]
 
 
-def _prepare_multiindex(joints, attribs):
+# def _prepare_multiindex(joints, attribs):
+#     '''Helper function that returns all attribs in two iterables.
+#        1st iterable will be used as the 1st level of the index.
+#        2nd iterable will be used as the 2nd level of the index.'''
+#     all_columns = it.izip(*it.product(joints, attribs))
+#     first_level = \
+#         it.chain(['header'] * len(header), all_columns.next(), ['pose', ])
+#     second_level = it.chain(header, all_columns.next(), ['pose', ])
+#     return first_level, second_level
+
+def _prepare_multiindex(joints, attribs, header=True):
     '''Helper function that returns all attribs in two iterables.
        1st iterable will be used as the 1st level of the index.
        2nd iterable will be used as the 2nd level of the index.'''
-    all_columns = it.izip(*it.product(joints, attribs))
+
+    combined = _combine_joints_attribs_for_multiindex(joints, attribs)
+    first_level = combined.next()
+    second_level = combined.next()
+    if header:
+        first_level, second_level = _add_header_to_multindex(first_level,
+                                                             second_level)
+    return first_level, second_level
+
+
+def _combine_joints_attribs_for_multiindex(joints, attribs):
+    ''' produces an iterator with two items. Each one is one level for
+        pandas.MultiIndex'''
+    return it.izip(*it.product(joints, attribs))
+
+
+def _add_header_to_multindex(level_1, level_2):
+    ''' Adds headder attributes to the output of
+        L{_combine_joints_attribs_for_multiindex(...)}'''
     first_level = \
-        it.chain(['header'] * len(header), all_columns.next(), ['pose', ])
-    second_level = it.chain(header, all_columns.next(), ['pose', ])
+        it.chain(['header'] * len(header), level_1, ['pose', ])
+    second_level = it.chain(header, level_2, ['pose', ])
     return first_level, second_level
 
 
-def make_multiindex(joints, attribs):
+def make_multiindex(joints, attribs, header=True):
     ''' Returns a pandas.MultiIndex from the entered attributes
         @param joints: first level of multiindex
         @type joints: list of strings
         @param joints: second level of multiindex
         @type joints: list of strings
         @return: the Multiindex itself
         @rtype: pandas.MultiIndex
+        @keyword header: wether to add the header cols to the MultiIndex or not
+        @type header: bool
     '''
-    multiind1, multiind2 = _prepare_multiindex(joints, attribs)
+    multiind1, multiind2 = _prepare_multiindex(joints, attribs, header)
     return pd.MultiIndex.from_arrays([list(multiind1), list(multiind2)],
                                      names=['joint', 'attrib'])
 
@@ -83,7 +113,7 @@ def load_user_file(file_):
 
 def clean_stand_prefix(x):
     '''Removes 'STAND_' prefix of a string'''
-    x.lstrip('STAND_')
+    return x.lstrip('STAND_')
 
 
 def prepare_df(df):