commit 97cd9a56c61bbb459ccb81f3baf916dce60e5c7a Author: 刘泽群 Date: Tue Jun 16 18:22:26 2026 +0800 Initial IMU EKF project diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..248d193 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +.venv/ +__pycache__/ +*.py[cod] + +output/ +*.csv +*.exe +*.log diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..2556558 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,15 @@ +# AGENTS.md + +## IMU CSV Notes + +1. Each CSV header records the IMU configuration used when collecting the data. The `odr` field and related IMU configuration fields are important. +2. For row-level data analysis, only use these columns: `sensor_uptime_s`, `temp_c`, `acc_x_g`, `acc_y_g`, `acc_z_g`, `gyro_x_dps`, `gyro_y_dps`, and `gyro_z_dps`. + +## Column Meanings + +- `sensor_uptime_s`: data arrival time in seconds. +- `temp_c`: IMU temperature, in degrees Celsius. +- `acc_x_g`, `acc_y_g`, `acc_z_g`: IMU acceleration on the x, y, and z axes, in `g`. +- `gyro_x_dps`, `gyro_y_dps`, `gyro_z_dps`: IMU angular offset/rate on the x, y, and z axes, in degrees per second. These values are not radians. + +Unit conversion is required for practical use. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8e6304c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,8 @@ +[project] +name = "imu" +version = "0.1.0" +requires-python = ">=3.13" +dependencies = [ + "allantools>=2024.6", + "numpy>=2.4.6", +] diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..acd7548 --- /dev/null +++ b/scripts/__init__.py @@ -0,0 +1 @@ +"""Local scripts for IMU processing.""" diff --git a/scripts/imu_ekf_core.py b/scripts/imu_ekf_core.py new file mode 100644 index 0000000..d421e86 --- /dev/null +++ b/scripts/imu_ekf_core.py @@ -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] diff --git a/scripts/run_imu_ekf.py b/scripts/run_imu_ekf.py new file mode 100644 index 0000000..f92de97 --- /dev/null +++ b/scripts/run_imu_ekf.py @@ -0,0 +1,519 @@ +"""Run the IMU attitude EKF over local CSV files and write a viewer.""" + +from __future__ import annotations + +import argparse +import csv +from dataclasses import dataclass +import json +import math +from pathlib import Path + +import numpy as np + +from scripts import imu_ekf_core as ekf + + +REQUIRED_COLUMNS = ( + "sensor_uptime_s", + "temp_c", + "acc_x_g", + "acc_y_g", + "acc_z_g", + "gyro_x_dps", + "gyro_y_dps", + "gyro_z_dps", +) +RESTART_PREVIOUS_MIN_S = 1.0 +RESTART_CURRENT_MAX_S = 0.01 + + +@dataclass(frozen=True) +class ImuRow: + sensor_uptime_s: float + temp_c: float + acc_g: tuple[float, float, float] + gyro_dps: tuple[float, float, float] + + +@dataclass(frozen=True) +class EkfFileResult: + input_csv: Path + output_csv: Path + input_rows: int + metadata: dict[str, str] + samples: list[dict[str, float]] + + +def read_imu_csv(path: Path): + metadata, _fieldnames = _read_metadata_and_header(path) + return metadata, iter_imu_rows(path) + + +def iter_imu_rows(path: Path): + reader = csv.DictReader(_iter_data_lines(path)) + _validate_columns(reader.fieldnames or [], path) + for data_row_index, record in enumerate(reader, start=1): + yield _row_from_record(record, path, data_row_index) + + +def process_file( + input_csv: Path, + output_dir: Path, + init_seconds: int = 3, + max_points: int = 2500, +) -> EkfFileResult: + init_seconds = _validate_init_seconds(init_seconds) + input_csv = Path(input_csv) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + metadata, _fieldnames = _read_metadata_and_header(input_csv) + + output_csv = output_dir / f"{input_csv.stem}_ekf.csv" + samples: list[dict[str, float]] = [] + input_rows = 0 + rows_seen = 0 + + with output_csv.open("w", encoding="utf-8", newline="") as handle: + fieldnames = [ + "sensor_uptime_s", + "segment_id", + "roll_deg", + "pitch_deg", + "yaw_deg", + "relative_yaw_deg", + "qw", + "qx", + "qy", + "qz", + "gyro_bias_x_dps", + "gyro_bias_y_dps", + "gyro_bias_z_dps", + "acc_residual_norm", + "dt_s", + "acc_update_used", + ] + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + + segment_id = 0 + state: ekf.ImuEkfState | None = None + init_buffer: list[ImuRow] = [] + previous_input_time: float | None = None + previous_output_time: float | None = None + relative_yaw_deg = 0.0 + previous_yaw_deg: float | None = None + + def reset_segment() -> None: + nonlocal state, init_buffer, previous_output_time, relative_yaw_deg, previous_yaw_deg + state = None + init_buffer = [] + previous_output_time = None + relative_yaw_deg = 0.0 + previous_yaw_deg = None + + def write_row(row: ImuRow) -> None: + nonlocal input_rows, previous_output_time, relative_yaw_deg, previous_yaw_deg + if state is None: + raise ValueError("EKF state is not initialized") + dt_s = 0.0 if previous_output_time is None else row.sensor_uptime_s - previous_output_time + previous_output_time = row.sensor_uptime_s + + used_update, residual_norm = ekf.step(state, dt_s, _acc_mps2(row), _gyro_rad_s(row)) + roll, pitch, yaw = ekf.quaternion_to_euler_deg(state.q) + if previous_yaw_deg is None: + relative_yaw_deg = 0.0 + previous_yaw_deg = yaw + else: + relative_yaw_deg += _unwrap_delta_deg(yaw - previous_yaw_deg) + previous_yaw_deg = yaw + bias_dps = tuple(math.degrees(value) for value in state.gyro_bias_rad_s) + out = { + "sensor_uptime_s": row.sensor_uptime_s, + "segment_id": segment_id, + "roll_deg": roll, + "pitch_deg": pitch, + "yaw_deg": yaw, + "relative_yaw_deg": relative_yaw_deg, + "qw": float(state.q[0]), + "qx": float(state.q[1]), + "qy": float(state.q[2]), + "qz": float(state.q[3]), + "gyro_bias_x_dps": bias_dps[0], + "gyro_bias_y_dps": bias_dps[1], + "gyro_bias_z_dps": bias_dps[2], + "acc_residual_norm": residual_norm, + "dt_s": dt_s, + "acc_update_used": int(used_update), + } + _validate_finite_mapping(out, "EKF output") + writer.writerow(out) + input_rows += 1 + _append_bounded(samples, _sample_for_html(out), max_points) + + def initialize_and_write_buffer() -> None: + nonlocal state, init_buffer + state = _initialize_state(init_buffer, init_seconds) + buffered_rows = init_buffer + init_buffer = [] + for buffered_row in buffered_rows: + write_row(buffered_row) + + for data_row_index, row in enumerate(iter_imu_rows(input_csv), start=1): + rows_seen += 1 + if previous_input_time is not None and row.sensor_uptime_s < previous_input_time: + if _is_device_restart(previous_input_time, row.sensor_uptime_s): + if state is None and init_buffer: + initialize_and_write_buffer() + segment_id += 1 + reset_segment() + else: + raise ValueError( + f"{input_csv} timestamp decreased at data row {data_row_index}: " + f"previous {previous_input_time}, current {row.sensor_uptime_s}" + ) + + if state is None: + init_buffer.append(row) + if init_seconds == 0 or row.sensor_uptime_s - init_buffer[0].sensor_uptime_s >= init_seconds: + initialize_and_write_buffer() + else: + write_row(row) + + previous_input_time = row.sensor_uptime_s + + if rows_seen == 0: + raise ValueError(f"{input_csv} has no IMU rows") + if state is None and init_buffer: + initialize_and_write_buffer() + + return EkfFileResult( + input_csv=input_csv, + output_csv=output_csv, + input_rows=input_rows, + metadata=metadata, + samples=samples, + ) + + +def write_html_report(results: list[EkfFileResult], html_path: Path, max_points: int = 2500) -> None: + html_path = Path(html_path) + html_path.parent.mkdir(parents=True, exist_ok=True) + payload = [] + for result in results: + payload.append( + { + "input_csv": str(result.input_csv.name), + "output_csv": str(result.output_csv), + "input_rows": result.input_rows, + "metadata": result.metadata, + "samples": result.samples[:max_points], + } + ) + + data_json = _json_for_script(payload) + html_text = f""" + + + + + IMU EKF Viewer + + + +
+

IMU EKF Viewer

+
Relative yaw only; no magnetometer heading reference is available.
+
+
+ + + + +""" + html_path.write_text(html_text, encoding="utf-8") + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Run attitude EKF over IMU CSV files.") + parser.add_argument("csv_files", nargs="*", type=Path, help="CSV files. Defaults to imu_*.csv.") + parser.add_argument("--output-dir", type=Path, default=Path("output") / "ekf") + parser.add_argument("--init-seconds", type=_parse_init_seconds, default=3) + parser.add_argument("--max-points", type=int, default=2500) + args = parser.parse_args(argv) + + csv_files = args.csv_files or sorted(Path(".").glob("imu_*.csv")) + if not csv_files: + raise SystemExit("No CSV files found.") + + results = [ + process_file(path, args.output_dir, init_seconds=args.init_seconds, max_points=args.max_points) + for path in csv_files + ] + write_html_report(results, args.output_dir / "ekf_viewer.html", max_points=args.max_points) + for result in results: + print(f"{result.input_csv}: {result.input_rows} rows -> {result.output_csv}") + print(f"viewer -> {args.output_dir / 'ekf_viewer.html'}") + return 0 + + +def _read_metadata_and_header(path: Path) -> tuple[dict[str, str], list[str]]: + metadata: dict[str, str] = {} + header_line: str | None = None + with Path(path).open("r", encoding="utf-8-sig", newline="") as handle: + for line in handle: + if line.startswith("#"): + key, value = _parse_metadata_line(line) + if key: + metadata[key] = value + continue + header_line = line + break + + if header_line is None: + raise ValueError(f"{path} missing CSV header") + + fieldnames = next(csv.reader([header_line])) + _validate_columns(fieldnames, path) + return metadata, fieldnames + + +def _iter_data_lines(path: Path): + with Path(path).open("r", encoding="utf-8-sig", newline="") as handle: + header_seen = False + for line in handle: + if not header_seen and line.startswith("#"): + continue + header_seen = True + yield line + + +def _parse_metadata_line(line: str) -> tuple[str | None, str]: + stripped = line[1:].strip() + if "=" not in stripped: + return None, "" + key, value = stripped.split("=", 1) + return key.strip(), value.strip() + + +def _validate_columns(fieldnames: list[str], path: Path) -> None: + missing = [column for column in REQUIRED_COLUMNS if column not in fieldnames] + if missing: + raise ValueError(f"{path} missing required columns: {', '.join(missing)}") + + +def _row_from_record(record: dict[str, str], path: Path, data_row_index: int) -> ImuRow: + return ImuRow( + sensor_uptime_s=_finite_float(record["sensor_uptime_s"], "sensor_uptime_s", path, data_row_index), + temp_c=_finite_float(record["temp_c"], "temp_c", path, data_row_index), + acc_g=( + _finite_float(record["acc_x_g"], "acc_x_g", path, data_row_index), + _finite_float(record["acc_y_g"], "acc_y_g", path, data_row_index), + _finite_float(record["acc_z_g"], "acc_z_g", path, data_row_index), + ), + gyro_dps=( + _finite_float(record["gyro_x_dps"], "gyro_x_dps", path, data_row_index), + _finite_float(record["gyro_y_dps"], "gyro_y_dps", path, data_row_index), + _finite_float(record["gyro_z_dps"], "gyro_z_dps", path, data_row_index), + ), + ) + + +def _validate_init_seconds(value) -> int: + if type(value) is not int or not 0 <= value <= 10: + raise ValueError("init_seconds must be a 0-10 second integer") + return value + + +def _parse_init_seconds(value: str) -> int: + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("init_seconds must be a 0-10 second integer") from exc + if str(parsed) != value: + raise argparse.ArgumentTypeError("init_seconds must be a 0-10 second integer") + try: + return _validate_init_seconds(parsed) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) from exc + + +def _initialize_state(rows: list[ImuRow], init_seconds: int) -> ekf.ImuEkfState: + if not rows: + raise ValueError("at least one IMU row is required for initialization") + gyro_samples = [np.zeros(3)] if init_seconds == 0 else [_gyro_rad_s(row) for row in rows] + return ekf.initialize_from_samples( + acc_mps2_samples=[_acc_mps2(row) for row in rows], + gyro_rad_s_samples=gyro_samples, + ) + + +def _acc_mps2(row: ImuRow) -> np.ndarray: + return np.array(row.acc_g, dtype=float) * ekf.GRAVITY_MPS2 + + +def _gyro_rad_s(row: ImuRow) -> np.ndarray: + return np.radians(np.array(row.gyro_dps, dtype=float)) + + +def _sample_for_html(row: dict[str, float]) -> dict[str, float]: + keys = ( + "sensor_uptime_s", + "roll_deg", + "pitch_deg", + "relative_yaw_deg", + "gyro_bias_x_dps", + "gyro_bias_y_dps", + "gyro_bias_z_dps", + "acc_residual_norm", + ) + sample = {key: float(row[key]) for key in keys} + _validate_finite_mapping(sample, "HTML sample") + return sample + + +def _unwrap_delta_deg(delta: float) -> float: + while delta > 180.0: + delta -= 360.0 + while delta <= -180.0: + delta += 360.0 + return delta + + +def _json_for_script(payload) -> str: + _validate_json_finite(payload) + return ( + json.dumps(payload, ensure_ascii=False, allow_nan=False) + .replace("&", "\\u0026") + .replace("<", "\\u003c") + .replace(">", "\\u003e") + ) + + +def _is_device_restart(previous_time: float, current_time: float) -> bool: + return previous_time > RESTART_PREVIOUS_MIN_S and current_time <= RESTART_CURRENT_MAX_S + + +def _finite_float(value: str, column: str, path: Path, data_row_index: int) -> float: + parsed = float(value) + if not math.isfinite(parsed): + raise ValueError(f"{path} column {column} must be finite at data row {data_row_index}") + return parsed + + +def _validate_finite_mapping(values: dict[str, float], label: str) -> None: + for key, value in values.items(): + if not math.isfinite(float(value)): + raise ValueError(f"{label} field {key} must be finite") + + +def _validate_json_finite(value, path: str = "$") -> None: + if isinstance(value, dict): + for key, child in value.items(): + _validate_json_finite(child, f"{path}.{key}") + elif isinstance(value, list): + for index, child in enumerate(value): + _validate_json_finite(child, f"{path}[{index}]") + elif isinstance(value, (float, np.floating)) and not math.isfinite(float(value)): + raise ValueError(f"JSON payload number at {path} must be finite") + + +def _append_bounded(samples: list[dict[str, float]], sample: dict[str, float], max_points: int) -> None: + if max_points <= 0: + return + samples.append(sample) + if len(samples) > max_points: + del samples[1::2] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_imu_ekf_core.py b/tests/test_imu_ekf_core.py new file mode 100644 index 0000000..95ec1d4 --- /dev/null +++ b/tests/test_imu_ekf_core.py @@ -0,0 +1,73 @@ +import math +import unittest + +import numpy as np + +from scripts import imu_ekf_core as ekf + + +G = 9.80665 + + +class ImuEkfCoreTests(unittest.TestCase): + def test_static_data_keeps_level_orientation(self): + state = ekf.initialize_from_samples( + acc_mps2_samples=[np.array([0.0, 0.0, G])] * 20, + gyro_rad_s_samples=[np.zeros(3)] * 20, + ) + + for _ in range(500): + ekf.step(state, 0.002, np.array([0.0, 0.0, G]), np.zeros(3)) + + roll, pitch, yaw = ekf.quaternion_to_euler_deg(state.q) + self.assertLess(abs(roll), 0.1) + self.assertLess(abs(pitch), 0.1) + self.assertLess(abs(yaw), 0.1) + + def test_initialization_estimates_gyro_bias_from_static_window(self): + bias = np.array([math.radians(1.5), math.radians(-0.5), math.radians(0.25)]) + + state = ekf.initialize_from_samples( + acc_mps2_samples=[np.array([0.0, 0.0, G])] * 50, + gyro_rad_s_samples=[bias] * 50, + ) + + np.testing.assert_allclose(state.gyro_bias_rad_s, bias, atol=1e-12) + + def test_quaternion_stays_normalized_after_many_steps(self): + state = ekf.initialize_from_samples( + acc_mps2_samples=[np.array([0.0, 0.0, G])] * 20, + gyro_rad_s_samples=[np.zeros(3)] * 20, + ) + + for _ in range(2000): + ekf.step( + state, + 0.001, + np.array([0.0, 0.0, G]), + np.array([0.01, -0.02, 0.03]), + ) + + self.assertAlmostEqual(float(np.linalg.norm(state.q)), 1.0, places=9) + + def test_accelerometer_gate_suppresses_large_linear_acceleration(self): + state = ekf.initialize_from_samples( + acc_mps2_samples=[np.array([0.0, 0.0, G])] * 20, + gyro_rad_s_samples=[np.zeros(3)] * 20, + ) + before = state.q.copy() + + used_update, residual_norm = ekf.step( + state, + 0.002, + np.array([4.0 * G, 0.0, 0.0]), + np.zeros(3), + ) + + self.assertFalse(used_update) + self.assertGreater(residual_norm, 0.0) + np.testing.assert_allclose(state.q, before, atol=1e-12) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_run_imu_ekf.py b/tests/test_run_imu_ekf.py new file mode 100644 index 0000000..4d8068b --- /dev/null +++ b/tests/test_run_imu_ekf.py @@ -0,0 +1,308 @@ +import csv +import math +import tempfile +import unittest +from pathlib import Path + +from scripts import run_imu_ekf + + +class RunImuEkfTests(unittest.TestCase): + def _write_sample_csv(self, path: Path, rows: list[tuple[float, float, float]]): + lines = [ + "# odr=0x0F - 500 Hz", + "sensor_uptime_s,temp_c,acc_x_g,acc_y_g,acc_z_g,gyro_x_dps,gyro_y_dps,gyro_z_dps", + ] + for sensor_time, gyro_z_dps, acc_z_g in rows: + lines.append(f"{sensor_time},28.0,0,0,{acc_z_g},0,0,{gyro_z_dps}") + path.write_text("\n".join(lines), encoding="utf-8-sig") + + def test_csv_parser_skips_metadata_and_reads_required_columns(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "sample.csv" + path.write_text( + "\n".join( + [ + "# odr=0x0F - 500 Hz", + "# gyro_bw=0x01 - ODR/4", + "sensor_uptime_s,temp_c,acc_x_g,acc_y_g,acc_z_g,gyro_x_dps,gyro_y_dps,gyro_z_dps,ignored", + "0.000,28.0,0,0,1,1,2,3,x", + "0.002,28.0,0,0,1,1,2,3,x", + ] + ), + encoding="utf-8-sig", + ) + + metadata, rows = run_imu_ekf.read_imu_csv(path) + self.assertNotIsInstance(rows, list) + rows = list(rows) + + self.assertEqual(metadata["odr"], "0x0F - 500 Hz") + self.assertEqual(metadata["gyro_bw"], "0x01 - ODR/4") + self.assertEqual(len(rows), 2) + self.assertEqual(rows[1].sensor_uptime_s, 0.002) + self.assertEqual(rows[0].gyro_dps[2], 3.0) + + def test_csv_parser_rejects_missing_required_columns(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bad.csv" + path.write_text( + "sensor_uptime_s,temp_c,acc_x_g,acc_y_g,acc_z_g,gyro_x_dps,gyro_y_dps\n" + "0.0,28,0,0,1,0,0\n", + encoding="utf-8-sig", + ) + + with self.assertRaisesRegex(ValueError, "gyro_z_dps"): + run_imu_ekf.read_imu_csv(path) + + def test_process_file_writes_expected_result_columns(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv(input_path, [(index * 0.002, 0.0, 1.0) for index in range(20)]) + + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=1) + + with result.output_csv.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + + self.assertEqual(len(rows), 20) + self.assertIn("roll_deg", rows[0]) + self.assertIn("relative_yaw_deg", rows[0]) + self.assertIn("segment_id", rows[0]) + self.assertIn("gyro_bias_z_dps", rows[0]) + self.assertEqual(result.input_rows, 20) + + def test_process_file_reads_imu_rows_once(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv(input_path, [(0.0, 0.0, 1.0)]) + rows = [ + run_imu_ekf.ImuRow(0.0, 28.0, (0.0, 0.0, 1.0), (0.0, 0.0, 0.0)), + run_imu_ekf.ImuRow(0.5, 28.0, (0.0, 0.0, 1.0), (0.0, 0.0, 0.0)), + run_imu_ekf.ImuRow(1.0, 28.0, (0.0, 0.0, 1.0), (0.0, 0.0, 0.0)), + ] + call_count = 0 + original_iter = run_imu_ekf.iter_imu_rows + + def single_use_iter(path): + nonlocal call_count + call_count += 1 + if call_count > 1: + raise AssertionError("process_file must stream iter_imu_rows once") + return iter(rows) + + run_imu_ekf.iter_imu_rows = single_use_iter + try: + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=1) + finally: + run_imu_ekf.iter_imu_rows = original_iter + + self.assertEqual(call_count, 1) + self.assertEqual(result.input_rows, 3) + + def test_init_seconds_zero_disables_gyro_bias_initialization(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv(input_path, [(0.0, 7.5, 1.0), (0.1, 7.5, 1.0)]) + + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=0) + + with result.output_csv.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + + self.assertEqual(float(rows[0]["gyro_bias_z_dps"]), 0.0) + self.assertEqual(float(rows[1]["gyro_bias_z_dps"]), 0.0) + + def test_init_seconds_rejects_values_outside_integer_range(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + self._write_sample_csv(input_path, [(0.0, 0.0, 1.0)]) + + with self.assertRaisesRegex(ValueError, "init_seconds.*0.*10.*integer"): + run_imu_ekf.process_file(input_path, Path(tmp) / "out", init_seconds=11) + + with self.assertRaisesRegex(ValueError, "init_seconds.*0.*10.*integer"): + run_imu_ekf.process_file(input_path, Path(tmp) / "out", init_seconds=1.5) + + def test_relative_yaw_is_unwrapped_in_csv_and_html_samples(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv(input_path, [(float(index), 100.0, 1.0) for index in range(5)]) + + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=0) + + with result.output_csv.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + + self.assertLess(float(rows[2]["yaw_deg"]), 0.0) + self.assertGreater(float(rows[2]["relative_yaw_deg"]), 180.0) + self.assertGreater(float(rows[4]["relative_yaw_deg"]), 360.0) + self.assertIn("relative_yaw_deg", result.samples[0]) + + def test_process_file_segments_device_restart_near_zero(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv( + input_path, + [(9.9, 0.0, 1.0), (10.0, 0.0, 1.0), (0.002, 0.0, 1.0), (0.004, 0.0, 1.0)], + ) + + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=0) + + with result.output_csv.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + + self.assertEqual([row["segment_id"] for row in rows], ["0", "0", "1", "1"]) + self.assertEqual(float(rows[2]["dt_s"]), 0.0) + + def test_process_file_flushes_short_uninitialized_segment_before_restart(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv( + input_path, + [(9.9, 0.0, 1.0), (10.0, 0.0, 1.0), (0.002, 20.0, 1.0), (0.502, 20.0, 1.0), (1.002, 20.0, 1.0)], + ) + + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=1) + + with result.output_csv.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + + self.assertEqual(len(rows), 5) + self.assertEqual([row["segment_id"] for row in rows], ["0", "0", "1", "1", "1"]) + self.assertEqual(float(rows[2]["dt_s"]), 0.0) + + def test_process_file_initializes_each_segment_from_its_own_window(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv( + input_path, + [ + (0.0, 0.0, 1.0), + (0.5, 0.0, 1.0), + (1.1, 0.0, 1.0), + (0.002, 20.0, 1.0), + (0.502, 20.0, 1.0), + (1.002, 20.0, 1.0), + ], + ) + + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=1) + + with result.output_csv.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + + segment0_bias = [float(row["gyro_bias_z_dps"]) for row in rows if row["segment_id"] == "0"] + segment1_bias = [float(row["gyro_bias_z_dps"]) for row in rows if row["segment_id"] == "1"] + self.assertTrue(segment0_bias) + self.assertTrue(segment1_bias) + self.assertTrue(all(abs(value) < 0.01 for value in segment0_bias)) + self.assertAlmostEqual(segment1_bias[-1], 20.0, delta=0.01) + + def test_process_file_uses_restart_initialization_window_not_single_row(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + output_dir = Path(tmp) / "out" + self._write_sample_csv( + input_path, + [ + (0.0, 0.0, 1.0), + (0.5, 0.0, 1.0), + (1.1, 0.0, 1.0), + (0.002, 0.0, 1.0), + (0.502, 20.0, 1.0), + (1.002, 20.0, 1.0), + ], + ) + + result = run_imu_ekf.process_file(input_path, output_dir, init_seconds=1) + + with result.output_csv.open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + + segment1_bias = [float(row["gyro_bias_z_dps"]) for row in rows if row["segment_id"] == "1"] + self.assertAlmostEqual(segment1_bias[-1], 40.0 / 3.0, delta=0.01) + + def test_process_file_rejects_timestamp_drop_not_near_zero(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + self._write_sample_csv(input_path, [(10.0, 0.0, 1.0), (9.5, 0.0, 1.0)]) + + with self.assertRaisesRegex(ValueError, "data row 2.*10.0.*9.5"): + run_imu_ekf.process_file(input_path, Path(tmp) / "out", init_seconds=0) + + def test_process_file_rejects_small_nonzero_timestamp_drop(self): + with tempfile.TemporaryDirectory() as tmp: + input_path = Path(tmp) / "imu_sample.csv" + self._write_sample_csv(input_path, [(0.91, 0.0, 1.0), (0.90, 0.0, 1.0)]) + + with self.assertRaisesRegex(ValueError, "data row 2.*0.91.*0.9"): + run_imu_ekf.process_file(input_path, Path(tmp) / "out", init_seconds=0) + + def test_iter_imu_rows_rejects_nan_and_infinity(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bad.csv" + path.write_text( + "\n".join( + [ + "sensor_uptime_s,temp_c,acc_x_g,acc_y_g,acc_z_g,gyro_x_dps,gyro_y_dps,gyro_z_dps", + "0.0,28.0,0,0,nan,0,0,0", + "0.1,28.0,0,0,1,0,0,inf", + ] + ), + encoding="utf-8-sig", + ) + + with self.assertRaisesRegex(ValueError, "acc_z_g.*finite.*data row 1"): + list(run_imu_ekf.iter_imu_rows(path)) + + def test_json_for_script_rejects_nan_and_infinity(self): + with self.assertRaisesRegex(ValueError, "finite"): + run_imu_ekf._json_for_script([{"bad": math.nan}]) + with self.assertRaisesRegex(ValueError, "finite"): + run_imu_ekf._json_for_script([{"bad": math.inf}]) + + def test_html_report_is_self_contained_and_bounded(self): + result = run_imu_ekf.EkfFileResult( + input_csv=Path("imu_sample.csv"), + output_csv=Path("out/.csv"), + input_rows=3, + metadata={"odr": "", "close": ""}, + samples=[ + { + "sensor_uptime_s": 0.0, + "roll_deg": 0.0, + "pitch_deg": 0.0, + "yaw_deg": 0.0, + "relative_yaw_deg": 0.0, + "gyro_bias_x_dps": 0.0, + "gyro_bias_y_dps": 0.0, + "gyro_bias_z_dps": 0.0, + "acc_residual_norm": 0.0, + } + ], + ) + + with tempfile.TemporaryDirectory() as tmp: + html_path = Path(tmp) / "viewer.html" + run_imu_ekf.write_html_report([result], html_path, max_points=1) + html = html_path.read_text(encoding="utf-8") + + self.assertIn("", html.lower()) + self.assertIn("application/json", html) + self.assertNotIn("innerHTML", html) + self.assertNotIn("