ekf2: migrate covariance prediction to SymForce

src/modules/ekf2/EKF/python/ekf_derivation/derivation.py

@@ -69,6 +69,53 @@ class VState(sf.Matrix):
 class MState(sf.Matrix):
     SHAPE = (State.n_states, State.n_states)
 
+def predict_covariance(
+        state: VState,
+        P: MState,
+        d_vel: sf.V3,
+        d_vel_var: sf.V3,
+        d_ang: sf.V3,
+        d_ang_var: sf.Scalar,
+        dt: sf.Scalar
+):
+    g = sf.Symbol("g") # does not appear in the jacobians
+
+    d_vel_b = sf.V3(state[State.d_vel_bx], state[State.d_vel_by], state[State.d_vel_bz])
+    d_vel_true = d_vel - d_vel_b
+
+    d_ang_b = sf.V3(state[State.d_ang_bx], state[State.d_ang_by], state[State.d_ang_bz])
+    d_ang_true = d_ang - d_ang_b
+
+    q = sf.V4(state[State.qw], state[State.qx], state[State.qy], state[State.qz])
+    R_to_earth = quat_to_rot_simplified(q)
+    v = sf.V3(state[State.vx], state[State.vy], state[State.vz])
+    p = sf.V3(state[State.px], state[State.py], state[State.pz])
+
+    q_new = quat_mult(q, sf.V4(1, 0.5 * d_ang_true[0],  0.5 * d_ang_true[1],  0.5 * d_ang_true[2]))
+    v_new = v + R_to_earth * d_vel_true + sf.V3(0 ,0 ,g) * dt
+    p_new = p + v * dt
+
+    # Predicted state vector at time t + dt
+    state_new = VState.block_matrix([[q_new], [v_new], [p_new], [sf.Matrix(state[State.d_ang_bx:State.n_states])]])
+
+    # State propagation jacobian
+    A = state_new.jacobian(state)
+    G = state_new.jacobian(sf.V6.block_matrix([[d_vel], [d_ang]]))
+
+    # Covariance propagation
+    var_u = sf.Matrix.diag([d_vel_var[0], d_vel_var[1], d_vel_var[2], d_ang_var, d_ang_var, d_ang_var])
+    P_new = A * P * A.T + G * var_u * G.T
+
+    # Generate the equations for the lower triangular matrix and the diagonal only
+    # Since the matrix is symmetric, the upper triangle does not need to be derived
+    # and can simply be copied in the implementation
+    for index in range(State.n_states):
+        for j in range(State.n_states):
+            if index > j:
+                P_new[index,j] = 0
+
+    return P_new
+
 def compute_airspeed_innov_and_innov_var(
         state: VState,
         P: MState,
@@ -148,8 +195,10 @@ def compute_sideslip_h_and_k(
 
     return (H.T, K)
 
+print("Derive EKF2 equations...")
 generate_px4_function(compute_airspeed_innov_and_innov_var, output_names=["innov", "innov_var"])
 generate_px4_function(compute_airspeed_h_and_k, output_names=["H", "K"])
 
 generate_px4_function(compute_sideslip_innov_and_innov_var, output_names=["innov", "innov_var"])
 generate_px4_function(compute_sideslip_h_and_k, output_names=["H", "K"])
+generate_px4_function(predict_covariance, output_names=["P_new"])

src/modules/ekf2/EKF/python/ekf_derivation/derivation_utils.py

@@ -34,9 +34,7 @@ Description:
     Common functions used for the derivation of most estimators
 """
 
-from symforce import symbolic as sm
-from symforce import geo
-from symforce import typing as T
+import symforce.symbolic as sf
 
 # q: quaternion describing rotation from frame 1 to frame 2
 # returns a rotation matrix derived form q which describes the same
@@ -47,17 +45,40 @@ def quat_to_rot(q):
     q2 = q[2]
     q3 = q[3]
 
-    Rot = geo.M33([[q0**2 + q1**2 - q2**2 - q3**2, 2*(q1*q2 - q0*q3), 2*(q1*q3 + q0*q2)],
+    Rot = sf.M33([[q0**2 + q1**2 - q2**2 - q3**2, 2*(q1*q2 - q0*q3), 2*(q1*q3 + q0*q2)],
                   [2*(q1*q2 + q0*q3), q0**2 - q1**2 + q2**2 - q3**2, 2*(q2*q3 - q0*q1)],
                    [2*(q1*q3-q0*q2), 2*(q2*q3 + q0*q1), q0**2 - q1**2 - q2**2 + q3**2]])
 
     return Rot
 
-def sign_no_zero(x) -> T.Scalar:
+def quat_to_rot_simplified(q):
+    q0 = q[0]
+    q1 = q[1]
+    q2 = q[2]
+    q3 = q[3]
+
+    # Use the simplified formula for unit quaternion to rotation matrix
+    # as it produces a simpler and more stable EKF derivation given
+    # the additional constraint: q0^2 + q1^2 + q2^2 + q3^2 = 1
+    Rot = sf.Matrix([[1 - 2*q2**2 - 2*q3**2, 2*(q1*q2 - q0*q3), 2*(q1*q3 + q0*q2)],
+                  [2*(q1*q2 + q0*q3), 1 - 2*q1**2 - 2*q3**2, 2*(q2*q3 - q0*q1)],
+                   [2*(q1*q3-q0*q2), 2*(q2*q3 + q0*q1), 1 - 2*q1**2 - 2*q2**2]])
+
+    return Rot
+
+def quat_mult(p,q):
+    r = sf.Matrix([p[0] * q[0] - p[1] * q[1] - p[2] * q[2] - p[3] * q[3],
+                p[0] * q[1] + p[1] * q[0] + p[2] * q[3] - p[3] * q[2],
+                p[0] * q[2] - p[1] * q[3] + p[2] * q[0] + p[3] * q[1],
+                p[0] * q[3] + p[1] * q[2] - p[2] * q[1] + p[3] * q[0]])
+
+    return r
+
+def sign_no_zero(x) -> sf.Scalar:
     """
     Returns -1 if x is negative, 1 if x is positive, and 1 if x is zero
     """
-    return 2 * sm.Min(sm.sign(x), 0) + 1
+    return 2 * sf.Min(sf.sign(x), 0) + 1
 
 def add_epsilon_sign(expr, var, eps):
     # Avoids a singularity at 0 while keeping the derivative correct
@@ -76,7 +97,6 @@ def generate_px4_function(function_name, output_names):
             output_dir="generated",
             skip_directory_nesting=True)
 
-    print("Files generated in {}:\n".format(metadata.output_dir))
     for f in metadata.generated_files:
         print("  |- {}".format(os.path.relpath(f, metadata.output_dir)))