Numba implementation of IMU orientation estimation with complementary filter

complementary_filter_numba.py

import math
import numpy as np
from numba import guvectorize
from scipy.signal import filtfilt, butter
from quaternion import \
    quaternion, from_rotation_vector, as_float_array, as_quat_array


# --- Some utility functions --- #
def multiply(q, r):
    return (
        r[0] * q[0] - r[1] * q[1] - r[2] * q[2] - r[3] * q[3],
        r[0] * q[1] + r[1] * q[0] - r[2] * q[3] + r[3] * q[2],
        r[0] * q[2] + r[1] * q[3] + r[2] * q[0] - r[3] * q[1],
        r[0] * q[3] - r[1] * q[2] + r[2] * q[1] + r[3] * q[0])


def conjugate(q):
    return q[0], -q[1], -q[2], -q[3]


def times(a, v):
    return a*v[0], a*v[1], a*v[2]


def dot(v, w):
    return v[0]*w[0] + v[1]*w[1] + v[2]*w[2]


def cross(v, w):
    return v[1]*w[2] - v[2]*w[1], v[2]*w[0] - v[0]*w[2], v[0]*w[1] - v[1]*w[0]


def normalize(v):
    norm = math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    return v[0]/norm, v[1]/norm, v[2]/norm


def rotate_vector(q, v):
    return multiply(multiply(q, np.hstack((0., v))), conjugate(q))[1:]


def from_axis_angle(v):
    angle = math.sqrt(v[0] ** 2 + v[1] ** 2 + v[2] ** 2)
    s = math.sin(angle/2)
    return math.cos(angle/2), v[0]*s/angle, v[1]*s/angle, v[2]*s/angle


# --- Numba implementation of complementary filter --- #
@guvectorize(
    ['(float64[:,:],float64[:,:],float64[:],float64,float64[:,:])'],
    '(n,p),(n,q),(p),()->(n,q)')
def rot(a, q_delta, g0, alpha, q):

    q[0] = [1., 0., 0., 0.]

    for i in range(1, a.shape[0]):

        # get rotation estimate from gyroscope
        q_w = multiply(q[i - 1], q_delta[i - 1])

        # get rotation estimate from accelerometer
        v_w = rotate_vector(q_w, a[i])
        n = normalize(cross(v_w, g0))
        phi = math.acos(dot(normalize(v_w), g0))
        q_a = from_axis_angle(times((1 - alpha) * phi, n))

        # fuse both estimates
        q[i] = multiply(q_a, q_w)


def estimate_orientation(a, w, t, alpha=0.9, g_ref=(0., 0., 1.),
                         theta_min=1e-6, highpass=.01, lowpass=.05):
    """ Estimate orientation with a complementary filter.

    Fuse linear acceleration and angular velocity measurements to obtain an
    estimate of orientation using a complementary filter as described in
    `Wetzstein 2017: 3-DOF Orientation Tracking with IMUs`_

    .. _Wetzstein 2017: 3-DOF Orientation Tracking with IMUs:
    https://pdfs.semanticscholar.org/5568/e2100cab0b573599accd2c77debd05ccf3b1.pdf

    Parameters
    ----------
    a : array-like, shape (N, 3)
        Acceleration measurements (in arbitrary units).

    w : array-like, shape (N, 3)
        Angular velocity measurements (in rad/s).

    t : array-like, shape (N,)
        Timestamps of the measurements (in s).

    alpha : float, default 0.9
        Weight of the angular velocity measurements in the estimate.

    g_ref : tuple, len 3, default (0., 0., 1.)
        Unit vector denoting direction of gravity.

    theta_min : float, default 1e-6
        Minimal angular velocity after filtering. Values smaller than this
        will be considered noise and are not used for the estimate.

    highpass : float, default .01
        Cutoff frequency of the high-pass filter for the angular velocity as
        fraction of Nyquist frequency.

    lowpass : float, default .05
        Cutoff frequency of the low-pass filter for the linear acceleration as
        fraction of Nyquist frequency.

    Returns
    -------
    q : array of quaternions, shape (N,)
        The estimated orientation for each measurement.
    """

    # initialize some things
    N = len(t)
    dt = np.diff(t)
    g_ref = np.array(g_ref)

    # get high-passed angular velocity
    w = filtfilt(*butter(5, highpass, btype='high'), w, axis=0)
    w[np.linalg.norm(w, axis=1) < theta_min] = 0
    q_delta = np.hstack((from_rotation_vector(w[1:] * dt[:, None]),
                         np.zeros(1, dtype=quaternion)))

    # get low-passed linear acceleration
    a = filtfilt(*butter(5, lowpass, btype='low'), a, axis=0)

    # run filter
    q = as_quat_array(rot(a, as_float_array(q_delta), g_ref, alpha))

    return q