rm3100/tests/synthetic.py

"""Write synthetic captures, so tests never need hardware or a recorded file.

Shared by test_capture.py, which uses it to break one guarantee at a time, and
test_compare.py, which uses it to plant a known answer -- a scale factor, a
rotation, a tone at a chosen fraction of the sample rate -- and check that the
analysis recovers it.

Defaults produce a clean, loadable capture; every argument exists to change one
thing about it.
"""

import numpy as np

import capture
import rm3100

CYCLE_COUNT = 100
DT = 1.0 / 250.0          # a deliberately non-nominal true period


def write_capture(path, rows=200, dt=DT, nominal_hz=282.0, header=None,
                  flags=None, drop_header_key=None, extra_lines=(),
                  index_from=0, index_step=1, start_time=1_700_000_000.0,
                  amplitude=1000.0, seed=0, cycle_count=CYCLE_COUNT,
                  counts=None, times=None):
    """Write a synthetic capture and return its path.

    `counts` overrides the generated signal with an (rows, 3) array of raw
    counts, which is how a test plants an exact answer. `cycle_count` moves the
    header's gain and is what makes a decimation pair possible: two captures
    whose cycle counts differ by an integer factor.

    `times` overrides the timestamps with absolute unix seconds, one per row.
    Without it the grid is exactly uniform, which is the right default but makes
    a whole class of question untestable: the chip's oscillator drifts through a
    real run, and the analysis that measures the drift needs a capture where the
    answer was planted. `drifting_times` below builds the usual case.
    """
    meta = {
        "rm3100_capture": 1,
        "nominal_rate_hz": nominal_hz,
        "tmrc_nominal_hz": 600.0,
        "tmrc": "0x92",
        "cycle_count": cycle_count,
        "tesla_per_count": repr(rm3100.tesla_per_count(cycle_count)),
        "i2c_address": "0x23",
        "bus_speed_khz": 750,
        "revid": "0x22",
        "calibrated_period_s": repr(dt),
    }
    meta.update(header or {})
    if drop_header_key:
        meta.pop(drop_header_key, None)

    rng = np.random.default_rng(seed)
    flags = flags or {}
    lines = [f"# {k}: {v}" for k, v in meta.items()]
    lines += list(extra_lines)
    lines.append("sample_index,system_time_unix,x_raw,y_raw,z_raw,warning")
    for i in range(rows):
        index = index_from + i * index_step
        warning = flags.get(i, "")
        if capture.WARN_MISSED in warning:
            x = y = z = 0
        elif counts is not None:
            x, y, z = (int(round(c)) for c in counts[i])
        else:
            x = int(amplitude + rng.normal(0, 3))
            y = int(2 * amplitude + rng.normal(0, 3))
            z = int(-amplitude + rng.normal(0, 3))
        stamp = start_time + index * dt if times is None else times[i]
        lines.append(f"{index},{stamp:.6f},{x},{y},{z},{warning}")
    path.write_text("\n".join(lines) + "\n")
    return str(path)


def drifting_times(rows, dt, ppm_per_second=0.0, jitter_s=0.0, seed=0,
                   start_time=1_700_000_000.0):
    """Timestamps for a capture whose sample period ramps linearly.

    The period at sample k is dt * (1 + ppm_per_second * 1e-6 * t), so the times
    are the integral of that -- quadratic in k, which is what a warming
    oscillator actually produces and what a straight-line fit of time against
    index cannot absorb.

    `jitter_s` adds independent noise to each timestamp without moving the
    underlying grid, standing in for host scheduling: it is what the drift has
    to be measured through, and it must not be mistaken for drift.
    """
    k = np.arange(rows)
    t = k * dt + 0.5 * ppm_per_second * 1e-6 * dt * k * (k - 1) * dt
    if jitter_s:
        t = t + np.random.default_rng(seed).normal(0, jitter_s, rows)
    return start_time + t


def field_counts(rows, mean_nt, cycle_count, noise_nt=0.0, seed=0, tones=()):
    """Raw counts for a field of a given mean, noise and planted tones.

    `mean_nt` is an (x, y, z) field in nanotesla; `tones` is a sequence of
    (axis_index, cycles_per_sample, amplitude_nT) added on top. Quantisation to
    integer counts is deliberate -- it is what the real file carries, and a test
    that skipped it would not exercise the dither the analysis relies on.
    """
    lsb = rm3100.tesla_per_count(cycle_count) * rm3100.NT_PER_TESLA
    rng = np.random.default_rng(seed)
    n = np.arange(rows)
    nt = np.tile(np.asarray(mean_nt, dtype=float), (rows, 1))
    if noise_nt:
        nt += rng.normal(0, noise_nt, size=(rows, 3))
    for axis, cycles_per_sample, amplitude in tones:
        nt[:, axis] += amplitude * np.cos(2 * np.pi * cycles_per_sample * n)
    return nt / lsb