IMU/scripts/imu_static_calibrator.py

"""Numeric-only stationary detector and yaw gyro bias estimator."""

from __future__ import annotations

from dataclasses import dataclass
import math

import numpy as np


MOVING = 0
CANDIDATE = 1
STATIC = 2
TIME_EPSILON_S = 1.0e-12


@dataclass(frozen=True)
class StaticCorrectionConfig:
    enter_seconds: float = 2.0
    gyro_threshold_dps: float = 0.5
    acc_norm_tolerance_g: float = 0.2
    acc_stability_threshold_g: float = 0.02


@dataclass
class StaticCorrectionState:
    config: StaticCorrectionConfig
    mode: int
    active_yaw_bias_z_dps: float
    candidate_elapsed_s: float
    candidate_count: int
    candidate_acc_mean_g: np.ndarray
    candidate_gyro_z_mean_dps: float
    static_count: int
    static_acc_mean_g: np.ndarray
    static_gyro_z_mean_dps: float


def initialize(
    config: StaticCorrectionConfig,
    initial_yaw_bias_z_dps: float,
) -> StaticCorrectionState:
    _validate_config(config)
    if not math.isfinite(initial_yaw_bias_z_dps):
        raise ValueError("initial_yaw_bias_z_dps must be finite")
    return StaticCorrectionState(
        config=config,
        mode=MOVING,
        active_yaw_bias_z_dps=float(initial_yaw_bias_z_dps),
        candidate_elapsed_s=0.0,
        candidate_count=0,
        candidate_acc_mean_g=np.zeros(3),
        candidate_gyro_z_mean_dps=0.0,
        static_count=0,
        static_acc_mean_g=np.zeros(3),
        static_gyro_z_mean_dps=0.0,
    )


def step(
    state: StaticCorrectionState,
    dt_s: float,
    acc_g: np.ndarray,
    gyro_dps: np.ndarray,
    gyro_bias_xy_dps: np.ndarray,
) -> bool:
    if not math.isfinite(dt_s) or dt_s < 0.0:
        raise ValueError("dt_s must be finite and non-negative")
    acc = _vector(acc_g, 3, "acc_g")
    gyro = _vector(gyro_dps, 3, "gyro_dps")
    bias_xy = _vector(gyro_bias_xy_dps, 2, "gyro_bias_xy_dps")

    gyro_residual = np.array(
        [
            gyro[0] - bias_xy[0],
            gyro[1] - bias_xy[1],
            gyro[2] - state.active_yaw_bias_z_dps,
        ]
    )
    absolute_gate_ok = (
        abs(float(np.linalg.norm(acc)) - 1.0) <= state.config.acc_norm_tolerance_g
        and float(np.linalg.norm(gyro_residual)) <= state.config.gyro_threshold_dps
    )

    if state.mode == STATIC:
        stable_acc = (
            float(np.linalg.norm(acc - state.static_acc_mean_g))
            <= state.config.acc_stability_threshold_g
        )
        if not absolute_gate_ok or not stable_acc:
            _reset_candidate(state)
            state.mode = MOVING
            return False
        state.static_count += 1
        state.static_acc_mean_g += (acc - state.static_acc_mean_g) / state.static_count
        state.static_gyro_z_mean_dps += (
            gyro[2] - state.static_gyro_z_mean_dps
        ) / state.static_count
        state.active_yaw_bias_z_dps = state.static_gyro_z_mean_dps
        return True

    if not absolute_gate_ok:
        _reset_candidate(state)
        state.mode = MOVING
        return False

    if state.mode == MOVING:
        _start_candidate(state, acc, gyro[2])
        return False

    stable_acc = (
        float(np.linalg.norm(acc - state.candidate_acc_mean_g))
        <= state.config.acc_stability_threshold_g
    )
    if not stable_acc:
        _start_candidate(state, acc, gyro[2])
        return False

    state.candidate_count += 1
    state.candidate_elapsed_s += dt_s
    state.candidate_acc_mean_g += (
        acc - state.candidate_acc_mean_g
    ) / state.candidate_count
    state.candidate_gyro_z_mean_dps += (
        gyro[2] - state.candidate_gyro_z_mean_dps
    ) / state.candidate_count
    if state.candidate_elapsed_s + TIME_EPSILON_S < state.config.enter_seconds:
        return False

    state.mode = STATIC
    state.static_count = state.candidate_count
    state.static_acc_mean_g = state.candidate_acc_mean_g.copy()
    state.static_gyro_z_mean_dps = state.candidate_gyro_z_mean_dps
    state.active_yaw_bias_z_dps = state.static_gyro_z_mean_dps
    return True


def _start_candidate(
    state: StaticCorrectionState,
    acc_g: np.ndarray,
    gyro_z_dps: float,
) -> None:
    state.mode = CANDIDATE
    state.candidate_elapsed_s = 0.0
    state.candidate_count = 1
    state.candidate_acc_mean_g = acc_g.copy()
    state.candidate_gyro_z_mean_dps = float(gyro_z_dps)


def _reset_candidate(state: StaticCorrectionState) -> None:
    state.candidate_elapsed_s = 0.0
    state.candidate_count = 0
    state.candidate_acc_mean_g.fill(0.0)
    state.candidate_gyro_z_mean_dps = 0.0


def _vector(value, size: int, name: str) -> np.ndarray:
    vector = np.asarray(value, dtype=float)
    if vector.shape != (size,) or not np.all(np.isfinite(vector)):
        raise ValueError(f"{name} must be a finite {size}-element vector")
    return vector


def _validate_config(config: StaticCorrectionConfig) -> None:
    values = (
        config.enter_seconds,
        config.gyro_threshold_dps,
        config.acc_norm_tolerance_g,
        config.acc_stability_threshold_g,
    )
    if not all(math.isfinite(value) for value in values):
        raise ValueError("static correction configuration must be finite")
    if config.enter_seconds <= 0.0:
        raise ValueError("enter_seconds must be positive")
    if config.gyro_threshold_dps <= 0.0:
        raise ValueError("gyro_threshold_dps must be positive")
    if config.acc_norm_tolerance_g <= 0.0:
        raise ValueError("acc_norm_tolerance_g must be positive")
    if config.acc_stability_threshold_g <= 0.0:
        raise ValueError("acc_stability_threshold_g must be positive")