Repository: tobac-project/tobac

Issue:
Calculate Bulk Statistics from Raw Data
Could you add the ability to calculate bulk statistics (maximum, mean, etc.) from the raw input data as opposed to tobac's smoothed data?

Discussion:
@freemansw1: @John-David-48 put this in at my request :) . He's meaning on the feature detection module; right now we calculate the bulk statistics on the post-smoothed data rather than the raw data. I think it would be worth doing as an option, but @JuliaKukulies what do you think?@JuliaKukulies: > @John-David-48 put this in at my request :) . He's meaning on the feature detection module; right now we calculate the bulk statistics on the post-smoothed data rather than the raw data. I think it would be worth doing as an option, but @JuliaKukulies what do you think?

That is a very good point @John-David-48 @freemansw1! I had previously not thought about this since I am not using the smoothing option too much. But yes, I think we should definitely make this possible. I can take care of this. 

@John-David-48 I also want to point out that you can use the postprocessing method to calculate the bulk statistics (`tobac.utils.get_statistics_from_mask`) and give the raw/unsmoothed data as the field to compute the statistics from. I know, it is not the same as for the feature detection method, because this would use the segmented objects rather than the detected features. It is also less efficient, but I just wanted to mention that out in case you were not aware of this option to calculate the statistics from. @JuliaKukulies: @freemansw1 I have added this to `v1.5.4` as of now because I do not think this would be too complicated. Anyhow, as we are getting closer to release the new version I dont want to stop us from this, so in that case we might push this to a later one. @JuliaKukulies: Hi @John-David-48, 

I have addressed your request and created a pull request for this: #449! Would you be willing to be the third reviewer for that pull request? It would be very helpful if you could test this with your data and see if it works well and gives you the desired result. If this is approved, we will have this new feature release in our next version (`tobac v1.5.4`). 

Let me know what you think!
Julia @w-k-jones: Resolved by #449 

Solution patch:
diff --git a/doc/index.rst b/doc/index.rst
index 1b215f20..d4f4b4e7 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -10,14 +10,14 @@ The software is set up in a modular way to include different algorithms for feat
 
 Individual features are identified as either maxima or minima in a 2D or 3D time-varying field (see :doc:`feature_detection_overview`). An associated volume can then be determined using these features with a separate (or identical) time-varying 2D or 3D field and a threshold value (see :doc:`segmentation`). The identified objects are linked into consistent trajectories representing the cloud over its lifecycle in the tracking step. Analysis and visualization methods provide a convenient way to use and display the tracking results.
 
-Version 1.2 of tobac and some example applications are described in a peer-reviewed article in Geoscientific Model Development as:
+**Version 1.2 of tobac** and some example applications are described in a peer-reviewed article in Geoscientific Model Development as:
 
 Heikenfeld, M., Marinescu, P. J., Christensen, M., Watson-Parris, D., Senf, F., van den Heever, S. C., and Stier, P.: tobac 1.2: towards a flexible framework for tracking and analysis of clouds in diverse datasets, Geosci. Model Dev., 12, 4551–4570, https://doi.org/10.5194/gmd-12-4551-2019, 2019.
 
-Version 1.5 of tobac and the major enhancements that came with that version are currently under review in Geoscientific Model Development:
+**Version 1.5 of tobac** and the major enhancements that came with that version are described in the following peer-reviewed article in Geoscientific Model Development:
 
+Sokolowsky, G. A., Freeman, S. W., Jones, W. K., Kukulies, J., Senf, F., Marinescu, P. J., Heikenfeld, M., Brunner, K. N., Bruning, E. C., Collis, S. M., Jackson, R. C., Leung, G. R., Pfeifer, N., Raut, B. A., Saleeby, S. M., Stier, P., and van den Heever, S. C.: tobac v1.5: Introducing Fast 3D Tracking, Splits and Mergers, and Other Enhancements for Identifying and Analysing Meteorological Phenomena. Geoscientific Model Development, 17(13), 5309-5330. https://doi.org/10.5194/gmd-17-5309-2024, 2024.
 
-Sokolowsky, G. A., Freeman, S. W., Jones, W. K., Kukulies, J., Senf, F., Marinescu, P. J., Heikenfeld, M., Brunner, K. N., Bruning, E. C., Collis, S. M., Jackson, R. C., Leung, G. R., Pfeifer, N., Raut, B. A., Saleeby, S. M., Stier, P., and van den Heever, S. C.: tobac v1.5: Introducing Fast 3D Tracking, Splits and Mergers, and Other Enhancements for Identifying and Analysing Meteorological Phenomena, EGUsphere [preprint], https://doi.org/10.5194/egusphere-2023-1722, 2023.
 
 The project is currently being extended by several contributors to include additional workflows and algorithms using the same structure, syntax, and data formats.
 
diff --git a/doc/publications.rst b/doc/publications.rst
index 0a0d4776..b22d7856 100644
--- a/doc/publications.rst
+++ b/doc/publications.rst
@@ -11,6 +11,10 @@ Refereed Publications
       :widths: 30
       :class: wy-table-responsive
 
+      * - Sokolowsky, G. A., Freeman, S. W., Jones, W. K., Kukulies, J., Senf, F., Marinescu, P. J., Heikenfeld, M., Brunner, K. N., Bruning, E. C., Collis, S. M., Jackson, R. C., Leung, G. R., Pfeifer, N., Raut, B. A., Saleeby, S. M., Stier, P., and van den Heever, S. C.: tobac v1.5 (2024). Introducing Fast 3D Tracking, Splits and Mergers, and Other Enhancements for Identifying and Analysing Meteorological Phenomena. Geoscientific Model Development, 17(13), 5309-5330. https://doi.org/10.5194/gmd-17-5309-2024. 
+
+      * - Heikenfeld, M., Marinescu, P. J., Christensen, M., Watson-Parris, D., Senf, F., van den Heever, S. C., and Stier, P. (2019). tobac 1.2: towards a flexible framework for tracking and analysis of clouds in diverse datasets, Geosci. Model Dev., 12, 4551–4570, https://doi.org/10.5194/gmd-12-4551-2019
+	      
       * - Bukowski, J., & van den Heever, S. C. (2021). Direct radiative effects in haboobs. *Journal of Geophysical Research: Atmospheres*, 126(21), e2021JD034814, doi:10.1029/2021JD034814.
 
       * - Bukowski, J. (2021). Mineral Dust Lofting and Interactions with Cold Pools (Doctoral dissertation, Colorado State University).
diff --git a/tobac/feature_detection.py b/tobac/feature_detection.py
index fc491ff0..075b8513 100644
--- a/tobac/feature_detection.py
+++ b/tobac/feature_detection.py
@@ -627,9 +627,9 @@ def feature_detection_threshold(
                         # find the updated label, and overwrite all of label_ind indices with
                         # updated label
                         labels_2_alt = labels_2[label_z, y_val_alt, x_val_alt]
-                        labels_2[
-                            label_locs_v, label_locs_h1, label_locs_h2
-                        ] = labels_2_alt
+                        labels_2[label_locs_v, label_locs_h1, label_locs_h2] = (
+                            labels_2_alt
+                        )
                         skip_list = np.append(skip_list, label_ind)
                         break
 
@@ -673,9 +673,9 @@ def feature_detection_threshold(
                         # find the updated label, and overwrite all of label_ind indices with
                         # updated label
                         labels_2_alt = labels_2[label_z, y_val_alt, label_x]
-                        labels_2[
-                            label_locs_v, label_locs_h1, label_locs_h2
-                        ] = labels_2_alt
+                        labels_2[label_locs_v, label_locs_h1, label_locs_h2] = (
+                            labels_2_alt
+                        )
                         new_label_ind = labels_2_alt
                         skip_list = np.append(skip_list, label_ind)
 
@@ -717,9 +717,9 @@ def feature_detection_threshold(
                         # find the updated label, and overwrite all of label_ind indices with
                         # updated label
                         labels_2_alt = labels_2[label_z, label_y, x_val_alt]
-                        labels_2[
-                            label_locs_v, label_locs_h1, label_locs_h2
-                        ] = labels_2_alt
+                        labels_2[label_locs_v, label_locs_h1, label_locs_h2] = (
+                            labels_2_alt
+                        )
                         new_label_ind = labels_2_alt
                         skip_list = np.append(skip_list, label_ind)
 
@@ -912,6 +912,7 @@ def feature_detection_multithreshold_timestep(
     wavelength_filtering: tuple[float] = None,
     strict_thresholding: bool = False,
     statistic: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None] = None,
+    statistics_unsmoothed: bool = False,
 ) -> pd.DataFrame:
     """Find features in each timestep.
 
@@ -984,6 +985,9 @@ def feature_detection_multithreshold_timestep(
             Default is None. Optional parameter to calculate bulk statistics within feature detection.
             Dictionary with callable function(s) to apply over the region of each detected feature and the name of the statistics to appear in the feature ou            tput dataframe. The functions should be the values and the names of the metric the keys (e.g. {'mean': np.mean})
 
+    statistics_unsmoothed: bool, optional
+            Default is False. If True, calculate the statistics on the raw data instead of the smoothed input data.
+
     Returns
     -------
     features_threshold : pandas DataFrame
@@ -1005,6 +1009,14 @@ def feature_detection_multithreshold_timestep(
     # get actual numpy array and make a copy so as not to change the data in the iris cube
     track_data = data_i.core_data().copy()
 
+    # keep a copy of the unsmoothed data (that can be used for calculating stats)
+    if statistics_unsmoothed:
+        if not statistic:
+            raise ValueError(
+                "Please provide the input parameter statistic to determine what statistics to calculate."
+            )
+    
+
     track_data = gaussian_filter(
         track_data, sigma=sigma_threshold
     )  # smooth data slightly to create rounded, continuous field
@@ -1117,14 +1129,24 @@ def feature_detection_multithreshold_timestep(
                 labels.ravel()[regions_old[key]] = key
                 # apply function to get statistics based on labeled regions and functions provided by the user
                 # the feature dataframe is updated by appending a column for each metric
-            features_thresholds = get_statistics(
-                features_thresholds,
-                labels,
-                track_data,
-                statistic=statistic,
-                index=np.unique(labels[labels > 0]),
-                id_column="idx",
-            )
+            if statistics_unsmoothed:
+                features_thresholds = get_statistics(
+                    features_thresholds,
+                    labels,
+                    data_i.core_data(),
+                    statistic=statistic,
+                    index=np.unique(labels[labels > 0]),
+                    id_column="idx",
+                )
+            else:
+                features_thresholds = get_statistics(
+                    features_thresholds,
+                    labels,
+                    track_data,
+                    statistic=statistic,
+                    index=np.unique(labels[labels > 0]),
+                    id_column="idx",
+                )
 
         logging.debug(
             "Finished feature detection for threshold "
@@ -1158,6 +1180,7 @@ def feature_detection_multithreshold(
     dz: Union[float, None] = None,
     strict_thresholding: bool = False,
     statistic: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None] = None,
+    statistics_unsmoothed: bool = False,
 ) -> pd.DataFrame:
     """Perform feature detection based on contiguous regions.
 
@@ -1370,6 +1393,7 @@ def feature_detection_multithreshold(
             wavelength_filtering=wavelength_filtering,
             strict_thresholding=strict_thresholding,
             statistic=statistic,
+            statistics_unsmoothed=statistics_unsmoothed,
         )
         # check if list of features is not empty, then merge features from different threshold
         # values into one DataFrame and append to list for individual timesteps:
@@ -1468,7 +1492,7 @@ def filter_min_distance(
     target: {'maximum', 'minimum'}, optional
         Flag to determine if tracking is targeting minima or maxima in
         the data. Default is 'maximum'.
-    PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both')
+    PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both'), optional
         Sets whether to use periodic boundaries, and if so in which directions.
         'none' means that we do not have periodic boundaries
         'hdim_1' means that we are periodic along hdim1
@@ -1540,8 +1564,8 @@ def filter_min_distance(
         feature_locations = features[
             [z_coordinate_name, y_coordinate_name, x_coordinate_name]
         ].to_numpy()
-        feature_locations[0] *= dz
-        feature_locations[1:] *= dxy
+        feature_locations[:, 0] *= dz
+        feature_locations[:, 1:] *= dxy
     else:
         feature_locations = (
             features[[y_coordinate_name, x_coordinate_name]].to_numpy() * dxy
diff --git a/tobac/merge_split.py b/tobac/merge_split.py
index 6e53eb97..012856b9 100644
--- a/tobac/merge_split.py
+++ b/tobac/merge_split.py
@@ -7,30 +7,91 @@
     
 """
 
-
-def merge_split_MEST(TRACK, dxy, distance=None, frame_len=5):
+import logging
+from typing import Optional
+from typing_extensions import Literal
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+import scipy.sparse
+from sklearn.neighbors import BallTree
+
+from tobac.utils.periodic_boundaries import build_distance_function
+from tobac.utils import internal as internal_utils
+
+
+def merge_split_MEST(
+    tracks: pd.DataFrame,
+    dxy: float,
+    dz: Optional[float] = None,
+    distance: Optional[float] = None,
+    frame_len: int = 5,
+    cell_number_unassigned: int = -1,
+    vertical_coord: Optional[str] | None = None,
+    PBC_flag: Literal["none", "hdim_1", "hdim_2", "both"] = None,
+    min_h1: Optional[int] = None,
+    max_h1: Optional[int] = None,
+    min_h2: Optional[int] = None,
+    max_h2: Optional[int] = None,
+) -> xr.Dataset:
     """
-    function to  postprocess tobac track data for merge/split cells using a minimum euclidian spanning tree
-
+    Search for merging splitting cells  in tobac tracking data using a minimum
+    euclidian spanning tree, and combine the merged cells into unique tracks.
 
     Parameters
     ----------
-    TRACK : pandas.core.frame.DataFrame
+    tracks : pandas.core.frame.DataFrame
         Pandas dataframe of tobac Track information
 
-    dxy : float, mandatory
+    dxy : float
         The x/y grid spacing of the data.
         Should be in meters.
 
+    dz : float, optional
+        Constant vertical grid spacing (m), default None. If None, the vertical
+        coord will be inferred automatically or from a specified coord given by
+        the vertical_coord parameter. An exception is raised if both dz and
+        vertical_coord are provided.
 
     distance : float, optional
-        Distance threshold determining how close two features must be in order to consider merge/splitting.
-        Default is 25x the x/y grid spacing of the data, given in dxy.
-        The distance should be in units of meters.
+        Distance threshold determining how close two features must be in order
+        to consider merge/splitting. Default is 25x the x/y grid spacing of the
+        data, given in dxy. The distance should be in units of meters.
 
     frame_len : float, optional
-        Threshold for the maximum number of frames that can separate the end of cell and the start of a related cell.
-        Default is five (5) frames.
+        Threshold for the maximum number of frames that can separate the end of
+        cell and the start of a related cell, by default 5 frames.
+
+    cell_number_unassigned: int, optional
+        Value given tp unassigned/non-tracked cells by tracking, by default -1.
+
+    vertical_coord: str, optional
+        Name of the vertical coordinate, default None. The vertical coordinate
+        used must have values in meters. If None, tries to auto-detect, or uses
+        constant vertical grid spacing if dz is specified. An exception is
+        raised if both dz and vertical_coord are provided.
+
+    PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both'), optional
+        Sets whether to use periodic boundaries, and if so in which directions.
+        'none' means that we do not have periodic boundaries
+        'hdim_1' means that we are periodic along hdim1
+        'hdim_2' means that we are periodic along hdim2
+        'both' means that we are periodic along both horizontal dimensions
+
+    min_h1: int, optional
+        Minimum real point in hdim_1, for use with periodic boundaries.
+
+    max_h1: int, optional
+        Maximum point in hdim_1, exclusive. max_h1-min_h1 should be the size of
+        hdim_1.
+
+    min_h2: int, optional
+        Minimum real point in hdim_2, for use with periodic boundaries.
+
+    max_h2: int, optional
+        Maximum point in hdim_2, exclusive. max_h2-min_h2 should be the size of
+        hdim_2.
 
     Returns
     -------
@@ -45,7 +106,6 @@ def merge_split_MEST(TRACK, dxy, distance=None, frame_len=5):
         - track_child_cell_count: The total number of features belonging to all child cells of a given track id.
         - cell_child_feature_count: The total number of features for each cell.
 
-
     Example usage:
         d = merge_split_MEST(Track)
         ds = tobac.utils.standardize_track_dataset(Track, refl_mask)
@@ -54,172 +114,205 @@ def merge_split_MEST(TRACK, dxy, distance=None, frame_len=5):
         both_ds.to_netcdf(os.path.join(savedir,'Track_features_merges.nc'))
 
     """
-    try:
-        import networkx as nx
-    except ImportError:
-        networkx = None
-
-    import logging
-    import numpy as np
-    from pandas.core.common import flatten
-    import xarray as xr
-    from scipy.spatial.distance import cdist
-
-    # Immediately convert pandas dataframe of track information to xarray:
-    TRACK = TRACK.to_xarray()
-    track_groups = TRACK.groupby("cell")
+
+    track_groups = tracks[tracks["cell"] != cell_number_unassigned].groupby("cell")
     first = track_groups.first()
     last = track_groups.last()
 
     if distance is None:
         distance = dxy * 25.0
 
-    a_names = list()
-    b_names = list()
-    dist = list()
-
-    # write all sets of points (a and b) as Nx2 arrays
-    l = len(last["hdim_2"].values)
-    cells = first["cell"].values
-    a_xy = np.zeros((l, 2))
-    a_xy[:, 0] = last["hdim_2"].values * dxy
-    a_xy[:, 1] = last["hdim_1"].values * dxy
-    b_xy = np.zeros((l, 2))
-    b_xy[:, 0] = first["hdim_2"].values * dxy
-    b_xy[:, 1] = first["hdim_1"].values * dxy
-    # Use cdist to find distance matrix
-    out = cdist(a_xy, b_xy)
-    # Find all cells under the distance threshold
-    j = np.where(out <= distance)
-
-    # Compile cells meeting the criteria to an array of both the distance and cell ids
-    a_names = cells[j[0]]
-    b_names = cells[j[1]]
-    dist = out[j]
-
-    # This is inputing data to the object which will perform the spanning tree.
-    g = nx.Graph()
-    for i in np.arange(len(dist)):
-        g.add_edge(a_names[i], b_names[i], weight=dist[i])
-
-    tree = nx.minimum_spanning_edges(g)
-    tree_list = list(tree)
-
-    new_tree = []
-
-    # Pruning the tree for time limits.
-    for i, j in enumerate(tree_list):
-        frame_a = np.nanmax(track_groups[j[0]].frame.values)
-        frame_b = np.nanmin(track_groups[j[1]].frame.values)
-        if np.abs(frame_a - frame_b) <= frame_len:
-            new_tree.append(tree_list[i][0:2])
-    new_tree_arr = np.array(new_tree)
-
-    TRACK["cell_parent_track_id"] = np.zeros(len(TRACK["cell"].values))
-    cell_id = np.unique(
-        TRACK.cell.values.astype(int)[~np.isnan(TRACK.cell.values.astype(int))]
-    )
-    track_id = dict()  # same size as number of total merged tracks
-
-    # Cleaning up tracks, combining tracks which contain the same cells.
-    arr = np.array([0])
-    for p in cell_id:
-        j = np.where(arr == int(p))
-        if len(j[0]) > 0:
-            continue
+    # As optional coordinate names are not yet implemented, set to defaults here:
+    y_coordinate_name = "hdim_1"
+    x_coordinate_name = "hdim_2"
+
+    # Check if we are 3D.
+    is_3D = "vdim" in tracks
+    if is_3D:
+        if dz is None:
+            # Find vertical coord name
+            z_coordinate_name = internal_utils.find_dataframe_vertical_coord(
+                variable_dataframe=tracks, vertical_coord=vertical_coord
+            )
+            dz = 1
         else:
-            k = np.where(new_tree_arr == p)
-            if len(k[0]) == 0:
-                track_id[p] = [p]
-                arr = np.append(arr, p)
+            # Use dz, raise error if both are set
+            if vertical_coord is None:
+                z_coordinate_name = "vdim"
             else:
-                temp1 = list(np.unique(new_tree_arr[k[0]]))
-                temp = list(np.unique(new_tree_arr[k[0]]))
-
-                for l in range(len(cell_id)):
-                    for i in temp1:
-                        k2 = np.where(new_tree_arr == i)
-                        temp.append(list(np.unique(new_tree_arr[k2[0]]).squeeze()))
-                        temp = list(flatten(temp))
-                        temp = list(np.unique(temp))
-
-                    if len(temp1) == len(temp):
-                        break
-                    temp1 = np.array(temp)
-
-                for i in temp1:
-                    k2 = np.where(new_tree_arr == i)
-                    temp.append(list(np.unique(new_tree_arr[k2[0]]).squeeze()))
-
-                temp = list(flatten(temp))
-                temp = list(np.unique(temp))
-                arr = np.append(arr, np.unique(temp))
-
-                track_id[np.nanmax(np.unique(temp))] = list(np.unique(temp))
+                raise ValueError(
+                    "dz and vertical_coord both set, vertical"
+                    " spacing is ambiguous. Set one to None."
+                )
+
+    # Calculate feature locations in cartesian coordinates
+    if is_3D:
+        cell_start_locations = np.stack(
+            [
+                first[var].values
+                for var in [z_coordinate_name, y_coordinate_name, x_coordinate_name]
+            ],
+            axis=-1,
+        )
+        cell_start_locations[:, 0] *= dz
+        cell_start_locations[:, 1:] *= dxy
+        cell_end_locations = np.stack(
+            [
+                last[var].values
+                for var in [z_coordinate_name, y_coordinate_name, x_coordinate_name]
+            ],
+            axis=-1,
+        )
+        cell_end_locations[0] *= dz
+        cell_end_locations[1:] *= dxy
+    else:
+        cell_start_locations = (
+            np.stack(
+                [first[var].values for var in [y_coordinate_name, x_coordinate_name]],
+                axis=-1,
+            )
+            * dxy
+        )
+        cell_end_locations = (
+            np.stack(
+                [last[var].values for var in [y_coordinate_name, x_coordinate_name]],
+                axis=-1,
+            )
+            * dxy
+        )
+
+    if PBC_flag in ["hdim_1", "hdim_2", "both"]:
+        # Note that we multiply by dxy to get the distances in spatial coordinates
+        dist_func = build_distance_function(
+            min_h1 * dxy, max_h1 * dxy, min_h2 * dxy, max_h2 * dxy, PBC_flag, is_3D
+        )
+        cell_start_tree = BallTree(
+            cell_start_locations, metric="pyfunc", func=dist_func
+        )
+
+    else:
+        cell_start_tree = BallTree(cell_start_locations, metric="euclidean")
+
+    neighbours, distances = cell_start_tree.query_radius(
+        cell_end_locations, r=distance, return_distance=True
+    )
 
-    cell_id = list(np.unique(TRACK.cell.values.astype(int)))
-    logging.debug("found cell ids")
+    # Input data to the graph which will perform the spanning tree.
+    nodes = np.repeat(
+        np.arange(len(neighbours), dtype=int), [len(n) for n in neighbours]
+    )
+    neighbours = np.concatenate(neighbours)
+    weights = np.concatenate(distances)
 
-    cell_parent_track_id = np.zeros(len(cell_id))
-    cell_parent_track_id[:] = -1
+    # Remove edges where the frame gap is greater than frame_len, and also remove connections to the same cell
+    wh_frame_len = (
+        np.abs(first["frame"].values[nodes] - last["frame"].values[neighbours])
+        <= frame_len
+    )
+    wh_valid_edge = np.logical_and(wh_frame_len, nodes != neighbours)
+    start_node_cells = first.index.values[nodes[wh_valid_edge]].astype(np.int32)
+    end_node_cells = last.index.values[neighbours[wh_valid_edge]].astype(np.int32)
+
+    cell_id = np.unique(tracks.cell.values)
+    cell_id = cell_id[cell_id != cell_number_unassigned].astype(int)
+    max_cell = np.max(cell_id)
+
+    if len(start_node_cells):
+        # We need to add a small value to the dists to prevent 0-length edges
+        cell_graph = scipy.sparse.coo_array(
+            (weights[wh_valid_edge] + 0.01, (start_node_cells, end_node_cells)),
+            shape=(max_cell + 1, max_cell + 1),
+        )
+        cell_graph = scipy.sparse.csgraph.minimum_spanning_tree(
+            cell_graph, overwrite=True
+        )
+        # Find remaining start/end nodes after calculating minimum spanning tree
+        start_node_cells, end_node_cells = cell_graph.nonzero()
+
+        cell_parent_track_id = scipy.sparse.csgraph.connected_components(cell_graph)[1][
+            cell_id
+        ]
+        cell_parent_track_id = (
+            np.unique(cell_parent_track_id, return_inverse=True)[1] + 1
+        )
+    else:
+        cell_parent_track_id = np.arange(cell_id.size, dtype=int) + 1
 
-    for i, id in enumerate(track_id, start=0):
-        for j in track_id[int(id)]:
-            cell_parent_track_id[cell_id.index(j)] = int(i)
+    track_dim = "track"
+    cell_dim = "cell"
+    feature_dim = "feature"
 
+    cell_parent_track_id = xr.DataArray(
+        cell_parent_track_id, dims=(cell_dim,), coords={cell_dim: cell_id}
+    )
     logging.debug("found cell parent track ids")
 
-    track_ids = np.array(np.unique(cell_parent_track_id))
+    track_id = np.unique(cell_parent_track_id)
     logging.debug("found track ids")
 
-    feature_parent_cell_id = list(TRACK.cell.values.astype(int))
-    logging.debug("found feature parent cell ids")
-
-    #     # This version includes all the feature regardless of if they are used in cells or not.
-    feature_id = list(TRACK.feature.values.astype(int))
+    # This version includes all the feature regardless of if they are used in cells or not.
+    feature_id = tracks.feature.values.astype(int)
     logging.debug("found feature ids")
 
-    feature_parent_track_id = []
-    feature_parent_track_id = np.zeros(len(feature_id))
-    for i, id in enumerate(feature_id):
-        cellid = feature_parent_cell_id[i]
-        if cellid < 0:
-            feature_parent_track_id[i] = -1
-        else:
-            feature_parent_track_id[i] = cell_parent_track_id[cell_id.index(cellid)]
+    feature_parent_cell_id = tracks.cell.values.astype(int)
+    feature_parent_cell_id = xr.DataArray(
+        feature_parent_cell_id,
+        dims=(feature_dim,),
+        coords={feature_dim: feature_id},
+    )
+    logging.debug("found feature parent cell ids")
 
-    track_child_cell_count = np.zeros(len(track_id))
-    for i, id in enumerate(track_id):
-        track_child_cell_count[i] = len(np.where(cell_parent_track_id == i)[0])
-    logging.debug("found track child cell count")
+    wh_feature_in_cell = (feature_parent_cell_id != cell_number_unassigned).values
+    feature_parent_track_id = np.full(wh_feature_in_cell.shape, cell_number_unassigned)
+    feature_parent_track_id[wh_feature_in_cell] = cell_parent_track_id.loc[
+        feature_parent_cell_id[wh_feature_in_cell]
+    ].values
+    feature_parent_track_id = xr.DataArray(
+        feature_parent_track_id,
+        dims=(feature_dim,),
+        coords={feature_dim: feature_id},
+    )
 
-    cell_child_feature_count = np.zeros(len(cell_id))
-    for i, id in enumerate(cell_id):
-        cell_child_feature_count[i] = len(track_groups[id].feature.values)
-    logging.debug("found cell child feature count")
+    track_child_cell_count = (
+        cell_parent_track_id.groupby(cell_parent_track_id).reduce(np.size).values
+    )
+    track_child_cell_count = xr.DataArray(
+        track_child_cell_count,
+        dims=(track_dim,),
+        coords={track_dim: track_id},
+    )
 
-    track_dim = "track"
-    cell_dim = "cell"
-    feature_dim = "feature"
+    cell_child_feature_count = (
+        feature_parent_cell_id[wh_feature_in_cell]
+        .groupby(feature_parent_cell_id[wh_feature_in_cell])
+        .reduce(np.size)
+        .values
+    )
+    cell_child_feature_count = xr.DataArray(
+        cell_child_feature_count, dims=(cell_dim), coords={cell_dim: cell_id}
+    )
 
-    d = xr.Dataset(
-        {
-            "track": (track_dim, track_ids),
-            "cell": (cell_dim, cell_id),
-            "cell_parent_track_id": (cell_dim, cell_parent_track_id),
-            "feature": (feature_dim, feature_id),
-            "feature_parent_cell_id": (feature_dim, feature_parent_cell_id),
-            "feature_parent_track_id": (feature_dim, feature_parent_track_id),
-            "track_child_cell_count": (track_dim, track_child_cell_count),
-            "cell_child_feature_count": (cell_dim, cell_child_feature_count),
-        }
+    cell_starts_with_split = np.isin(cell_id, start_node_cells)
+    cell_starts_with_split = xr.DataArray(
+        cell_starts_with_split, dims=(cell_dim), coords={cell_dim: cell_id}
     )
 
-    d = d.set_coords(["feature", "cell", "track"])
+    cell_ends_with_merge = np.isin(cell_id, end_node_cells)
+    cell_ends_with_merge = xr.DataArray(
+        cell_ends_with_merge, dims=(cell_dim), coords={cell_dim: cell_id}
+    )
 
-    #     assert len(cell_id) == len(cell_parent_track_id)
-    #     assert len(feature_id) == len(feature_parent_cell_id)
-    #     assert sum(track_child_cell_count) == len(cell_id)
-    #     assert sum(cell_child_feature_count) == len(feature_id)
+    merge_split_ds = xr.Dataset(
+        data_vars={
+            "cell_parent_track_id": cell_parent_track_id,
+            "feature_parent_cell_id": feature_parent_cell_id,
+            "feature_parent_track_id": feature_parent_track_id,
+            "track_child_cell_count": track_child_cell_count,
+            "cell_child_feature_count": cell_child_feature_count,
+            "cell_starts_with_split": cell_starts_with_split,
+            "cell_ends_with_merge": cell_ends_with_merge,
+        },
+        coords={feature_dim: feature_id, cell_dim: cell_id, track_dim: track_id},
+    )
 
-    return d
+    return merge_split_ds
diff --git a/tobac/segmentation.py b/tobac/segmentation.py
index fe2eda2e..4697a25d 100644
--- a/tobac/segmentation.py
+++ b/tobac/segmentation.py
@@ -824,15 +824,15 @@ def segmentation_timestep(
                 )
 
                 # edit value in buddy_features dataframe
-                buddy_features.hdim_1.values[
-                    buddy_looper
-                ] = pbc_utils.transfm_pbc_point(
-                    float(buddy_feat.hdim_1), hdim1_min, hdim1_max
+                buddy_features.hdim_1.values[buddy_looper] = (
+                    pbc_utils.transfm_pbc_point(
+                        float(buddy_feat.hdim_1), hdim1_min, hdim1_max
+                    )
                 )
-                buddy_features.hdim_2.values[
-                    buddy_looper
-                ] = pbc_utils.transfm_pbc_point(
-                    float(buddy_feat.hdim_2), hdim2_min, hdim2_max
+                buddy_features.hdim_2.values[buddy_looper] = (
+                    pbc_utils.transfm_pbc_point(
+                        float(buddy_feat.hdim_2), hdim2_min, hdim2_max
+                    )
                 )
 
                 buddy_looper = buddy_looper + 1
@@ -1010,9 +1010,9 @@ def segmentation_timestep(
                                 segmentation_mask_3[z_val_o, y_val_o, x_val_o]
                                 != segmentation_mask_4.data[z_seg, y_seg, x_seg]
                             ):
-                                segmentation_mask_3[
-                                    z_val_o, y_val_o, x_val_o
-                                ] = segmentation_mask_4.data[z_seg, y_seg, x_seg]
+                                segmentation_mask_3[z_val_o, y_val_o, x_val_o] = (
+                                    segmentation_mask_4.data[z_seg, y_seg, x_seg]
+                                )
         if not is_3D_seg:
             segmentation_mask_3 = segmentation_mask_3[0]
 
diff --git a/tobac/utils/decorators.py b/tobac/utils/decorators.py
index 90e600b5..8d304e71 100644
--- a/tobac/utils/decorators.py
+++ b/tobac/utils/decorators.py
@@ -75,9 +75,7 @@ def _conv_kwargs_irispandas_to_xarray(conv_kwargs: dict):
         key: (
             convert_cube_to_dataarray(arg)
             if isinstance(arg, iris.cube.Cube)
-            else arg.to_xarray()
-            if isinstance(arg, pd.DataFrame)
-            else arg
+            else arg.to_xarray() if isinstance(arg, pd.DataFrame) else arg
         )
         for key, arg in zip(conv_kwargs.keys(), conv_kwargs.values())
     }
@@ -123,9 +121,7 @@ def _conv_kwargs_xarray_to_irispandas(conv_kwargs: dict):
         key: (
             xr.DataArray.to_iris(arg)
             if isinstance(arg, xr.DataArray)
-            else arg.to_dataframe()
-            if isinstance(arg, xr.Dataset)
-            else arg
+            else arg.to_dataframe() if isinstance(arg, xr.Dataset) else arg
         )
         for key, arg in zip(conv_kwargs.keys(), conv_kwargs.values())
     }
@@ -340,9 +336,7 @@ def wrapper(*args, **kwargs):
                         (
                             convert_cube_to_dataarray(arg)
                             if type(arg) == iris.cube.Cube
-                            else arg.to_xarray()
-                            if type(arg) == pd.DataFrame
-                            else arg
+                            else arg.to_xarray() if type(arg) == pd.DataFrame else arg
                         )
                         for arg in args
                     ]


Test patch:
diff --git a/tobac/tests/test_merge_split.py b/tobac/tests/test_merge_split.py
index 83200192..8b3ecdb4 100644
--- a/tobac/tests/test_merge_split.py
+++ b/tobac/tests/test_merge_split.py
@@ -3,6 +3,7 @@
 
 import pandas as pd
 import numpy as np
+import pytest
 
 import datetime
 
@@ -78,7 +79,7 @@ def test_merge_split_MEST():
 
     # Test a successful merger
     mergesplit_output_merged = mergesplit.merge_split_MEST(
-        output, dxy=1, distance=dist_between + 50
+        output, dxy=10, distance=(dist_between + 50) * 10
     )
 
     # These cells should have merged together.
@@ -86,9 +87,256 @@ def test_merge_split_MEST():
 
     # Test an unsuccessful merger.
     mergesplit_output_unmerged = mergesplit.merge_split_MEST(
-        output, dxy=1, distance=dist_between - 50
+        output, dxy=10, distance=(dist_between - 50) * 10
     )
 
     # These cells should NOT have merged together.
     print(mergesplit_output_unmerged["track"])
     assert len(mergesplit_output_unmerged["track"]) == 2
+
+
+def test_merge_split_MEST_PBC():
+    """
+    Test PBC handling for merge_split_MEST
+    """
+    test_features = pd.DataFrame(
+        {
+            "feature": [1, 2, 3, 4],
+            "hdim_1": [1, 89, 1, 99],
+            "hdim_2": [50, 50, 50, 50],
+            "cell": [1, 2, 1, 2],
+            "frame": [0, 0, 1, 1],
+            "time": [
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1, 0, 5),
+                datetime.datetime(2000, 1, 1, 0, 5),
+            ],
+        }
+    )
+    # Test without PBCs
+    mergesplit_output_no_pbc = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=25,
+    )
+
+    assert len(mergesplit_output_no_pbc["track"]) == 2
+
+    # Test with PBC in hdim_1, cells should merge
+    mergesplit_output_hdim_1_pbc = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=25,
+        PBC_flag="hdim_1",
+        min_h1=0,
+        max_h1=100,
+        min_h2=0,
+        max_h2=100,
+    )
+
+    assert len(mergesplit_output_hdim_1_pbc["track"]) == 1
+
+    # Test with PBC in hdim_2, cells should not merge
+    mergesplit_output_hdim_2_pbc = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=25,
+        PBC_flag="hdim_2",
+        min_h1=0,
+        max_h1=100,
+        min_h2=0,
+        max_h2=100,
+    )
+
+    assert len(mergesplit_output_hdim_2_pbc["track"]) == 2
+
+
+def test_merge_split_MEST_frame_len():
+    """
+    Test the frame_len parameter of merge_split_MEST
+    """
+    test_features = pd.DataFrame(
+        {
+            "feature": [1, 2, 3, 4],
+            "hdim_1": [50, 50, 50, 50],
+            "hdim_2": [50, 50, 50, 50],
+            "cell": [1, 1, 2, 2],
+            "frame": [0, 1, 3, 4],
+            "time": [
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1, 0, 5),
+                datetime.datetime(2000, 1, 1, 0, 5),
+            ],
+        }
+    )
+
+    # Test with short frame_len, expect no link
+    mergesplit_output = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=25,
+        frame_len=1,
+    )
+    assert len(mergesplit_output["track"]) == 2
+
+    # Test with longer frame_len, expect link
+    mergesplit_output = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=25,
+        frame_len=2,
+    )
+    assert len(mergesplit_output["track"]) == 1
+
+
+def test_merge_split_MEST_no_cell():
+    """
+    Test merge/split in cases with features with no cell
+    """
+    test_features = pd.DataFrame(
+        {
+            "feature": [1, 2, 3],
+            "hdim_1": [25, 30, 50],
+            "hdim_2": [25, 30, 50],
+            "cell": [1, -1, 1],
+            "frame": [0, 0, 1],
+            "time": [
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1, 0, 5),
+            ],
+        }
+    )
+
+    mergesplit_output = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=25,
+    )
+
+    assert len(mergesplit_output["track"]) == 1
+
+    assert mergesplit_output["feature_parent_cell_id"].values[1] == -1
+
+
+def test_merge_split_MEST_3D():
+    """
+    Test merge/split support for 3D tracks and dz input
+    """
+
+    test_features = pd.DataFrame(
+        {
+            "feature": [1, 2, 3, 4],
+            "vdim": [1, 2, 1, 2],
+            "hdim_1": [50, 40, 50, 40],
+            "hdim_2": [50, 50, 50, 50],
+            "cell": [1, 2, 1, 2],
+            "frame": [0, 0, 1, 1],
+            "altitude": [500, 750, 1500, 2000],
+            "time": [
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1, 0, 5),
+                datetime.datetime(2000, 1, 1, 0, 5),
+            ],
+        }
+    )
+
+    # Test with dz=10, expect merge
+    mergesplit_output_3d_merge = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        dz=10,
+        distance=20,
+    )
+
+    assert len(mergesplit_output_3d_merge["track"]) == 1
+
+    # Test with dz=25, expect no merge
+    mergesplit_output_3d_nomerge = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        dz=25,
+        distance=20,
+    )
+
+    assert len(mergesplit_output_3d_nomerge["track"]) == 2
+
+    # Test providing vertical_coord
+    mergesplit_output_3d_coord_merge = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        vertical_coord="altitude",
+        distance=1100,
+    )
+    mergesplit_output_3d_coord_nomerge = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        vertical_coord="altitude",
+        distance=20,
+    )
+
+    assert len(mergesplit_output_3d_coord_merge["track"]) == 1
+    assert len(mergesplit_output_3d_coord_nomerge["track"]) == 2
+
+    # Test auto find vertical_coord
+    mergesplit_output_3d_coord_merge = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=1100,
+    )
+    mergesplit_output_3d_coord_nomerge = mergesplit.merge_split_MEST(
+        test_features,
+        dxy=1,
+        distance=20,
+    )
+
+    assert len(mergesplit_output_3d_coord_merge["track"]) == 1
+    assert len(mergesplit_output_3d_coord_nomerge["track"]) == 2
+
+    # Test error if both dz and coord are provided
+    with pytest.raises(ValueError):
+        mergesplit.merge_split_MEST(
+            test_features,
+            dxy=1,
+            dz=1,
+            vertical_coord="auto",
+            distance=20,
+        )
+
+    # Test that wrong vertical coord name causes an error
+    with pytest.raises(ValueError):
+        mergesplit.merge_split_MEST(
+            test_features,
+            dxy=1,
+            vertical_coord="invalid_coord_name",
+            distance=20,
+        )
+
+    # Test that auto search fails if coordinate name is not in default list
+    test_features = pd.DataFrame(
+        {
+            "feature": [1, 2, 3, 4],
+            "vdim": [1, 2, 1, 2],
+            "hdim_1": [50, 40, 50, 40],
+            "hdim_2": [50, 50, 50, 50],
+            "cell": [1, 2, 1, 2],
+            "frame": [0, 0, 1, 1],
+            "invalid_coord_name": [500, 1500, 1000, 2000],
+            "time": [
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1),
+                datetime.datetime(2000, 1, 1, 0, 5),
+                datetime.datetime(2000, 1, 1, 0, 5),
+            ],
+        }
+    )
+
+    with pytest.raises(ValueError):
+        mergesplit.merge_split_MEST(
+            test_features,
+            dxy=1,
+            distance=20,
+        )
diff --git a/tobac/tests/test_utils_bulk_statistics.py b/tobac/tests/test_utils_bulk_statistics.py
index 1db036b0..c62ee821 100644
--- a/tobac/tests/test_utils_bulk_statistics.py
+++ b/tobac/tests/test_utils_bulk_statistics.py
@@ -8,6 +8,37 @@
 import tobac.testing as tb_test
 
 
+@pytest.mark.parametrize("statistics_unsmoothed", [(False), (True)])
+def test_bulk_statistics_fd(statistics_unsmoothed):
+    """
+    Assure that bulk statistics in feature detection work, both on smoothed and raw data
+    """
+    ### Test 2D data with time dimension
+    test_data = tb_test.make_simple_sample_data_2D().core_data()
+    common_dset_opts = {
+        "in_arr": test_data,
+        "data_type": "iris",
+    }
+    test_data_iris = tb_test.make_dataset_from_arr(
+        time_dim_num=0, y_dim_num=1, x_dim_num=2, **common_dset_opts
+    )
+    stats = {"feature_max": np.max}
+
+    # detect features
+    threshold = 7
+    fd_output = tobac.feature_detection.feature_detection_multithreshold(
+        test_data_iris,
+        dxy=1000,
+        threshold=[threshold],
+        n_min_threshold=100,
+        target="maximum",
+        statistic=stats,
+        statistics_unsmoothed=statistics_unsmoothed,
+    )
+
+    assert "feature_max" in fd_output.columns
+
+
 @pytest.mark.parametrize(
     "id_column, index", [("feature", [1]), ("feature_id", [1]), ("cell", [1])]
 )