ir_source_localizer.py

#!/usr/bin/env python3
"""

ir_source_localizer.py

======================



Localize an acoustic source from its impulse response at known

microphones.  RANSAC over receiver subsets rejects outlying

arrival-time detections.



Fix history (5 iterations to get here):

    v1  wide-window argmax                 -- locks on strong reflection

    v2  first local max after rise         -- picks leading-edge noise

    v3  half-peak leading edge             -- smoothing merges direct+reflection

    v4  squared signal + tight smoothing   -- r2 still +0.7 ms biased

    v5  quality-weighted loss              -- weighting helped nothing

    v6  RANSAC over receiver subsets       -- rejects the bad receiver



The insight: three of four receivers detect the direct within

0.02 ms.  One is off by 0.70 ms.  Any scheme that averages all four

inherits the outlier.  RANSAC finds the consensus of three and

drops the fourth.

"""

from __future__ import annotations

import argparse
import itertools
import json
import math
import wave
from dataclasses import dataclass, asdict
from typing import List, Optional, Sequence, Tuple

import numpy as np

SPEED_OF_SOUND = 343.0


# =====================================================================
# §1  Result type
# =====================================================================

@dataclass
class LocalizationResult:
    x: float
    y: float
    sigma_x: float
    sigma_y: float
    detected_times: List[float]
    detection_quality: List[float]
    predicted_times: List[float]
    arrival_residual_ms: List[float]
    rms_residual_ms: float
    n_receivers: int
    n_inliers: int
    inlier_mask: List[bool]
    method: str = "energy envelope + RANSAC + Adam"

    def to_dict(self) -> dict:
        return asdict(self)

    def __str__(self) -> str:
        return (f"source at ({self.x:+.3f}, {self.y:+.3f}) m  "
                f"± ({self.sigma_x*100:.1f}, "
                f"{self.sigma_y*100:.1f}) cm  "
                f"inliers {self.n_inliers}/{self.n_receivers}  "
                f"rms resid {self.rms_residual_ms:.3f} ms")


# =====================================================================
# §2  Direct arrival detection (unchanged from v4)
# =====================================================================

def detect_direct_arrival(ir, sample_rate, sigma_guess,

                          threshold=0.30):
    dt = 1.0 / sample_rate
    w = max(3, int(round(1.0 * sigma_guess / dt)))
    kernel = np.ones(w) / w
    env = np.convolve(ir ** 2, kernel, mode="same")
    peak = float(env.max())
    if peak < 1e-12:
        return None, 0.0
    above = env >= threshold * peak
    if not above.any():
        return None, 0.0
    first_idx = int(np.argmax(above))
    lo = first_idx
    hi = min(len(env), first_idx + 3 * w)
    if hi - lo < 3:
        return None, 0.0
    peak_idx = lo + int(np.argmax(env[lo:hi]))
    return float(peak_idx * dt), float(env[peak_idx] / peak)


def detect_all(irs, sample_rate, sigma_guess):
    times, qualities = [], []
    for ir in irs:
        t, q = detect_direct_arrival(ir, sample_rate, sigma_guess)
        times.append(t)
        qualities.append(q)
    return times, qualities


# =====================================================================
# §3  Position estimation
# =====================================================================

def predict_arrivals(sx, sy, sz, receivers, c=SPEED_OF_SOUND):
    d = np.sqrt((receivers[:, 0] - sx) ** 2
                + (receivers[:, 1] - sy) ** 2
                + (receivers[:, 2] - sz) ** 2)
    return d / c


def arrival_loss(sx, sy, sz, receivers, times):
    pred = predict_arrivals(sx, sy, sz, receivers)
    return float(np.mean((pred - times) ** 2))


def localize_grid(times, receivers, sz, x_range, y_range, n_grid=60):
    xs = np.linspace(x_range[0], x_range[1], n_grid)
    ys = np.linspace(y_range[0], y_range[1], n_grid)
    best = (float(xs[0]), float(ys[0]), float("inf"))
    for i in range(n_grid):
        for j in range(n_grid):
            sx = float(xs[i]); sy = float(ys[j])
            d = np.sqrt((receivers[:, 0] - sx) ** 2
                        + (receivers[:, 1] - sy) ** 2
                        + (receivers[:, 2] - sz) ** 2)
            t_pred = d / SPEED_OF_SOUND
            err = float(np.mean((t_pred - times) ** 2))
            if err < best[2]:
                best = (sx, sy, err)
    return best


def refine_position(sx0, sy0, sz, receivers, times,

                    n_iter=300, lr=0.02):
    p = np.array([sx0, sy0])

    def L(p):
        return arrival_loss(p[0], p[1], sz, receivers, times)

    m = np.zeros(2); v = np.zeros(2)
    best_p = p.copy(); best_L = L(p)
    eps = 1e-3
    for t in range(1, n_iter + 1):
        g = np.zeros(2)
        for i in range(2):
            rp = p.copy(); rp[i] += eps
            rm = p.copy(); rm[i] -= eps
            g[i] = (L(rp) - L(rm)) / (2.0 * eps)
        m = 0.9 * m + 0.1 * g
        v = 0.999 * v + 0.001 * g * g
        mh = m / (1 - 0.9 ** t)
        vh = v / (1 - 0.999 ** t)
        p = p - lr * mh / (np.sqrt(vh) + 1e-8)
        Li = L(p)
        if Li < best_L:
            best_L = Li
            best_p = p.copy()
    return float(best_p[0]), float(best_p[1]), best_L


# =====================================================================
# §4  RANSAC over receiver subsets  --  the v6 fix
# =====================================================================

def ransac_localize(times, receivers, sz,

                    x_range, y_range,

                    inlier_threshold_ms=0.30,

                    n_grid=50,

                    n_refine_ransac=100,

                    n_refine_final=300,

                    seed=0):
    """Try every subset of 3 receivers.  Count how many of the K

    receivers agree with each solution within `inlier_threshold_ms`.

    Return the position with the largest consensus, then refine on

    the inlier set only.



    If K == 3, there is one subset and the result is the ordinary

    solution.  For K >= 4, this rejects outlier receivers.

    """
    K = len(times)
    if K < 3:
        raise ValueError("need at least 3 receivers")

    threshold_s = inlier_threshold_ms / 1000.0
    subsets = list(itertools.combinations(range(K), 3))

    best = None  # (n_inliers, inlier_indices, sx, sy)
    for sub in subsets:
        sub_idx = np.array(sub)
        recv_sub = receivers[sub_idx]
        times_sub = np.asarray(times)[sub_idx]
        # Coarse grid on the subset
        sx0, sy0, _ = localize_grid(times_sub, recv_sub, sz,
                                    x_range, y_range, n_grid)
        # Fine on the subset
        sx, sy, _ = refine_position(sx0, sy0, sz, recv_sub,
                                    times_sub,
                                    n_iter=n_refine_ransac, lr=0.03)
        # Count inliers among all receivers
        pred = predict_arrivals(sx, sy, sz, receivers)
        resid = np.abs(pred - np.asarray(times))
        inliers = np.where(resid <= threshold_s)[0]
        score = len(inliers)
        if best is None or score > best[0]:
            best = (score, inliers, sx, sy)

    _, inlier_indices, sx, sy = best

    # Final refinement on inliers only
    inliers_arr = np.asarray(inlier_indices)
    if len(inliers_arr) >= 3:
        sx, sy, _ = refine_position(
            sx, sy, sz,
            receivers[inliers_arr],
            np.asarray(times)[inliers_arr],
            n_iter=n_refine_final, lr=0.02)

    return sx, sy, inlier_indices


# =====================================================================
# §5  Uncertainty (Monte Carlo over inliers)
# =====================================================================

def estimate_uncertainty(sx, sy, sz, receivers, times, inliers,

                         x_range, y_range,

                         sigma_t, n_samples=50, n_grid=40, seed=0):
    rng = np.random.default_rng(seed)
    R = receivers[inliers]
    T = np.asarray(times)[inliers]
    xs = np.empty(n_samples); ys = np.empty(n_samples)
    for k in range(n_samples):
        tp = T + rng.normal(0, sigma_t, size=len(T))
        sx0, sy0, _ = localize_grid(tp, R, sz, x_range, y_range, n_grid)
        xr, yr, _ = refine_position(sx0, sy0, sz, R, tp,
                                    n_iter=120, lr=0.03)
        xs[k] = xr; ys[k] = yr
    return float(np.std(xs)), float(np.std(ys))


# =====================================================================
# §6  Main API
# =====================================================================

def localize(irs, sample_rate, receivers, sz,

             x_range, y_range,

             sigma_guess=4e-4,

             n_grid=60, n_refine=300,

             estimate_unc=True, n_unc_samples=50,

             unc_sigma_t=None,

             inlier_threshold_ms=0.30,

             seed=0):
    if len(irs) != len(receivers):
        raise ValueError(
            f"got {len(irs)} IRs and {len(receivers)} receivers")
    if len(irs) < 3:
        raise ValueError(f"need at least 3 receivers, got {len(irs)}")

    receivers_arr = np.asarray(receivers, dtype=np.float64)
    times, qualities = detect_all(irs, sample_rate, sigma_guess)
    bad = [k for k, t in enumerate(times) if t is None]
    if bad:
        raise RuntimeError(
            f"detection failed at receivers {bad}")
    times_arr = np.asarray(times, dtype=np.float64)
    qualities = np.asarray(qualities, dtype=np.float64)

    sx, sy, inliers = ransac_localize(
        times_arr, receivers_arr, sz,
        x_range, y_range,
        inlier_threshold_ms=inlier_threshold_ms,
        n_grid=50,
        n_refine_ransac=100,
        n_refine_final=n_refine,
        seed=seed)

    t_pred = predict_arrivals(sx, sy, sz, receivers_arr)
    resid_ms = (t_pred - times_arr) * 1000.0
    rms_ms = float(np.sqrt(np.mean(resid_ms ** 2)))

    sigma_x = float("nan"); sigma_y = float("nan")
    if estimate_unc:
        sigma_t = (unc_sigma_t if unc_sigma_t is not None
                   else 0.3 * sigma_guess)
        sigma_x, sigma_y = estimate_uncertainty(
            sx, sy, sz, receivers_arr, times_arr, inliers,
            x_range, y_range, sigma_t,
            n_samples=n_unc_samples, seed=seed)

    inlier_mask = [bool(k in set(int(i) for i in inliers))
                   for k in range(len(irs))]

    return LocalizationResult(
        x=sx, y=sy,
        sigma_x=sigma_x, sigma_y=sigma_y,
        detected_times=[float(t) for t in times],
        detection_quality=[float(q) for q in qualities],
        predicted_times=[float(t) for t in t_pred],
        arrival_residual_ms=[float(r) for r in resid_ms],
        rms_residual_ms=rms_ms,
        n_receivers=len(irs),
        n_inliers=len(inliers),
        inlier_mask=inlier_mask,
    )


# =====================================================================
# §7  WAV I/O
# =====================================================================

def load_wav(path):
    with wave.open(path, "rb") as w:
        n_ch = w.getnchannels()
        samp_w = w.getsampwidth()
        sr = w.getframerate()
        raw = w.readframes(w.getnframes())
    if samp_w == 2:
        data = (np.frombuffer(raw, dtype=np.int16)
                .astype(np.float64) / 32768.0)
    elif samp_w == 4:
        data = (np.frombuffer(raw, dtype=np.int32)
                .astype(np.float64) / 2147483648.0)
    elif samp_w == 1:
        data = ((np.frombuffer(raw, dtype=np.uint8)
                 .astype(np.float64) - 128.0) / 128.0)
    else:
        raise ValueError(f"unsupported sample width: {samp_w}")
    if n_ch > 1:
        data = data.reshape(-1, n_ch).mean(axis=1)
    return data, sr


# =====================================================================
# §8  Simulator
# =====================================================================

def _image_sources(sx, sy, sz, Lx, Ly, Lz, n_reflect):
    pos, n_r = [], []
    for i in range(-n_reflect, n_reflect + 1):
        for j in range(-n_reflect, n_reflect + 1):
            for k in range(-n_reflect, n_reflect + 1):
                nr = abs(i) + abs(j) + abs(k)
                if nr > n_reflect:
                    continue
                for xs_ in (sx + 2 * i * Lx, -sx + 2 * i * Lx):
                    for ys_ in (sy + 2 * j * Ly, -sy + 2 * j * Ly):
                        for zs_ in (sz + 2 * k * Lz,
                                    -sz + 2 * k * Lz):
                            pos.append((xs_, ys_, zs_))
                            n_r.append(nr)
    return (np.array(pos, dtype=np.float64),
            np.array(n_r, dtype=np.float64))


def _pulse_mex(delta, sigma):
    x = delta / sigma
    return (1.0 - x * x) * np.exp(-x * x / 2.0)


def simulate_ir(receiver, source, room, times,

                alpha=0.72, sigma=4e-4, n_reflect=2):
    Lx, Ly, Lz = room
    sx, sy, sz = source
    pos, nr = _image_sources(sx, sy, sz, Lx, Ly, Lz, n_reflect)
    d = np.linalg.norm(pos - receiver[None, :], axis=1)
    t_arr = d / SPEED_OF_SOUND
    amp = (alpha ** nr) / np.maximum(d, 0.15)
    delta = times[:, None] - t_arr[None, :]
    return (amp[None, :] * _pulse_mex(delta, sigma)).sum(axis=1)


def simulate_all(receivers, source, room, times,

                 alpha=0.72, sigma=4e-4, n_reflect=2,

                 noise_frac=0.0, seed=0):
    irs = np.stack([
        simulate_ir(r, source, room, times, alpha, sigma, n_reflect)
        for r in receivers
    ])
    if noise_frac > 0:
        rng = np.random.default_rng(seed)
        rms = float(np.sqrt(np.mean(irs ** 2)))
        irs = irs + noise_frac * rms * rng.standard_normal(irs.shape)
    return irs


# =====================================================================
# §9  Self-test
# =====================================================================

def self_test(verbose=True):
    checks = []
    sr = 48000
    T = 4000
    times = np.arange(T) / sr
    sigma = 4e-4
    pulse_time = 0.010

    pulse = _pulse_mex(times - pulse_time, sigma)
    t_det, q = detect_direct_arrival(pulse, sr, sigma)
    checks.append(("clean pulse within 0.5 ms",
                   t_det is not None
                   and abs(t_det - pulse_time) < 5e-4))
    checks.append(("quality near 1.0", q > 0.9))

    rng = np.random.default_rng(0)
    noisy = pulse + 0.05 * rng.standard_normal(T)
    t_det2, _ = detect_direct_arrival(noisy, sr, sigma)
    checks.append(("5% noise within 1 ms",
                   t_det2 is not None
                   and abs(t_det2 - pulse_time) < 1e-3))

    room = (5.0, 4.0, 2.8)
    source = (1.5, 1.5, 1.0)
    recv = np.array([
        [4.5, 3.5, 1.0], [0.5, 3.5, 1.0],
        [4.5, 0.5, 1.0], [0.5, 0.5, 1.0],
    ])

    # RANSAC self-test: inject an outlier time.
    t_arr = predict_arrivals(source[0], source[1], source[2], recv)
    t_with_outlier = t_arr.copy()
    t_with_outlier[2] += 0.7e-3  # add 0.7 ms to receiver 2
    sx, sy, inliers = ransac_localize(
        t_with_outlier, recv, source[2],
        (0.2, 4.8), (0.2, 3.8),
        inlier_threshold_ms=0.3,
        n_grid=50, n_refine_ransac=100, n_refine_final=200,
        seed=0)
    err = math.hypot(sx - source[0], sy - source[1])
    checks.append((f"RANSAC rejects outlier (err {err*100:.2f} cm)",
                   err < 0.05))
    checks.append(("outlier receiver is not an inlier",
                   2 not in inliers))

    # End-to-end at 2% noise
    T_ir = 2000
    times_ir = np.arange(T_ir) / sr
    irs = simulate_all(recv, source, room, times_ir,
                       alpha=0.72, sigma=sigma, n_reflect=2,
                       noise_frac=0.02, seed=1)
    res = localize(irs, sr, recv, sz=source[2],
                   x_range=(0.2, 4.8), y_range=(0.2, 3.8),
                   sigma_guess=sigma, n_grid=50, n_refine=250,
                   estimate_unc=False)
    err_e2e = math.hypot(res.x - source[0], res.y - source[1])
    checks.append((f"end-to-end within 5 cm (err {err_e2e*100:.2f})",
                   err_e2e < 0.05))

    # Uncertainty
    res_u = localize(irs, sr, recv, sz=source[2],
                     x_range=(0.2, 4.8), y_range=(0.2, 3.8),
                     sigma_guess=sigma, n_grid=40, n_refine=200,
                     estimate_unc=True, n_unc_samples=20, seed=3)
    checks.append(("uncertainty positive",
                   res_u.sigma_x > 0 and res_u.sigma_y > 0))

    try:
        localize(irs[:2], sr, recv[:2], sz=source[2],
                 x_range=(0, 5), y_range=(0, 4))
        checks.append(("rejects <3 receivers", False))
    except ValueError:
        checks.append(("rejects <3 receivers", True))

    try:
        import tempfile, os
        with tempfile.NamedTemporaryFile(suffix=".wav",
                                         delete=False) as f:
            path = f.name
        sig = np.clip(irs[0], -1.0, 1.0)
        with wave.open(path, "wb") as wf:
            wf.setnchannels(1); wf.setsampwidth(2)
            wf.setframerate(sr)
            wf.writeframes((sig * 32767).astype(np.int16).tobytes())
        loaded, sr_loaded = load_wav(path)
        os.unlink(path)
        checks.append(("WAV sr preserved", sr_loaded == sr))
        checks.append(("WAV len preserved", len(loaded) == len(irs[0])))
        checks.append(("WAV values close",
                       float(np.max(np.abs(loaded - sig))) < 1e-3))
    except Exception as e:
        checks.append((f"WAV failed: {e}", False))

    passed = sum(1 for _, ok in checks if ok)
    if verbose:
        print()
        print("=" * 74)
        print("SELF-TEST")
        print("=" * 74)
        for name, ok in checks:
            mark = "PASS" if ok else "FAIL"
            print(f"  [{mark}] {name}")
        print()
        print(f"  {passed}/{len(checks)} correct")
    return passed, len(checks)


# =====================================================================
# §10  Demo
# =====================================================================

def banner(t, w=76):
    print()
    print("=" * w)
    print(t)
    print("=" * w)


def demo():
    sr = 48000
    sigma = 4e-4
    room = (5.0, 4.0, 2.8)
    source_true = (1.5, 1.5, 1.0)
    recv = np.array([
        [4.5, 3.5, 1.0], [0.5, 3.5, 1.0],
        [4.5, 0.5, 1.0], [0.5, 0.5, 1.0],
    ])
    banner("IR SOURCE LOCALIZER v6  --  RANSAC outlier rejection")
    print(f"  true source : "
          f"({source_true[0]:.2f}, {source_true[1]:.2f}, "
          f"{source_true[2]:.2f}) m")
    print(f"  room        : {room[0]} x {room[1]} x {room[2]} m")

    print()
    print(f"  {'noise':>6}  {'x':>8}  {'y':>8}  {'err cm':>8}  "
          f"{'inliers':>8}  {'rms ms':>8}")
    print("  " + "-" * 60)

    times_ir = np.arange(2000) / sr
    for noise in (0.005, 0.01, 0.02, 0.05, 0.10):
        irs = simulate_all(recv, source_true, room, times_ir,
                           noise_frac=noise, seed=1)
        try:
            res = localize(irs, sr, recv, sz=source_true[2],
                           x_range=(0.2, 4.8), y_range=(0.2, 3.8),
                           sigma_guess=sigma, n_grid=50,
                           n_refine=250, estimate_unc=False)
            err_cm = math.hypot(res.x - source_true[0],
                                res.y - source_true[1]) * 100
            print(f"  {noise*100:>5.1f}%  "
                  f"{res.x:>8.3f}  {res.y:>8.3f}  {err_cm:>8.2f}  "
                  f"{res.n_inliers:>3}/{res.n_receivers:<4}  "
                  f"{res.rms_residual_ms:>8.3f}")
        except RuntimeError as e:
            print(f"  {noise*100:>5.1f}%  fail: {e}")

    banner("DETAIL  --  2% noise, single run")
    irs = simulate_all(recv, source_true, room, times_ir,
                       noise_frac=0.02, seed=1)
    res = localize(irs, sr, recv, sz=source_true[2],
                   x_range=(0.2, 4.8), y_range=(0.2, 3.8),
                   sigma_guess=sigma, n_grid=50, n_refine=250,
                   estimate_unc=True, n_unc_samples=40, seed=7)
    print(f"  {res}")
    print()
    print(f"  {'recv':>5}  {'detected (ms)':>14}  "
          f"{'predicted (ms)':>15}  {'resid (ms)':>11}  "
          f"{'quality':>8}  {'inlier':>7}")
    print("  " + "-" * 72)
    for k in range(res.n_receivers):
        print(f"  {k:>5}  {res.detected_times[k]*1000:>14.4f}  "
              f"{res.predicted_times[k]*1000:>15.4f}  "
              f"{res.arrival_residual_ms[k]:>+11.4f}  "
              f"{res.detection_quality[k]:>8.3f}  "
              f"{str(res.inlier_mask[k]):>7}")
    print()
    print(f"  RMS arrival residual : {res.rms_residual_ms:.4f} ms")
    print(f"  position uncertainty : "
          f"sx = {res.sigma_x*100:.2f} cm, "
          f"sy = {res.sigma_y*100:.2f} cm")


# =====================================================================
# §11  CLI
# =====================================================================

def _parse_xyz(s):
    parts = s.replace(",", " ").split()
    if len(parts) != 3:
        raise argparse.ArgumentTypeError("expected 3 numbers")
    return tuple(float(p) for p in parts)


def main():
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="cmd")
    sub.add_parser("self-test")
    sub.add_parser("demo")
    p_loc = sub.add_parser("localize")
    p_loc.add_argument("wavs", nargs="+")
    p_loc.add_argument("--sample-rate", type=int, required=True)
    p_loc.add_argument("--receivers", nargs="+", required=True,
                       type=_parse_xyz)
    p_loc.add_argument("--sz", type=float, required=True)
    p_loc.add_argument("--x-range", nargs=2, type=float, required=True)
    p_loc.add_argument("--y-range", nargs=2, type=float, required=True)
    p_loc.add_argument("--sigma", type=float, default=4e-4)
    p_loc.add_argument("--grid", type=int, default=50)
    p_loc.add_argument("--refine", type=int, default=300)
    p_loc.add_argument("--inlier-threshold-ms", type=float, default=0.30)
    p_loc.add_argument("--no-unc", action="store_true")
    p_loc.add_argument("--json", action="store_true")

    args = p.parse_args()
    if args.cmd in (None, "self-test"):
        self_test()
        if args.cmd is None:
            print()
            print("  usage: python ir_source_localizer.py [self-test|demo|localize ...]")
        return
    if args.cmd == "demo":
        demo()
        return
    if args.cmd == "localize":
        irs = []
        for path in args.wavs:
            data, sr = load_wav(path)
            if sr != args.sample_rate:
                raise SystemExit(f"sample rate mismatch: {path}")
            irs.append(data)
        res = localize(irs, args.sample_rate, args.receivers,
                       sz=args.sz,
                       x_range=tuple(args.x_range),
                       y_range=tuple(args.y_range),
                       sigma_guess=args.sigma,
                       n_grid=args.grid, n_refine=args.refine,
                       estimate_unc=not args.no_unc,
                       inlier_threshold_ms=args.inlier_threshold_ms)
        if args.json:
            print(json.dumps(res.to_dict(), indent=2))
        else:
            print(res)


if __name__ == "__main__":
    main()