# Path: imusim/maths/kalman.py
# class UnscentedKalmanFilter(KalmanFilter):
#     """
#     Implementation of the unscented Kalman filter algorithm for additive noise.
# 
#     Implementation equations are taken from Wan and Van Der Merwe, "The
#     Unscented Kalman Filter", table 7.3.2.
# 
#     @ivar stateUpdateFunction: Non-linear function to propagate state into the
#         next timestep. The function is called with two arguments, the current
#         Nx1 state vector and the current Mx1 control input vector. It should
#         return a new Nx1 state vector. N and M are the numbers of states and
#         control inputs.
#     @ivar measurementFunction: Non-linear function relating state variables
#         to expected measurements. The function is called with the current Nx1
#         state vector as an argument, and should return an Mx1 measurement
#         vector. N and M are the numbers of states and measurement outputs.
#     @ivar state: Current state estimate.
#     @ivar stateCovariance: State covariance matrix.
#     @ivar processCovariance: Process noise covariance matrix.
#     @ivar measurementCovariance: Measurement noise covariance matrix.
#     """
# 
#     def __init__(self, stateUpdateFunction, measurementFunction, state,
#             stateCovariance, processCovariance, measurementCovariance):
#         """
#         Construct unscented Kalman filter.
# 
#         @param stateUpdateFunction: Non-linear function to propagate state into
#             the next timestep. The function is called with two arguments, the
#             current Nx1 state vector and the current Mx1 control input vector.
#             It should return a new Nx1 state vector. N and M are the numbers of
#             states and control inputs.
#         @param measurementFunction: Non-linear function relating state
#             variables to expected measurements. The function is called with the
#             current Nx1 state vector as an argument, and should return an Mx1
#             measurement vector. N and M are the numbers of states and measurement
#             outputs.
#         @param state: Initial state estimate.
#         @param stateCovariance: Initial state covariance matrix.
#         @param processCovariance: Process noise covariance matrix.
#         @param measurementCovariance: Measurement noise covariance matrix.
#         """
# 
#         self.stateUpdateFunction = stateUpdateFunction
#         self.measurementFunction = measurementFunction
# 
#         self.state = state
#         self.stateCovariance = stateCovariance
#         self.processCovariance = processCovariance
#         self.measurementCovariance = measurementCovariance
# 
#     @property
#     def stateUpdateFunction(self):
#         return self._stateUpdateUT._function
# 
#     @stateUpdateFunction.setter
#     def stateUpdateFunction(self, function):
#         self._stateUpdateUT = UnscentedTransform(function)
# 
#     @property
#     def measurementFunction(self):
#         return self._measurementUT._function
# 
#     @measurementFunction.setter
#     def measurementFunction(self, function):
#         self._measurementUT = UnscentedTransform(function)
# 
#     def predict(self, control=None):
#         state, stateCovariance = self._stateUpdateUT(self._x, self._P, *[control])
#         self._x = state
#         self._P = stateCovariance + self._Q
#         assert self._x.shape == (self._states,1)
#         assert np.all(np.isfinite(self._x))
#         assert self._P.shape == (self._states, self._states)
#         assert np.all(np.isfinite(self._P))
# 
#     def innovation(self, measurement):
#         measurement = np.asanyarray(measurement).reshape(-1,1)
#         predictedMeasurement, predictionCovariance, \
#             stateSigmas, measurementSigmas, weights = \
#                 self._measurementUT(self._x, self._P, returnSigmas=True)
#         x,y = self._x, predictedMeasurement
#         sigmaPoints = izip(weights, stateSigmas, measurementSigmas)
#         innovation = measurement - predictedMeasurement
#         innovationCovariance = self._R + predictionCovariance
#         crossCovariance = np.sum((w*(X-x)*(Y-y).T for (w,X,Y) in sigmaPoints), axis=0)
#         return innovation, innovationCovariance, crossCovariance
# 
#     def correct(self, measurement):
#         measurement = np.asanyarray(measurement).reshape(-1,1)
#         innovation, innovationCovariance, crossCovariance = \
#                 self.innovation(measurement)
#         self._K = np.dot(crossCovariance, inv(innovationCovariance))
#         self._x = self._x + self._K * innovation
#         assert self._x.shape == (self._states,1)
#         assert np.all(np.isfinite(self._x))
#         self._P = self._P - self._K * innovationCovariance * self._K.T
#         assert self._P.shape == (self._states, self._states)
#         assert np.all(np.isfinite(self._P))


import numpy as np
from imusim.maths.kalman import UnscentedKalmanFilter
        estimateCovariances[i] = ukf.stateCovariance

    # Allow ten samples for estimates to start settling
    errors = trueStates[:,10:]-estimates[:,10:]
    assert np.mean(errors[0]**2) < 0.01
    assert np.mean(errors[1]**2) < 0.01
    assert np.var(errors[0]) < measurementCovariance[0,0]
    assert np.var(errors[1]) < measurementCovariance[1,1]

def testNonLinearStateUpdate():
    """
    Test filter with a non-linear state update
    """

    SAMPLES = 1000

    def stateUpdate(state, control):
        state = np.asarray(state)
        newState = np.empty_like(state)
        newState[0] = state[0] + 0.001*state[1]**2
        newState[1] = state[1] + 0.01
        return newState

    measurementFunc = lambda state: state

    initialState = np.matrix('0;0')
    initialCovariance = np.matrix('1,0.5;0.5,1')
    processCovariance = np.matrix('0.01,0.01;0.01,0.01')
    measurementCovariance = np.matrix('0.25,0;0,0.25')


# Path: imusim/maths/kalman.py
# class UnscentedKalmanFilter(KalmanFilter):
#     """
#     Implementation of the unscented Kalman filter algorithm for additive noise.
# 
#     Implementation equations are taken from Wan and Van Der Merwe, "The
#     Unscented Kalman Filter", table 7.3.2.
# 
#     @ivar stateUpdateFunction: Non-linear function to propagate state into the
#         next timestep. The function is called with two arguments, the current
#         Nx1 state vector and the current Mx1 control input vector. It should
#         return a new Nx1 state vector. N and M are the numbers of states and
#         control inputs.
#     @ivar measurementFunction: Non-linear function relating state variables
#         to expected measurements. The function is called with the current Nx1
#         state vector as an argument, and should return an Mx1 measurement
#         vector. N and M are the numbers of states and measurement outputs.
#     @ivar state: Current state estimate.
#     @ivar stateCovariance: State covariance matrix.
#     @ivar processCovariance: Process noise covariance matrix.
#     @ivar measurementCovariance: Measurement noise covariance matrix.
#     """
# 
#     def __init__(self, stateUpdateFunction, measurementFunction, state,
#             stateCovariance, processCovariance, measurementCovariance):
#         """
#         Construct unscented Kalman filter.
# 
#         @param stateUpdateFunction: Non-linear function to propagate state into
#             the next timestep. The function is called with two arguments, the
#             current Nx1 state vector and the current Mx1 control input vector.
#             It should return a new Nx1 state vector. N and M are the numbers of
#             states and control inputs.
#         @param measurementFunction: Non-linear function relating state
#             variables to expected measurements. The function is called with the
#             current Nx1 state vector as an argument, and should return an Mx1
#             measurement vector. N and M are the numbers of states and measurement
#             outputs.
#         @param state: Initial state estimate.
#         @param stateCovariance: Initial state covariance matrix.
#         @param processCovariance: Process noise covariance matrix.
#         @param measurementCovariance: Measurement noise covariance matrix.
#         """
# 
#         self.stateUpdateFunction = stateUpdateFunction
#         self.measurementFunction = measurementFunction
# 
#         self.state = state
#         self.stateCovariance = stateCovariance
#         self.processCovariance = processCovariance
#         self.measurementCovariance = measurementCovariance
# 
#     @property
#     def stateUpdateFunction(self):
#         return self._stateUpdateUT._function
# 
#     @stateUpdateFunction.setter
#     def stateUpdateFunction(self, function):
#         self._stateUpdateUT = UnscentedTransform(function)
# 
#     @property
#     def measurementFunction(self):
#         return self._measurementUT._function
# 
#     @measurementFunction.setter
#     def measurementFunction(self, function):
#         self._measurementUT = UnscentedTransform(function)
# 
#     def predict(self, control=None):
#         state, stateCovariance = self._stateUpdateUT(self._x, self._P, *[control])
#         self._x = state
#         self._P = stateCovariance + self._Q
#         assert self._x.shape == (self._states,1)
#         assert np.all(np.isfinite(self._x))
#         assert self._P.shape == (self._states, self._states)
#         assert np.all(np.isfinite(self._P))
# 
#     def innovation(self, measurement):
#         measurement = np.asanyarray(measurement).reshape(-1,1)
#         predictedMeasurement, predictionCovariance, \
#             stateSigmas, measurementSigmas, weights = \
#                 self._measurementUT(self._x, self._P, returnSigmas=True)
#         x,y = self._x, predictedMeasurement
#         sigmaPoints = izip(weights, stateSigmas, measurementSigmas)
#         innovation = measurement - predictedMeasurement
#         innovationCovariance = self._R + predictionCovariance
#         crossCovariance = np.sum((w*(X-x)*(Y-y).T for (w,X,Y) in sigmaPoints), axis=0)
#         return innovation, innovationCovariance, crossCovariance
# 
#     def correct(self, measurement):
#         measurement = np.asanyarray(measurement).reshape(-1,1)
#         innovation, innovationCovariance, crossCovariance = \
#                 self.innovation(measurement)
#         self._K = np.dot(crossCovariance, inv(innovationCovariance))
#         self._x = self._x + self._K * innovation
#         assert self._x.shape == (self._states,1)
#         assert np.all(np.isfinite(self._x))
#         self._P = self._P - self._K * innovationCovariance * self._K.T
#         assert self._P.shape == (self._states, self._states)
#         assert np.all(np.isfinite(self._P))


# Path: imusim/tests/maths/ukf_test.py
import numpy as np
from imusim.maths.kalman import UnscentedKalmanFilter

        estimateCovariances[i] = ukf.stateCovariance

    # Allow ten samples for estimates to start settling
    errors = trueStates[:,10:]-estimates[:,10:]
    assert np.mean(errors[0]**2) < 0.01
    assert np.mean(errors[1]**2) < 0.01
    assert np.var(errors[0]) < measurementCovariance[0,0]
    assert np.var(errors[1]) < measurementCovariance[1,1]

def testNonLinearStateUpdate():
    """
    Test filter with a non-linear state update
    """

    SAMPLES = 1000

    def stateUpdate(state, control):
        state = np.asarray(state)
        newState = np.empty_like(state)
        newState[0] = state[0] + 0.001*state[1]**2
        newState[1] = state[1] + 0.01
        return newState

    measurementFunc = lambda state: state

    initialState = np.matrix('0;0')
    initialCovariance = np.matrix('1,0.5;0.5,1')
    processCovariance = np.matrix('0.01,0.01;0.01,0.01')
    measurementCovariance = np.matrix('0.25,0;0,0.25')


    ukf = UnscentedKalmanFilter(stateUpdate, measurementFunc, initialState,