logWMSE, an audio quality metric with support for digital silence target. Useful for evaluating audio source separation systems, even when there are many audio tracks or stems.
import os
import pickle
from pathlib import Path
from typing import Callable

import numpy as np
import scipy.signal
from audiomentations import (
    Compose,
    HighPassFilter,
    HighShelfFilter,
    PeakingFilter,
    LowPassFilter,
)
from numpy.typing import NDArray

from log_wmse_audio_quality.constants import N_FFT, INTERNAL_SAMPLE_RATE


def get_human_hearing_sensitivity_filter_set() -> (
    Callable[[NDArray[np.float32], int], NDArray[np.float32]]
):
    """Compose a set of filters that together form a frequency response that resembles
    human hearing sensitivity to different frequencies. Inspired by Fenton & Lee (2017).
    """
    return Compose(
        [
            HighShelfFilter(
                min_center_freq=1500.0,
                max_center_freq=1500.0,
                min_gain_db=5.0,
                max_gain_db=5.0,
                min_q=1 / np.sqrt(2),
                max_q=1 / np.sqrt(2),
                p=1.0,
            ),
            HighPassFilter(
                min_cutoff_freq=120.0,
                max_cutoff_freq=120.0,
                min_rolloff=12,
                max_rolloff=12,
                zero_phase=False,
                p=1.0,
            ),
            PeakingFilter(
                min_center_freq=500.0,
                max_center_freq=500.0,
                min_gain_db=2.5,
                max_gain_db=2.5,
                min_q=1.5 / np.sqrt(2),
                max_q=1.5 / np.sqrt(2),
                p=1.0,
            ),
            LowPassFilter(
                min_cutoff_freq=10_000,
                max_cutoff_freq=10_000,
                min_rolloff=12,
                max_rolloff=12,
                zero_phase=True,
                p=1.0,
            ),
        ]
    )


def get_zero_phase_equivalent_filter_impulse_response(
    filter_func: Callable[[NDArray[np.float32], int], NDArray[np.float32]],
    sample_rate: int,
    n_fft: int = 2 * N_FFT,
) -> NDArray:
    """Extract the target response from the given filter_func (which may be not
    zero-phase). The idea is to construct a zero-phase equivalent of the given
    filter_func. Credits: mmxgn (2023)"""
    # Get the impulse response of the filter
    delta = np.zeros(n_fft, dtype=np.float32)
    delta[len(delta) // 2] = 1.0
    impulse_response = filter_func(delta, sample_rate)

    w, h = scipy.signal.freqz(impulse_response, worN=n_fft // 2 + 1)
    linear_target_response = np.abs(h)

    # Compute impulse response
    impulse_response = np.fft.fftshift(np.fft.irfft(linear_target_response, n_fft))

    # Make it symmetric
    center_sample = len(impulse_response) // 2 + 1
    return impulse_response[center_sample - 2000 : center_sample + 2000]


if __name__ == "__main__":
    """Calculate and write the filter impulse response"""
    ir = get_zero_phase_equivalent_filter_impulse_response(
        get_human_hearing_sensitivity_filter_set(), sample_rate=INTERNAL_SAMPLE_RATE
    )
    here = Path(os.path.abspath(os.path.dirname(__file__)))
    base_dir = here.parent
    package_dir = base_dir / "log_wmse_audio_quality"
    with open(package_dir / "filter_ir.pkl", "wb") as f:
        pickle.dump(ir, f)

dev/create_freq_weighting_filter_ir.py
nomonosound/log-wmse-audio-quality
import matplotlib.pyplot as plt

from dev.create_freq_weighting_filter_ir import (
    get_human_hearing_sensitivity_filter_set,
    get_zero_phase_equivalent_filter_impulse_response,
)
from log_wmse_audio_quality.freq_weighting_filter import (
    HumanHearingSensitivityFilter,
)

if __name__ == "__main__":
    # This script plots the phase response of the human hearing sensitivity filter.
    # Ideally, the filter should be zero-phase, so it does not introduce any offsets
    # that may negatively affect the precision of the metric.

    sample_rate = 44100
    filters = get_human_hearing_sensitivity_filter_set()
    linear_phase_filter = HumanHearingSensitivityFilter(
        get_zero_phase_equivalent_filter_impulse_response(
            filters, sample_rate=sample_rate
        )
    )

    # Plot
    plt.figure()
    plt.plot(linear_phase_filter.impulse_response)
    plt.title("Filter impulse response")
    plt.grid()
    plt.show()

dev/plot_filter_impulse_response.py
nomonosound/log-wmse-audio-quality
import math
from typing import Callable, Iterable, List

import fast_align_audio
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm

from dev.create_freq_weighting_filter_ir import (
    get_human_hearing_sensitivity_filter_set,
    get_zero_phase_equivalent_filter_impulse_response,
)
from log_wmse_audio_quality.freq_weighting_filter import (
    HumanHearingSensitivityFilter,
)


def get_phase_shifts_deg(
    filter_func: Callable, frequencies: Iterable
) -> List[np.float64]:
    T = 1 / 44100
    t = np.arange(0, 1, T)  # 1 second worth of samples

    phase_shifts = []

    for freq in tqdm(frequencies):
        # Generate sinusoid
        sinusoid = np.sin(2 * np.pi * freq * t).astype(np.float32)

        # Pass through filters
        output = filter_func(sinusoid)

        # Normalize the output for compatibility with the MSE-based offset calculation
        output /= np.amax(output)

        # Find lag
        max_offset_samples = 1 + int(math.ceil(0.5 * 44100 / freq))
        lag, _ = fast_align_audio.find_best_alignment_offset(
            reference_signal=output,
            delayed_signal=sinusoid,
            max_offset_samples=max_offset_samples,
        )

        # Convert lag to phase shift in degrees
        phase_shift = lag * T * freq * 360  # in degrees

        phase_shifts.append(phase_shift)

    return phase_shifts


if __name__ == "__main__":
    # This script plots the phase response of the human hearing sensitivity filter.
    # Ideally, the filter should be zero-phase, so it does not introduce any offsets
    # that may negatively affect the precision of the metric.

    sample_rate = 44100
    filters = get_human_hearing_sensitivity_filter_set()
    filters_simple_callable = lambda x: filters(x, sample_rate)
    linear_phase_filter = HumanHearingSensitivityFilter(
        get_zero_phase_equivalent_filter_impulse_response(
            filters, sample_rate=sample_rate
        )
    )

    # Frequencies to analyze
    frequencies = np.linspace(20, 18000, 600)

    phase_shifts = get_phase_shifts_deg(filters_simple_callable, frequencies)
    phase_shifts2 = get_phase_shifts_deg(linear_phase_filter, frequencies)

    # Plot
    plt.figure()
    plt.semilogx(frequencies, phase_shifts)
    plt.semilogx(frequencies, phase_shifts2)
    plt.title("Filter phase response")
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Phase Shift (degrees)")
    plt.grid()
    plt.show()

dev/plot_filter_phase_shift.py
nomonosound/log-wmse-audio-quality
from typing import Callable, Iterable, List

import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray
from tqdm import tqdm

from dev.create_freq_weighting_filter_ir import (
    get_human_hearing_sensitivity_filter_set,
    get_zero_phase_equivalent_filter_impulse_response,
)
from log_wmse_audio_quality.freq_weighting_filter import (
    HumanHearingSensitivityFilter,
)


def calculate_rms_db(samples: NDArray):
    """
    Given a numpy array of audio samples, return its RMS power level in decibels.
    Avoids NaNs when provided with zeros as input.
    """
    return (
        20 * np.log10(1000000 * np.sqrt(np.mean(np.square(samples), axis=-1)) + 1)
        - 120.0
    )


def get_frequency_response(
    filter_func: Callable, frequencies: Iterable, sample_rate: int
) -> List[np.float64]:
    T = 1 / sample_rate
    t = np.arange(0, 1, T)  # 1 second worth of samples
    loudness_differences = []

    for freq in tqdm(frequencies):
        # Generate sinusoid
        sinusoid = np.sin(2 * np.pi * freq * t).astype(np.float32)
        unfiltered_rms = calculate_rms_db(sinusoid)

        # Pass through filter
        output = filter_func(sinusoid)

        filtered_rms = calculate_rms_db(output)

        loudness_differences.append(filtered_rms - unfiltered_rms)

    return loudness_differences


if __name__ == "__main__":
    # This script plots the frequency response of the human hearing sensitivity filter.
    sample_rate = 44100

    filters = get_human_hearing_sensitivity_filter_set()
    filters_simple_callable = lambda x: filters(x, sample_rate)
    filters2 = HumanHearingSensitivityFilter(
        get_zero_phase_equivalent_filter_impulse_response(
            filters, sample_rate=sample_rate
        )
    )

    # Frequencies to analyze
    frequencies = np.linspace(20, 20000, 400)

    loudness_differences = get_frequency_response(
        filters_simple_callable, frequencies, sample_rate
    )
    loudness_differences2 = get_frequency_response(filters2, frequencies, sample_rate)

    # Plot
    plt.figure(figsize=(9, 4))
    plt.semilogx(frequencies, loudness_differences, label="audiomentations filter")
    # plt.semilogx(frequencies, loudness_differences2, label="zero phase alternative")
    plt.title("Frequency response")
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("dB")
    xticks = [25, 50, 100, 200, 400, 800, 1500, 3000, 6000, 10000, 20000]
    xlabels = [25, 50, 100, 200, 400, 800, 1500, "3k", "6k", "10k", "20k"]
    plt.xticks(xticks, xlabels)
    plt.grid()
    # plt.legend()
    plt.show()

dev/plot_frequency_response.py
nomonosound/log-wmse-audio-quality
from .metric import calculate_log_wmse

__version__ = "0.1.0"

log_wmse_audio_quality/__init__.py
nomonosound/log-wmse-audio-quality
from log_wmse_audio_quality.utils import convert_decibels_to_amplitude_ratio

# N_FFT relates to the calculation of the impulse response of the human hearing
# sensitivity filter
N_FFT = 4096

# ERROR_TOLERANCE_THRESHOLD is relative to 0 dB RMS
ERROR_TOLERANCE_THRESHOLD = convert_decibels_to_amplitude_ratio(-68.0)

# This scaler makes the scale of values closer to SDR, where an increase
# in the tenths place is a meaningful improvement. I believe this makes it easier to
# compare numbers at a glance, e.g. when numbers are presented in a table.
SCALER = -4.0

# The main purpose of EPS is to avoid taking the log of a zero value
EPS = 1e-8

# INTERNAL_SAMPLE_RATE is the sample rate at which the metric operates internally,
# mainly due to the human hearing sensitivity filter. When audio with different
# sample rate gets passed, it gets resampled to INTERNAL_SAMPLE_RATE.
INTERNAL_SAMPLE_RATE = 44100

log_wmse_audio_quality/constants.py
nomonosound/log-wmse-audio-quality
import pickle
from typing import Optional
import importlib.resources as resources
import numpy as np
import scipy.signal
from numpy.typing import NDArray


class HumanHearingSensitivityFilter:
    def __init__(self, impulse_response: Optional[NDArray] = None):
        if impulse_response is None:
            with resources.open_binary("log_wmse_audio_quality", "filter_ir.pkl") as f:
                self.impulse_response = pickle.load(f)
        else:
            self.impulse_response = impulse_response

    def __call__(self, audio: NDArray[np.float32]):
        """Apply the filter to the given audio. The sample rate of the audio
        should be the same as the impulse response."""
        if audio.ndim == 2:
            placeholder = np.zeros(shape=audio.shape, dtype=np.float32)
            for chn_idx in range(audio.shape[0]):
                placeholder[chn_idx, :] = scipy.signal.oaconvolve(
                    audio[chn_idx], self.impulse_response, "same"
                ).astype(np.float32)
            return placeholder
        else:
            return scipy.signal.oaconvolve(audio, self.impulse_response, "same").astype(
                np.float32
            )

log_wmse_audio_quality/freq_weighting_filter.py
nomonosound/log-wmse-audio-quality
from typing import Callable

import numpy as np
import soxr
from numpy.typing import NDArray

from log_wmse_audio_quality.constants import (
    ERROR_TOLERANCE_THRESHOLD,
    SCALER,
    EPS,
    INTERNAL_SAMPLE_RATE,
)
from log_wmse_audio_quality.freq_weighting_filter import HumanHearingSensitivityFilter
from log_wmse_audio_quality.utils import calculate_rms


def _calculate_single_channel_log_wmse(
    input_rms: np.float32,
    filters: Callable,
    processed_audio: NDArray[np.float32],
    target_audio: NDArray[np.float32],
):
    """
    Calculate the logWMSE metric for a single channel of audio. This is considered an
    internal function used by calculate_wlog_mse.
    """
    if np.sum(input_rms) == 0:
        return np.log(EPS) * SCALER
    scaling_factor = 1 / input_rms

    differences = filters(processed_audio * scaling_factor) - filters(
        target_audio * scaling_factor
    )

    # Tolerate minor errors, because those can't be heard anyway
    differences[np.abs(differences) < ERROR_TOLERANCE_THRESHOLD] = 0.0

    return np.log((differences**2).mean() + EPS) * SCALER


def calculate_log_wmse(
    unprocessed_audio: NDArray,
    processed_audio: NDArray,
    target_audio: NDArray,
    sample_rate: int,
):
    """
    A custom audio metric, logWMSE, that tries to fix a few shortcomings of common metrics.
    The most important property is the support for digital silence in the target, which
    (SI-)SDR, SIR, SAR, ISR, VISQOL_audio, STOI, CDPAM and VISQOL do not support.

    MSE is well-defined for digital silence targets, but has a bunch of other issues:
    * The values are commonly ridiculously small, like between 1e-8 and 1e-3, which
        makes number formatting and sight-reading hard
    * It’s not tailored for audio
    * It’s not scale-invariant
    * It doesn’t align with frequency sensitivity of human hearing
    * It’s not invariant to tiny errors that don’t matter because humans can’t hear
        those errors anyway
    * It’s not logarithmic, like human hearing is

    So this custom metric attempts to solve all the problems mentioned above.
    It's essentially the log of a frequency-weighted MSE, with a few bells and whistles.
    """
    # A custom audio quality metric that
    assert unprocessed_audio.ndim in (1, 2)
    assert processed_audio.ndim in (1, 2)
    assert target_audio.ndim in (1, 2)
    assert processed_audio.shape == target_audio.shape
    is_unprocessed_audio_mono = (
        unprocessed_audio.ndim == 1 or unprocessed_audio.shape[0] == 1
    )
    is_processed_audio_mono = processed_audio.ndim == 1 or processed_audio.shape[0] == 1
    assert is_unprocessed_audio_mono == is_processed_audio_mono

    if sample_rate != INTERNAL_SAMPLE_RATE:
        # Resample to a common sample rate that the filters are designed for
        unprocessed_audio = soxr.resample(
            unprocessed_audio.T,
            sample_rate,
            INTERNAL_SAMPLE_RATE,
            quality="HQ",
        ).T
        processed_audio = soxr.resample(
            processed_audio.T,
            sample_rate,
            INTERNAL_SAMPLE_RATE,
            quality="HQ",
        ).T
        target_audio = soxr.resample(
            target_audio.T,
            sample_rate,
            INTERNAL_SAMPLE_RATE,
            quality="HQ",
        ).T

    # Get filter for human hearing sensitivity
    filters = HumanHearingSensitivityFilter()

    input_rms = calculate_rms(filters(unprocessed_audio))

    if processed_audio.ndim == 2:
        values = np.zeros(shape=input_rms.shape, dtype=np.float32)
        for chn_idx in range(input_rms.shape[0]):
            if processed_audio.ndim == 2:
                values[chn_idx] = _calculate_single_channel_log_wmse(
                    input_rms[chn_idx],
                    filters,
                    processed_audio[chn_idx],
                    target_audio[chn_idx],
                )

        return np.mean(values)
    else:
        return _calculate_single_channel_log_wmse(
            input_rms, filters, processed_audio, target_audio
        )

log_wmse_audio_quality/metric.py
nomonosound/log-wmse-audio-quality
import numpy as np
from numpy.typing import NDArray


def calculate_rms(samples: NDArray):
    """Given a numpy array of audio samples, return its RMS power level."""
    return np.sqrt(np.mean(np.square(samples), axis=-1))


def convert_decibels_to_amplitude_ratio(decibels):
    return 10 ** (decibels / 20)

log_wmse_audio_quality/utils.py
nomonosound/log-wmse-audio-quality
import codecs
import os
import re

from setuptools import setup, find_packages

with open("README.md", "r") as readme_file:
    long_description = readme_file.read()

here = os.path.abspath(os.path.dirname(__file__))


def read(*parts):
    with codecs.open(os.path.join(here, *parts), "r") as fp:
        return fp.read()


def find_version(*file_paths):
    version_file = read(*file_paths)
    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", version_file, re.M)
    if version_match:
        return version_match.group(1)
    raise RuntimeError("Unable to find version string.")


setup(
    name="log-wmse-audio-quality",
    version=find_version("log_wmse_audio_quality", "__init__.py"),
    author="Iver Jordal",
    description=(
        "logWMSE is an audio quality metric with support for digital silence target."
        " Useful for evaluating audio source separation systems, even when there are"
        " many audio tracks or stems."
    ),
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/nomonosound/log-wmse-audio-quality",
    packages=find_packages(exclude=["dev", "tests"]),
    package_data={"log_wmse_audio_quality": ["filter_ir.pkl"]},
    license="Apache 2.0",
    install_requires=[
        "numpy>=1.21,<2",
        "soxr>=0.3.2,<1",
        "scipy>=1.3,<2",
    ],
    python_requires=">=3.8",
    classifiers=[
        "Programming Language :: Python :: 3",
        "Operating System :: OS Independent",
        "Development Status :: 3 - Alpha",
        "Intended Audience :: Developers",
        "Intended Audience :: Science/Research",
        "Topic :: Multimedia",
        "Topic :: Multimedia :: Sound/Audio",
        "Topic :: Scientific/Engineering",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
    project_urls={
        "Issue Tracker": "https://github.com/nomonosound/log-wmse-audio-quality/issues",
    },
)

setup.py
nomonosound/log-wmse-audio-quality
import numpy as np
import pytest
import soxr
from numpy.testing import assert_array_equal
from numpy.typing import NDArray

from dev.create_freq_weighting_filter_ir import get_zero_phase_equivalent_filter_impulse_response, \
    get_human_hearing_sensitivity_filter_set
from log_wmse_audio_quality import calculate_log_wmse
from log_wmse_audio_quality.freq_weighting_filter import (
    HumanHearingSensitivityFilter
)
from dev.plot_filter_phase_shift import get_phase_shifts_deg
from dev.plot_frequency_response import get_frequency_response


def generate_sine_wave(frequency: float, length: int, sample_rate: int) -> NDArray:
    t = np.linspace(0, length / sample_rate, length, endpoint=False)
    return np.sin(2 * np.pi * frequency * t).astype(np.float32)


class TestMetrics:
    def test_no_side_effect(self):
        # Test that calculate_log_wmse does not change the numpy arrays passed to it
        np.random.seed(42)
        sample_rate = 44100
        input_sound = np.random.uniform(low=-1.0, high=1.0, size=(sample_rate,)).astype(
            np.float32
        )
        input_sound_copy = np.copy(input_sound)
        est_sound = input_sound * 0.1
        est_sound_copy = np.copy(est_sound)
        target_sound = np.zeros((sample_rate,), dtype=np.float32)
        target_sound_copy = np.copy(target_sound)

        calculate_log_wmse(input_sound, est_sound, target_sound, sample_rate)

        assert_array_equal(input_sound, input_sound_copy)
        assert_array_equal(est_sound, est_sound_copy)
        assert_array_equal(target_sound, target_sound_copy)

    def test_1d_2d_shape(self):
        np.random.seed(42)
        sample_rate = 44100
        input_sound_1d = np.random.uniform(
            low=-1.0, high=1.0, size=(sample_rate,)
        ).astype(np.float32)
        est_sound_1d = input_sound_1d * 0.1
        target_sound_1d = np.zeros((sample_rate,), dtype=np.float32)

        log_mse_1d = calculate_log_wmse(
            input_sound_1d, est_sound_1d, target_sound_1d, sample_rate
        )

        input_sound_2d_mono = input_sound_1d.reshape((1, -1))
        est_sound_2d_mono = est_sound_1d.reshape((1, -1))
        target_sound_2d_mono = target_sound_1d.reshape((1, -1))
        log_mse_2d_mono = calculate_log_wmse(
            input_sound_2d_mono, est_sound_2d_mono, target_sound_2d_mono, sample_rate
        )

        input_sound_2d_stereo = np.vstack((input_sound_1d, input_sound_1d))
        est_sound_2d_stereo = np.vstack((est_sound_1d, est_sound_1d))
        target_sound_2d_stereo = np.vstack((target_sound_1d, target_sound_1d))
        log_mse_2d_stereo = calculate_log_wmse(
            input_sound_2d_stereo,
            est_sound_2d_stereo,
            target_sound_2d_stereo,
            sample_rate,
        )

        assert log_mse_2d_mono == pytest.approx(log_mse_1d)
        assert log_mse_2d_stereo == pytest.approx(log_mse_1d)

    def test_digital_silence_target(self):
        np.random.seed(42)
        sample_rate = 44100
        input_sound = np.random.uniform(low=-1.0, high=1.0, size=(sample_rate,)).astype(
            np.float32
        )
        est_sound = input_sound * 0.1
        target_sound = np.zeros((sample_rate,), dtype=np.float32)

        log_mse = calculate_log_wmse(input_sound, est_sound, target_sound, sample_rate)
        assert log_mse == pytest.approx(18.42, abs=0.01)

    def test_digital_silence_input_and_target(self):
        np.random.seed(42)
        sample_rate = 44100
        input_sound = np.zeros((sample_rate,), dtype=np.float32)
        est_sound = np.zeros_like(input_sound)
        target_sound = np.zeros_like(input_sound)

        log_mse = calculate_log_wmse(input_sound, est_sound, target_sound, sample_rate)
        assert log_mse == pytest.approx(73.68, abs=0.01)

    def test_exact_match(self):
        np.random.seed(42)
        sample_rate = 44100
        target_sound = np.random.uniform(
            low=-0.1, high=0.1, size=(sample_rate,)
        ).astype(np.float32)
        input_sound = np.copy(target_sound)
        est_sound = np.copy(target_sound)

        log_mse = calculate_log_wmse(input_sound, est_sound, target_sound, sample_rate)
        assert log_mse == pytest.approx(73.68, abs=0.01)

    def test_greater_is_better(self):
        np.random.seed(42)
        sample_rate = 44100
        target_sound = np.random.uniform(
            low=-0.1, high=0.1, size=(sample_rate,)
        ).astype(np.float32)
        input_sound = target_sound + np.random.uniform(
            low=-0.5, high=0.5, size=(sample_rate,)
        ).astype(np.float32)
        est_sound1 = target_sound + np.random.uniform(
            low=-0.01, high=0.01, size=(sample_rate,)
        ).astype(np.float32)
        est_sound12 = target_sound + np.random.uniform(
            low=-0.001, high=0.001, size=(sample_rate,)
        ).astype(np.float32)

        log_mse1 = calculate_log_wmse(
            input_sound, est_sound1, target_sound, sample_rate
        )
        log_mse2 = calculate_log_wmse(
            input_sound, est_sound12, target_sound, sample_rate
        )
        assert log_mse2 > log_mse1

    def test_scale_invariance(self):
        np.random.seed(42)
        sample_rate = 44100
        target_sound = np.random.uniform(
            low=-0.1, high=0.1, size=(sample_rate,)
        ).astype(np.float32)
        input_sound = target_sound + np.random.uniform(
            low=-0.5, high=0.5, size=(sample_rate,)
        ).astype(np.float32)
        est_sound = target_sound + np.random.uniform(
            low=-0.01, high=0.01, size=(sample_rate,)
        ).astype(np.float32)

        log_mse1 = calculate_log_wmse(input_sound, est_sound, target_sound, sample_rate)

        target_sound *= 0.01
        est_sound *= 0.01
        input_sound *= 0.01

        log_mse2 = calculate_log_wmse(input_sound, est_sound, target_sound, sample_rate)
        assert log_mse2 == pytest.approx(log_mse1)

    def test_small_error_tolerance(self):
        np.random.seed(42)
        sample_rate = 44100
        target_sound = np.random.uniform(
            low=-0.1, high=0.1, size=(sample_rate,)
        ).astype(np.float32)
        input_sound = target_sound + np.random.uniform(
            low=-1.0, high=1.0, size=(44100,)
        ).astype(np.float32)
        est_sound1 = np.copy(target_sound)

        weak_hiss = np.random.uniform(
            low=-0.0001, high=0.0001, size=(sample_rate,)
        ).astype(np.float32)
        est_sound2 = est_sound1 + weak_hiss

        log_mse1 = calculate_log_wmse(
            input_sound, est_sound1, target_sound, sample_rate
        )
        log_mse2 = calculate_log_wmse(
            input_sound, est_sound2, target_sound, sample_rate
        )
        assert log_mse2 == pytest.approx(log_mse1, abs=0.01)

    def test_frequency_weighting_filter_response(self):
        np.random.seed(42)
        sample_rate = 44100
        filters = HumanHearingSensitivityFilter(
            get_zero_phase_equivalent_filter_impulse_response(
                get_human_hearing_sensitivity_filter_set(), sample_rate
            )
        )

        frequencies = np.linspace(20, 20000, 2000)
        frequency_response = get_frequency_response(
            filter_func=filters, frequencies=frequencies, sample_rate=sample_rate
        )

        freq_50hz_index = np.argmin(np.abs(frequencies - 50))
        freq_1000hz_index = np.argmin(np.abs(frequencies - 1000))
        freq_3000hz_index = np.argmin(np.abs(frequencies - 3000))
        freq_16000hz_index = np.argmin(np.abs(frequencies - 16000))

        assert frequency_response[freq_50hz_index] == pytest.approx(-16.0, abs=1.0)
        assert frequency_response[freq_1000hz_index] == pytest.approx(1.3, abs=1.0)
        assert frequency_response[freq_3000hz_index] == pytest.approx(4.2, abs=1.0)
        assert frequency_response[freq_16000hz_index] == pytest.approx(-11.8, abs=1.0)

    def test_frequency_weighting_filter_zero_phase(self):
        np.random.seed(42)
        sample_rate = 44100
        filters = HumanHearingSensitivityFilter(
            get_zero_phase_equivalent_filter_impulse_response(
                get_human_hearing_sensitivity_filter_set(), sample_rate
            )
        )
        frequencies = np.linspace(20, 18000, 400)
        phase_shifts_deg = get_phase_shifts_deg(filters, frequencies)
        # Check that the phase is no more than 1 degree off
        assert np.amax(phase_shifts_deg) < 1.0
        assert np.amin(phase_shifts_deg) > -1.0

    def test_human_hearing_frequency_sensitivity(self):
        np.random.seed(42)

        sample_rate = 44100

        target_sound = np.random.uniform(low=-0.03, high=0.3, size=(44100,)).astype(
            np.float32
        )
        low_freq_sine = 0.3 * generate_sine_wave(
            50.0, target_sound.shape[-1], sample_rate
        )
        high_freq_sine = 0.3 * generate_sine_wave(
            3000.0, target_sound.shape[-1], sample_rate
        )
        input_sound = np.copy(target_sound) + low_freq_sine + high_freq_sine
        # The 1st est reduced the low freq sine quite well, but not the high freq sine
        est_sound_enhanced_only_low_freq = (
            np.copy(target_sound) + 0.05 * low_freq_sine + high_freq_sine
        )
        # The 2nd est reduced the high freq sine quite well, but not the low freq sine
        est_sound_enhanced_only_high_freq = (
            np.copy(target_sound) + low_freq_sine + 0.05 * high_freq_sine
        )

        # Because human hearing is much more sensitive to 3000 Hz than 50 Hz, the
        # estimate that reduced the 3200 Hz sine should get a significantly higher score
        log_mse1 = calculate_log_wmse(
            input_sound, est_sound_enhanced_only_low_freq, target_sound, sample_rate
        )
        log_mse2 = calculate_log_wmse(
            input_sound, est_sound_enhanced_only_high_freq, target_sound, sample_rate
        )
        # We check that log_mse1 is much lower than log_mse2
        assert log_mse1 == pytest.approx(log_mse2 - 17.2, abs=0.5)

    def test_various_sample_rates(self):
        # Check that the metric returns similar scores for similar-sounding audio at
        # different sample rates
        np.random.seed(42)

        sample_rate = 44100

        target_sound = np.random.uniform(low=-0.03, high=0.3, size=(44100,)).astype(
            np.float32
        )
        low_freq_sine = 0.3 * generate_sine_wave(
            50.0, target_sound.shape[-1], sample_rate
        )
        high_freq_sine = 0.3 * generate_sine_wave(
            3000.0, target_sound.shape[-1], sample_rate
        )
        input_sound = np.copy(target_sound) + low_freq_sine + high_freq_sine
        # The 1st est reduced the low freq sine quite well, but not the high freq sine
        est_sound = np.copy(target_sound) + 0.05 * low_freq_sine + 0.05 * high_freq_sine

        # low-pass filter the sounds by downsampling them
        low_sr = 36000
        input_sound = soxr.resample(
            input_sound,
            sample_rate,
            low_sr,
            quality="HQ",
        )
        est_sound = soxr.resample(
            est_sound,
            sample_rate,
            low_sr,
            quality="HQ",
        )
        target_sound = soxr.resample(
            target_sound,
            sample_rate,
            low_sr,
            quality="HQ",
        )

        scores = []
        for sample_rate in [36000, 44100, 48000]:
            resampled_input_sound = soxr.resample(
                input_sound,
                low_sr,
                sample_rate,
                quality="HQ",
            )
            resampled_est_sound = soxr.resample(
                est_sound,
                low_sr,
                sample_rate,
                quality="HQ",
            )
            resampled_target_sound = soxr.resample(
                target_sound,
                low_sr,
                sample_rate,
                quality="HQ",
            )
            log_mse = calculate_log_wmse(
                resampled_input_sound,
                resampled_est_sound,
                resampled_target_sound,
                sample_rate,
            )
            scores.append(log_mse)

        # Check that all scores are equal
        assert scores[0] == pytest.approx(scores[1], abs=0.01)
        assert scores[0] == pytest.approx(scores[2], abs=0.01)

tests/test_metric.py
nomonosound/log-wmse-audio-quality