dither tests

Author: CamiWilliams

test/assets/sinewave_soxdither.wav

@@ -268,6 +268,21 @@ def test_scale_to_interval(self):
         self.assertTrue(torch.max(waveform_scaled) <= scaled)
         self.assertTrue(torch.min(waveform_scaled) >= -scaled)
 
+    def test_dither(self):
+        waveform, sample_rate = torchaudio.load(self.test_filepath)
+        waveform_dithered = F.dither(waveform)
+        waveform_dithered_noiseshaped = F.dither(waveform, noise_shaping=True)
+
+        test_filepath_sox_dither = os.path.join(self.test_dirpath, "assets", "sinewave_soxdither.wav")
+        sox_dither_waveform, sox_sr = torchaudio.load(test_filepath_sox_dither)
+
+        self.assertTrue(torch.allclose(waveform_dithered, sox_dither_waveform, rtol=1e-03, atol=1e-03))
+
+        test_filepath_sox_dither_ns = os.path.join(self.test_dirpath, "assets", "sinewave_soxdither_noiseshaping.wav")
+        sox_dither_waveform_ns, sox_sr = torchaudio.load(test_filepath_sox_dither_ns)
+
+        self.assertTrue(torch.allclose(waveform_dithered_noiseshaped, sox_dither_waveform_ns, rtol=1e-03, atol=1e-03))
+
 
 def _num_stft_bins(signal_len, fft_len, hop_length, pad):
     return (signal_len + 2 * pad - fft_len + hop_length) // hop_length

test/assets/sinewave_soxdither_noiseshaping.wav

@@ -842,21 +842,20 @@ def scale_to_interval(waveform, interval=1.0):
 
 
 def _add_noise_shaping(dithered_waveform, waveform):
-  r"""Noise shaping is calculated by error:
-  error[n] = dithered[n] - original[n]
-  noise_shaped_waveform[n] = dithered[n] + error[n-1]
-  """
-  error = dithered_waveform - waveform
-
-  # add error[n-1] to dithered_waveform[n], so offset the error by 1 index
-  error0_offset = torch.cat((torch.zeros(1), error[0]))
-  error[0] = error0_offset[:waveform.size()[1]]
-  error1_offset = torch.cat((torch.zeros(1), error[1]))
-  error[1] = error1_offset[:waveform.size()[1]]
+    r"""Noise shaping is calculated by error:
+    error[n] = dithered[n] - original[n]
+    noise_shaped_waveform[n] = dithered[n] + error[n-1]
+    """
+    error = dithered_waveform - waveform
 
-  noise_shaped = dithered_waveform + error
+    # add error[n-1] to dithered_waveform[n], so offset the error by 1 index
+    for index in range(error.size()[0]):
+        err = error[index]
+        error_offset = torch.cat((torch.zeros(1), err))
+        error[index] = error_offset[:waveform.size()[1]]
 
-  return noise_shaped
+    noise_shaped = dithered_waveform + error
+    return noise_shaped
 
 
 def dither(waveform, probability_density_function="TPDF", noise_shaping=False, filter=None):
@@ -890,42 +889,43 @@ def dither(waveform, probability_density_function="TPDF", noise_shaping=False, f
 
     number_of_bits = 16
 
-    up_scaling = 2 ** (number_of_bits-1) - 2
+    up_scaling = 2 ** (number_of_bits - 1) - 2
     signal_scaled = waveform * up_scaling
     down_scaling = 2 ** (number_of_bits - 1)
 
+    dithered = waveform
+
     if (probability_density_function == "RPDF"):
         RPDF_dither = waveform[0][random.randint(1, wave_size)] - 0.5
 
         signal_scaled_RPDF_dithered = signal_scaled + RPDF_dither
         quantised_signal_scaled_RPDF_dithered = torch.round(signal_scaled_RPDF_dithered)
         quantised_signal_RPDF_dithered = quantised_signal_scaled_RPDF_dithered / down_scaling
 
-        if noise_shaped: return _add_noise_shaping(quantised_signal_RPDF_dithered, waveform)
-
-        return quantised_signal_RPDF_dithered
+        dithered = quantised_signal_RPDF_dithered
     elif (probability_density_function == "GPDF"):
         gaussian_dither = (waveform[0][random.randint(1, wave_size)]
-            + waveform[0][random.randint(1, wave_size)]
-            + waveform[0][random.randint(1, wave_size)]
-            + waveform[0][random.randint(1, wave_size)]
-            + waveform[0][random.randint(1, wave_size)]
-            + waveform[0][random.randint(1, wave_size)]) / 6
+                           + waveform[0][random.randint(1, wave_size)]
+                           + waveform[0][random.randint(1, wave_size)]
+                           + waveform[0][random.randint(1, wave_size)]
+                           + waveform[0][random.randint(1, wave_size)]
+                           + waveform[0][random.randint(1, wave_size)]) / 6
 
         signal_scaled_gaussian_dithered = signal_scaled + gaussian_dither
         quantised_signal_scaled_gaussian_dithered = torch.round(signal_scaled_gaussian_dithered)
         quantised_signal_gaussian_dithered = quantised_signal_scaled_gaussian_dithered / down_scaling
 
-        if noise_shaped: return _add_noise_shaping(quantised_signal_gaussian_dithered, waveform)
-
-        return quantised_signal_gaussian_dithered
+        dithered = quantised_signal_gaussian_dithered
     else:
         TPDF_dither = waveform[0][random.randint(1, wave_size)] - waveform[0][random.randint(1, wave_size)]
 
         signal_scaled_TPDF_dithered = signal_scaled + TPDF_dither
         quantised_signal_scaled_TPDF_dithered = torch.round(signal_scaled_TPDF_dithered)
         quantised_signal_TPDF_dithered = quantised_signal_scaled_TPDF_dithered / down_scaling
 
-        if noise_shaped: return _add_noise_shaping(quantised_signal_TPDF_dithered, waveform)
+        dithered = quantised_signal_TPDF_dithered
 
-        return quantised_signal_TPDF_dithered
+    if noise_shaping:
+        return _add_noise_shaping(dithered, waveform)
+    else:
+        return dithered