midmid/tempo_map.py

"""Tempo map derivation from beat tracker output."""

import numpy as np

from midmid.beat_tracker import BeatData


def derive_tempo_map(
    beat_data: BeatData, change_threshold: float = 0.08,
) -> list[tuple[float, float]]:
    """Derive a tempo map from beat data.

    Returns list of (time_seconds, bpm) tuples, sorted by time.
    """
    beats = beat_data.beats
    if len(beats) < 2:
        return [(0.0, 120.0)]

    intervals = np.diff(beats)
    bpms = 60.0 / intervals

    median_bpm = np.median(bpms)
    valid = (bpms > median_bpm * 0.6) & (bpms < median_bpm * 1.6)
    if not np.any(valid):
        return [(0.0, float(median_bpm))]

    valid_bpms = bpms[valid]
    if np.std(valid_bpms) / np.mean(valid_bpms) < change_threshold:
        avg_bpm = float(np.mean(valid_bpms))
        return [(0.0, _round_bpm(avg_bpm))]

    tempo_map = []
    current_bpm = float(bpms[0]) if valid[0] else float(median_bpm)
    tempo_map.append((0.0, _round_bpm(current_bpm)))

    window = 4
    for i in range(window, len(bpms) - window + 1, window):
        chunk = bpms[i : i + window]
        chunk_valid = chunk[(chunk > median_bpm * 0.6) & (chunk < median_bpm * 1.6)]
        if len(chunk_valid) == 0:
            continue
        local_bpm = float(np.mean(chunk_valid))
        if abs(local_bpm - current_bpm) / current_bpm > change_threshold:
            current_bpm = local_bpm
            tempo_map.append((float(beats[i]), _round_bpm(current_bpm)))

    return tempo_map


def get_median_bpm(beat_data: BeatData) -> float:
    if len(beat_data.beats) < 2:
        return 120.0
    intervals = np.diff(beat_data.beats)
    bpms = 60.0 / intervals
    return float(_round_bpm(np.median(bpms)))


def estimate_time_signature(beat_data: BeatData) -> int:
    if len(beat_data.downbeats) < 2:
        return 4

    beats = beat_data.beats
    downbeats = beat_data.downbeats

    counts = []
    for i in range(len(downbeats) - 1):
        start, end = downbeats[i], downbeats[i + 1]
        n = np.sum((beats >= start) & (beats < end))
        if 2 <= n <= 7:
            counts.append(n)

    if not counts:
        return 4

    values, freq = np.unique(counts, return_counts=True)
    return int(values[np.argmax(freq)])


def _round_bpm(bpm: float) -> float:
    return round(float(bpm), 2)