IMU/scripts/imu_ekf_core.py

"""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]