text: import matplotlib.pyplot as plt
  import numpy as np
  np.set_printoptions(precision=2)
  import scipy.linalg
  from termcolor import cprint
  from utils import *
  
  
  class NUMPYIEKF:
      Id2 = np.eye(2)
      Id3 = np.eye(3)
      Id6 = np.eye(6)
      IdP = np.eye(21)
  
      def __init__(self, parameter_class=None):
  
          # variables to initialize with `filter_parameters`
          self.g = None
          self.cov_omega = None
          self.cov_acc = None
          self.cov_b_omega = None
          self.cov_b_acc = None
          self.cov_Rot_c_i = None
          self.cov_t_c_i = None
          self.cov_lat = None
          self.cov_up = None
          self.cov_b_omega0 = None
          self.cov_b_acc0 = None
          self.cov_Rot0 = None
          self.cov_v0 = None
          self.cov_Rot_c_i0 = None
          self.cov_t_c_i0 = None
          self.Q = None
          self.Q_dim = None
          self.n_normalize_rot = None
          self.n_normalize_rot_c_i = None
          self.P_dim = None
          self.verbose = None
  
          # set the parameters
          if parameter_class is None:
              filter_parameters = NUMPYIEKF.Parameters()
          else:
              filter_parameters = parameter_class()
          self.filter_parameters = filter_parameters
          self.set_param_attr()
  
      class Parameters:
          g = np.array([0, 0, -9.80665])
          """gravity vector"""
  
          P_dim = 21
          """covariance dimension"""
  
          Q_dim = 18
          """process noise covariance dimension"""
  
          # Process noise covariance
          cov_omega = 1e-3
          """gyro covariance"""
          cov_acc = 1e-2
          """accelerometer covariance"""
          cov_b_omega = 6e-9
          """gyro bias covariance"""
          cov_b_acc = 2e-4
          """accelerometer bias covariance"""
          cov_Rot_c_i = 1e-9
          """car to IMU orientation covariance"""
          cov_t_c_i = 1e-9
          """car to IMU translation covariance"""
  
          cov_lat = 0.2
          """Zero lateral velocity covariance"""
          cov_up = 300
          """Zero lateral velocity covariance"""
  
          cov_Rot0 = 1e-3
          """initial pitch and roll covariance"""
          cov_b_omega0 = 6e-3
          """initial gyro bias covariance"""
          cov_b_acc0 = 4e-3
          """initial accelerometer bias covariance"""
          cov_v0 = 1e-1
          """initial velocity covariance"""
          cov_Rot_c_i0 = 1e-6
          """initial car to IMU pitch and roll covariance"""
          cov_t_c_i0 = 5e-3
          """initial car to IMU translation covariance"""
  
          # numerical parameters
          n_normalize_rot = 100
          """timestamp before normalizing orientation"""
          n_normalize_rot_c_i = 1000
          """timestamp before normalizing car to IMU orientation"""
  
          def __init__(self, **kwargs):
              self.set(**kwargs)
  
          def set(self, **kwargs):
              for key, value in kwargs.items():
                  setattr(self, key, value)
  
      def set_param_attr(self):
  
          # get a list of attribute only
          attr_list = [a for a in dir(self.filter_parameters) if not a.startswith('__')
                       and not callable(getattr(self.filter_parameters, a))]
          for attr in attr_list:
              setattr(self, attr, getattr(self.filter_parameters, attr))
  
          self.Q = np.diag([self.cov_omega, self.cov_omega, self. cov_omega,
                             self.cov_acc, self.cov_acc, self.cov_acc,
                             self.cov_b_omega, self.cov_b_omega, self.cov_b_omega,
                             self.cov_b_acc, self.cov_b_acc, self.cov_b_acc,
                             self.cov_Rot_c_i, self.cov_Rot_c_i, self.cov_Rot_c_i,
                             self.cov_t_c_i, self.cov_t_c_i, self.cov_t_c_i])
  
      def run(self, t, u, measurements_covs, v_mes, p_mes, N, ang0):
          dt = t[1:] - t[:-1]  # (s)
          if N is None:
              N = u.shape[0]
          Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P = self.init_run(dt, u, p_mes, v_mes,
                                         ang0, N)
  
          for i in range(1, N):
              Rot[i], v[i], p[i], b_omega[i], b_acc[i], Rot_c_i[i], t_c_i[i], P = \
                  self.propagate(Rot[i-1], v[i-1], p[i-1], b_omega[i-1], b_acc[i-1], Rot_c_i[i-1],
                                 t_c_i[i-1], P, u[i], dt[i-1])
  
              Rot[i], v[i], p[i], b_omega[i], b_acc[i], Rot_c_i[i], t_c_i[i], P = \
                  self.update(Rot[i], v[i], p[i], b_omega[i], b_acc[i], Rot_c_i[i], t_c_i[i], P, u[i],
                              i, measurements_covs[i])
              # correct numerical error every second
              if i % self.n_normalize_rot == 0:
                  Rot[i] = self.normalize_rot(Rot[i])
              # correct numerical error every 10 seconds
              if i % self.n_normalize_rot_c_i == 0:
                  Rot_c_i[i] = self.normalize_rot(Rot_c_i[i])
          return Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i
  
      def init_run(self, dt, u, p_mes, v_mes, ang0, N):
          Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i = self.init_saved_state(dt, N, ang0)
          Rot[0] = self.from_rpy(ang0[0], ang0[1], ang0[2])
          v[0] = v_mes[0]
          P = self.init_covariance()
          return Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P
  
      def init_covariance(self):
          P = np.zeros((self.P_dim, self.P_dim))
          P[:2, :2] = self.cov_Rot0*self.Id2  # no yaw error
          P[3:5, 3:5] = self.cov_v0*self.Id2
          P[9:12, 9:12] = self.cov_b_omega0*self.Id3
          P[12:15, 12:15] = self.cov_b_acc0*self.Id3
          P[15:18, 15:18] = self.cov_Rot_c_i0*self.Id3
          P[18:21, 18:21] = self.cov_t_c_i0*self.Id3
          return P
  
      def init_saved_state(self, dt, N, ang0):
          Rot = np.zeros((N, 3, 3))
          v = np.zeros((N, 3))
          p = np.zeros((N, 3))
          b_omega = np.zeros((N, 3))
          b_acc = np.zeros((N, 3))
          Rot_c_i = np.zeros((N, 3, 3))
          t_c_i = np.zeros((N, 3))
          Rot_c_i[0] = np.eye(3)
          return Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i
  
      def propagate(self, Rot_prev, v_prev, p_prev, b_omega_prev, b_acc_prev, Rot_c_i_prev,
                    t_c_i_prev, P_prev, u, dt):
          acc = Rot_prev.dot(u[3:6] - b_acc_prev) + self.g
          v = v_prev + acc * dt
          p = p_prev + v_prev*dt + 1/2 * acc * dt**2
          omega = u[:3] - b_omega_prev
          Rot = Rot_prev.dot(self.so3exp(omega * dt))
          b_omega = b_omega_prev
          b_acc = b_acc_prev
          Rot_c_i = Rot_c_i_prev
          t_c_i = t_c_i_prev
          P = self.propagate_cov(P_prev, Rot_prev, v_prev, p_prev, b_omega_prev,
                                 b_acc_prev, u, dt)
          return Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P
  
      def propagate_cov(self, P_prev, Rot_prev, v_prev, p_prev, b_omega_prev,
                        b_acc_prev, u, dt):
          F = np.zeros((self.P_dim, self.P_dim))
          G = np.zeros((self.P_dim, self.Q_dim))
          v_skew_rot = self.skew(v_prev).dot(Rot_prev)
          p_skew_rot = self.skew(p_prev).dot(Rot_prev)
  
          F[3:6, :3] = self.skew(self.g)
          F[6:9, 3:6] = self.Id3
          G[3:6, 3:6] = Rot_prev
          F[3:6, 12:15] = -Rot_prev
          G[:3, :3] = Rot_prev
          G[3:6, :3] = v_skew_rot
          G[6:9, :3] = p_skew_rot
          F[:3, 9:12] = -Rot_prev
          F[3:6, 9:12] = -v_skew_rot
          F[6:9, 9:12] = -p_skew_rot
          G[9:15, 6:12] = self.Id6
          G[15:18, 12:15] = self.Id3
          G[18:21, 15:18] = self.Id3
  
          F = F * dt
          G = G * dt
          F_square = F.dot(F)
          F_cube = F_square.dot(F)
          Phi = self.IdP + F + 1/2*F_square + 1/6*F_cube
          P = Phi.dot(P_prev + G.dot(self.Q).dot(G.T)).dot(Phi.T)
          return P
  
      def update(self, Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P, u, i, measurement_cov):
          # orientation of body frame
          Rot_body = Rot.dot(Rot_c_i)
          # velocity in imu frame
          v_imu = Rot.T.dot(v)
          # velocity in body frame
          v_body = Rot_c_i.T.dot(v_imu)
          # velocity in body frame in the vehicle axis
          v_body += self.skew(t_c_i).dot(u[:3] - b_omega)
          Omega = self.skew(u[:3] - b_omega)
  
          # Jacobian w.r.t. car frame
          H_v_imu = Rot_c_i.T.dot(self.skew(v_imu))
          H_t_c_i = -self.skew(t_c_i)
  
          H = np.zeros((2, self.P_dim))
          H[:, 3:6] = Rot_body.T[1:]
          H[:, 15:18] = H_v_imu[1:]
          H[:, 9:12] = H_t_c_i[1:]
          H[:, 18:21] = -Omega[1:]
          r = - v_body[1:]
          R = np.diag(measurement_cov)
          Rot_up, v_up, p_up, b_omega_up, b_acc_up, Rot_c_i_up, t_c_i_up, P_up = \
              self.state_and_cov_update(Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P, H, r, R)
          return Rot_up, v_up, p_up, b_omega_up, b_acc_up, Rot_c_i_up, t_c_i_up, P_up
  
      @staticmethod
      def state_and_cov_update(Rot, v, p, b_omega, b_acc, Rot_c_i, t_c_i, P, H, r, R):
          S = H.dot(P).dot(H.T) + R
          K = (np.linalg.solve(S, P.dot(H.T).T)).T
          dx = K.dot(r)
  
          dR, dxi = NUMPYIEKF.sen3exp(dx[:9])
          dv = dxi[:, 0]
          dp = dxi[:, 1]
          Rot_up = dR.dot(Rot)
          v_up = dR.dot(v) + dv
          p_up = dR.dot(p) + dp
  
          b_omega_up = b_omega + dx[9:12]
          b_acc_up = b_acc + dx[12:15]
  
          dR = NUMPYIEKF.so3exp(dx[15:18])
          Rot_c_i_up = dR.dot(Rot_c_i)
          t_c_i_up = t_c_i + dx[18:21]
  
          I_KH = NUMPYIEKF.IdP - K.dot(H)
          P_up = I_KH.dot(P).dot(I_KH.T) + K.dot(R).dot(K.T)
          P_up = (P_up + P_up.T)/2
          return Rot_up, v_up, p_up, b_omega_up, b_acc_up, Rot_c_i_up, t_c_i_up, P_up
  
      @staticmethod
      def skew(x):
          X = np.array([[0, -x[2], x[1]],
                        [x[2], 0, -x[0]],
                        [-x[1], x[0], 0]])
          return X
  
      @staticmethod
      def rot_from_2_vectors(v1, v2):
          v1 = v1/np.linalg.norm(v1)
          v2 = v2/np.linalg.norm(v2)
          v = np.cross(v1, v2)
          cosang = np.dot(v1, v2)
          sinang = np.linalg.norm(v)
          Rot = NUMPYIEKF.Id3 + NUMPYIEKF.skew(v) + \
                NUMPYIEKF.skew(v).dot(NUMPYIEKF.skew(v))*(1-cosang)/(sinang**2)
          Rot = NUMPYIEKF.normalize_rot(Rot)
          return Rot
  
      @staticmethod
      def sen3exp(xi):
          phi = xi[:3]
  
          angle = np.linalg.norm(phi)
  
          # Near |phi|==0, use first order Taylor expansion
          if np.abs(angle) < 1e-8:
              skew_phi = np.array([[0, -phi[2], phi[1]],
                          [phi[2], 0, -phi[0]],
                          [-phi[1], phi[0], 0]])
              J = NUMPYIEKF.Id3 + 0.5 * skew_phi
              Rot = NUMPYIEKF.Id3 + skew_phi
          else:
              axis = phi / angle
              skew_axis = np.array([[0, -axis[2], axis[1]],
                          [axis[2], 0, -axis[0]],
                          [-axis[1], axis[0], 0]])
              s = np.sin(angle)
              c = np.cos(angle)
              J = (s / angle) * NUMPYIEKF.Id3 \
                     + (1 - s / angle) * np.outer(axis, axis) + ((1 - c) / angle) * skew_axis
              Rot = c * NUMPYIEKF.Id3 + (1 - c) * np.outer(axis, axis) + s * skew_axis
  
          x = J.dot(xi[3:].reshape(-1, 3).T)
          return Rot, x
  
      @staticmethod
      def so3exp(phi):
          angle = np.linalg.norm(phi)
  
          # Near phi==0, use first order Taylor expansion
          if np.abs(angle) < 1e-8:
              skew_phi = np.array([[0, -phi[2], phi[1]],
                        [phi[2], 0, -phi[0]],
                        [-phi[1], phi[0], 0]])
              return np.identity(3) + skew_phi
  
          axis = phi / angle
          skew_axis = np.array([[0, -axis[2], axis[1]],
                        [axis[2], 0, -axis[0]],
                        [-axis[1], axis[0], 0]])
          s = np.sin(angle)
          c = np.cos(angle)
  
          return c * NUMPYIEKF.Id3 + (1 - c) * np.outer(axis, axis) + s * skew_axis
  
      @staticmethod
      def so3left_jacobian(phi):
          """
  
          :param phi:
          :return:
          """
  
          angle = np.linalg.norm(phi)
  
          # Near |phi|==0, use first order Taylor expansion
          if np.abs(angle) < 1e-8:
              skew_phi = np.array([[0, -phi[2], phi[1]],
                        [phi[2], 0, -phi[0]],
                        [-phi[1], phi[0], 0]])
              return NUMPYIEKF.Id3 + 0.5 * skew_phi
  
          axis = phi / angle
          skew_axis = np.array([[0, -axis[2], axis[1]],
                        [axis[2], 0, -axis[0]],
                        [-axis[1], axis[0], 0]])
          s = np.sin(angle)
          c = np.cos(angle)
  
          return (s / angle) * NUMPYIEKF.Id3 \
                 + (1 - s / angle) * np.outer(axis, axis) + ((1 - c) / angle) * skew_axis
  
      @staticmethod
      def normalize_rot(Rot):
  
          # The SVD is commonly written as a = U S V.H.
          # The v returned by this function is V.H and u = U.
          U, _, V = np.linalg.svd(Rot, full_matrices=False)
  
          S = np.eye(3)
          S[2, 2] = np.linalg.det(U) * np.linalg.det(V)
          return U.dot(S).dot(V)
  
      @staticmethod
      def from_rpy(roll, pitch, yaw):
          return NUMPYIEKF.rotz(yaw).dot(NUMPYIEKF.roty(pitch).dot(NUMPYIEKF.rotx(roll)))
  
      @staticmethod
      def rotx(t):
          c = np.cos(t)
          s = np.sin(t)
          return np.array([[1,  0,  0],
                           [0,  c, -s],
                           [0,  s,  c]])
  
      @staticmethod
      def roty(t):
          c = np.cos(t)
          s = np.sin(t)
          return np.array([[c,  0,  s],
                           [0,  1,  0],
                           [-s, 0,  c]])
  
      @staticmethod
      def rotz(t):
          c = np.cos(t)
          s = np.sin(t)
          return np.array([[c, -s,  0],
                           [s,  c,  0],
                           [0,  0,  1]])
  
      @staticmethod
      def to_rpy(Rot):
          pitch = np.arctan2(-Rot[2, 0], np.sqrt(Rot[0, 0]**2 + Rot[1, 0]**2))
  
          if np.isclose(pitch, np.pi / 2.):
              yaw = 0.
              roll = np.arctan2(Rot[0, 1], Rot[1, 1])
          elif np.isclose(pitch, -np.pi / 2.):
              yaw = 0.
              roll = -np.arctan2(Rot[0, 1], Rot[1, 1])
          else:
              sec_pitch = 1. / np.cos(pitch)
              yaw = np.arctan2(Rot[1, 0] * sec_pitch,
                               Rot[0, 0] * sec_pitch)
              roll = np.arctan2(Rot[2, 1] * sec_pitch,
                                Rot[2, 2] * sec_pitch)
          return roll, pitch, yaw
  
      def set_learned_covariance(self, torch_iekf):
          torch_iekf.set_Q()
          self.Q = torch_iekf.Q.cpu().detach().numpy()
  
          beta = torch_iekf.initprocesscov_net.init_cov(torch_iekf)\
              .detach().cpu().numpy()
  
          self.cov_Rot0 *= beta[0]
          self.cov_v0 *= beta[1]
          self.cov_b_omega0 *= beta[2]
          self.cov_b_acc0 *= beta[3]
          self.cov_Rot_c_i0 *= beta[4]
          self.cov_t_c_i0 *= beta[5]
  