RosbagToPandas Reindexing

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
from __future__ import division
import numpy as np
import pandas as pd       # pip install pandas
import rosbag_pandas      # pip install rosbag_pandas

# Fast Introduction to pandas:
# http://www.gregreda.com/2013/10/26/intro-to-pandas-data-structures/
#
# tl;dr: pandas extends numpy array with named columns and a time index


# Example usage:
# dataframe = rosbag_pandas.bag_to_dataframe('2014-08-20-17-37-47.bag')
# ridf=reindex_df(dataframe, dataframe.uC_time__data)
# print ridf
# plt.plot(ridf.index, ridf.sensors_imu_data_raw__angular_velocity_z)

def datetime_index_to_sec(dti):
    '''Converts pandas DatetimeIndex into seconds.
    This may only work for DatetimeIndex from rosbag_pandas. Untested on other.'''
    return dti.map(lambda x:x.value/1e9)

def notnull_values(series):
    '''Returns all not null values in the panda Series'''
    return series[series.notnull()].values

def notnull_keys(series):
    '''Returns all not null keys in the panda Series'''
    return series[series.notnull()].keys()

def notnull_keys_as_sec(series):
    '''Returns all not null keys in the panda Series as seconds'''
    return datetime_index_to_sec(notnull_keys(series))

def get_closest_values(desired_timestamps, timestamps, data):
    '''Returns the values in data whose timestamps are closest to the desired_timestamps.
    Assumes sorted desired_timestamps and uniform dt.'''
    assert len(desired_timestamps.shape)==1
    assert len(timestamps.shape)==1
    assert len(data.shape)==1
    assert timestamps.shape[0] == data.shape[0]

    n = timestamps.shape[0]
    dt = np.diff(timestamps).mean()
    start = timestamps[0]

    values = np.empty_like(desired_timestamps)
    idx = (desired_timestamps-start)/dt
    idx = np.round(idx)
    idx[idx<0]=0
    idx[idx>n-1]=n-1
    idx=idx.astype('int')
    values = data[idx]

    return values

def get_closest_values_from_df(reference_series, series):
    '''Returns the values in series whose timestamps are closest to the timestamps in the reference_series.
    Assumes uniform dt in reference_series.'''
    return get_closest_values(
        notnull_keys_sec(reference_series),
        notnull_keys_sec(series),
        notnull_values(series))

def reindex_dataframe(dataframe, reference_series, idx_start_from_zero=True):
    '''Reindex the pandas dataframe to only contain the timestamps from reference_series.
    The closest values (time-wise) for each column are used.
    idx_start_from_zero specifies whether the reindexed timestamps shall start from zero.
    Only columns with float64 values are include in the reindexed dataframe.
    '''
    if type(reference_series) is str:
        reference_series = dataframe[reference_series]

    idx = notnull_keys_sec(reference_series)
    n = idx.shape[0]
    data = dict()
    for col in dataframe.columns:
        col = df.columns[i]
        data[col] = np.empty(shape=(n))
        series = df[col]
        series = series[series.notnull()]
        values = series.values
        keys = datetime_index_to_sec(series.keys())
        values = get_closest_values(idx,keys,values)
        data[col] = values

    if idx_start_from_zero:
        idx -= idx[0]

    reindexed = pd.DataFrame(data,idx)

    return reindexed