Initial IMU EKF project
This commit is contained in:
@@ -0,0 +1,208 @@
|
||||
"""Numeric-only attitude EKF core for IMU accelerometer/gyroscope samples.
|
||||
|
||||
The wrapper layer owns CSV parsing, units, files, and visualization. This file
|
||||
keeps explicit numeric state so the algorithm can be ported to fixed-size C.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
GRAVITY_MPS2 = 9.80665
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImuEkfState:
|
||||
q: np.ndarray
|
||||
gyro_bias_rad_s: np.ndarray
|
||||
p: np.ndarray
|
||||
gyro_noise_var: float
|
||||
bias_noise_var: float
|
||||
acc_noise_var: float
|
||||
acc_gate_mps2: float
|
||||
|
||||
|
||||
def initialize_from_samples(
|
||||
acc_mps2_samples,
|
||||
gyro_rad_s_samples,
|
||||
gyro_noise_var: float = 1.0e-5,
|
||||
bias_noise_var: float = 1.0e-8,
|
||||
acc_noise_var: float = 2.5e-3,
|
||||
acc_gate_mps2: float = 2.0,
|
||||
) -> ImuEkfState:
|
||||
acc_mean = _mean_vector(acc_mps2_samples)
|
||||
gyro_mean = _mean_vector(gyro_rad_s_samples)
|
||||
q = _quaternion_from_two_vectors(_normalize3(acc_mean), np.array([0.0, 0.0, 1.0]))
|
||||
p = np.diag([1.0e-3, 1.0e-3, 1.0e-3, 1.0e-4, 1.0e-4, 1.0e-4])
|
||||
return ImuEkfState(
|
||||
q=_quat_normalize(q),
|
||||
gyro_bias_rad_s=gyro_mean.copy(),
|
||||
p=p,
|
||||
gyro_noise_var=gyro_noise_var,
|
||||
bias_noise_var=bias_noise_var,
|
||||
acc_noise_var=acc_noise_var,
|
||||
acc_gate_mps2=acc_gate_mps2,
|
||||
)
|
||||
|
||||
|
||||
def predict(state: ImuEkfState, dt_s: float, gyro_rad_s: np.ndarray) -> None:
|
||||
if dt_s <= 0.0:
|
||||
return
|
||||
|
||||
omega = np.asarray(gyro_rad_s, dtype=float) - state.gyro_bias_rad_s
|
||||
state.q = _quat_normalize(_quat_multiply(state.q, _quat_from_rotvec(omega * dt_s)))
|
||||
|
||||
f = np.eye(6)
|
||||
f[0:3, 0:3] -= _skew(omega) * dt_s
|
||||
f[0:3, 3:6] = -np.eye(3) * dt_s
|
||||
|
||||
q_noise = np.zeros((6, 6))
|
||||
q_noise[0:3, 0:3] = np.eye(3) * state.gyro_noise_var * dt_s * dt_s
|
||||
q_noise[3:6, 3:6] = np.eye(3) * state.bias_noise_var * dt_s
|
||||
state.p = f @ state.p @ f.T + q_noise
|
||||
|
||||
|
||||
def update_accel(state: ImuEkfState, acc_mps2: np.ndarray) -> tuple[bool, float]:
|
||||
acc = np.asarray(acc_mps2, dtype=float)
|
||||
acc_norm = float(np.linalg.norm(acc))
|
||||
if acc_norm <= 1.0e-12:
|
||||
return False, 0.0
|
||||
|
||||
z_meas = acc / acc_norm
|
||||
z_pred = _rotate_world_to_body(state.q, np.array([0.0, 0.0, 1.0]))
|
||||
residual = z_meas - z_pred
|
||||
residual_norm = float(np.linalg.norm(residual))
|
||||
|
||||
if abs(acc_norm - GRAVITY_MPS2) > state.acc_gate_mps2:
|
||||
return False, residual_norm
|
||||
|
||||
h = np.zeros((3, 6))
|
||||
h[:, 0:3] = _skew(z_pred)
|
||||
r = np.eye(3) * state.acc_noise_var
|
||||
s = h @ state.p @ h.T + r
|
||||
k = state.p @ h.T @ np.linalg.inv(s)
|
||||
dx = k @ residual
|
||||
|
||||
state.q = _quat_normalize(_quat_multiply(state.q, _quat_from_rotvec(dx[0:3])))
|
||||
state.gyro_bias_rad_s += dx[3:6]
|
||||
|
||||
i = np.eye(6)
|
||||
kh = k @ h
|
||||
state.p = (i - kh) @ state.p @ (i - kh).T + k @ r @ k.T
|
||||
return True, residual_norm
|
||||
|
||||
|
||||
def step(
|
||||
state: ImuEkfState,
|
||||
dt_s: float,
|
||||
acc_mps2: np.ndarray,
|
||||
gyro_rad_s: np.ndarray,
|
||||
) -> tuple[bool, float]:
|
||||
predict(state, dt_s, gyro_rad_s)
|
||||
return update_accel(state, acc_mps2)
|
||||
|
||||
|
||||
def quaternion_to_euler_deg(q: np.ndarray) -> tuple[float, float, float]:
|
||||
w, x, y, z = _quat_normalize(q)
|
||||
|
||||
sinr_cosp = 2.0 * (w * x + y * z)
|
||||
cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
|
||||
roll = math.atan2(sinr_cosp, cosr_cosp)
|
||||
|
||||
sinp = 2.0 * (w * y - z * x)
|
||||
if abs(sinp) >= 1.0:
|
||||
pitch = math.copysign(math.pi / 2.0, sinp)
|
||||
else:
|
||||
pitch = math.asin(sinp)
|
||||
|
||||
siny_cosp = 2.0 * (w * z + x * y)
|
||||
cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
|
||||
yaw = math.atan2(siny_cosp, cosy_cosp)
|
||||
|
||||
return math.degrees(roll), math.degrees(pitch), math.degrees(yaw)
|
||||
|
||||
|
||||
def _mean_vector(samples) -> np.ndarray:
|
||||
vectors = [np.asarray(sample, dtype=float) for sample in samples]
|
||||
if not vectors:
|
||||
raise ValueError("at least one sample is required")
|
||||
return np.mean(np.vstack(vectors), axis=0)
|
||||
|
||||
|
||||
def _normalize3(v: np.ndarray) -> np.ndarray:
|
||||
norm = float(np.linalg.norm(v))
|
||||
if norm <= 1.0e-12:
|
||||
raise ValueError("cannot normalize a zero vector")
|
||||
return np.asarray(v, dtype=float) / norm
|
||||
|
||||
|
||||
def _skew(v: np.ndarray) -> np.ndarray:
|
||||
x, y, z = v
|
||||
return np.array(
|
||||
[
|
||||
[0.0, -z, y],
|
||||
[z, 0.0, -x],
|
||||
[-y, x, 0.0],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _quat_normalize(q: np.ndarray) -> np.ndarray:
|
||||
q = np.asarray(q, dtype=float)
|
||||
norm = float(np.linalg.norm(q))
|
||||
if norm <= 1.0e-12:
|
||||
raise ValueError("cannot normalize a zero quaternion")
|
||||
out = q / norm
|
||||
if out[0] < 0.0:
|
||||
out = -out
|
||||
return out
|
||||
|
||||
|
||||
def _quat_multiply(a: np.ndarray, b: np.ndarray) -> np.ndarray:
|
||||
aw, ax, ay, az = a
|
||||
bw, bx, by, bz = b
|
||||
return np.array(
|
||||
[
|
||||
aw * bw - ax * bx - ay * by - az * bz,
|
||||
aw * bx + ax * bw + ay * bz - az * by,
|
||||
aw * by - ax * bz + ay * bw + az * bx,
|
||||
aw * bz + ax * by - ay * bx + az * bw,
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _quat_conjugate(q: np.ndarray) -> np.ndarray:
|
||||
return np.array([q[0], -q[1], -q[2], -q[3]])
|
||||
|
||||
|
||||
def _quat_from_rotvec(rotvec: np.ndarray) -> np.ndarray:
|
||||
angle = float(np.linalg.norm(rotvec))
|
||||
if angle <= 1.0e-12:
|
||||
return _quat_normalize(np.array([1.0, rotvec[0] / 2.0, rotvec[1] / 2.0, rotvec[2] / 2.0]))
|
||||
axis = rotvec / angle
|
||||
half = angle / 2.0
|
||||
return np.array([math.cos(half), *(math.sin(half) * axis)])
|
||||
|
||||
|
||||
def _quaternion_from_two_vectors(source: np.ndarray, target: np.ndarray) -> np.ndarray:
|
||||
source = _normalize3(source)
|
||||
target = _normalize3(target)
|
||||
dot = float(np.dot(source, target))
|
||||
if dot < -0.999999:
|
||||
axis = _normalize3(np.cross(np.array([1.0, 0.0, 0.0]), source))
|
||||
if float(np.linalg.norm(axis)) <= 1.0e-12:
|
||||
axis = _normalize3(np.cross(np.array([0.0, 1.0, 0.0]), source))
|
||||
return np.array([0.0, *axis])
|
||||
cross = np.cross(source, target)
|
||||
return _quat_normalize(np.array([1.0 + dot, cross[0], cross[1], cross[2]]))
|
||||
|
||||
|
||||
def _rotate_world_to_body(q_body_to_world: np.ndarray, v_world: np.ndarray) -> np.ndarray:
|
||||
q_conj = _quat_conjugate(_quat_normalize(q_body_to_world))
|
||||
v_quat = np.array([0.0, v_world[0], v_world[1], v_world[2]])
|
||||
rotated = _quat_multiply(_quat_multiply(q_conj, v_quat), q_body_to_world)
|
||||
return rotated[1:4]
|
||||
Reference in New Issue
Block a user