diff --git a/tobac/segmentation/watershed_segmentation.py b/tobac/segmentation/watershed_segmentation.py index 5d40cb30..eb130d75 100644 --- a/tobac/segmentation/watershed_segmentation.py +++ b/tobac/segmentation/watershed_segmentation.py @@ -286,6 +286,7 @@ def segmentation_3D( PBC_flag="none", seed_3D_flag="column", statistic=None, + suppress_warnings=False, ): """Wrapper for the segmentation()-function.""" @@ -301,6 +302,7 @@ def segmentation_3D( PBC_flag=PBC_flag, seed_3D_flag=seed_3D_flag, statistic=statistic, + suppress_warnings=suppress_warnings, ) @@ -316,6 +318,7 @@ def segmentation_2D( PBC_flag="none", seed_3D_flag="column", statistic=None, + suppress_warnings=False, ): """Wrapper for the segmentation()-function.""" return segmentation( @@ -330,6 +333,7 @@ def segmentation_2D( PBC_flag=PBC_flag, seed_3D_flag=seed_3D_flag, statistic=statistic, + suppress_warnings=suppress_warnings, ) @@ -1136,6 +1140,7 @@ def segmentation( segment_number_unassigned: int = 0, statistic: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None] = None, time_padding: Optional[datetime.timedelta] = datetime.timedelta(seconds=0.5), + suppress_warnings: Optional[bool] = False, return_cells: bool = False, stubs: Optional[int] = None, ) -> tuple[xr.DataArray, pd.DataFrame]: @@ -1215,6 +1220,8 @@ def segmentation( timestep that is time_padding off of the feature. Extremely useful when converting between micro- and nanoseconds, as is common when using Pandas dataframes. + suppress_warnings: bool, optional + If True, suppresses warnings. Default is False. return_cells: bool, optional (default: False) If True, the segmentation mask returned will use the cell values of the input dataframe, rather than the feature values. This requires the @@ -1287,7 +1294,7 @@ def segmentation( features_out_list = [] - if len(field.coords[time_var_name]) == 1: + if not suppress_warnings and len(field.coords[time_var_name]) == 1: warnings.warn( "As of v1.6.0, segmentation with time length 1 will return time as a coordinate" " instead of dropping it (i.e., output will now be 1xMxN instead of MxN). ", diff --git a/tobac/tests/segmentation_tests/test_segmentation.py b/tobac/tests/segmentation_tests/test_segmentation.py index 9702ee05..19263a37 100644 --- a/tobac/tests/segmentation_tests/test_segmentation.py +++ b/tobac/tests/segmentation_tests/test_segmentation.py @@ -1,3 +1,7 @@ +import warnings +import pytest +import tobac.segmentation as seg +import tobac.segmentation.watershed_segmentation as watershed_segmentation from datetime import datetime import numpy as np import pandas as pd @@ -1347,3 +1351,55 @@ def test_segmentation_returns_early_for_empty_features(): assert np.all(sef_mask_none.data == 0) assert features_out_none is None + + +@pytest.mark.parametrize( + "seg_func", + [ + watershed_segmentation.segmentation, + watershed_segmentation.segmentation_2D, + watershed_segmentation.segmentation_3D, + ], +) +@pytest.mark.parametrize( + "suppress_warnings, expect_warning", + [ + (False, True), + (True, False), + ("default", True), + ], +) +def test_segmentation_suppress_warnings(seg_func, suppress_warnings, expect_warning): + """Test that the 1xMxN warning is issued when suppress_warnings=False + and suppressed when suppress_warnings=True + """ + test_dset_size = (1, 50, 50) + test_data = np.zeros(test_dset_size) + test_data_xarray = testing.make_dataset_from_arr( + test_data, data_type="xarray", time_dim_num=0, y_dim_num=1, x_dim_num=2 + ) + features = testing.generate_single_feature( + start_h1=10, start_h2=10, max_h1=50, max_h2=50 + ) + + kwargs = {"dxy": 1000.0, "threshold": 1.5} + if suppress_warnings != "default": + kwargs["suppress_warnings"] = suppress_warnings + + warning_msg = ( + "As of v1.6.0, segmentation with time length 1 will return time as a coordinate" + ) + + if expect_warning: + with pytest.warns(UserWarning, match=warning_msg): + seg_func(features, test_data_xarray, **kwargs) + else: + with warnings.catch_warnings(record=True) as record: + warnings.simplefilter("always") + seg_func(features, test_data_xarray, **kwargs) + matching_warnings = [ + w + for w in record + if issubclass(w.category, UserWarning) and warning_msg in str(w.message) + ] + assert len(matching_warnings) == 0