diff --git a/mlx_audio/dsp.py b/mlx_audio/dsp.py index 8f3895f9c..105a98640 100644 --- a/mlx_audio/dsp.py +++ b/mlx_audio/dsp.py @@ -94,6 +94,16 @@ def bartlett(size, periodic=False): } +# K-weighting stage parameters. These are the analogue prototypes behind the +# BS.1770 Table 1 / Table 2 coefficients, so a redesign at any sampling rate +# still matches the tabulated 48 kHz response. +_K_WEIGHT_SHELF_FREQ = 1681.974450955533 +_K_WEIGHT_SHELF_Q = 0.7071752369554196 +_K_WEIGHT_SHELF_GAIN_DB = 3.999843853973347 +_K_WEIGHT_HIGHPASS_FREQ = 38.13547087602444 +_K_WEIGHT_HIGHPASS_Q = 0.5003270373238773 + + def _validate_loudness_audio(data: np.ndarray, rate: int, block_size: float) -> None: if not isinstance(data, np.ndarray): raise ValueError("Data must be of type numpy.ndarray.") @@ -115,44 +125,43 @@ def _biquad_coefficients( rate: int, filter_type: str, ) -> tuple[np.ndarray, np.ndarray]: - amplitude = 10 ** (gain_db / 40.0) - omega = 2.0 * math.pi * (center_freq / rate) - alpha = math.sin(omega) / (2.0 * q_factor) + """Build a K-weighting stage as specified by ITU-R BS.1770. + + The Recommendation tabulates both stages only at 48 kHz (Tables 1 and 2) and + requires other rates to "provide the same frequency response". These designs + reproduce the tabulated coefficients exactly at 48 kHz and rescale correctly. + """ + k = math.tan(math.pi * center_freq / rate) + denominator = 1.0 + k / q_factor + k * k + a = np.array( + [ + 1.0, + 2.0 * (k * k - 1.0) / denominator, + (1.0 - k / q_factor + k * k) / denominator, + ] + ) if filter_type == "high_shelf": - b0 = amplitude * ( - (amplitude + 1) - + (amplitude - 1) * math.cos(omega) - + 2 * math.sqrt(amplitude) * alpha - ) - b1 = -2 * amplitude * ((amplitude - 1) + (amplitude + 1) * math.cos(omega)) - b2 = amplitude * ( - (amplitude + 1) - + (amplitude - 1) * math.cos(omega) - - 2 * math.sqrt(amplitude) * alpha - ) - a0 = ( - (amplitude + 1) - - (amplitude - 1) * math.cos(omega) - + 2 * math.sqrt(amplitude) * alpha - ) - a1 = 2 * ((amplitude - 1) - (amplitude + 1) * math.cos(omega)) - a2 = ( - (amplitude + 1) - - (amplitude - 1) * math.cos(omega) - - 2 * math.sqrt(amplitude) * alpha + shelf_gain = 10.0 ** (gain_db / 20.0) + # Exponent as used by the BS.1770 reference designs (libebur128, and + # pyloudnorm's "DeMan" filter class); it is what lands stage 1 on Table 1. + mid_gain = shelf_gain**0.4996667741545416 + b = np.array( + [ + (shelf_gain + mid_gain * k / q_factor + k * k) / denominator, + 2.0 * (k * k - shelf_gain) / denominator, + (shelf_gain - mid_gain * k / q_factor + k * k) / denominator, + ] ) elif filter_type == "high_pass": - b0 = (1 + math.cos(omega)) / 2 - b1 = -(1 + math.cos(omega)) - b2 = (1 + math.cos(omega)) / 2 - a0 = 1 + alpha - a1 = -2 * math.cos(omega) - a2 = 1 - alpha + # Table 2 gives the numerator as exactly [1, -2, 1]. Normalising it by + # the denominator (as the RBJ cookbook high-pass does) costs 0.043 dB + # across the whole band above ~500 Hz. + b = np.array([1.0, -2.0, 1.0]) else: raise ValueError(f"Unsupported filter type: {filter_type}") - return np.array([b0, b1, b2]) / a0, np.array([a0, a1, a2]) / a0 + return b, a def lfilter(b: np.ndarray, a: np.ndarray, data: np.ndarray) -> np.ndarray: @@ -206,9 +215,15 @@ def _k_weight_audio(data: np.ndarray, rate: int) -> np.ndarray: weighted = np.array(data, dtype=np.float64, copy=True) high_shelf_b, high_shelf_a = _biquad_coefficients( - 4.0, 1 / math.sqrt(2), 1500.0, rate, "high_shelf" + _K_WEIGHT_SHELF_GAIN_DB, + _K_WEIGHT_SHELF_Q, + _K_WEIGHT_SHELF_FREQ, + rate, + "high_shelf", + ) + high_pass_b, high_pass_a = _biquad_coefficients( + 0.0, _K_WEIGHT_HIGHPASS_Q, _K_WEIGHT_HIGHPASS_FREQ, rate, "high_pass" ) - high_pass_b, high_pass_a = _biquad_coefficients(0.0, 0.5, 38.0, rate, "high_pass") for channel in range(weighted.shape[1]): weighted[:, channel] = _apply_lfilter( @@ -242,19 +257,20 @@ def integrated_loudness( absolute_threshold = -70.0 step = 1.0 - overlap - duration_seconds = num_samples / rate - num_blocks = int( - np.round(((duration_seconds - block_size) / (block_size * step))) + 1 - ) + # BS.1770: a gating block is 400 ms "to the nearest sample", and "incomplete + # gating blocks at the end of the measurement interval are not used". + block_samples = int(round(block_size * rate)) + hop_samples = int(round(block_size * step * rate)) + num_blocks = 1 + (num_samples - block_samples) // hop_samples block_indices = np.arange(0, num_blocks) mean_square = np.zeros((num_channels, num_blocks), dtype=np.float64) for channel in range(num_channels): for block_index in block_indices: - lower = int(block_size * (block_index * step) * rate) - upper = int(block_size * (block_index * step + 1) * rate) - mean_square[channel, block_index] = (1.0 / (block_size * rate)) * np.sum( - np.square(input_data[lower:upper, channel]) + lower = block_index * hop_samples + upper = lower + block_samples + mean_square[channel, block_index] = ( + np.sum(np.square(input_data[lower:upper, channel])) / block_samples ) with warnings.catch_warnings(): @@ -276,7 +292,7 @@ def integrated_loudness( gated_blocks = [ block_index for block_index, loudness in enumerate(block_loudness) - if loudness >= absolute_threshold + if loudness > absolute_threshold ] with warnings.catch_warnings(): diff --git a/mlx_audio/tests/test_dsp.py b/mlx_audio/tests/test_dsp.py index 8a40adbaf..1039ad8b6 100644 --- a/mlx_audio/tests/test_dsp.py +++ b/mlx_audio/tests/test_dsp.py @@ -93,46 +93,146 @@ def test_utils_lazy_imports(): assert result.returncode == 0, f"Lazy import failed: {result.stderr}" -def test_integrated_loudness_matches_reference_values(): - """Verify BS.1770 loudness matches fixed reference outputs.""" - from mlx_audio.dsp import integrated_loudness +def _sine(peak_dbfs, seconds, rate, freq=997.0): + n = int(seconds * rate) + return 10.0 ** (peak_dbfs / 20.0) * np.sin(2 * np.pi * freq * np.arange(n) / rate) + + +def test_k_weighting_matches_bs1770_coefficients(): + """The two K-weighting stages must equal the coefficients ITU-R BS.1770 + tabulates for 48 kHz in Table 1 (spherical head) and Table 2 (RLB).""" + from mlx_audio.dsp import ( + _K_WEIGHT_HIGHPASS_FREQ, + _K_WEIGHT_HIGHPASS_Q, + _K_WEIGHT_SHELF_FREQ, + _K_WEIGHT_SHELF_GAIN_DB, + _K_WEIGHT_SHELF_Q, + _biquad_coefficients, + ) - rng = np.random.default_rng(0) - mono = (rng.standard_normal(24000) * 0.02).astype(np.float64) - stereo = (rng.standard_normal((24000, 2)) * 0.015).astype(np.float64) + shelf_b, shelf_a = _biquad_coefficients( + _K_WEIGHT_SHELF_GAIN_DB, + _K_WEIGHT_SHELF_Q, + _K_WEIGHT_SHELF_FREQ, + 48000, + "high_shelf", + ) + np.testing.assert_allclose( + shelf_b, [1.53512485958697, -2.69169618940638, 1.19839281085285], atol=1e-12 + ) + np.testing.assert_allclose( + shelf_a, [1.0, -1.69065929318241, 0.73248077421585], atol=1e-12 + ) - assert integrated_loudness(mono, 24000) == pytest.approx( - -31.147497580698033, abs=1e-12 + pass_b, pass_a = _biquad_coefficients( + 0.0, _K_WEIGHT_HIGHPASS_Q, _K_WEIGHT_HIGHPASS_FREQ, 48000, "high_pass" ) - assert integrated_loudness(stereo, 24000) == pytest.approx( - -30.587340400145717, abs=1e-12 + np.testing.assert_allclose(pass_b, [1.0, -2.0, 1.0], atol=1e-12) + np.testing.assert_allclose( + pass_a, [1.0, -1.99004745483398, 0.99007225036621], atol=1e-12 ) + # BS.1770 Note 1: the -0.691 constant cancels the K-weighting gain at 997 Hz, + # so that gain has to be 0.691 dB. + z = np.exp(-2j * np.pi * 997.0 / 48000) + + def gain_db(b, a): + num = b[0] + b[1] * z + b[2] * z**2 + den = a[0] + a[1] * z + a[2] * z**2 + return 20.0 * np.log10(np.abs(num / den)) + + total = gain_db(shelf_b, shelf_a) + gain_db(pass_b, pass_a) + assert total == pytest.approx(0.691, abs=1e-3) + + +def test_integrated_loudness_matches_bs1770_997hz_anchor(): + """BS.1770: a 0 dB FS 997 Hz sine on one channel reads -3.01 LKFS, and the + scale is 1 LKFS per dB.""" + from mlx_audio.dsp import integrated_loudness + + for peak_dbfs in (0.0, -20.0, -40.0): + measured = integrated_loudness(_sine(peak_dbfs, 2.0, 48000), 48000) + assert measured == pytest.approx(peak_dbfs - 3.01, abs=0.01) + + +def test_integrated_loudness_block_hop_11025hz(): + """At 44.1 kHz / 4 the 400 ms block is 4410 samples and the hop is + round(1102.5) = 1102, so block starts are exact integer multiples of the + hop (uniform spacing, no drift), and the 997 Hz anchor still holds.""" + from mlx_audio.dsp import integrated_loudness + + rate = 11025 + assert int(round(0.4 * rate)) == 4410 + assert int(round(0.4 * 0.25 * rate)) == 1102 -def test_normalize_loudness_matches_reference_values(): - """Verify loudness normalization matches fixed reference outputs.""" + # -23 dB FS 997 Hz reads -23 - 3.01 LKFS; the K-weighting is redesigned + # per rate from the analogue prototypes, which costs < 0.05 dB (48 kHz + # reads -26.010, 44.1 kHz -26.007). + measured = integrated_loudness(_sine(-23.0, 2.0, rate), rate) + assert measured == pytest.approx(-23.0 - 3.01, abs=0.05) + + +def test_integrated_loudness_ignores_incomplete_final_block(): + """BS.1770: "Incomplete gating blocks at the end of the measurement interval + are not used", so trailing samples that do not complete a block cannot move + the reading and a steady tone measures the same at any length.""" + from mlx_audio.dsp import integrated_loudness + + rate = 24000 + tone = _sine(-23.0, 0.5, rate) + reference = integrated_loudness(tone, rate) + hop = int(0.4 * 0.25 * rate) + for extra in (1, hop // 2, hop - 1): + padded = np.concatenate([tone, _sine(-23.0, 0.5, rate)[:extra]]) + assert integrated_loudness(padded, rate) == pytest.approx(reference, abs=1e-12) + + readings = [ + integrated_loudness(_sine(-23.0, duration, rate), rate) + for duration in (0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70) + ] + assert max(readings) - min(readings) < 0.01 + + +def test_integrated_loudness_gates_quiet_passages(): + """The two-stage -70 LKFS / -10 LU gating keeps a loud programme's reading + from being dragged down by quiet passages around it.""" + from mlx_audio.dsp import integrated_loudness + + rate = 24000 + loud = _sine(-23.0, 8.0, rate) + + # -60 dBFS sits above the -70 LKFS absolute threshold, so it is the relative + # threshold that has to exclude it. + quiet = _sine(-60.0, 2.0, rate) + assert integrated_loudness( + np.concatenate([quiet, loud, quiet]), rate + ) == pytest.approx(integrated_loudness(loud, rate), abs=0.25) + + # Below the absolute threshold, pushing a passage even further down cannot + # change the result at all. + readings = [ + integrated_loudness( + np.concatenate([_sine(peak, 2.0, rate), loud, _sine(peak, 2.0, rate)]), + rate, + ) + for peak in (-100.0, -140.0) + ] + assert readings[0] == pytest.approx(readings[1], abs=1e-7) + + +def test_normalize_loudness_reaches_target(): + """Normalizing to a target LUFS has to actually land on that target.""" from mlx_audio.dsp import integrated_loudness, normalize_loudness rng = np.random.default_rng(0) mono = (rng.standard_normal(24000) * 0.02).astype(np.float64) - measured = integrated_loudness(mono, 24000, block_size=0.4) + measured = integrated_loudness(mono, 24000) normalized = normalize_loudness(mono, measured, -18.0) - assert np.max(np.abs(normalized)) == pytest.approx(0.4083656963780373, abs=1e-12) + assert integrated_loudness(normalized, 24000) == pytest.approx(-18.0, abs=1e-9) np.testing.assert_allclose( - normalized[:5], - np.array( - [ - 0.011424693401069328, - -0.01200393625946315, - 0.058193108743361296, - 0.009531930078445609, - -0.04867452152305382, - ] - ), - atol=1e-12, - rtol=0.0, + normalized, mono * 10.0 ** ((-18.0 - measured) / 20.0), rtol=1e-12 )