rm3100/logger.py

#!/usr/bin/env python3
"""Log RM3100 magnetometer data over a CH347 USB-I2C adapter.

Run ./setup.sh first to install the udev rule and create the virtualenv, then:

    ./.venv/bin/python logger.py --duration 10
    ./.venv/bin/python logger.py --scan-only

The sampling loop is kept clear of everything that is not I2C traffic: raw
counts go onto a queue, and a writer thread does the unit conversion, CSV
formatting and console display. At the rates the sensor can reach (~540 Hz at
cycle count 50) that formatting work is otherwise the bottleneck, not the bus.
"""

import argparse
import csv
import math
import queue
import sys
import threading
import time
from datetime import datetime, timedelta, timezone

import ch347
import rm3100

CSV_FIELDS = [
    "timestamp_iso", "elapsed_s",
    "x_raw", "y_raw", "z_raw",
    "x_uT", "y_uT", "z_uT",
    "magnitude_uT",
]

BUS_SPEEDS = {
    20: ch347.SPEED_20KHZ,
    100: ch347.SPEED_100KHZ,
    400: ch347.SPEED_400KHZ,
    750: ch347.SPEED_750KHZ,
}

CONSOLE_REFRESH_S = 0.05
FLUSH_INTERVAL_S = 0.5

# Bounded so a stalled writer degrades predictably instead of exhausting memory
# on a long run. Far above the depth a healthy writer ever reaches.
QUEUE_MAX = 200_000

_SENTINEL = object()


def parse_args():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--cycle-count", type=int, default=rm3100.DEFAULT_CYCLE_COUNT,
                   help="cycle count per axis; lower is faster but noisier. "
                        "50 allows ~540 Hz, 200 allows ~150 Hz (default: %(default)s)")
    p.add_argument("--tmrc", type=lambda s: int(s, 0), default=0x96,
                   help="continuous-mode rate register, 0x92 (fastest) to 0x9D "
                        "(slowest) (default: 0x96, ~37 Hz)")
    p.add_argument("--duration", type=float, default=0.0,
                   help="seconds to log, or 0 to run until Ctrl-C (default: %(default)s)")
    p.add_argument("--output", default=None,
                   help="CSV output path (default: rm3100_<timestamp>.csv)")
    p.add_argument("--address", type=lambda s: int(s, 0), default=None,
                   help="I2C address, skipping the scan (default: autodetect)")
    p.add_argument("--bus-speed", type=int, choices=[20, 100, 400, 750], default=100,
                   help="I2C bus speed in kHz; 400 measured best for high sample "
                        "rates (default: %(default)s)")
    p.add_argument("--scan-only", action="store_true",
                   help="scan the bus, report what responded, and exit")
    return p.parse_args()


def find_sensor(bus, address):
    """Locate the RM3100, or exit with wiring guidance."""
    if address is not None:
        print(f"Using I2C address 0x{address:02x} (scan skipped)")
        return address

    candidates = [a for a in rm3100.RM3100.ADDRESSES if bus.probe(a)]
    if not candidates:
        print("No RM3100 responded at 0x20-0x23.", file=sys.stderr)
        others = bus.scan()
        if others:
            print("Other devices on the bus: "
                  + ", ".join(f"0x{a:02x}" for a in others), file=sys.stderr)
        else:
            print("Nothing responded anywhere on the bus. Check, in order:\n"
                  "  - I2CEN tied high (otherwise the chip stays in SPI mode)\n"
                  "  - the adapter's voltage jumper (the RM3100 wants ~3.3 V)\n"
                  "  - SDA/SCL not swapped\n"
                  "  - SDA and SCL pull-up resistors present", file=sys.stderr)
        sys.exit(1)

    if len(candidates) > 1:
        print("Multiple devices responded at "
              + ", ".join(f"0x{a:02x}" for a in candidates)
              + "; use --address to pick one.", file=sys.stderr)
        sys.exit(1)

    print(f"Found a device at 0x{candidates[0]:02x}")
    return candidates[0]


def start_sensor(sensor, cycle_count, tmrc):
    """Identify and configure the sensor, printing each step."""
    revid = sensor.revid()
    if revid == rm3100.EXPECTED_REVID:
        print(f"REVID 0x{revid:02x} -- RM3100 confirmed")
    else:
        print(f"WARNING: REVID 0x{revid:02x}, expected "
              f"0x{rm3100.EXPECTED_REVID:02x}", file=sys.stderr)

    sensor.set_cycle_counts(cycle_count)
    readback = sensor.get_cycle_counts()
    if readback != (cycle_count,) * 3:
        print(f"ERROR: cycle count read back as {readback}, expected "
              f"{(cycle_count,) * 3}", file=sys.stderr)
        sys.exit(1)
    gain = rm3100.gain_lsb_per_ut(cycle_count)
    print(f"Cycle counts set to {readback} -- gain {gain:.1f} LSB/uT")

    sensor.configure()
    sensor.set_rate(tmrc)
    rate = rm3100.TMRC_RATES[tmrc]
    print(f"Rate register 0x{tmrc:02x} -- requesting about {rate:g} Hz")

    sensor.start_cmm()
    print("Continuous measurement mode started")
    return rate


def writer_thread(q, path, gain, start_wall, stats):
    """Drain raw samples: convert, format, write CSV, drive the console.

    Everything here is deliberately off the sampling thread.
    """
    with open(path, "w", newline="") as handle:
        out = csv.writer(handle)
        out.writerow(CSV_FIELDS)
        last_print = 0.0
        last_flush = time.monotonic()

        while True:
            item = q.get()
            if item is _SENTINEL:
                break
            elapsed, cx, cy, cz = item

            ux, uy, uz = cx / gain, cy / gain, cz / gain
            magnitude = math.sqrt(ux * ux + uy * uy + uz * uz)
            # Wall-clock is reconstructed from the monotonic offset rather than
            # sampled per reading: one fewer syscall in the hot path, and immune
            # to NTP steps mid-capture.
            stamp = (start_wall + timedelta(seconds=elapsed)).isoformat()

            out.writerow([stamp, f"{elapsed:.5f}", cx, cy, cz,
                          f"{ux:.4f}", f"{uy:.4f}", f"{uz:.4f}",
                          f"{magnitude:.4f}"])
            stats["rows"] += 1

            now = time.monotonic()
            if now - last_flush >= FLUSH_INTERVAL_S:
                handle.flush()
                last_flush = now
            if now - last_print >= CONSOLE_REFRESH_S:
                last_print = now
                print(f"\r{elapsed:8.2f}s  X {ux:+9.3f}  Y {uy:+9.3f}  "
                      f"Z {uz:+9.3f}  |B| {magnitude:8.3f} uT  "
                      f"({stats['rows']} samples, "
                      f"{stats['rows'] / max(elapsed, 1e-9):.0f} Hz)",
                      end="", flush=True)
        handle.flush()


def sample_loop(sensor, q, duration, rate, stats):
    """Read the sensor as fast as it produces data. I2C and nothing else."""
    sample_timeout = max(2.0, 5.0 / rate)
    read_raw = sensor.read_raw
    data_ready = sensor.data_ready
    monotonic = time.monotonic
    put = q.put_nowait

    start = monotonic()
    timeouts = 0
    deadline = start + duration if duration > 0 else float("inf")

    while monotonic() < deadline:
        if not data_ready():
            if monotonic() - start > sample_timeout and stats["rows"] == 0:
                timeouts += 1
                print(f"\nWARNING: no data ready within {sample_timeout:.1f} s",
                      file=sys.stderr)
                if timeouts >= 3:
                    print("Giving up after 3 consecutive timeouts.", file=sys.stderr)
                    return
            continue

        cx, cy, cz = read_raw()
        try:
            put((monotonic() - start, cx, cy, cz))
        except queue.Full:
            stats["dropped"] += 1


def main():
    args = parse_args()

    if args.cycle_count < 1 or args.cycle_count > 0xFFFF:
        sys.exit(f"--cycle-count {args.cycle_count} outside 1..65535")
    if args.tmrc not in rm3100.TMRC_RATES:
        sys.exit(f"--tmrc 0x{args.tmrc:02x} is not a valid rate register value")

    try:
        bus = ch347.CH347I2C(BUS_SPEEDS[args.bus_speed])
    except IOError as exc:
        sys.exit(str(exc))

    with bus:
        print(f"CH347 adapter opened, I2C at {args.bus_speed} kHz")

        if args.scan_only:
            found = bus.scan()
            print("Devices found: " + ", ".join(f"0x{a:02x}" for a in found)
                  if found else "No devices responded on the bus.")
            return

        address = find_sensor(bus, args.address)
        sensor = rm3100.RM3100(bus, address)
        rate = start_sensor(sensor, args.cycle_count, args.tmrc)

        path = args.output or datetime.now().strftime("rm3100_%Y%m%d_%H%M%S.csv")
        print(f"Logging to {path} -- Ctrl-C to stop\n")

        q = queue.Queue(maxsize=QUEUE_MAX)
        stats = {"rows": 0, "dropped": 0}
        writer = threading.Thread(
            target=writer_thread,
            args=(q, path, rm3100.gain_lsb_per_ut(args.cycle_count),
                  datetime.now(timezone.utc), stats),
            daemon=True)
        writer.start()

        try:
            sample_loop(sensor, q, args.duration, rate, stats)
        except KeyboardInterrupt:
            pass
        finally:
            q.put(_SENTINEL)
            writer.join(timeout=30)
            print()
            try:
                sensor.stop_cmm()
                print("Continuous measurement mode stopped")
            except IOError as exc:
                print(f"WARNING: could not stop CMM: {exc}", file=sys.stderr)
            if stats["dropped"]:
                print(f"WARNING: dropped {stats['dropped']} samples -- the writer "
                      "could not keep up", file=sys.stderr)
            print(f"Wrote {stats['rows']} samples to {path}")


if __name__ == "__main__":
    main()