"""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
from scripts import imu_static_calibrator as static_calibrator
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,
yaw_bias_seconds: int = 60,
static_correction_seconds: float = 2.0,
static_gyro_threshold_dps: float = 0.5,
static_acc_norm_tolerance_g: float = 0.2,
static_acc_stability_threshold_g: float = 0.02,
max_points: int = 2500,
) -> EkfFileResult:
init_seconds = _validate_init_seconds(init_seconds)
yaw_bias_seconds = _validate_yaw_bias_seconds(yaw_bias_seconds)
static_correction_seconds = _validate_static_correction_seconds(static_correction_seconds)
static_gyro_threshold_dps = _validate_positive_float(
static_gyro_threshold_dps, "static_gyro_threshold_dps"
)
static_acc_norm_tolerance_g = _validate_positive_float(
static_acc_norm_tolerance_g, "static_acc_norm_tolerance_g"
)
static_acc_stability_threshold_g = _validate_positive_float(
static_acc_stability_threshold_g, "static_acc_stability_threshold_g"
)
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",
"fixed_yaw_bias_z_dps",
"active_yaw_bias_z_dps",
"is_static",
"acc_residual_norm",
"dt_s",
"acc_update_used",
]
writer = csv.DictWriter(handle, fieldnames=fieldnames)
writer.writeheader()
segment_id = 0
state: ekf.ImuEkfState | None = None
static_state: static_calibrator.StaticCorrectionState | None = None
yaw_bias_buffer: list[ImuRow] = []
init_buffer: list[ImuRow] = []
fixed_yaw_bias_enabled = yaw_bias_seconds > 0
fixed_yaw_bias_z_dps: float | None = 0.0 if yaw_bias_seconds == 0 else None
active_yaw_bias_z_dps: float | None = fixed_yaw_bias_z_dps
static_correction_enabled = static_correction_seconds > 0.0
wrapper_yaw_bias_enabled = fixed_yaw_bias_enabled or static_correction_enabled
static_config = None
if static_correction_enabled:
static_config = static_calibrator.StaticCorrectionConfig(
enter_seconds=static_correction_seconds,
gyro_threshold_dps=static_gyro_threshold_dps,
acc_norm_tolerance_g=static_acc_norm_tolerance_g,
acc_stability_threshold_g=static_acc_stability_threshold_g,
)
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, static_state, yaw_bias_buffer, init_buffer
nonlocal fixed_yaw_bias_z_dps, active_yaw_bias_z_dps
nonlocal previous_output_time, relative_yaw_deg, previous_yaw_deg
state = None
static_state = None
yaw_bias_buffer = []
init_buffer = []
fixed_yaw_bias_z_dps = 0.0 if yaw_bias_seconds == 0 else None
active_yaw_bias_z_dps = fixed_yaw_bias_z_dps
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
nonlocal active_yaw_bias_z_dps
if state is None:
raise ValueError("EKF state is not initialized")
if fixed_yaw_bias_z_dps is None:
raise ValueError("fixed yaw bias is not initialized")
if active_yaw_bias_z_dps is None:
raise ValueError("active yaw bias is not initialized")
if wrapper_yaw_bias_enabled:
state.gyro_bias_rad_s[2] = 0.0
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
is_static = False
if static_state is not None:
bias_xy_dps = np.degrees(state.gyro_bias_rad_s[0:2])
is_static = static_calibrator.step(
static_state,
dt_s,
np.array(row.acc_g, dtype=float),
np.array(row.gyro_dps, dtype=float),
bias_xy_dps,
)
active_yaw_bias_z_dps = static_state.active_yaw_bias_z_dps
corrected_row = _row_with_yaw_bias(row, active_yaw_bias_z_dps)
used_update, residual_norm = ekf.step(
state, dt_s, _acc_mps2(corrected_row), _gyro_rad_s(corrected_row)
)
if wrapper_yaw_bias_enabled:
state.gyro_bias_rad_s[2] = 0.0
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
elif is_static:
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],
"fixed_yaw_bias_z_dps": fixed_yaw_bias_z_dps,
"active_yaw_bias_z_dps": active_yaw_bias_z_dps,
"is_static": int(is_static),
"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, static_state, init_buffer
if active_yaw_bias_z_dps is None:
raise ValueError("active yaw bias is not initialized")
corrected_init_rows = [
_row_with_yaw_bias(row, active_yaw_bias_z_dps) for row in init_buffer
]
state = _initialize_state(corrected_init_rows, init_seconds)
if wrapper_yaw_bias_enabled:
state.gyro_bias_rad_s[2] = 0.0
if static_config is not None:
static_state = static_calibrator.initialize(
static_config,
initial_yaw_bias_z_dps=active_yaw_bias_z_dps,
)
buffered_rows = init_buffer
init_buffer = []
for buffered_row in buffered_rows:
write_row(buffered_row)
def process_row_with_yaw_bias(row: ImuRow) -> None:
if active_yaw_bias_z_dps is None:
raise ValueError("active yaw bias is not initialized")
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)
def initialize_fixed_yaw_bias_from_buffer() -> None:
nonlocal yaw_bias_buffer, fixed_yaw_bias_z_dps, active_yaw_bias_z_dps
if fixed_yaw_bias_z_dps is not None:
return
if not yaw_bias_buffer:
raise ValueError("at least one IMU row is required for fixed yaw bias initialization")
fixed_yaw_bias_z_dps = _fixed_yaw_bias_z_dps(yaw_bias_buffer)
active_yaw_bias_z_dps = fixed_yaw_bias_z_dps
buffered_rows = yaw_bias_buffer
yaw_bias_buffer = []
for buffered_row in buffered_rows:
process_row_with_yaw_bias(buffered_row)
def flush_segment() -> None:
if fixed_yaw_bias_z_dps is None and yaw_bias_buffer:
initialize_fixed_yaw_bias_from_buffer()
if state is None and init_buffer:
initialize_and_write_buffer()
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):
flush_segment()
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 fixed_yaw_bias_z_dps is None:
if yaw_bias_buffer and row.sensor_uptime_s - yaw_bias_buffer[0].sensor_uptime_s >= yaw_bias_seconds:
initialize_fixed_yaw_bias_from_buffer()
process_row_with_yaw_bias(row)
else:
yaw_bias_buffer.append(row)
else:
process_row_with_yaw_bias(row)
previous_input_time = row.sensor_uptime_s
if rows_seen == 0:
raise ValueError(f"{input_csv} has no IMU rows")
flush_segment()
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
"""
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("--yaw-bias-seconds", type=_parse_yaw_bias_seconds, default=60)
parser.add_argument("--static-correction-seconds", type=_parse_static_correction_seconds, default=2.0)
parser.add_argument("--static-gyro-threshold-dps", type=_parse_positive_float, default=0.5)
parser.add_argument("--static-acc-norm-tolerance-g", type=_parse_positive_float, default=0.2)
parser.add_argument("--static-acc-stability-threshold-g", type=_parse_positive_float, default=0.02)
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,
yaw_bias_seconds=args.yaw_bias_seconds,
static_correction_seconds=args.static_correction_seconds,
static_gyro_threshold_dps=args.static_gyro_threshold_dps,
static_acc_norm_tolerance_g=args.static_acc_norm_tolerance_g,
static_acc_stability_threshold_g=args.static_acc_stability_threshold_g,
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 _validate_yaw_bias_seconds(value) -> int:
if type(value) is not int or value < 0:
raise ValueError("yaw_bias_seconds must be a non-negative integer")
return value
def _validate_static_correction_seconds(value) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError("static_correction_seconds must be a 0-10 second number")
value = float(value)
if not math.isfinite(value) or not 0.0 <= value <= 10.0:
raise ValueError("static_correction_seconds must be a 0-10 second number")
return value
def _validate_positive_float(value, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be a finite positive number")
value = float(value)
if not math.isfinite(value) or value <= 0.0:
raise ValueError(f"{name} must be a finite positive number")
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 _parse_yaw_bias_seconds(value: str) -> int:
try:
parsed = int(value)
except ValueError as exc:
raise argparse.ArgumentTypeError("yaw_bias_seconds must be a non-negative integer") from exc
if str(parsed) != value:
raise argparse.ArgumentTypeError("yaw_bias_seconds must be a non-negative integer")
try:
return _validate_yaw_bias_seconds(parsed)
except ValueError as exc:
raise argparse.ArgumentTypeError(str(exc)) from exc
def _parse_static_correction_seconds(value: str) -> float:
try:
parsed = float(value)
return _validate_static_correction_seconds(parsed)
except ValueError as exc:
raise argparse.ArgumentTypeError(str(exc)) from exc
def _parse_positive_float(value: str) -> float:
try:
parsed = float(value)
return _validate_positive_float(parsed, "value")
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 _fixed_yaw_bias_z_dps(rows: list[ImuRow]) -> float:
return sum(row.gyro_dps[2] for row in rows) / len(rows)
def _row_with_yaw_bias(row: ImuRow, yaw_bias_z_dps: float) -> ImuRow:
return ImuRow(
sensor_uptime_s=row.sensor_uptime_s,
temp_c=row.temp_c,
acc_g=row.acc_g,
gyro_dps=(row.gyro_dps[0], row.gyro_dps[1], row.gyro_dps[2] - yaw_bias_z_dps),
)
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",
"fixed_yaw_bias_z_dps",
"active_yaw_bias_z_dps",
"is_static",
"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())