Stub Inscopix .imu (9 channels, first 100 samples per stream; session timingInfo numTimes patched)

generate_imu_stub.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "numpy",
# ]
# ///
"""
Stub an Inscopix raw .imu file down to its first N samples.

Raw .imu layout (isxcore `IMUFile`, structs in `api/isxIMUFile.h`):

    [ IMUHeader (112 B) ]
    [ AccPayload[accCount] @ accOffset ]
    [ MagPayload[magCount] @ magOffset ]   (often empty)
    [ OriPayload[oriCount] @ oriOffset ]
    [ inscopix-session JSON @ sessionOffset ]   (no length trailer)

IMUHeader is packed (no padding): two uint64 epoch num/den, int32 utcOffset,
uint16 fileFormat + reserved, then {count, offset, size} uint64 triples for
acc / mag / ori, then sessionOffset + sessionSize. Each payload is 16 bytes:
uint64 timeStamp (ms), int16 data[3] (xyz / yaw-pitch-roll), int16 reserved.

Two things must be edited together to make a valid stub:
  1. the header counts/offsets/sizes, and
  2. the session JSON `timingInfo` per-stream `numTimes` and `end` time.
The reader (`IMUFile::parse`) takes each trace's length from the JSON
`timingInfo.<stream>.numTimes`, NOT from the header, so truncating the data
without fixing the JSON makes the read abort on an internal assertion.

Verified by round-tripping the output through `isx.GpioSet.read`.

Usage:
    uv run generate_imu_stub.py SOURCE.imu OUTPUT.imu [--n 100]

See vault notes: inscopix_gpio_imu_format, inscopix_gpio_imu_format_evidence.
"""
import argparse
import json

import numpy as np

IMU_HEADER = np.dtype(
    [
        ("epochNum", "<u8"),
        ("epochDen", "<u8"),
        ("utcOffset", "<i4"),
        ("fileFormat", "<u2"),
        ("reserved", "<u2"),
        ("accCount", "<u8"),
        ("accOffset", "<u8"),
        ("accSize", "<u8"),
        ("magCount", "<u8"),
        ("magOffset", "<u8"),
        ("magSize", "<u8"),
        ("oriCount", "<u8"),
        ("oriOffset", "<u8"),
        ("oriSize", "<u8"),
        ("sessionOffset", "<u8"),
        ("sessionSize", "<u8"),
    ]
)  # itemsize 112

IMU_PAYLOAD = np.dtype([("timeStamp", "<u8"), ("data", "<i2", 3), ("reserved", "<i2")])  # itemsize 16

# JSON `timingInfo` key for each header stream
STREAM_JSON_KEY = {"acc": "accelerometer", "mag": "magnetometer", "ori": "orientation"}


def _read_block(raw, header, stream, n):
    offset = int(header[f"{stream}Offset"])
    size = int(header[f"{stream}Size"])
    block = np.frombuffer(raw[offset : offset + size], dtype=IMU_PAYLOAD)
    return np.array(block[:n])  # copy so it survives reuse


def stub_imu(source_path, output_path, n):
    raw = open(source_path, "rb").read()
    header = np.frombuffer(raw[: IMU_HEADER.itemsize], dtype=IMU_HEADER)[0]

    acc = _read_block(raw, header, "acc", n)
    mag = _read_block(raw, header, "mag", n)
    ori = _read_block(raw, header, "ori", n)

    # Patch the session JSON timing to match the truncated counts.
    session = json.loads(raw[int(header["sessionOffset"]) :])
    timing = session["timingInfo"]
    start_ms = timing["start"]["secsSinceEpoch"]["num"]  # denominator is 1000
    longest = 0
    for stream, block in (("acc", acc), ("mag", mag), ("ori", ori)):
        key = STREAM_JSON_KEY[stream]
        if key not in timing:
            continue
        count = len(block)
        timing[key]["numTimes"] = count
        timing[key]["dropped"] = [d for d in timing[key].get("dropped", []) if d < count]
        longest = max(longest, count)
    # Bring the global end time back to the truncated duration (acc/ori share the base rate).
    period_ms = timing[STREAM_JSON_KEY["acc"]]["periodMs"]
    step_ms = period_ms["num"] / period_ms["den"]
    timing["end"]["secsSinceEpoch"] = {"num": int(round(start_ms + longest * step_ms)), "den": 1000}
    session_bytes = json.dumps(session, separators=(",", ":")).encode("utf-8")

    # Rebuild the header with new counts/offsets/sizes; blocks laid out acc, mag, ori.
    new_header = np.zeros(1, dtype=IMU_HEADER)[0]
    for field in ("epochNum", "epochDen", "utcOffset", "fileFormat"):
        new_header[field] = header[field]
    payload_size = IMU_PAYLOAD.itemsize
    cursor = IMU_HEADER.itemsize
    for stream, block in (("acc", acc), ("mag", mag), ("ori", ori)):
        new_header[f"{stream}Count"] = len(block)
        new_header[f"{stream}Offset"] = cursor
        new_header[f"{stream}Size"] = len(block) * payload_size
        cursor += len(block) * payload_size
    new_header["sessionOffset"] = cursor
    new_header["sessionSize"] = len(session_bytes)

    with open(output_path, "wb") as out:
        out.write(new_header.tobytes())
        out.write(acc.tobytes())
        out.write(mag.tobytes())
        out.write(ori.tobytes())
        out.write(session_bytes)
    return output_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", help="path to the full raw .imu file")
    parser.add_argument("output", help="path to write the stubbed .imu file")
    parser.add_argument("--n", type=int, default=100, help="number of samples per stream to keep")
    args = parser.parse_args()
    path = stub_imu(args.source, args.output, args.n)
    print(f"wrote {path}")