import numpy as np
import pytest
from augsig import augment, Augment
from augsig.noisifier import noisify, burstify, burst_mask
from augsig.warper import (
rand_knots,
twarp_bezier, twarp_pchip,
adrift_bezier, adrift_pchip,
amod_bezier, amod_pchip,
drift_linear,
)
from augsig.utils import normalize, freqfilt
# Shared fixtures
N = 500
RNG = np.random.default_rng(0)
@pytest.fixture
def signal():
rng = np.random.default_rng(42)
return rng.standard_normal(N)
# ---------------------------------------------------------------------------
# utils
# ---------------------------------------------------------------------------
class TestNormalize:
def test_output_range(self, signal):
out = normalize(signal)
assert out.min() >= 0.0
assert out.max() <= 1.0 + 1e-5
def test_shape_preserved(self, signal):
assert normalize(signal).shape == signal.shape
def test_constant_signal_no_nan(self):
out = normalize(np.ones(100))
assert not np.any(np.isnan(out))
class TestFreqfilt:
def test_lowpass_attenuates_high_freq(self, signal):
filtered = freqfilt(signal, 0.0, 0.1)
assert np.std(filtered) < np.std(signal)
def test_passthrough_full_band(self, signal):
out = freqfilt(signal, 0.0, 1.0)
np.testing.assert_array_equal(out, signal)
def test_shape_preserved(self, signal):
assert freqfilt(signal, 0.0, 0.5).shape == signal.shape
def test_invalid_cutoffs_raises(self, signal):
with pytest.raises(ValueError):
freqfilt(signal, 0.8, 0.2)
def test_out_of_range_raises(self, signal):
with pytest.raises(ValueError):
freqfilt(signal, -0.1, 0.5)
# ---------------------------------------------------------------------------
# noisifier
# ---------------------------------------------------------------------------
class TestNoisify:
def test_output_shape(self, signal):
noisy, noise = noisify(signal, snr_db=20, rng=np.random.default_rng(0))
assert noisy.shape == signal.shape
assert noise.shape == signal.shape
def test_inf_snr_returns_original(self, signal):
noisy, noise = noisify(signal, snr_db=float('inf'))
np.testing.assert_array_equal(noisy, signal)
assert np.all(noise == 0)
def test_snr_approximately_correct(self, signal):
snr_db = 20
noisy, noise = noisify(signal, snr_db=snr_db, rng=np.random.default_rng(1))
signal_power = np.mean(signal ** 2)
noise_power = np.mean(noise ** 2)
measured_snr_db = 10 * np.log10(signal_power / noise_power)
assert abs(measured_snr_db - snr_db) < 2.0
@pytest.mark.parametrize("color", ["white", "pink", "brown", "blue", "violet"])
def test_noise_colors(self, signal, color):
noisy, noise = noisify(signal, snr_db=15, color=color, rng=np.random.default_rng(2))
assert noisy.shape == signal.shape
assert not np.any(np.isnan(noise))
@pytest.mark.parametrize("dist", ["gauss", "uniform", "laplace"])
def test_noise_distributions(self, signal, dist):
noisy, noise = noisify(signal, snr_db=15, dist=dist, rng=np.random.default_rng(3))
assert noisy.shape == signal.shape
def test_resample_dist(self, signal):
noisy, noise = noisify(signal, snr_db=15, dist="resample",
resample_pool="self", rng=np.random.default_rng(4))
assert noisy.shape == signal.shape
def test_minus_inf_snr_unit_power(self, signal):
noise, _ = noisify(signal, snr_db=float('-inf'), rng=np.random.default_rng(5))
assert abs(np.mean(noise ** 2) - 1.0) < 0.1
def test_reproducibility(self, signal):
n1, _ = noisify(signal, snr_db=20, rng=np.random.default_rng(99))
n2, _ = noisify(signal, snr_db=20, rng=np.random.default_rng(99))
np.testing.assert_array_equal(n1, n2)
class TestBurstify:
def test_output_shape(self, signal):
noisy, burst = burstify(signal, snr_db=float('-inf'), rng=np.random.default_rng(0))
assert noisy.shape == signal.shape
assert burst.shape == signal.shape
def test_inf_snr_returns_original(self, signal):
noisy, burst = burstify(signal, snr_db=float('inf'))
np.testing.assert_array_equal(noisy, signal)
assert np.all(burst == 0)
def test_burst_is_sparse(self, signal):
_, burst = burstify(signal, snr_db=float('-inf'), n_bursts=3,
burst_width=10, rng=np.random.default_rng(7))
nonzero_ratio = np.count_nonzero(burst) / len(burst)
assert nonzero_ratio < 0.5
# ---------------------------------------------------------------------------
# warper
# ---------------------------------------------------------------------------
class TestTimeWarp:
@pytest.mark.parametrize("fn", [twarp_bezier, twarp_pchip])
def test_output_shape(self, signal, fn):
out = fn(signal, rng=np.random.default_rng(0))
assert out.shape == signal.shape
@pytest.mark.parametrize("fn", [twarp_bezier, twarp_pchip])
def test_value_range_preserved(self, signal, fn):
out = fn(signal, rng=np.random.default_rng(0))
assert out.min() >= signal.min() - 1e-6
assert out.max() <= signal.max() + 1e-6
class TestAmpDrift:
@pytest.mark.parametrize("fn", [adrift_bezier, adrift_pchip])
def test_output_shape(self, signal, fn):
out = fn(signal, variance=0.05, rng=np.random.default_rng(0))
assert out.shape == signal.shape
@pytest.mark.parametrize("fn", [adrift_bezier, adrift_pchip])
def test_small_variance_stays_close(self, signal, fn):
out = fn(signal, variance=0.01, rng=np.random.default_rng(0))
assert np.max(np.abs(out - signal)) < 0.1
class TestAmpMod:
@pytest.mark.parametrize("fn", [amod_bezier, amod_pchip])
def test_output_shape(self, signal, fn):
out = fn(signal, variance=0.05, rng=np.random.default_rng(0))
assert out.shape == signal.shape
@pytest.mark.parametrize("fn", [amod_bezier, amod_pchip])
def test_small_variance_stays_close(self, signal, fn):
out = fn(signal, variance=0.01, rng=np.random.default_rng(0))
np.testing.assert_allclose(out, signal, atol=0.1)
class TestDriftLinear:
def test_output_shape(self, signal):
out = drift_linear(signal, a=0.1, b=0.0, rng=np.random.default_rng(0))
assert out.shape == signal.shape
def test_zero_drift_unchanged(self, signal):
out = drift_linear(signal, a=0.0, b=0.0, rng=np.random.default_rng(0))
np.testing.assert_array_almost_equal(out, signal)
def test_interval_sampling(self, signal):
out = drift_linear(signal, a=[-0.5, 0.5], b=[-0.1, 0.1], rng=np.random.default_rng(0))
assert out.shape == signal.shape
def test_invalid_a_raises(self, signal):
with pytest.raises(ValueError):
drift_linear(signal, a=[1, 2, 3], b=0.0)
# ---------------------------------------------------------------------------
# augmenter
# ---------------------------------------------------------------------------
class TestAugment:
def test_output_shape_single_aug(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 20}}
out = augment(signal, config, seed=0)
assert out.shape == (N, 2) # original + 1 variant
def test_original_preserved_in_col0(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 20}}
out = augment(signal, config, seed=0)
np.testing.assert_array_equal(out[:, 0], signal)
def test_num_copies(self, signal):
config = {"a1": {"num_copies": 3, "Add_noise": True, "SNRdb": 20}}
out = augment(signal, config, seed=0)
assert out.shape == (N, 4) # original + 3 copies
def test_multiple_recipes(self, signal):
config = {
"a1": {"Flip": True},
"a2": {"Invert": True},
"a3": {"Add_noise": True, "SNRdb": 15},
}
out = augment(signal, config, seed=0)
assert out.shape == (N, 4)
def test_normalize_output_true(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 20}}
out = augment(signal, config, seed=0, normalize_output=True)
variant = out[:, 1]
assert variant.min() >= 0.0
assert variant.max() <= 1.0 + 1e-5
def test_normalize_output_false(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 5}}
out_norm = augment(signal, config, seed=0, normalize_output=True)
out_raw = augment(signal, config, seed=0, normalize_output=False)
assert not np.allclose(out_norm[:, 1], out_raw[:, 1])
def test_invalid_input_raises(self):
with pytest.raises(ValueError):
augment(np.ones((10, 2)), {})
def test_reproducibility(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 20}}
out1 = augment(signal, config, seed=42)
out2 = augment(signal, config, seed=42)
np.testing.assert_array_equal(out1, out2)
def test_flip(self, signal):
config = {"a1": {"Flip": True}}
out = augment(signal, config, seed=0, normalize_output=False)
np.testing.assert_array_equal(out[:, 1], np.flip(signal))
def test_invert(self, signal):
config = {"a1": {"Invert": True}}
out = augment(signal, config, seed=0, normalize_output=False)
np.testing.assert_array_almost_equal(out[:, 1], np.max(signal) - signal)
@pytest.mark.parametrize("key,extra", [
("Bezier_time_warp", {}),
("PCHIP_time_warp", {}),
("Bezier_amp_drift", {}),
("PCHIP_amp_drift", {}),
("Bezier_amp_mod", {}),
("PCHIP_amp_mod", {}),
("lf_noise", {"lf_noise_amplitude": 0.05, "lf_noise_frequency": 0.05}),
("powerline", {"powerline_amplitude": 0.05, "powerline_frequency": 0.1}),
("burst", {"burst_amplitude": 0.1, "burst_number": 3}),
("drift", {"drift_a": [-0.1, 0.1], "drift_b": [-0.05, 0.05]}),
])
def test_each_augmentation_runs(self, signal, key, extra):
config = {"a1": {key: True, **extra}}
out = augment(signal, config, seed=0)
assert out.shape == (N, 2)
assert not np.any(np.isnan(out))
class TestAugmentClass:
def test_callable(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 20}}
aug = Augment(config, seed=0)
out = aug(signal)
assert out.shape == (N, 2)
def test_seed_override(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 20}}
aug = Augment(config, seed=0)
out1 = aug(signal, seed=99)
out2 = aug(signal, seed=99)
np.testing.assert_array_equal(out1, out2)
def test_normalize_output_false(self, signal):
config = {"a1": {"Add_noise": True, "SNRdb": 5}}
aug = Augment(config, seed=0, normalize_output=False)
out = aug(signal)
assert out.shape == (N, 2)
# ---------------------------------------------------------------------------
# Regression tests for PhysioNet review comments (submission vuCy1q8QeG9EZqJLvMB1)
#
# Each class below pins one reviewer item so a later change cannot silently
# reintroduce the defect. Tags match PhysioNet_materials/reviewers-comments-s1.txt
# ---------------------------------------------------------------------------
WARP_FNS = [twarp_bezier, twarp_pchip, adrift_bezier,
adrift_pchip, amod_bezier, amod_pchip]
RAND_KNOTS_FNS = [twarp_bezier, twarp_pchip, adrift_pchip, amod_pchip]
BEZIER_AMP_FNS = [adrift_bezier, amod_bezier]
class TestSignalShapeHandling:
"""B3c: bare squeeze() replaced by explicit 1D validation."""
@pytest.mark.parametrize("fn", WARP_FNS)
@pytest.mark.parametrize("shape", [(N,), (N, 1), (1, N)])
def test_singleton_axes_accepted(self, fn, shape):
x = np.random.default_rng(0).standard_normal(shape)
out = fn(x, k=4, variance=0.01, rng=np.random.default_rng(0))
assert out.shape == (N,)
@pytest.mark.parametrize("fn", WARP_FNS)
def test_true_2d_rejected(self, fn):
x = np.zeros((3, 100))
with pytest.raises(ValueError, match="1D array"):
fn(x, k=4, variance=0.01, rng=np.random.default_rng(0))
def test_error_reports_original_shape(self):
with pytest.raises(ValueError, match=r"\(1, 2, 100\)"):
twarp_bezier(np.zeros((1, 2, 100)), k=4, rng=np.random.default_rng(0))
@pytest.mark.parametrize("shape", [(N,), (N, 1), (1, N)])
def test_drift_linear_singleton_axes(self, shape):
x = np.random.default_rng(0).standard_normal(shape)
out = drift_linear(x, a=[-0.3, 0.3], b=[-0.1, 0.1], rng=np.random.default_rng(0))
assert out.shape == (N,)
def test_burst_mask_singleton_axes(self):
noise = np.random.default_rng(0).standard_normal((1, N))
out = burst_mask(noise, n_bursts=2, burst_width=10, burst_base=0.0, burst_onset=5)
assert out.shape == (N,)
def test_burst_mask_rejects_true_2d(self):
with pytest.raises(ValueError, match="1D array"):
burst_mask(np.zeros((3, 100)), n_bursts=1, burst_width=5,
burst_base=0.0, burst_onset=0)
class TestAugmentShapeHandling:
"""B3c / A2: augment() agrees with the functions it calls."""
@pytest.mark.parametrize("shape", [(N,), (N, 1), (1, N)])
def test_singleton_axes_accepted(self, shape):
x = np.random.default_rng(0).standard_normal(shape)
out = augment(x, {"a": {"Add_noise": True, "SNRdb": 20}}, seed=0)
assert out.shape == (N, 2)
def test_true_2d_rejected(self):
with pytest.raises(ValueError, match=r"\(3, 100\)"):
augment(np.zeros((3, 100)), {"a": {"Flip": True}}, seed=0)
def test_list_input_accepted(self):
x = list(np.random.default_rng(0).standard_normal(N))
out = augment(x, {"a": {"Flip": True}}, seed=0)
assert out.shape == (N, 2)
class TestKValidation:
"""B3a / B3b: k is validated, with the bound each group actually needs."""
def test_rand_knots_returns_k_total_knots(self):
for k in (3, 4, 5, 6):
x_vals, y_vals = rand_knots(k=k, variance=1e-9, rng=np.random.default_rng(0))
assert len(x_vals) == k
assert len(y_vals) == k
@pytest.mark.parametrize("fn", RAND_KNOTS_FNS)
@pytest.mark.parametrize("k", [0, 1, 2])
def test_rand_knots_group_requires_k_ge_3(self, signal, fn, k):
with pytest.raises(ValueError, match="k must be"):
fn(signal, k=k, variance=0.01, rng=np.random.default_rng(0))
@pytest.mark.parametrize("fn", BEZIER_AMP_FNS)
def test_bezier_amp_rejects_k_zero(self, signal, fn):
# k=0 produced an empty envelope: amod_bezier zeroed the signal outright
with pytest.raises(ValueError, match="k must be"):
fn(signal, k=0, variance=0.05, rng=np.random.default_rng(0))
@pytest.mark.parametrize("fn", BEZIER_AMP_FNS)
@pytest.mark.parametrize("k", [1, 2])
def test_bezier_amp_allows_small_positive_k(self, signal, fn, k):
out = fn(signal, k=k, variance=0.05, rng=np.random.default_rng(0))
assert out.shape == signal.shape
assert np.all(np.isfinite(out))
@pytest.mark.parametrize("fn", WARP_FNS)
def test_non_integer_k_rejected(self, signal, fn):
with pytest.raises(ValueError, match="k must be"):
fn(signal, k=3.5, variance=0.01, rng=np.random.default_rng(0))
class TestFreqfiltShortSignal:
"""B3d: short signals get a package-level error, not a SciPy internal one."""
def test_short_signal_raises(self):
with pytest.raises(ValueError, match="too short"):
freqfilt(np.zeros(20), 0.1, 0.8)
def test_message_avoids_scipy_internals(self):
with pytest.raises(ValueError) as exc:
freqfilt(np.zeros(20), 0.1, 0.8)
message = str(exc.value)
assert "padlen" not in message
assert "20" in message
def test_lowpass_minimum_is_lower_than_bandpass(self):
# an order-4 lowpass needs > 15 samples; the bandpass needs > 27
lowpassed = freqfilt(np.random.default_rng(0).standard_normal(20), 0.0, 0.8)
assert lowpassed.shape == (20,)
with pytest.raises(ValueError, match="too short"):
freqfilt(np.zeros(20), 0.1, 0.8)
def test_full_band_bypasses_length_check(self):
x = np.zeros(5)
np.testing.assert_array_equal(freqfilt(x, 0, 1), x)
def test_long_signal_filters(self):
out = freqfilt(np.random.default_rng(0).standard_normal(200), 0.1, 0.8)
assert out.shape == (200,)
class TestBurstCount:
"""B2: n_bursts=0 means no bursts, not full-length broadband noise."""
def test_zero_bursts_adds_nothing(self, signal):
noisy, burst = burstify(signal, snr_db=10, n_bursts=0, rng=np.random.default_rng(0))
assert np.count_nonzero(burst) == 0
np.testing.assert_array_equal(noisy, signal)
def test_negative_bursts_rejected(self, signal):
with pytest.raises(ValueError, match="n_bursts"):
burstify(signal, snr_db=10, n_bursts=-1, rng=np.random.default_rng(0))
def test_coverage_increases_with_burst_count(self, signal):
_, b1 = burstify(signal, snr_db=10, n_bursts=1, burst_width=10,
rng=np.random.default_rng(3))
_, b5 = burstify(signal, snr_db=10, n_bursts=5, burst_width=10,
rng=np.random.default_rng(3))
assert 0 < np.count_nonzero(b1) < np.count_nonzero(b5) < len(signal)
class TestResamplePool:
"""B1: resample_pool accepts an external array, as the README documents."""
def test_ndarray_pool(self, signal):
pool = np.random.default_rng(1).standard_normal(1000)
noisy, noise = noisify(signal, snr_db=15, dist="resample", resample_pool=pool,
rng=np.random.default_rng(0))
assert noisy.shape == signal.shape
assert noise.shape == signal.shape
def test_list_pool(self, signal):
pool = list(np.random.default_rng(1).standard_normal(100))
noisy, _ = noisify(signal, snr_db=15, dist="resample", resample_pool=pool,
rng=np.random.default_rng(0))
assert noisy.shape == signal.shape
def test_self_pool_still_works(self, signal):
noisy, _ = noisify(signal, snr_db=15, dist="resample", resample_pool="self",
rng=np.random.default_rng(0))
assert noisy.shape == signal.shape
def test_missing_pool_rejected(self, signal):
with pytest.raises(ValueError, match="resample_pool"):
noisify(signal, snr_db=15, dist="resample", resample_pool=None,
rng=np.random.default_rng(0))
def test_samples_are_drawn_from_the_pool(self, signal):
# a two-valued pool must yield exactly two distinct noise levels
pool = np.array([-1.0, 1.0])
_, noise = noisify(signal, snr_db=15, dist="resample", resample_pool=pool,
zero_mean=False, rng=np.random.default_rng(0))
assert len(np.unique(np.round(noise, 9))) == 2