Preprocessing API¶

def align(
    envelopes: dict[str, np.ndarray],
    *,
    hop: float = LANDMARK_HOP,
    min_quality: float = LANDMARK_MIN_VOTES,
    reference: str | None = None,
    pairwise: Callable[
        [np.ndarray, np.ndarray, float], tuple[float, float]
    ] = landmark_offset,
) -> Alignment:
    """Solve every input's offset from the envelopes, using all pairs at once.

    ``reference`` is the input left at offset zero, defaulting to the first;
    which one is chosen only shifts the whole timeline, it does not change the
    inputs' positions relative to each other.
    """
    ids = list(envelopes)
    return solve_offsets(
        ids,
        measure_pairs(envelopes, hop, pairwise),
        min_quality=min_quality,
        reference=reference,
    )

def align_media(
    paths: dict[str, str | Path],
    *,
    reference: str | None = None,
    progress: Callable[[float], bool] | None = None,
) -> Alignment:
    """Offsets for a set of recordings, keyed the way the inputs are."""
    features = {}
    for index, (name, path) in enumerate(paths.items()):
        values = landmark_features(path)
        if len(values) == 0:
            logger.warning("input %r has no audio to align on; skipping", name)
        else:
            features[name] = values
        # Reading is nearly all of the work, so it gets nearly all of the bar.
        if progress is not None and progress(0.95 * (index + 1) / len(paths)) is False:
            return Alignment(offsets={})
    missing = [name for name in paths if name not in features]
    if reference is not None and reference not in features:
        logger.warning("reference input %r has no audio to align on", reference)
        return Alignment(offsets={}, unaligned=list(paths))
    result = align(features, reference=reference)
    result.unaligned.extend(name for name in missing if name not in result.unaligned)
    if progress is not None:
        progress(1.0)
    return result

def landmark_features(
    media_path: str | Path,
    sample_rate: int = LANDMARK_SAMPLE_RATE,
) -> np.ndarray:
    """Sparse ``(frame, hash)`` fingerprints for blind alignment."""
    from scipy.ndimage import maximum_filter

    samples = audio_samples(media_path, sample_rate)
    if len(samples) < LANDMARK_FFT:
        return np.empty((0, 2), dtype=np.int64)

    frame_count = 1 + (len(samples) - LANDMARK_FFT) // LANDMARK_HOP_SAMPLES
    window = np.hanning(LANDMARK_FFT).astype(np.float32)
    peaks: list[tuple[int, int, float]] = []
    chunk_frames = 4096
    time_radius = 2
    low_bin, high_bin = 3, LANDMARK_FFT // 2 - 2

    for core_start in range(0, frame_count, chunk_frames):
        core_end = min(core_start + chunk_frames, frame_count)
        ext_start = max(0, core_start - time_radius)
        ext_end = min(frame_count, core_end + time_radius)
        starts = np.arange(ext_start, ext_end) * LANDMARK_HOP_SAMPLES
        indices = starts[:, None] + np.arange(LANDMARK_FFT)
        frames = samples[indices] * window
        magnitude = np.abs(np.fft.rfft(frames, axis=1)).astype(np.float32)
        log_magnitude = np.log1p(1000.0 * magnitude)
        local_max = maximum_filter(log_magnitude, size=(5, 7), mode="nearest")
        candidates = log_magnitude == local_max
        candidates[:, :low_bin] = False
        candidates[:, high_bin + 1 :] = False

        first_group = core_start // _LANDMARK_PEAK_BLOCK
        last_group = (core_end + _LANDMARK_PEAK_BLOCK - 1) // _LANDMARK_PEAK_BLOCK
        for group in range(first_group, last_group):
            lo = max(group * _LANDMARK_PEAK_BLOCK, core_start)
            hi = min((group + 1) * _LANDMARK_PEAK_BLOCK, core_end)
            local_lo, local_hi = lo - ext_start, hi - ext_start
            sample_lo = lo * LANDMARK_HOP_SAMPLES
            sample_hi = min(
                len(samples),
                (hi - 1) * LANDMARK_HOP_SAMPLES + LANDMARK_FFT,
            )
            if np.sqrt(np.mean(samples[sample_lo:sample_hi] ** 2)) < 1e-5:
                continue
            where = np.argwhere(candidates[local_lo:local_hi])
            if len(where) == 0:
                continue
            values = log_magnitude[
                where[:, 0] + local_lo,
                where[:, 1],
            ]
            baseline = float(np.median(log_magnitude[local_lo:local_hi]))
            useful = np.flatnonzero(values > baseline + 1.0)
            if len(useful) == 0:
                continue
            useful = useful[np.argsort(values[useful])[-_LANDMARK_PEAKS_PER_BLOCK:]]
            for selected in useful:
                frame = int(where[selected, 0] + lo)
                frequency = int(where[selected, 1])
                peaks.append((frame, frequency, float(values[selected])))

    if len(peaks) < 2:
        return np.empty((0, 2), dtype=np.int64)
    peaks.sort()
    peak_times = np.asarray([p[0] for p in peaks])
    fingerprints: list[tuple[int, int]] = []
    for anchor_time, anchor_frequency, _ in peaks:
        lo = int(np.searchsorted(peak_times, anchor_time + _LANDMARK_MIN_DELTA))
        hi = int(np.searchsorted(peak_times, anchor_time + _LANDMARK_MAX_DELTA + 1))
        if hi <= lo:
            continue
        targets = sorted(peaks[lo:hi], key=lambda p: p[2], reverse=True)[
            :_LANDMARK_TARGETS
        ]
        for target_time, target_frequency, _ in targets:
            delta = target_time - anchor_time
            fingerprint = (anchor_frequency << 16) | (target_frequency << 8) | delta
            fingerprints.append((anchor_time, fingerprint))
    return np.asarray(fingerprints, dtype=np.int64).reshape(-1, 2)

def landmark_offset(
    a: np.ndarray, b: np.ndarray, hop: float = LANDMARK_HOP
) -> tuple[float, float]:
    """How many seconds to add to b's clock to align with a.

    Confidence is the number of independent hash matches agreeing within four
    landmark frames. Hashes occurring very often are discarded because they describe
    repetitive tones rather than distinctive acoustic events.
    """
    if len(a) == 0 or len(b) == 0:
        return 0.0, 0.0
    a_by_hash = _hash_index(a)
    b_by_hash = _hash_index(b)

    matches: list[tuple[int, int]] = []
    for fingerprint in a_by_hash.keys() & b_by_hash.keys():
        a_times = a_by_hash[fingerprint]
        b_times = b_by_hash[fingerprint]
        if (
            len(a_times) > _LANDMARK_MAX_OCCURRENCES
            or len(b_times) > _LANDMARK_MAX_OCCURRENCES
        ):
            continue
        matches.extend((tb, ta) for ta in a_times for tb in b_times)
    if not matches:
        return 0.0, 0.0

    matched = np.asarray(matches, dtype=float)
    lags = matched[:, 1] - matched[:, 0]
    order = np.argsort(lags)
    ordered = lags[order]
    # find largest collection of matches with offsets that agree within four landmark frames
    starts = np.searchsorted(ordered, ordered - _LANDMARK_LAG_TOLERANCE, side="left")
    best_hi = int(np.argmax(np.arange(len(ordered)) - starts))
    inliers = order[starts[best_hi] : best_hi + 1]

    agreeing = lags[inliers]
    return float(np.median(agreeing) * hop), float(len(agreeing))

def measure_pairs(
    envelopes: dict[str, np.ndarray],
    hop: float = LANDMARK_HOP,
    pairwise: Callable[
        [np.ndarray, np.ndarray, float], tuple[float, float]
    ] = landmark_offset,
) -> list[PairOffset]:
    """Measure every pair of envelopes against each other."""
    return [
        PairOffset(a, b, *pairwise(envelopes[a], envelopes[b], hop))
        for a, b in itertools.combinations(envelopes, 2)
    ]

def solve_offsets(
    ids: list[str],
    pairs: list[PairOffset],
    *,
    min_quality: float = LANDMARK_MIN_VOTES,
    reference: str | None = None,
) -> Alignment:
    """Least-squares offsets from pair measurements, ignoring ones that failed.

    Each locked pair contributes ``offset(b) - offset(a) = lag``, weighted by
    its quality, and the reference is pinned to zero. With more pairs than
    unknowns the leftover disagreement becomes :attr:`Alignment.residual`.
    """
    if not ids:
        return Alignment(offsets={})
    reference = reference or ids[0]
    if reference not in ids:
        raise ValueError(f"reference input {reference!r} is not available")

    locked: list[PairOffset] = []
    for pair in pairs:
        if pair.quality < min_quality:
            logger.warning(
                "inputs %r and %r did not lock (quality %.1f); "
                "they may not overlap in time",
                pair.a,
                pair.b,
                pair.quality,
            )
            continue
        locked.append(pair)

    neighbours = {name: set() for name in ids}
    for pair in locked:
        neighbours[pair.a].add(pair.b)
        neighbours[pair.b].add(pair.a)
    connected = {reference}
    frontier = [reference]
    while frontier:
        name = frontier.pop()
        for neighbour in neighbours[name] - connected:
            connected.add(neighbour)
            frontier.append(neighbour)
    solved_ids = [name for name in ids if name in connected]
    unaligned = [name for name in ids if name not in connected]
    index = {name: i for i, name in enumerate(solved_ids)}

    rows: list[np.ndarray] = []
    values: list[float] = []
    weights: list[float] = []
    for pair in locked:
        if pair.a not in connected or pair.b not in connected:
            continue
        row = np.zeros(len(solved_ids))
        row[index[pair.b]] = 1.0
        row[index[pair.a]] = -1.0
        rows.append(row)
        values.append(pair.lag)
        # Weighted least squares multiplies rows by sqrt(weight).
        weights.append(np.sqrt(pair.quality))

    # Pin the reference to zero, weighted so the solve cannot trade it away.
    pin = np.zeros(len(solved_ids))
    pin[index[reference]] = 1.0
    rows.append(pin)
    values.append(0.0)
    weights.append(max(weights, default=1.0) * 100)

    design = np.asarray(rows)
    measured = np.asarray(values)
    weight = np.asarray(weights)
    solution, *_ = np.linalg.lstsq(
        design * weight[:, None], measured * weight, rcond=None
    )
    # Residual over the pair equations only; the pin is a constraint, not data.
    leftover = design[:-1] @ solution - measured[:-1]
    residual = float(np.sqrt((leftover**2).mean())) if len(leftover) else 0.0
    return Alignment(
        offsets={name: float(solution[index[name]]) for name in solved_ids},
        unaligned=unaligned,
        residual=residual,
    )

def analyse_clock_rates(
    paths: dict[str, str | Path],
    offsets: dict[str, float],
    *,
    window: float = DEFAULT_WINDOW,
    search: float = DEFAULT_SEARCH,
    min_quality: float = SPECTRAL_MIN_QUALITY,
    min_drift_ppm: float = MIN_DRIFT_PPM,
    progress: Callable[[float], bool] | None = None,
) -> ClockRateAnalysis:
    """Measure local lags and fit a timeline for every usable input.

    The reference is selected automatically as the recording with the greatest
    total overlap with the others on the existing alignment timeline. Its
    clock is the one the others are measured against, so it keeps a rate of
    one; all returned fits are expressed on the existing experiment clock and
    can therefore be applied directly to the inputs.

    A clock difference smaller than ``min_drift_ppm`` is left uncorrected.
    """
    if set(paths) != set(offsets):
        raise ValueError("paths and offsets must describe the same inputs")
    if len(paths) < 2:
        raise ValueError("clock-rate analysis needs at least two inputs")

    features = {}
    unavailable = []
    for index, (name, path) in enumerate(paths.items()):
        values = spectral_features(path)
        if len(values):
            features[name] = values
        else:
            unavailable.append(name)
        _continue(progress, 0.55 * (index + 1) / len(paths))
    if len(features) < 2:
        raise ValueError("fewer than two inputs have usable audio")

    reference = _reference_with_most_overlap(features, offsets)
    reference_offset = offsets[reference]
    points: dict[str, list[OffsetPoint]] = {reference: []}
    fits: dict[str, Timeline] = {}
    others = [name for name in features if name != reference]
    for index, name in enumerate(others):
        start = 0.55 + 0.45 * index / len(others)
        extent = 0.45 / len(others)

        def curve_progress(value: float, start=start, extent=extent) -> bool:
            if progress is None:
                return True
            return progress(start + extent * value)

        measured = offset_curve(
            features[reference],
            features[name],
            offsets[name] - reference_offset,
            window=window,
            search=search,
            min_quality=min_quality,
            progress=curve_progress,
        )
        _continue(progress, start + extent)
        if not measured:
            unavailable.append(name)
            continue
        points[name] = [
            replace(
                point,
                time=point.time + reference_offset,
                offset=point.offset + reference_offset,
            )
            for point in measured
        ]
        relative_fit = fit_timeline(measured, min_drift_ppm=min_drift_ppm)
        if relative_fit is not None:
            fits[name] = replace(
                relative_fit,
                offset=relative_fit.offset + reference_offset,
            )

    _continue(progress, 1.0)
    return ClockRateAnalysis(reference, points, fits, unavailable)

def fit_timeline(
    points: list[OffsetPoint],
    *,
    min_drift_ppm: float = MIN_DRIFT_PPM,
    min_span: float = MIN_DRIFT_SPAN,
    confidence: float = DRIFT_CONFIDENCE,
) -> Timeline | None:
    """Fit where a drifting recording starts and how fast its clock runs.

    The fit is a straight line through the measured offsets, so its slope is
    the difference between the two devices' clocks. It is a Theil-Sen line —
    the median of the slopes between every pair of points — because a window
    that locks onto the wrong lag misses by a second where the others agree to
    a few milliseconds, and least squares would follow it.

    ``None`` unless the slope clears three separate bars: measured over at least
    ``min_span`` of recording, so it is not extrapolated across a session from
    a moment of it; a ``confidence`` interval that excludes no difference at
    all, so it is not noise; and at least ``min_drift_ppm``, so it is worth
    correcting.
    """
    if not points:
        return None
    local = np.asarray([point.time - point.offset for point in points])
    experiment = np.asarray([point.time for point in points])
    rate = _fitted_rate(local, experiment, min_drift_ppm, min_span, confidence)
    if rate is None:
        return None
    # also use the median for the offset to reduce effect of outliers
    offset = float(np.median(experiment - local * rate))
    return Timeline(offset=offset, rate=rate)

def offset_curve(
    reference: np.ndarray,
    other: np.ndarray,
    offset: float,
    *,
    window: float = 10.0,
    search: float = 12.0,
    min_quality: float = SPECTRAL_MIN_QUALITY,
    progress: Callable[[float], bool] | None = None,
) -> list[OffsetPoint]:
    """Measure ``other``'s offset repeatedly across the experiment.

    ``reference`` and ``other`` are arrays returned by :func:`spectral_features`.

    ``offset`` seeds the search and only has to be close enough that the true
    lag falls within ``search`` of it.

    ``window`` is how much audio each measurement correlates. Windows do not
    overlap, so their errors are independent and the spread of the points is
    an honest measure of how well the lag is known.

    ``min_quality`` is the lock threshold for each window.
    """
    points: list[OffsetPoint] = []
    if len(reference) == 0 or len(other) == 0:
        return points
    span = max(len(reference), len(other) + int(offset / SPECTRAL_HOP)) * SPECTRAL_HOP
    starts = np.arange(0.0, span, window)
    for index, start in enumerate(starts):
        a0, a1 = int(start / SPECTRAL_HOP), int((start + window) / SPECTRAL_HOP)
        b0 = int((start - offset - search) / SPECTRAL_HOP)
        b1 = int((start + window - offset + search) / SPECTRAL_HOP)
        if progress is not None and progress((index + 1) / len(starts)) is False:
            return points
        if min(a0, b0) < 0 or a1 > len(reference) or b1 > len(other):
            continue
        lag, quality = pairwise_offset(reference[a0:a1], other[b0:b1])
        if quality < min_quality:
            continue
        points.append(OffsetPoint(start + window / 2, (a0 - b0) * SPECTRAL_HOP + lag))
    return points

def pairwise_offset(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
    """Seconds to add to ``b``'s clock to reach ``a``'s, and the lock quality.

    Each spectral feature is mean-subtracted and correlated independently, then
    the correlations are combined by the length of their vector.

    Quality is how many standard deviations the best lag stands above the rest
    of the curve, so it says whether one lag is singled out rather than how
    strongly the recordings resemble each other.
    """
    if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[1]:
        raise ValueError("feature arrays must be two-dimensional with matching columns")
    if len(a) == 0 or len(b) == 0:
        return 0.0, 0.0
    size = 1 << int(np.ceil(np.log2(len(a) + len(b))))
    a = a - a.mean(axis=0)
    b = b - b.mean(axis=0)
    squares = np.zeros(size)
    for column in range(a.shape[1]):
        band = np.fft.irfft(
            np.fft.rfft(a[:, column], size) * np.conj(np.fft.rfft(b[:, column], size)),
            size,
        )
        squares += band**2
    correlation = np.sqrt(squares)
    # rearrange from FFT order into lags running -(len(b)-1) .. len(a)-1.
    correlation = np.concatenate((correlation[-(len(b) - 1) :], correlation[: len(a)]))
    lags = np.arange(-(len(b) - 1), len(a))
    peak = int(np.argmax(correlation))
    spread = correlation.std()
    quality = (
        float((correlation[peak] - correlation.mean()) / spread) if spread else 0.0
    )
    return float(lags[peak] * SPECTRAL_HOP), quality

def spectral_features(media_path: str | Path) -> np.ndarray:
    """Cepstral features of one recording, several values per frame.

    The log-mel spectrogram comes from the faster-whisper FeatureExtractor.
    """
    from faster_whisper.feature_extractor import FeatureExtractor
    from scipy.fft import dct

    samples = audio_samples(media_path, SAMPLE_RATE)
    if len(samples) == 0:
        return np.zeros((0, SPECTRAL_COEFFICIENTS), dtype=np.float32)
    mel = np.asarray(FeatureExtractor()(samples, padding=False)).T
    cepstra = dct(mel, axis=1, norm="ortho")[:, :SPECTRAL_COEFFICIENTS]
    return (cepstra - cepstra.mean(axis=0)) / (cepstra.std(axis=0) + 1e-9)

def audio_samples(path: str | Path, sample_rate: int) -> np.ndarray:
    """Decode audio when possible, returning an empty array otherwise."""
    try:
        return load_audio(path, sample_rate)
    except Exception:
        return np.zeros(0, dtype=np.float32)

def load_audio(audio_path: str | Path, sample_rate: int = SAMPLE_RATE) -> np.ndarray:
    """Decode a recording onto its own timeline, as mono samples.

    Every decoded frame is placed at the position its timestamp gives it, so
    stretches a recorder never wrote are represented as silence.
    """
    import av

    with av.open(str(audio_path)) as container:
        if not container.streams.audio:
            return np.zeros(0, dtype=np.float32)
        stream = container.streams.audio[0]
        resampler = av.AudioResampler(format="s16", layout="mono", rate=sample_rate)
        pieces: list[tuple[int, np.ndarray]] = []
        # a trailing None flushes whatever the resampler is still holding.
        for frame in chain(container.decode(stream), [None]):
            for resampled in resampler.resample(frame):
                samples = resampled.to_ndarray().reshape(-1)
                if resampled.pts is not None and samples.size:
                    pieces.append((resampled.pts, samples))
        if not pieces:
            return np.zeros(0, dtype=np.float32)
        origin = _origin(container, sample_rate, pieces[0][0])

    length = max(pts - origin + len(samples) for pts, samples in pieces)
    buffer = np.zeros(max(length, 0), dtype=np.int16)
    for pts, samples in pieces:
        start = pts - origin
        if start + len(samples) <= 0:
            continue
        buffer[max(start, 0) : start + len(samples)] = samples[max(-start, 0) :]
    return buffer.astype(np.float32) / 32768.0