Katherine preproc touch data

aopy/preproc/bmi3d.py

@@ -161,7 +161,7 @@ def _correct_tracking_task_data(data, metadata, contains_hand=True):
     user_screen = np.dot(user_world, exp_mapping) # intuitive screen coords (x: right/left, y: up/down, z: into/out of the screen)
 
     # incremental perturbations
-    if b'incremental_rotation' in metadata['features']:
+    if 'incremental_rotation' in metadata['features']:
         x_fixed = metadata['init_rotation_x']==metadata['final_rotation_x']
         y_fixed = metadata['init_rotation_y']==metadata['final_rotation_y']
         z_fixed = metadata['init_rotation_z']==metadata['final_rotation_z']
@@ -226,6 +226,42 @@ def _correct_tracking_task_data(data, metadata, contains_hand=True):
     user_bounded = np.array([np.clip(task['user_screen'][:,i], bounds[i*2], bounds[i*2+1]) for i in range(3)]).T
     task['disturbance'] = cursor_bounded - user_bounded # only bmi3d z axis has non-zero values but all axes retain existing NaN values from user_screen
 
+    print('...correcting tracking task frame shift')
+    return task
+
+def _correct_touch_app_data(original_task):
+    '''
+    This function replaces NaN values in cursor data saved by an older version of the tablet touch app (where the cursor would disappear whenever
+    there was no touch).
+    Before the first touch input, cursor values remain NaN but the 'plant_visible' field is updated to False to indicate that the cursor
+    was hidden. After the first touch input, NaN cursor values are replaced with the last valid cursor position and 'plant_visible' is set to False.
+    This function does not affect the 'user_screen' field, which is NaN when there is no touch input.  
+
+    Args:
+        original_task (nt,): original array of task data with specified dtypes
+
+    Returns:
+        task (nt,): corrected array of task data with specified dtypes
+    '''
+    # list of task data fields to keep
+    keys = list(original_task.dtype.names)
+    dtypes = [(key, original_task.dtype.fields[key][0]) for key in keys]
+
+    # construct corrected task data
+    task = np.zeros(len(original_task), dtype=dtypes)
+    for key in keys:
+        task[key] = original_task[key]
+
+    # indicate that cursor was hidden when NaN
+    original_cursor = pd.DataFrame(original_task['cursor'])
+    nan_idx = np.where(np.isnan(original_cursor).any(axis=1))[0]
+    task['plant_visible'][nan_idx] = False
+
+    # replace NaNs in cursor with last valid value
+    corrected_cursor = original_cursor.ffill()
+    task['cursor'] = np.array(corrected_cursor)
+
+    print('...correcting touch app data')
     return task
 
 def parse_bmi3d(data_dir, files):
@@ -565,6 +601,10 @@ def _prepare_bmi3d_v0(data, metadata):
             if 'current_target' in task.dtype.names:
                 # task data from bmi3d has bugs in saved reference and disturbance
                 task = _correct_tracking_task_data(data, metadata, contains_hand=False)
+
+        # special handling for an early version of tablet touch app
+        if 'tablet_touch' in metadata['features']:
+            task = _correct_touch_app_data(task)
 
         data['task'] = task
     else: