def test_2d_tv_denoise():
rng = np.random.RandomState(123)
data = ascent().astype(float)
data_noisy = data + data.std() * rng.randn(*data.shape)
data_clean = tv_denoise(data, weight=60)
norm_noisy = np.linalg.norm(data - data_noisy) / np.linalg.norm(data)
norm_clean = np.linalg.norm(data - data_clean) / np.linalg.norm(data)
np.testing.assert_allclose(norm_noisy, 0.48604971)
np.testing.assert_allclose(norm_clean, 0.10888393)
def test_3d_tv_denoise():
rng = np.random.RandomState(123)
x, y, z = np.ogrid[0:40, 0:40, 0:40]
data = (x - 22)**2 + (y - 20)**2 + (z - 17)**2 < 8**2
data = 255 * data.astype(float)
data_noisy = data + data.std() * rng.randn(*data.shape)
data_clean = tv_denoise(data_noisy, weight=100)
norm_noisy = np.linalg.norm(data - data_noisy) / np.linalg.norm(data)
norm_clean = np.linalg.norm(data - data_clean) / np.linalg.norm(data)
print(norm_clean)
np.testing.assert_allclose(norm_noisy, 0.98151071)
np.testing.assert_allclose(norm_clean, 0.12519535)
def test_tv_denoise_error():
with pytest.raises(ValueError, match="may be denoised"):
_ = tv_denoise(np.array([1, 2, 3]))
