Fix remaining issues with datetime conversions and coordinate naming for iris/xarray compatibility

tobac/feature_detection.py

@@ -1180,7 +1180,8 @@ def feature_detection_multithreshold(
     strict_thresholding: bool = False,
     statistic: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None] = None,
     statistics_unsmoothed: bool = False,
-    preserve_iris_datetime_types: bool = True,
+    use_standard_names: Optional[bool] = None,
+    converted_from_iris: bool = False,
     **kwargs,
 ) -> pd.DataFrame:
     """Perform feature detection based on contiguous regions.
@@ -1264,6 +1265,11 @@ def feature_detection_multithreshold(
         If True, a feature can only be detected if all previous thresholds have been met.
         Default is False.
 
+    use_standard_names: bool
+        If true, when interpolating a coordinate, it looks for a standard_name
+        and uses that to name the output coordinate, to mimic iris functionality.
+        If false, uses the actual name of the coordinate to output.
+
     preserve_iris_datetime_types: bool, optional, default: True
         If True, for iris input, preserve the original datetime type (typically
         `cftime.DatetimeGregorian`) where possible. For xarray input, this parameter has no
@@ -1409,18 +1415,23 @@ def feature_detection_multithreshold(
     if any([not x.empty for x in list_features_timesteps]):
         features = pd.concat(list_features_timesteps, ignore_index=True)
         features["feature"] = features.index + feature_number_start
-        #    features_filtered = features.drop(features[features['num'] < min_num].index)
-        #    features_filtered.drop(columns=['idx','num','threshold_value'],inplace=True)
+
+        if use_standard_names is None:
+            use_standard_names = True if converted_from_iris else False
+
         if "vdim" in features:
             features = add_coordinates_3D(
                 features,
                 field_in,
                 vertical_coord=vertical_coord,
-                preserve_iris_datetime_types=kwargs["converted_from_iris"]
-                & preserve_iris_datetime_types,
+                use_standard_names=use_standard_names,
             )
         else:
-            features = add_coordinates(features, field_in)
+            features = add_coordinates(
+                features,
+                field_in,
+                use_standard_names=use_standard_names,
+            )
 
         # Loop over DataFrame to remove features that are closer than distance_min to each
         # other:
@@ -1452,12 +1463,6 @@ def feature_detection_multithreshold(
                 )
             features = pd.concat(filtered_features, ignore_index=True)
 
-            features = add_coordinates(
-                features,
-                field_in,
-                preserve_iris_datetime_types=kwargs["converted_from_iris"]
-                & preserve_iris_datetime_types,
-            )
     else:
         features = None
         logging.debug("No features detected")

tobac/segmentation/watershed_segmentation.py

@@ -49,6 +49,7 @@
 from tobac.utils import internal as internal_utils
 from tobac.utils import get_statistics
 from tobac.utils import decorators
+from tobac.utils.generators import field_and_features_over_time
 
 
 def add_markers(
@@ -1257,32 +1258,21 @@ def segmentation(
     )
     features_out_list = []
 
-    # Iris workaround: convert cftime to datetime64
-
-    if np.issubdtype(features["time"].dtype, np.datetime64):
-        # we are (likely) a numpy datetime
-        all_times = features["time"]
-    else:
-        all_times = features["time"].map(np.datetime64)
-
     if 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). ",
             UserWarning,
         )
 
-    for time_iteration_number, time_iteration_value in enumerate(
-        field.coords[time_var_name]
+    for (
+        time_iteration_number,
+        time_iteration_value,
+        field_at_time,
+        features_i,
+    ) in field_and_features_over_time(
+        field, features, time_var_name=time_var_name, time_padding=time_padding
     ):
-        field_at_time = field.isel({time_var_name: time_iteration_number})
-        if time_padding is not None:
-            padded_conv = pd.Timedelta(time_padding).to_timedelta64()
-            min_time = time_iteration_value.values - padded_conv
-            max_time = time_iteration_value.values + padded_conv
-            features_i = features.loc[all_times.between(min_time, max_time)]
-        else:
-            features_i = features.loc[all_times == time_iteration_value.values]
         segmentation_out_i, features_out_i = segmentation_timestep(
             field_at_time,
             features_i,
@@ -1304,10 +1294,7 @@ def segmentation(
             segmentation_out_i
         )
         features_out_list.append(features_out_i)
-        logging.debug(
-            "Finished segmentation for "
-            + pd.to_datetime(time_iteration_value.values).strftime("%Y-%m-%d %H:%M:%S")
-        )
+        logging.debug(f"Finished segmentation for {time_iteration_value.values}")
 
     # Merge output from individual timesteps:
     features_out = pd.concat(features_out_list)

tobac/tests/test_datetime.py

@@ -0,0 +1,128 @@
+from datetime import datetime
+
+import numpy as np
+import pandas as pd
+import cftime
+
+import tobac.utils.datetime as datetime_utils
+
+
+def test_to_cftime():
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_cftime(date, "standard") == cftime.datetime(2000, 1, 1)
+        assert datetime_utils.to_cftime(date, "gregorian") == cftime.DatetimeGregorian(
+            2000, 1, 1
+        )
+        assert datetime_utils.to_cftime(date, "360_day") == cftime.Datetime360Day(
+            2000, 1, 1
+        )
+        assert datetime_utils.to_cftime(date, "365_day") == cftime.DatetimeNoLeap(
+            2000, 1, 1
+        )
+
+
+def test_to_timestamp():
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_timestamp(date) == pd.to_datetime("2000-01-01")
+
+
+def test_to_datetime():
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_datetime(date) == datetime(2000, 1, 1)
+
+
+def test_to_datetime64():
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert datetime_utils.to_datetime64(date) == np.datetime64(
+            "2000-01-01 00:00:00.000000000"
+        )
+
+
+def test_to_datestr():
+    test_dates = [
+        "2000-01-01",
+        "2000-01-01 00:00:00",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        np.datetime64("2000-01-01 00:00:00"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for date in test_dates:
+        assert (
+            datetime_utils.to_datestr(date) == "2000-01-01T00:00:00.000000000"
+            or datetime_utils.to_datestr(date) == "2000-01-01T00:00:00"
+        )
+
+
+def test_match_datetime_format():
+    test_dates = [
+        "2000-01-01T00:00:00.000000000",
+        datetime(2000, 1, 1),
+        np.datetime64("2000-01-01 00:00:00.000000000"),
+        pd.to_datetime("2000-01-01"),
+        cftime.datetime(2000, 1, 1),
+        cftime.DatetimeGregorian(2000, 1, 1),
+        cftime.Datetime360Day(2000, 1, 1),
+        cftime.DatetimeNoLeap(2000, 1, 1),
+    ]
+
+    for target in test_dates:
+        for date in test_dates:
+            assert datetime_utils.match_datetime_format(date, target) == target