Repository: iver56/audiomentations
Pull request #387 (2025-04-20 04:09:33)

Add optional weights argument to OneOf for weighted random selection
## Description

This addition should leave all existing scripts working the same, but allow for an optional argument for `OneOf` that allows weighted selection of the transform, instead of uniform random. 

If `weights` are not supplied, behavior will default to the way it has always been (uniform random), so this is fully backwards compatible :)

## Motivation

I have a large number of impulse response (IR) files that I wish to use at random in some augmentations for a model I am training, for example from my repo: [impulse-response-audio-files](https://github.com/worldveil/impulse-response-audio-files).

However, not all categories of these IRs are equally as important to our usecase, and as such, we'd like to upsample certain classes (ie: venues, large rooms), and downsample others (ie: small rooms, misc other IRs). 

## Example script

Let's envision a scenario in which we'd like to upsample one category 4x vs another (0.8 vs 0.2):

```python
from pathlib import Path
from collections import Counter

from audiomentations import Compose, ApplyImpulseResponse, OneOf
import numpy as np

# load some impulse responses from: https://github.com/worldveil/impulse-response-audio-files
IR_DIR = "./impulse-response-audio-files"
venues_dir1 = Path(IR_DIR) / "echo_thief" / "Venues"
misc_dir2 = Path(IR_DIR) / "voxengo"
weights = [0.8, 0.2]

augment = Compose([
    OneOf([
        ApplyImpulseResponse(ir_path=venues_dir1, p=1.0),
        ApplyImpulseResponse(ir_path=misc_dir2, p=1.0)
    ], p=0.5, weights=weights),
])

# Generate 2 seconds of dummy audio for the sake of example
samples = np.random.uniform(low=-0.2, high=0.2, size=(32000,)).astype(np.float32)

# let's count the number of times each transform is chosen to verify it's working
index_chosen = Counter()
one_of = augment.transforms[-1]

for i in range(1000):
    # Augment/transform/perturb the audio data
    augmented_samples = augment(samples=samples, sample_rate=44100)

    # roll the dice
    for i, transform in enumerate(augment.transforms):
        if i == len(augment.transforms) - 1:
            #print(f"We chose transform: {one_of.transforms[transform.transform_index].__class__}, which was index={transform.transform_index} (weights={one_of.weights})")
            index_chosen[transform.transform_index] += 1

# we should expect to see zero 4x as many times as 1
print(index_chosen)
```

For me, the output is:

```
Counter({0: 791, 1: 209})
```

which aligns with our expectation.

## Tests

Are all passing, with the exception of the two that depend on `fast-align-audio` which I wasn't able to install in the 30 min I had to make this PR :) but from quick glance this PR should not affect that functionality at all.


Patch:
diff --git a/audiomentations/core/composition.py b/audiomentations/core/composition.py
index a0d9861..2103ef4 100644
--- a/audiomentations/core/composition.py
+++ b/audiomentations/core/composition.py
@@ -1,5 +1,5 @@
 import random
-from typing import Any
+from typing import Any, Optional
 
 import numpy as np
 from numpy.typing import NDArray
@@ -7,6 +7,7 @@ from numpy.typing import NDArray
 from audiomentations.core.transforms_interface import BaseSpectrogramTransform
 from audiomentations.core.serialization import get_shortest_class_fullname
 from audiomentations.core.utils import format_args
+from audiomentations.core.sampling import WeightedChoiceSampler
 
 REPR_INDENT_STEP = 2
 
@@ -259,31 +260,48 @@ class SomeOf(BaseCompose):
 class OneOf(BaseCompose):
     """
     OneOf randomly picks one of the given transforms when called, and applies that
-    transform.
+    transform. Optional `weights` can be supplied to guide the probability of each transform being chosen.
     Usage example:
     ```
     augment = OneOf([
         TimeStretch(min_rate=0.8, max_rate=1.25, p=1.0),
         PitchShift(min_semitones=-4, max_semitones=4, p=1.0),
-    ])
+    ], weights=[0.2, 0.8])
     # Generate 2 seconds of dummy audio for the sake of example
     samples = np.random.uniform(low=-0.2, high=0.2, size=(32000,)).astype(np.float32)
     # Augment/transform/perturb the audio data
     augmented_samples = augment(samples=samples, sample_rate=16000)
-    # Result: The audio was either time-stretched or pitch-shifted, but not both
+    # Result: The audio was either time-stretched (less likely) or pitch-shifted (4x more likely), but not both
     ```
     """
 
-    def __init__(self, transforms, p: float = 1.0):
+    def __init__(self, transforms, p: float = 1.0, weights: Optional[list[float]] = None):
         super().__init__(transforms, p)
         self.transform_index = 0
         self.should_apply = True
 
+        # Handle weights initialization and validation
+        if weights is not None:
+            if len(weights) != len(transforms):
+                raise ValueError("Length of weights must match length of transforms")
+            if any(w < 0 for w in weights):
+                raise ValueError("All weights must be non-negative")
+            if sum(weights) == 0:
+                raise ValueError("Sum of weights must be greater than 0")
+            # Normalize weights to sum to 1
+            self.weights = np.array(weights)
+            self.weights /= np.sum(self.weights)
+        else:
+            # If no weights provided, use uniform distribution
+            self.weights = np.ones(len(transforms)) / len(transforms)
+
+        self.sampler = WeightedChoiceSampler(self.weights)
+
     def randomize_parameters(self, *args, **kwargs):
         super().randomize_parameters(*args, **kwargs)
         self.should_apply = random.random() < self.p
         if self.should_apply:
-            self.transform_index = random.randint(0, len(self.transforms) - 1)
+            self.transform_index = self.sampler.sample(size=1)[0]
 
     def __call__(self, *args, **kwargs):
         if not self.are_parameters_frozen:
diff --git a/audiomentations/core/sampling.py b/audiomentations/core/sampling.py
new file mode 100644
index 0000000..ea27689
--- /dev/null
+++ b/audiomentations/core/sampling.py
@@ -0,0 +1,30 @@
+import numpy as np
+
+
+class WeightedChoiceSampler:
+    """
+    A class for sampling from a weighted distribution using the cumulative distribution function (CDF).
+
+    Usage:
+    >>> sampler = WeightedChoiceSampler(weights=[0.2, 0.8])
+    >>> samples = sampler.sample(size=10)
+    """
+    def __init__(self, weights: list[float] = None):
+        assert weights is not None and len(weights) > 0, "Weights must be provided"
+
+        # ensure weights are normalized
+        weights = np.array(weights)
+        weights /= np.sum(weights)
+        assert np.all(weights >= 0), "Weights must be non-negative"
+
+        # initialize the sampler
+        self.num_transforms = len(weights)
+        self.cdf = np.cumsum(weights)
+        # Ensure the last element is exactly 1.0 due to potential floating point errors
+        self.cdf[-1] = 1.0
+
+    # Modified sample method to take size
+    def sample(self, size: int = 1) -> np.ndarray:
+        random_vals = np.random.rand(size)
+        # Use searchsorted with the array of random values
+        return np.searchsorted(self.cdf, random_vals, side='right')
diff --git a/docs/index.md b/docs/index.md
index 2c40f75..aff0656 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -101,7 +101,9 @@ Same as Compose, but for spectrogram transforms
 
 ## `OneOf`
 
-OneOf randomly picks one of the given transforms when called, and applies that transform.
+OneOf randomly picks one of the given transforms when called, and applies that transform. 
+
+An optional `weights` list of floats may be given to guide the probability of each transform for being chosen. If not specified, a transform is chosen uniformly at random.
 
 ## `SomeOf`
 


Test patch:
diff --git a/tests/test_one_of.py b/tests/test_one_of.py
index 83b20da..6a388f8 100644
--- a/tests/test_one_of.py
+++ b/tests/test_one_of.py
@@ -1,7 +1,8 @@
 import os
+import pytest
 
 import numpy as np
-from numpy.testing import assert_array_equal, assert_array_almost_equal
+from numpy.testing import assert_array_equal, assert_array_almost_equal, assert_allclose
 
 from audiomentations import (
     ClippingDistortion,
@@ -194,3 +195,81 @@ def test_one_of_spectrogram_magnitude_with_p_0():
     assert augmented_spectrogram.shape == spectrogram.shape
     assert augmented_spectrogram.dtype == spectrogram.dtype
     assert_array_equal(augmented_spectrogram, spectrogram)
+
+def test_one_of_weights():
+    samples = np.array([1.0, 0.5, -0.25, -0.125, 0.0], dtype=np.float32)
+    sample_rate = 16000
+    
+    transforms = [
+        Gain(min_gain_db=6.0, max_gain_db=6.0, p=1.0),  # Gain approx 2.0
+        Gain(min_gain_db=-6.0, max_gain_db=-6.0, p=1.0), # Gain approx 0.5
+    ]
+    weights = [0.8, 0.2]
+    
+    augmenter = OneOf(transforms=transforms, p=1.0, weights=weights)
+    
+    counts = [0, 0]
+    num_runs = 2000
+    for _ in range(num_runs):
+        perturbed_samples = augmenter(samples=samples, sample_rate=sample_rate)
+        # Check which transform was applied based on the gain effect
+        if np.allclose(perturbed_samples[0], samples[0] * 2.0, atol=1e-5):
+            counts[0] += 1
+        elif np.allclose(perturbed_samples[0], samples[0] * 0.5, atol=1e-5):
+            counts[1] += 1
+        else:
+             # This should not happen if p=1.0 and gains are fixed
+             raise AssertionError("Unexpected output sample value")
+
+    observed_proportions = np.array(counts) / num_runs
+    assert_allclose(observed_proportions, weights, atol=0.05)
+
+
+def test_one_of_weights_normalization():
+    samples = np.array([1.0], dtype=np.float32)
+    sample_rate = 16000
+    
+    transforms = [
+        Gain(min_gain_db=6.0, max_gain_db=6.0, p=1.0), 
+        Gain(min_gain_db=-6.0, max_gain_db=-6.0, p=1.0),
+    ]
+    # Weights don't sum to 1, should be normalized to [0.8, 0.2]
+    weights = [4, 1] 
+    expected_normalized_weights = [0.8, 0.2]
+
+    augmenter = OneOf(transforms=transforms, p=1.0, weights=weights)
+    
+    # Check internal weights are normalized
+    assert_allclose(augmenter.weights, expected_normalized_weights)
+
+    counts = [0, 0]
+    num_runs = 2000
+    for _ in range(num_runs):
+        perturbed_samples = augmenter(samples=samples, sample_rate=sample_rate)
+        if np.allclose(perturbed_samples[0], samples[0] * 2.0, atol=1e-5):
+            counts[0] += 1
+        elif np.allclose(perturbed_samples[0], samples[0] * 0.5, atol=1e-5):
+            counts[1] += 1
+
+    observed_proportions = np.array(counts) / num_runs
+    # Check proportions match the *normalized* weights
+    assert_allclose(observed_proportions, expected_normalized_weights, atol=0.05)
+
+
+def test_one_of_invalid_weights():
+    transforms = [
+        Gain(p=1.0),
+        PolarityInversion(p=1.0),
+    ]
+
+    # Length mismatch
+    with pytest.raises(ValueError, match="Length of weights must match"):
+        OneOf(transforms=transforms, weights=[0.5])
+        
+    # Negative weights
+    with pytest.raises(ValueError, match="All weights must be non-negative"):
+        OneOf(transforms=transforms, weights=[0.5, -0.1])
+
+    # Weights sum to zero
+    with pytest.raises(ValueError, match="Sum of weights must be greater than 0"):
+        OneOf(transforms=transforms, weights=[0.0, 0.0])
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
new file mode 100644
index 0000000..8f45d6d
--- /dev/null
+++ b/tests/test_sampling.py
@@ -0,0 +1,139 @@
+import pytest
+import numpy as np
+from numpy.testing import assert_allclose, assert_array_equal, assert_equal
+
+from audiomentations.core.sampling import WeightedChoiceSampler
+
+
+class TestWeightedChoiceSampler:
+    def test_init_valid_weights(self):
+        weights = [0.2, 0.8]
+        sampler = WeightedChoiceSampler(weights=weights)
+        assert sampler.num_transforms == 2
+        assert_allclose(sampler.cdf, [0.2, 1.0])
+        assert sampler.cdf[-1] == 1.0  # Ensure last element is exactly 1.0
+
+    def test_init_valid_weights_needs_normalization(self):
+        weights = [1, 4] # Sums to 5
+        sampler = WeightedChoiceSampler(weights=weights)
+        assert sampler.num_transforms == 2
+        assert_allclose(sampler.cdf, [0.2, 1.0])
+        assert sampler.cdf[-1] == 1.0
+
+    def test_init_invalid_weights_none(self):
+        with pytest.raises(AssertionError, match="Weights must be provided"):
+            WeightedChoiceSampler(weights=None)
+
+    def test_init_invalid_weights_empty(self):
+        with pytest.raises(AssertionError, match="Weights must be provided"):
+            WeightedChoiceSampler(weights=[])
+
+    def test_init_invalid_weights_negative(self):
+        # WeightedChoiceSampler uses numpy's assert_all which raises AssertionError
+        with pytest.raises(AssertionError, match="Weights must be non-negative"):
+            WeightedChoiceSampler(weights=[0.5, -0.1])
+
+    def test_init_weights_all_zero(self):
+        # WeightedChoiceSampler checks for sum > 0 via division, leading to RuntimeWarning
+        # and then NaN values, which fail the non-negative check.
+        with pytest.raises(AssertionError, match="Weights must be non-negative"):
+            WeightedChoiceSampler(weights=[0.0, 0.0])
+
+    def test_sample_single(self):
+        weights = [0.1, 0.2, 0.7]
+        sampler = WeightedChoiceSampler(weights=weights)
+        sample_index = sampler.sample(size=1)[0]
+        assert 0 <= sample_index < 3
+
+    def test_sample_multiple(self):
+        weights = [0.1, 0.2, 0.7]
+        sampler = WeightedChoiceSampler(weights=weights)
+        sample_indices = sampler.sample(size=100)
+        assert len(sample_indices) == 100
+        assert all(0 <= idx < 3 for idx in sample_indices)
+
+    def test_sample_distribution(self):
+        weights = [0.1, 0.3, 0.6]
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 20000
+        sample_indices = sampler.sample(size=num_samples)
+
+        counts = np.bincount(sample_indices, minlength=len(weights))
+        observed_proportions = counts / num_samples
+
+        normalized_weights = np.array(weights) / np.sum(weights)
+
+        # Use a tolerance suitable for statistical tests
+        assert_allclose(observed_proportions, normalized_weights, atol=0.02)
+
+    def test_sample_edge_case_one_hot_first(self):
+        weights = [1.0, 0.0, 0.0]
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 100
+        sample_indices = sampler.sample(size=num_samples)
+        assert all(idx == 0 for idx in sample_indices)
+
+    def test_sample_edge_case_one_hot_middle(self):
+        weights = [0.0, 1.0, 0.0]
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 100
+        sample_indices = sampler.sample(size=num_samples)
+        assert all(idx == 1 for idx in sample_indices)
+
+    def test_sample_edge_case_uniform(self):
+        weights = [1.0, 1.0, 1.0]
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 30000
+        sample_indices = sampler.sample(size=num_samples)
+
+        counts = np.bincount(sample_indices, minlength=len(weights))
+        observed_proportions = counts / num_samples
+
+        expected_proportions = np.array([1/3, 1/3, 1/3])
+
+        assert_allclose(observed_proportions, expected_proportions, atol=0.02)
+
+    def test_sample_distribution_uniform(self):
+        # Use normalized weights for testing distribution
+        weights = np.array([0.5, 0.5])
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 10000
+        sample_indices = sampler.sample(size=num_samples)
+        
+        counts = np.bincount(sample_indices, minlength=sampler.num_transforms)
+        observed_proportions = counts / num_samples
+        
+        assert_allclose(observed_proportions, weights, atol=0.05)
+
+    def test_sample_distribution_non_uniform(self):
+         # Use normalized weights for testing distribution
+        weights = np.array([0.1, 0.7, 0.2])
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 20000
+        sample_indices = sampler.sample(size=num_samples)
+
+        counts = np.bincount(sample_indices, minlength=sampler.num_transforms)
+        observed_proportions = counts / num_samples
+        
+        assert_allclose(observed_proportions, weights, atol=0.05)
+
+    def test_sample_distribution_with_zero_weight(self):
+         # Use normalized weights for testing distribution
+        weights = np.array([0.5, 0.0, 0.5])
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 10000
+        sample_indices = sampler.sample(size=num_samples)
+
+        counts = np.bincount(sample_indices, minlength=sampler.num_transforms)
+        observed_proportions = counts / num_samples
+        
+        assert counts[1] == 0 # Index 1 should never be chosen
+        assert_allclose(observed_proportions, weights, atol=0.05)
+
+    def test_sample_single_weight(self):
+        weights = [1.0]
+        sampler = WeightedChoiceSampler(weights=weights)
+        num_samples = 100
+        sample_indices = sampler.sample(size=num_samples)
+        assert_array_equal(sample_indices, np.zeros(num_samples, dtype=int))
+        assert sampler.cdf[-1] == 1.0 
\ No newline at end of file