From ee1aec02531fccbea4b6629bdac7b8ee4ce6f094 Mon Sep 17 00:00:00 2001 From: Sean Freeman Date: Wed, 4 Oct 2023 22:49:07 -0500 Subject: [PATCH 1/3] adding type hints for segmentation --- tobac/segmentation.py | 77 ++++++++++++++++++++++++------------------- 1 file changed, 43 insertions(+), 34 deletions(-) diff --git a/tobac/segmentation.py b/tobac/segmentation.py index b173f48c..cb70d975 100644 --- a/tobac/segmentation.py +++ b/tobac/segmentation.py @@ -31,10 +31,14 @@ """ import copy import logging + +import iris.cube import numpy as np +import pandas as pd +from typing_extensions import Literal +from typing import Union import skimage -import numpy as np from . import utils as tb_utils from .utils import periodic_boundaries as pbc_utils @@ -42,8 +46,13 @@ def add_markers( - features, marker_arr, seed_3D_flag, seed_3D_size=5, level=None, PBC_flag="none" -): + features: pd.DataFrame, + marker_arr: np.array, + seed_3D_flag: Literal["column", "box"], + seed_3D_size: Union[int, tuple[int]] = 5, + level: Union[None, slice] = None, + PBC_flag: Literal["none", "hdim_1", "hdim_2", "both"] = "none", +) -> np.array: """Adds markers for watershedding using the `features` dataframe to the marker_arr. @@ -309,21 +318,21 @@ def segmentation_2D( def segmentation_timestep( - field_in, - features_in, - dxy, - threshold=3e-3, - target="maximum", - level=None, - method="watershed", - max_distance=None, - vertical_coord=None, - PBC_flag="none", - seed_3D_flag="column", - seed_3D_size=5, - segment_number_below_threshold=0, - segment_number_unassigned=0, -): + field_in: iris.cube.Cube, + features_in: pd.DataFrame, + dxy: float, + threshold: float = 3e-3, + target: Literal["maximum", "minimum"] = "maximum", + level: Union[None, slice] = None, + method: Literal["watershed"] = "watershed", + max_distance: Union[None, float] = None, + vertical_coord: Union[str, None] = None, + PBC_flag: Literal["none", "hdim_1", "hdim_2", "both"] = "none", + seed_3D_flag: Literal["column", "box"] = "column", + seed_3D_size: Union[int, tuple[int]] = 5, + segment_number_below_threshold: int = 0, + segment_number_unassigned: int = 0, +) -> tuple[iris.cube.Cube, pd.DataFrame]: """Perform watershedding for an individual time step of the data. Works for both 2D and 3D data @@ -1023,7 +1032,7 @@ def check_add_unseeded_across_bdrys( border_max: int, markers_arr: np.array, inplace: bool = True, -): +) -> np.array: """Add new markers to unseeded but eligible regions when they are bordering an appropriate boundary. @@ -1080,21 +1089,21 @@ def check_add_unseeded_across_bdrys( def segmentation( - features, - field, - dxy, - threshold=3e-3, - target="maximum", - level=None, - method="watershed", - max_distance=None, - vertical_coord=None, - PBC_flag="none", - seed_3D_flag="column", - seed_3D_size=5, - segment_number_below_threshold=0, - segment_number_unassigned=0, -): + features: pd.DataFrame, + field: iris.cube.Cube, + dxy: float, + threshold: float = 3e-3, + target: Literal["maximum", "minimum"] = "maximum", + level: Union[None, slice] = None, + method: Literal["watershed"] = "watershed", + max_distance: Union[None, float] = None, + vertical_coord: Union[str, None] = None, + PBC_flag: Literal["none", "hdim_1", "hdim_2", "both"] = "none", + seed_3D_flag: Literal["column", "box"] = "column", + seed_3D_size: Union[int, tuple[int]] = 5, + segment_number_below_threshold: int = 0, + segment_number_unassigned: int = 0, +) -> tuple[iris.cube.Cube, pd.DataFrame]: """Use watershedding to determine region above a threshold value around initial seeding position for all time steps of the input data. Works both in 2D (based on single seeding From 818b48f118193137332ce5af7552ae439481af77 Mon Sep 17 00:00:00 2001 From: Sean Freeman Date: Thu, 5 Oct 2023 11:35:55 -0500 Subject: [PATCH 2/3] adding typing_extensions as a requirement --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 2f79a67c..22b4a2b4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ iris xarray cartopy trackpy +typing_extensions \ No newline at end of file From a635af8d928ef3a09c44c4f4bb3b8cc365a64509 Mon Sep 17 00:00:00 2001 From: Sean Freeman Date: Wed, 8 Nov 2023 09:05:22 -0600 Subject: [PATCH 3/3] formatting --- tobac/segmentation.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tobac/segmentation.py b/tobac/segmentation.py index b9ea5a29..f0011b78 100644 --- a/tobac/segmentation.py +++ b/tobac/segmentation.py @@ -344,7 +344,7 @@ def segmentation_timestep( seed_3D_size: Union[int, tuple[int]] = 5, segment_number_below_threshold: int = 0, segment_number_unassigned: int = 0, - statistics: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None]=None, + statistics: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None] = None, ) -> tuple[iris.cube.Cube, pd.DataFrame]: """Perform watershedding for an individual time step of the data. Works for both 2D and 3D data @@ -1126,8 +1126,7 @@ def segmentation( seed_3D_size: Union[int, tuple[int]] = 5, segment_number_below_threshold: int = 0, segment_number_unassigned: int = 0, - statistics: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None]=None, - + statistics: Union[dict[str, Union[Callable, tuple[Callable, dict]]], None] = None, ) -> tuple[iris.cube.Cube, pd.DataFrame]: """Use watershedding to determine region above a threshold value around initial seeding position for all time steps of