diff --git a/tobac/segmentation.py b/tobac/segmentation.py index 2f9d6599..60555295 100644 --- a/tobac/segmentation.py +++ b/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 @@ -148,12 +148,13 @@ def segmentation_timestep( if len(ndim_vertical) > 1: raise ValueError("please specify 1 dimensional vertical coordinate") for index, row in features_in.iterrows(): + print("type: ", type(level)) 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" diff --git a/tobac/testing.py b/tobac/testing.py index 4b7d8fc8..36338a44 100644 --- a/tobac/testing.py +++ b/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,6 +440,14 @@ 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 ------- @@ -440,13 +455,26 @@ def make_dataset_from_arr(in_arr, data_type="xarray"): """ 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'") diff --git a/tobac/tests/test_segmentation.py b/tobac/tests/test_segmentation.py new file mode 100644 index 00000000..e851cb91 --- /dev/null +++ b/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)) + )