Basic 1D Kalman filter implementation with `x_{t+1} = v_t + x_0` dynamics

kalmanfilt.py

"""
Adapted from work by Ian Glover (https://gist.github.com/manicai/922976)
"""
import numpy as np

__all__ = ["kalmanfilt1d"]


def kalmanfilt1d(x, time_step=0.1, n_init_samples=10, mesurement_noise=30, process_variance=0.01):
    """1-D Kalman filter with x1 = vt + x0 dynamics.

    Parameters:
    x (np.array)            : Input data.
    time_step (float)       : Time step length.
    n_init_samples (int)    : How many samples to include initial state estimation.

    Returns:
    np.array() of size Nx2 estimates for position and velocity [[position, velocity], ...]
    """
    ##########################################################################
    # Model
    # Initial state estimates
    half_n = n_init_samples // 2
    estimate_position = np.median(x[0:n_init_samples])
    delta_position = np.median(
        x[half_n:half_n+n_init_samples]) - estimate_position
    estimate_velocity = (delta_position) / (max(half_n, 1) * time_step)

    # Covariance matrix
    P_xx = 0.1  # Variance of the position
    P_xv = 0.1  # Covariance of position and velocity
    P_vv = 0.1  # Variance of the velocity

    ##########################################################################
    # Model parameters
    position_process_variance = process_variance
    velocity_process_variance = process_variance
    R = mesurement_noise  # Measurement noise variance

    N = x.shape[0]
    output = np.zeros([N, 2])

    for i, measurement in enumerate(x):
        ##################################################################
        # Temporal update (predictive)
        estimate_position += estimate_velocity * time_step

        # Update covariance
        P_xx += time_step * (2.0 * P_xv + time_step * P_vv)
        P_xv += time_step * P_vv

        P_xx += time_step * position_process_variance
        P_vv += time_step * velocity_process_variance

        ##################################################################
        # Observational update (reactive)
        vi = 1.0 / (P_xx + R)
        kx = P_xx * vi
        kv = P_xv * vi

        if np.isnan(measurement):
            measurement = estimate_position

        estimate_position += (measurement - estimate_position) * kx
        estimate_velocity += (measurement - estimate_position) * kv

        P_xx *= (1 - kx)
        P_xv *= (1 - kx)
        P_vv -= kv * P_xv

        output[i, :] = [estimate_position, estimate_velocity]

    return output