Merge pull request #92 from freemansw1/fix_level_segmentation

Fixed `level` parameter in segmentation to work and added segmentation tests

tobac/segmentation.py

@@ -69,7 +69,7 @@ def segmentation_timestep(
     target:     string
                 switch to determine if algorithm looks strating from maxima or minima in input field (maximum: starting from maxima (default), minimum: starting from minima)
     level       slice
-                levels at which to seed the cells for the watershedding algorithm
+                vertical levels at which to seed the cells for the watershedding algorithm
     method:     string
                 flag determining the algorithm to use (currently watershedding implemented)
     max_distance: float
@@ -149,11 +149,11 @@ def segmentation_timestep(
             raise ValueError("please specify 1 dimensional vertical coordinate")
         for index, row in features_in.iterrows():
             if ndim_vertical[0] == 0:
-                markers[:, int(row["hdim_1"]), int(row["hdim_2"])] = row["feature"]
+                markers[level, int(row["hdim_1"]), int(row["hdim_2"])] = row["feature"]
             elif ndim_vertical[0] == 1:
-                markers[int(row["hdim_1"]), :, int(row["hdim_2"])] = row["feature"]
+                markers[int(row["hdim_1"]), level, int(row["hdim_2"])] = row["feature"]
             elif ndim_vertical[0] == 2:
-                markers[int(row["hdim_1"]), int(row["hdim_2"]), :] = row["feature"]
+                markers[int(row["hdim_1"]), int(row["hdim_2"]), level] = row["feature"]
     else:
         raise ValueError(
             "Segmentations routine only possible with 2 or 3 spatial dimensions"

tobac/testing.py

@@ -423,7 +423,14 @@ def make_sample_data_3D_3blobs(data_type="iris", invert_xy=False):
     return sample_data
 
 
-def make_dataset_from_arr(in_arr, data_type="xarray"):
+def make_dataset_from_arr(
+    in_arr,
+    data_type="xarray",
+    time_dim_num=None,
+    z_dim_num=None,
+    y_dim_num=0,
+    x_dim_num=1,
+):
     """Makes a dataset (xarray or iris) for feature detection/segmentation from
     a raw numpy/dask/etc. array.
 
@@ -433,20 +440,41 @@ def make_dataset_from_arr(in_arr, data_type="xarray"):
         The input array to convert to iris/xarray
     data_type: str('xarray' or 'iris')
         Type of the dataset to return
+    time_dim_num: int or None
+        What axis is the time dimension on, None for a single timestep
+    z_dim_num: int or None
+        What axis is the z dimension on, None for a 2D array
+    y_dim_num: int
+        What axis is the y dimension on, typically 0 for a 2D array
+    x_dim_num: int
+        What axis is the x dimension on, typically 1 for a 2D array
 
     Returns
     -------
     Iris or xarray dataset with everything we need for feature detection/tracking.
 
     """
     import xarray as xr
+    import iris
 
+    if time_dim_num is not None:
+        raise NotImplementedError("Time dimension not yet implemented in this function")
+
+    is_3D = z_dim_num is not None
     output_arr = xr.DataArray(in_arr)
+    if is_3D:
+        z_max = in_arr.shape[z_dim_num]
 
     if data_type == "xarray":
         return output_arr
     elif data_type == "iris":
-        return output_arr.to_iris()
+        out_arr_iris = output_arr.to_iris()
+        if is_3D:
+            out_arr_iris.add_dim_coord(
+                iris.coords.DimCoord(np.arange(0, z_max), standard_name="altitude"),
+                z_dim_num,
+            )
+        return out_arr_iris
     else:
         raise ValueError("data_type must be 'xarray' or 'iris'")
 

tobac/tests/test_segmentation.py

@@ -0,0 +1,114 @@
+import tobac.testing as testing
+import tobac.segmentation as seg
+
+
+def test_segmentation_timestep_level():
+    """Tests `tobac.segmentation.segmentation_timestep` with a 2D
+    input feature and a 3D segmentation array, specifying the `level` parameter.
+    """
+    # Before we can run segmentation, we must run feature detection.
+
+    # start by building a simple dataset with a single feature
+    import numpy as np
+
+    test_dset_size = (20, 50, 50)
+    test_hdim_1_pt = 20.0
+    test_hdim_2_pt = 20.0
+    test_vdim_pt = 2
+    test_hdim_1_sz = 5
+    test_hdim_2_sz = 5
+    test_vdim_sz = 3
+    test_dxy = 1000
+
+    vdim_start_feat = int(np.ceil(test_vdim_pt - test_vdim_sz / 2))
+    vdim_end_feat = int(np.ceil(test_vdim_pt + test_vdim_sz / 2))
+    hdim_1_start_feat = int(np.ceil(test_hdim_1_pt - test_hdim_1_sz / 2))
+    hdim_1_end_feat = int(np.ceil(test_hdim_1_pt + test_hdim_1_sz / 2))
+    hdim_2_start_feat = int(np.ceil(test_hdim_2_pt - test_hdim_2_sz / 2))
+    hdim_2_end_feat = int(np.ceil(test_hdim_2_pt + test_hdim_2_sz / 2))
+
+    test_amp = 2
+
+    test_data = np.zeros(test_dset_size)
+    test_data = testing.make_feature_blob(
+        test_data,
+        test_hdim_1_pt,
+        test_hdim_2_pt,
+        test_vdim_pt,
+        h1_size=test_hdim_1_sz,
+        h2_size=test_hdim_2_sz,
+        v_size=test_vdim_sz,
+        amplitude=test_amp,
+    )
+
+    # Make a second feature, above the first.
+
+    delta_height = 8
+    test_data = testing.make_feature_blob(
+        test_data,
+        test_hdim_1_pt,
+        test_hdim_2_pt,
+        test_vdim_pt + delta_height,
+        h1_size=test_hdim_1_sz,
+        h2_size=test_hdim_2_sz,
+        v_size=test_vdim_sz,
+        amplitude=test_amp,
+    )
+
+    test_data_iris = testing.make_dataset_from_arr(
+        test_data, data_type="iris", z_dim_num=0, y_dim_num=1, x_dim_num=2
+    )
+    # Generate dummy feature dataset
+    test_feature_ds = testing.generate_single_feature(start_h1=20.0, start_h2=20.0)
+
+    out_seg_mask, out_df = seg.segmentation_timestep(
+        field_in=test_data_iris,
+        features_in=test_feature_ds,
+        dxy=test_dxy,
+        threshold=1.5,
+    )
+    out_seg_mask_arr = out_seg_mask.core_data()
+    # Make sure that all labeled points are segmented, before setting specific levels
+    assert np.all(
+        out_seg_mask_arr[
+            vdim_start_feat:vdim_end_feat,
+            hdim_1_start_feat:hdim_1_end_feat,
+            hdim_2_start_feat:hdim_2_end_feat,
+        ]
+        == np.ones((test_vdim_sz, test_hdim_1_sz, test_hdim_2_sz))
+    )
+    assert np.all(
+        out_seg_mask_arr[
+            vdim_start_feat + delta_height : vdim_end_feat + delta_height,
+            hdim_1_start_feat:hdim_1_end_feat,
+            hdim_2_start_feat:hdim_2_end_feat,
+        ]
+        == np.ones((test_vdim_sz, test_hdim_1_sz, test_hdim_2_sz))
+    )
+
+    # now set specific levels
+    out_seg_mask, out_df = seg.segmentation_timestep(
+        field_in=test_data_iris,
+        features_in=test_feature_ds,
+        dxy=test_dxy,
+        level=slice(vdim_start_feat, vdim_end_feat),
+        threshold=1.5,
+    )
+    out_seg_mask_arr = out_seg_mask.core_data()
+    # Make sure that all labeled points are segmented, before setting specific levels
+    assert np.all(
+        out_seg_mask_arr[
+            vdim_start_feat:vdim_end_feat,
+            hdim_1_start_feat:hdim_1_end_feat,
+            hdim_2_start_feat:hdim_2_end_feat,
+        ]
+        == np.ones((test_vdim_sz, test_hdim_1_sz, test_hdim_2_sz))
+    )
+    assert np.all(
+        out_seg_mask_arr[
+            vdim_start_feat + delta_height : vdim_end_feat + delta_height,
+            hdim_1_start_feat:hdim_1_end_feat,
+            hdim_2_start_feat:hdim_2_end_feat,
+        ]
+        == np.zeros((test_vdim_sz, test_hdim_1_sz, test_hdim_2_sz))
+    )