Revised IMM algorithm to use kf book's terminology

The variable names were borrowed from several sources, and not
consistent with my book.

filterpy/kalman/imm.py

@@ -46,13 +46,13 @@ def __init__(self, filters, mu, M):
             List of N filters. filters[i] is the ith Kalman filter in the
             IMM estimator.
 
-        mu : (N,N) ndarray of float
+        mu : (N,) ndarray of float
             mode probability: mu[i] is the probability that
             filter i is the correct one.
 
         M : (N,N) ndarray of float
             Markov chain transition matrix. M[i,j] is the probability of
-            switching from filter i to filter j.
+            switching from filter j to filter i.
         """
 
         assert len(filters) > 1
@@ -61,20 +61,20 @@ def __init__(self, filters, mu, M):
         self.mu = mu
         self.M = M
 
+        # compute # random variables in the state
         x_shape = filters[0].x.shape
         try:
-            self.N = x_shape[0]
+            n_states = x_shape[0]
         except:
-            self.N = x_shape
+            n_states = x_shape
 
         self.x = np.zeros(x_shape)
-        self.P = np.zeros((self.N, self.N))
+        self.P = np.zeros((n_states, n_states))
 
-        self.cbar = dot(self.M.T, self.mu)
-        self.n = len(filters)
+        self.N = len(filters) # number of filters
 
 
-    def update(self, z):
+    def update(self, z, u=None):
         """
         Add a new measurement (z) to the Kalman filter. If z is None, nothing
         is changed.
@@ -84,41 +84,33 @@ def update(self, z):
 
         z : np.array
             measurement for this update.
+
+        u : np.array, optional
+            u[i] contains the control input for the ith filter
         """
 
+        # run update on each filter, and save the likelihood in L
         L = zeros(len(self.filters))
         for i, f in enumerate(self.filters):
             f.update(z)
-            L[i] = f.likelihood    # prior * likelihood
-
-        # compute mode probabilities for this step
-        self.mu =  self.cbar * L
-        self.mu /= sum(self.mu) # normalize
-
-
-        # compute mixed IMM state and covariance
-        self.x.fill(0.)
-        self.P.fill(0.)
-
-        for f, w in zip(self.filters, self.mu):
-            self.x += f.x*w
-
-        for f, w in zip(self.filters, self.mu):
-            y = f.x - self.x
-            self.P += w*(np.outer(y, y) + f.P)
-
+            L[i] = f.likelihood
 
         # initial condition IMM state, covariance
         xs, Ps = [], []
-        self.cbar = dot(self.M.T, self.mu)
-
-        omega = np.zeros((self.n, self.n))
-        for i in range(self.n):
-            omega[i, 0] = self.M[i, 0]*self.mu[i] / self.cbar[0]
-            omega[i, 1] = self.M[i, 1]*self.mu[i] / self.cbar[1]
-
-
-        # compute initial states
+        # each element j = sum M_ij * mu_i
+
+        # cbar is the total probability, after interaction, 
+        # that the target is in state j. We use it as the
+        # normalization constant.
+        self.cbar = dot(self.mu, self.M) 
+
+        # compute mixing probabilities
+        omega = np.zeros((self.N, self.N))
+        for i in range(self.N):
+            for j in range(self.N):
+                omega[i, j] = (self.M[i, j] * self.mu[i]) / self.cbar[j]
+
+        # compute mixed initial conditions
         for i, (f, w) in enumerate(zip(self.filters, omega.T)):
             x = np.zeros(self.x.shape)
             for kf, wj in zip(self.filters, w):
@@ -128,11 +120,29 @@ def update(self, z):
             P = np.zeros(self.P.shape)
             for kf, wj in zip(self.filters, w):
                 y = kf.x - x
-                P += wj*(np.outer(y, y) + kf.P)
+                P += wj * (np.outer(y, y) + kf.P)
             Ps.append(P)
 
-        for i in range(len(xs)):
-            f = self.filters[i]
+        # perform predict step using the mixed initial conditions
+        for i, f in enumerate(self.filters):
             # propagate using the mixed state estimate and covariance
             f.x = dot(f.F, xs[i])
+            if u is not None:
+                f.x += dot(f.B, u[i])
             f.P = dot3(f.F, Ps[i], f.F.T) + f.Q
+
+
+        # compute mixed IMM state and covariance
+        self.x.fill(0.)
+        self.P.fill(0.)
+
+        for f, w in zip(self.filters, self.mu):
+            self.x += f.x * w
+
+        for f, w in zip(self.filters, self.mu):
+            y = f.x - self.x
+            self.P += w * (np.outer(y, y) + f.P)
+
+        # update mode probabilities from total probability * likelihood
+        self.mu =  self.cbar * L
+        self.mu /= sum(self.mu) # normalize

filterpy/kalman/tests/test_imm.py

@@ -193,13 +193,14 @@ def test_imm():
         pass
     ca = KalmanFilter(6, 2)
     cano = KalmanFilter(6, 2)
+    dt2 = (dt**2)/2
     ca.F = np.array(
-        [[1, dt, dt**2/2, 0, 0, 0],
-         [0, 1, dt, 0, 0, 0],
-         [0, 0, 1, 0, 0, 0],
-         [0, 0, 0, 1, dt, dt**2/2],
-         [0, 0, 0, 0, 1, dt],
-         [0, 0, 0, 0, 0, 1]])
+        [[1, dt, dt2, 0, 0,  0],
+         [0, 1,  dt,  0, 0,  0],
+         [0, 0,   1,  0, 0,  0],
+         [0, 0,   0,  1, dt, dt2],
+         [0, 0,   0,  0,  1, dt],
+         [0, 0,   0,  0,  0,  1]])
     cano.F = ca.F.copy()
 
     ca.x = np.array([[2000., 0, 0, 10000, -15, 0]]).T
@@ -210,9 +211,9 @@ def test_imm():
     ca.R *= r**2
     cano.R *= r**2
     cano.Q *= 0
-    q = np.array([[.05, .125, .16666666666666666666666666667],
-         [.125, .333333333333333333333333333333333333, .5],
-         [.166666666666666666666666667, .5, 1]])*1.e-3
+    q = np.array([[.05, .125, 1/6],
+         [.125, 1/3, .5],
+         [1/6, .5, 1]])*1.e-3
 
     ca.Q[0:3, 0:3] = q
     ca.Q[3:6, 3:6] = q
@@ -223,17 +224,14 @@ def test_imm():
 
     filters = [ca, cano]
 
-    trans = np.array([[0.8, 0.2],
-                      [0.05, 0.95]])
-
     trans = np.array([[0.97, 0.03],
                       [0.03, 0.97]])
 
     bank = IMMEstimator(filters, (0.5, 0.5), trans)
 
     xs, probs = [], []
     cvxs, caxs = [], []
-    for i, z in enumerate(zs):
+    for i, z in enumerate(zs[0:10]):
         z = np.array([z]).T
         bank.update(z)
         #print(ca.likelihood, cano.likelihood)