Added suppress warning option

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). ",

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