diff --git a/meteor/tv.py b/meteor/tv.py index 56bd3ec..81118f0 100644 --- a/meteor/tv.py +++ b/meteor/tv.py @@ -63,7 +63,7 @@ def tv_denoise_difference_map( full_output: bool = False, difference_map_amplitude_column: str = "DF", difference_map_phase_column: str = "PHIC", - lambda_values_to_scan: Sequence[float] | None = None, + lambda_values_to_scan: Sequence[float] | np.ndarray | None = None, ) -> rs.DataSet | tuple[rs.DataSet, TvDenoiseResult]: """Single-pass TV denoising of a difference map. diff --git a/test/unit/test_tv.py b/test/unit/test_tv.py index 34e7a73..2a672bb 100644 --- a/test/unit/test_tv.py +++ b/test/unit/test_tv.py @@ -111,6 +111,34 @@ def noisy_map() -> rs.DataSet: return displaced_single_atom_difference_map_coefficients(noise_sigma=0.03) +@pytest.mark.parametrize( + "lambda_values_to_scan", + [ + None, + [ + 1.0, + ], + ], +) +@pytest.mark.parametrize("full_output", [False, True]) +def test_tv_denoise_difference_map_smoke( + lambda_values_to_scan: None | Sequence[float], + full_output: bool, + noisy_map: rs.DataSet, +) -> None: + output = tv.tv_denoise_difference_map( + difference_map_coefficients=noisy_map, + lambda_values_to_scan=lambda_values_to_scan, + full_output=full_output, + ) # type: ignore + if full_output: + assert len(output) == 2 + assert isinstance(output[0], rs.DataSet) + assert isinstance(output[1], tv.TvDenoiseResult) + else: + assert isinstance(output, rs.DataSet) + + @pytest.mark.parametrize("lambda_values_to_scan", [None, np.logspace(-3, 2, 100)]) def test_tv_denoise_difference_map( lambda_values_to_scan: None | Sequence[float],