209 lines
6.1 KiB
Python
209 lines
6.1 KiB
Python
"""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]
|