From 17c55eef1ce232736af99d7fe643f1e4d4ca5aca Mon Sep 17 00:00:00 2001 From: William Jones Date: Thu, 7 Dec 2023 01:25:59 +0000 Subject: [PATCH 1/7] Update build_distance_function to always provide integer values to calc_distance_coords_pbc, even if given None --- tobac/tracking.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tobac/tracking.py b/tobac/tracking.py index 6e2b8028..362ead05 100644 --- a/tobac/tracking.py +++ b/tobac/tracking.py @@ -643,9 +643,9 @@ def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag): return functools.partial( pbc_utils.calc_distance_coords_pbc, - min_h1=min_h1, - max_h1=max_h1, - min_h2=min_h2, - max_h2=max_h2, + min_h1=min_h1 if min_h1 is not None else 0, + max_h1=max_h1 if max_h1 is not None else 0, + min_h2=min_h2 if min_h2 is not None else 0, + max_h2=max_h2 if max_h2 is not None else 0, PBC_flag=PBC_flag, ) From 17e12d27efd69a5e0e991dc6074a920c6bcefc1d Mon Sep 17 00:00:00 2001 From: William Jones Date: Thu, 7 Dec 2023 01:31:38 +0000 Subject: [PATCH 2/7] Add additional tests for build_distance_function with different PBC options --- tobac/tests/test_tracking.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tobac/tests/test_tracking.py b/tobac/tests/test_tracking.py index d21b1254..1266b210 100644 --- a/tobac/tests/test_tracking.py +++ b/tobac/tests/test_tracking.py @@ -285,6 +285,17 @@ def test_build_distance_function(): 1.4142135 ) + test_func = tobac.tracking.build_distance_function(0, 10, None, None, "hdim_1") + assert test_func(np.array((0, 9, 9)), np.array((0, 0, 9))) == pytest.approx(1) + + test_func = tobac.tracking.build_distance_function(None, None, 0, 10, "hdim_2") + assert test_func(np.array((0, 9, 9)), np.array((0, 9, 0))) == pytest.approx(1) + + test_func = tobac.tracking.build_distance_function(None, None, None, None, "none") + assert test_func(np.array((0, 9, 9)), np.array((0, 0, 0))) == pytest.approx( + (2 * 81) ** 0.5 + ) + @pytest.mark.parametrize( "point_init, speed, dxy, actual_dz, v_max, use_dz, features_connected", From 22495082fe11af76d82c14a04ed70056eb413fa2 Mon Sep 17 00:00:00 2001 From: William Jones Date: Thu, 7 Dec 2023 12:46:46 +0000 Subject: [PATCH 3/7] Improved validation of inputs in build_distance_function --- tobac/tests/test_tracking.py | 23 +++++++++- tobac/tracking.py | 84 ++++++++++++++++++++++++++++++++++-- 2 files changed, 102 insertions(+), 5 deletions(-) diff --git a/tobac/tests/test_tracking.py b/tobac/tests/test_tracking.py index 1266b210..631cf327 100644 --- a/tobac/tests/test_tracking.py +++ b/tobac/tests/test_tracking.py @@ -274,7 +274,7 @@ def test_linking_trackpy(): ) -def test_build_distance_function(): +def test_build_distance_function_3D(): """Tests ```tobac.tracking.build_distance_function``` Currently tests: that this produces an object that is suitable to call from trackpy @@ -297,6 +297,27 @@ def test_build_distance_function(): ) +def test_build_distance_function_2D(): + """Tests ```tobac.tracking.build_distance_function``` + Currently tests: + that this produces an object that is suitable to call from trackpy + """ + + test_func = tobac.tracking.build_distance_function(0, 10, 0, 10, "both") + assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx(1.4142135) + + test_func = tobac.tracking.build_distance_function(0, 10, None, None, "hdim_1") + assert test_func(np.array((9, 9)), np.array((0, 9))) == pytest.approx(1) + + test_func = tobac.tracking.build_distance_function(None, None, 0, 10, "hdim_2") + assert test_func(np.array((9, 9)), np.array((9, 0))) == pytest.approx(1) + + test_func = tobac.tracking.build_distance_function(None, None, None, None, "none") + assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx( + (2 * 81) ** 0.5 + ) + + @pytest.mark.parametrize( "point_init, speed, dxy, actual_dz, v_max, use_dz, features_connected", [ diff --git a/tobac/tracking.py b/tobac/tracking.py index 362ead05..549f4e3a 100644 --- a/tobac/tracking.py +++ b/tobac/tracking.py @@ -641,11 +641,87 @@ def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag): """ import functools + h1_size, h2_size = validate_pbc_dims(min_h1, max_h1, min_h2, max_h2, PBC_flag) + return functools.partial( pbc_utils.calc_distance_coords_pbc, - min_h1=min_h1 if min_h1 is not None else 0, - max_h1=max_h1 if max_h1 is not None else 0, - min_h2=min_h2 if min_h2 is not None else 0, - max_h2=max_h2 if max_h2 is not None else 0, + min_h1=0, + max_h1=h1_size, + min_h2=0, + max_h2=h2_size, PBC_flag=PBC_flag, ) + + +def validate_pbc_dims( + min_h1: int, max_h1: int, min_h2: int, max_h2: int, PBC_flag: str +) -> tuple[int, int]: + """Validate the input parameters for build_distance_function and return size of each axis + + Parameters + ---------- + min_h1: int + Minimum point in hdim_1 + max_h1: int + Maximum point in hdim_1 + min_h2: int + Minimum point in hdim_2 + max_h2: int + Maximum point in hdim_2 + PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both') + 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 + + Returns + ------- + tuple[int, int] + size of domain in hdim1 and hdim2 + """ + if PBC_flag == "none": + return (0, 0) + if PBC_flag == "both": + invalid_dim_limits = invalid_limit_names( + min_h1=min_h1, max_h1=max_h1, min_h2=min_h2, max_h2=max_h2 + ) + if invalid_dim_limits: + raise PBCLimitError(invalid_dim_limits, PBC_flag) + return (max_h1 - min_h1, max_h2 - min_h2) + if PBC_flag == "hdim_1": + invalid_dim_limits = invalid_limit_names(min_h1=min_h1, max_h1=max_h1) + if invalid_dim_limits: + raise PBCLimitError(invalid_dim_limits, PBC_flag) + return (max_h1 - min_h1, 0) + if PBC_flag == "hdim_2": + invalid_dim_limits = invalid_limit_names(min_h2=min_h2, max_h2=max_h2) + if invalid_dim_limits: + raise PBCLimitError(invalid_dim_limits, PBC_flag) + return (0, max_h2 - min_h2) + # if PBC_flag not in ('none', 'hdim_1', 'hdim_2', 'both'): + raise PBCflagError() + + +def invalid_limit_names(**limits) -> list[str]: + """Return the names of keywords if their value is None + + Returns + ------- + list[str] + List of provided keywords with value None + """ + return [k for k, v in limits.items() if v is None] + + +class PBCflagError(ValueError): + def __init__(self): + super.init( + "PBC_flag keyword is not valid, must be one of ['none', 'hdim_1', 'hdim_2', 'both']" + ) + + +class PBCLimitError(ValueError): + def __init__(self, invalid_limits, PBC_flag): + self.message = f"Keyword parameters {invalid_limits} must be provided for PBC_flag {PBC_flag}" + super.init(self.message) From f6d3555fc47062f6e5a9a5512efdc36e2b987cc9 Mon Sep 17 00:00:00 2001 From: William Jones Date: Thu, 7 Dec 2023 13:33:04 +0000 Subject: [PATCH 4/7] Move build_distance_func to pbc_utils --- tobac/feature_detection.py | 3 +- tobac/tests/test_pbc_utils.py | 44 +++++++++++ tobac/tests/test_tracking.py | 53 ++----------- tobac/tracking.py | 120 +---------------------------- tobac/utils/periodic_boundaries.py | 116 ++++++++++++++++++++++++++++ 5 files changed, 169 insertions(+), 167 deletions(-) diff --git a/tobac/feature_detection.py b/tobac/feature_detection.py index a1263cd4..d5baeeb1 100644 --- a/tobac/feature_detection.py +++ b/tobac/feature_detection.py @@ -30,7 +30,6 @@ import iris import xarray as xr -from tobac.tracking import build_distance_function from tobac.utils import internal as internal_utils from tobac.utils import periodic_boundaries as pbc_utils from tobac.utils.general import spectral_filtering @@ -1541,7 +1540,7 @@ def filter_min_distance( # Check if we have PBCs. 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( + dist_func = pbc_utils.build_distance_function( min_h1 * dxy, max_h1 * dxy, min_h2 * dxy, max_h2 * dxy, PBC_flag ) features_tree = BallTree(feature_locations, metric="pyfunc", func=dist_func) diff --git a/tobac/tests/test_pbc_utils.py b/tobac/tests/test_pbc_utils.py index c88ac2e6..95564048 100644 --- a/tobac/tests/test_pbc_utils.py +++ b/tobac/tests/test_pbc_utils.py @@ -286,6 +286,50 @@ def test_get_pbc_coordinates(): ) +def test_build_distance_function_3D(): + """Tests ```pbc_utils.build_distance_function``` + Currently tests: + that this produces an object that is suitable to call from trackpy + """ + + test_func = pbc_utils.build_distance_function(0, 10, 0, 10, "both") + assert test_func(np.array((0, 9, 9)), np.array((0, 0, 0))) == pytest.approx( + 1.4142135 + ) + + test_func = pbc_utils.build_distance_function(0, 10, None, None, "hdim_1") + assert test_func(np.array((0, 9, 9)), np.array((0, 0, 9))) == pytest.approx(1) + + test_func = pbc_utils.build_distance_function(None, None, 0, 10, "hdim_2") + assert test_func(np.array((0, 9, 9)), np.array((0, 9, 0))) == pytest.approx(1) + + test_func = pbc_utils.build_distance_function(None, None, None, None, "none") + assert test_func(np.array((0, 9, 9)), np.array((0, 0, 0))) == pytest.approx( + (2 * 81) ** 0.5 + ) + + +def test_build_distance_function_2D(): + """Tests ```pbc_utils.build_distance_function``` + Currently tests: + that this produces an object that is suitable to call from trackpy + """ + + test_func = pbc_utils.build_distance_function(0, 10, 0, 10, "both") + assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx(1.4142135) + + test_func = pbc_utils.build_distance_function(0, 10, None, None, "hdim_1") + assert test_func(np.array((9, 9)), np.array((0, 9))) == pytest.approx(1) + + test_func = pbc_utils.build_distance_function(None, None, 0, 10, "hdim_2") + assert test_func(np.array((9, 9)), np.array((9, 0))) == pytest.approx(1) + + test_func = pbc_utils.build_distance_function(None, None, None, None, "none") + assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx( + (2 * 81) ** 0.5 + ) + + def test_weighted_circmean() -> None: """ Test that weighted_circmean gives the expected results compared to scipy.stats.circmean diff --git a/tobac/tests/test_tracking.py b/tobac/tests/test_tracking.py index 631cf327..e810ce54 100644 --- a/tobac/tests/test_tracking.py +++ b/tobac/tests/test_tracking.py @@ -2,16 +2,17 @@ Test for the trackpy tracking functions """ import datetime +import copy import pytest -import tobac.testing -import tobac.tracking -import copy +import numpy as np import pandas as pd from pandas.testing import assert_frame_equal -import numpy as np import trackpy as tp +import tobac.testing +import tobac.tracking + def convert_cell_dtype_if_appropriate(output, expected_output): """Helper function to convert datatype of output if @@ -274,50 +275,6 @@ def test_linking_trackpy(): ) -def test_build_distance_function_3D(): - """Tests ```tobac.tracking.build_distance_function``` - Currently tests: - that this produces an object that is suitable to call from trackpy - """ - - test_func = tobac.tracking.build_distance_function(0, 10, 0, 10, "both") - assert test_func(np.array((0, 9, 9)), np.array((0, 0, 0))) == pytest.approx( - 1.4142135 - ) - - test_func = tobac.tracking.build_distance_function(0, 10, None, None, "hdim_1") - assert test_func(np.array((0, 9, 9)), np.array((0, 0, 9))) == pytest.approx(1) - - test_func = tobac.tracking.build_distance_function(None, None, 0, 10, "hdim_2") - assert test_func(np.array((0, 9, 9)), np.array((0, 9, 0))) == pytest.approx(1) - - test_func = tobac.tracking.build_distance_function(None, None, None, None, "none") - assert test_func(np.array((0, 9, 9)), np.array((0, 0, 0))) == pytest.approx( - (2 * 81) ** 0.5 - ) - - -def test_build_distance_function_2D(): - """Tests ```tobac.tracking.build_distance_function``` - Currently tests: - that this produces an object that is suitable to call from trackpy - """ - - test_func = tobac.tracking.build_distance_function(0, 10, 0, 10, "both") - assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx(1.4142135) - - test_func = tobac.tracking.build_distance_function(0, 10, None, None, "hdim_1") - assert test_func(np.array((9, 9)), np.array((0, 9))) == pytest.approx(1) - - test_func = tobac.tracking.build_distance_function(None, None, 0, 10, "hdim_2") - assert test_func(np.array((9, 9)), np.array((9, 0))) == pytest.approx(1) - - test_func = tobac.tracking.build_distance_function(None, None, None, None, "none") - assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx( - (2 * 81) ** 0.5 - ) - - @pytest.mark.parametrize( "point_init, speed, dxy, actual_dz, v_max, use_dz, features_connected", [ diff --git a/tobac/tracking.py b/tobac/tracking.py index 549f4e3a..648bac3d 100644 --- a/tobac/tracking.py +++ b/tobac/tracking.py @@ -329,7 +329,9 @@ def linking_trackpy( # which we need for PBCs, neighbor_strategy must be 'BTree'. # I think this shouldn't change results, but it will degrade performance. neighbor_strategy = "BTree" - dist_func = build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag) + dist_func = pbc_utils.build_distance_function( + min_h1, max_h1, min_h2, max_h2, PBC_flag + ) else: neighbor_strategy = "KDTree" @@ -609,119 +611,3 @@ def remap_particle_to_cell_nv(particle_cell_map, input_particle): """ return particle_cell_map[input_particle] - - -def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag): - """Function to build a partial ```calc_distance_coords_pbc``` function - suitable for use with trackpy - - Parameters - ---------- - min_h1: int - Minimum point in hdim_1 - max_h1: int - Maximum point in hdim_1 - min_h2: int - Minimum point in hdim_2 - max_h2: int - Maximum point in hdim_2 - PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both') - 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 - - Returns - ------- - function object - A version of calc_distance_coords_pbc suitable to be called by - just f(coords_1, coords_2) - - """ - import functools - - h1_size, h2_size = validate_pbc_dims(min_h1, max_h1, min_h2, max_h2, PBC_flag) - - return functools.partial( - pbc_utils.calc_distance_coords_pbc, - min_h1=0, - max_h1=h1_size, - min_h2=0, - max_h2=h2_size, - PBC_flag=PBC_flag, - ) - - -def validate_pbc_dims( - min_h1: int, max_h1: int, min_h2: int, max_h2: int, PBC_flag: str -) -> tuple[int, int]: - """Validate the input parameters for build_distance_function and return size of each axis - - Parameters - ---------- - min_h1: int - Minimum point in hdim_1 - max_h1: int - Maximum point in hdim_1 - min_h2: int - Minimum point in hdim_2 - max_h2: int - Maximum point in hdim_2 - PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both') - 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 - - Returns - ------- - tuple[int, int] - size of domain in hdim1 and hdim2 - """ - if PBC_flag == "none": - return (0, 0) - if PBC_flag == "both": - invalid_dim_limits = invalid_limit_names( - min_h1=min_h1, max_h1=max_h1, min_h2=min_h2, max_h2=max_h2 - ) - if invalid_dim_limits: - raise PBCLimitError(invalid_dim_limits, PBC_flag) - return (max_h1 - min_h1, max_h2 - min_h2) - if PBC_flag == "hdim_1": - invalid_dim_limits = invalid_limit_names(min_h1=min_h1, max_h1=max_h1) - if invalid_dim_limits: - raise PBCLimitError(invalid_dim_limits, PBC_flag) - return (max_h1 - min_h1, 0) - if PBC_flag == "hdim_2": - invalid_dim_limits = invalid_limit_names(min_h2=min_h2, max_h2=max_h2) - if invalid_dim_limits: - raise PBCLimitError(invalid_dim_limits, PBC_flag) - return (0, max_h2 - min_h2) - # if PBC_flag not in ('none', 'hdim_1', 'hdim_2', 'both'): - raise PBCflagError() - - -def invalid_limit_names(**limits) -> list[str]: - """Return the names of keywords if their value is None - - Returns - ------- - list[str] - List of provided keywords with value None - """ - return [k for k, v in limits.items() if v is None] - - -class PBCflagError(ValueError): - def __init__(self): - super.init( - "PBC_flag keyword is not valid, must be one of ['none', 'hdim_1', 'hdim_2', 'both']" - ) - - -class PBCLimitError(ValueError): - def __init__(self, invalid_limits, PBC_flag): - self.message = f"Keyword parameters {invalid_limits} must be provided for PBC_flag {PBC_flag}" - super.init(self.message) diff --git a/tobac/utils/periodic_boundaries.py b/tobac/utils/periodic_boundaries.py index 52c93c4b..121dddc7 100644 --- a/tobac/utils/periodic_boundaries.py +++ b/tobac/utils/periodic_boundaries.py @@ -255,6 +255,122 @@ def calc_distance_coords_pbc( return np.sqrt(np.sum(deltas**2)) +def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag): + """Function to build a partial ```calc_distance_coords_pbc``` function + suitable for use with trackpy + + Parameters + ---------- + min_h1: int + Minimum point in hdim_1 + max_h1: int + Maximum point in hdim_1 + min_h2: int + Minimum point in hdim_2 + max_h2: int + Maximum point in hdim_2 + PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both') + 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 + + Returns + ------- + function object + A version of calc_distance_coords_pbc suitable to be called by + just f(coords_1, coords_2) + + """ + import functools + + h1_size, h2_size = validate_pbc_dims(min_h1, max_h1, min_h2, max_h2, PBC_flag) + + return functools.partial( + calc_distance_coords_pbc, + min_h1=0, + max_h1=h1_size, + min_h2=0, + max_h2=h2_size, + PBC_flag=PBC_flag, + ) + + +def validate_pbc_dims( + min_h1: int, max_h1: int, min_h2: int, max_h2: int, PBC_flag: str +) -> tuple[int, int]: + """Validate the input parameters for build_distance_function and return size of each axis + + Parameters + ---------- + min_h1: int + Minimum point in hdim_1 + max_h1: int + Maximum point in hdim_1 + min_h2: int + Minimum point in hdim_2 + max_h2: int + Maximum point in hdim_2 + PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both') + 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 + + Returns + ------- + tuple[int, int] + size of domain in hdim1 and hdim2 + """ + if PBC_flag == "none": + return (0, 0) + if PBC_flag == "both": + invalid_dim_limits = invalid_limit_names( + min_h1=min_h1, max_h1=max_h1, min_h2=min_h2, max_h2=max_h2 + ) + if invalid_dim_limits: + raise PBCLimitError(invalid_dim_limits, PBC_flag) + return (max_h1 - min_h1, max_h2 - min_h2) + if PBC_flag == "hdim_1": + invalid_dim_limits = invalid_limit_names(min_h1=min_h1, max_h1=max_h1) + if invalid_dim_limits: + raise PBCLimitError(invalid_dim_limits, PBC_flag) + return (max_h1 - min_h1, 0) + if PBC_flag == "hdim_2": + invalid_dim_limits = invalid_limit_names(min_h2=min_h2, max_h2=max_h2) + if invalid_dim_limits: + raise PBCLimitError(invalid_dim_limits, PBC_flag) + return (0, max_h2 - min_h2) + # if PBC_flag not in ('none', 'hdim_1', 'hdim_2', 'both'): + raise PBCflagError() + + +def invalid_limit_names(**limits) -> list[str]: + """Return the names of keywords if their value is None + + Returns + ------- + list[str] + List of provided keywords with value None + """ + return [k for k, v in limits.items() if v is None] + + +class PBCflagError(ValueError): + def __init__(self): + super.init( + "PBC_flag keyword is not valid, must be one of ['none', 'hdim_1', 'hdim_2', 'both']" + ) + + +class PBCLimitError(ValueError): + def __init__(self, invalid_limits, PBC_flag): + self.message = f"Keyword parameters {invalid_limits} must be provided for PBC_flag {PBC_flag}" + super.init(self.message) + + def weighted_circmean( values: np.ndarray, weights: np.ndarray, From 009995890c98edec3e17a63f3ea069cc5b7f30d1 Mon Sep 17 00:00:00 2001 From: William Jones Date: Thu, 7 Dec 2023 15:57:07 +0000 Subject: [PATCH 5/7] Refactor and reorganise PBC distance functions --- tobac/feature_detection.py | 2 +- tobac/segmentation.py | 60 ++++---- tobac/tests/test_pbc_utils.py | 239 ++++++++++++++++++++++------- tobac/tracking.py | 2 +- tobac/utils/periodic_boundaries.py | 68 +++----- 5 files changed, 237 insertions(+), 134 deletions(-) diff --git a/tobac/feature_detection.py b/tobac/feature_detection.py index d5baeeb1..5dbb0702 100644 --- a/tobac/feature_detection.py +++ b/tobac/feature_detection.py @@ -1541,7 +1541,7 @@ def filter_min_distance( if PBC_flag in ["hdim_1", "hdim_2", "both"]: # Note that we multiply by dxy to get the distances in spatial coordinates dist_func = pbc_utils.build_distance_function( - min_h1 * dxy, max_h1 * dxy, min_h2 * dxy, max_h2 * dxy, PBC_flag + min_h1 * dxy, max_h1 * dxy, min_h2 * dxy, max_h2 * dxy, PBC_flag, is_3D ) features_tree = BallTree(feature_locations, metric="pyfunc", func=dist_func) neighbours = features_tree.query_radius(feature_locations, r=min_distance) diff --git a/tobac/segmentation.py b/tobac/segmentation.py index b59d7975..a6b945db 100644 --- a/tobac/segmentation.py +++ b/tobac/segmentation.py @@ -179,6 +179,10 @@ def add_markers( h2_end_coord=hdim_2_max, PBC_flag=PBC_flag, ) + # Build distance function ahead of time, 3D always true as we then reduce + dist_func = pbc_utils.build_distance_function( + 0, h1_len, 0, h2_len, PBC_flag, True + ) for seed_box in all_seed_boxes: # Need to see if there are any other points seeded # in this seed box first. @@ -198,32 +202,28 @@ def add_markers( # Get its global index so that we can calculate # distance and set the array. local_index = it.multi_index - global_index = ( - local_index[0] + z_seed_start, - local_index[1] + seed_box[0], - local_index[2] + seed_box[2], + global_index = np.array( + ( + local_index[0] + z_seed_start, + local_index[1] + seed_box[0], + local_index[2] + seed_box[2], + ) ) # If it's a background marker, we can just set it # with the feature we're working on. if curr_box_pt == bg_marker: - marker_arr[global_index] = row["feature"] + marker_arr[*global_index] = row["feature"] continue # it has another feature in it. Calculate the distance # from its current set feature and the new feature. if is_3D: - curr_coord = (row["vdim"], row["hdim_1"], row["hdim_2"]) + curr_coord = np.array( + (row["vdim"], row["hdim_1"], row["hdim_2"]) + ) else: - curr_coord = (0, row["hdim_1"], row["hdim_2"]) - - dist_from_curr_pt = pbc_utils.calc_distance_coords_pbc( - np.array(global_index), - np.array(curr_coord), - min_h1=0, - max_h1=h1_len, - min_h2=0, - max_h2=h2_len, - PBC_flag=PBC_flag, - ) + curr_coord = np.array((0, row["hdim_1"], row["hdim_2"])) + + dist_from_curr_pt = dist_func(global_index, curr_coord) # This is technically an O(N^2) operation, but # hopefully performance isn't too bad as this should @@ -232,29 +232,25 @@ def add_markers( features["feature"] == curr_box_pt ].iloc[0] if is_3D: - orig_coord = ( - orig_row["vdim"], - orig_row["hdim_1"], - orig_row["hdim_2"], + orig_coord = np.array( + ( + orig_row["vdim"], + orig_row["hdim_1"], + orig_row["hdim_2"], + ) ) else: - orig_coord = (0, orig_row["hdim_1"], orig_row["hdim_2"]) - dist_from_orig_pt = pbc_utils.calc_distance_coords_pbc( - np.array(global_index), - np.array(orig_coord), - min_h1=0, - max_h1=h1_len, - min_h2=0, - max_h2=h2_len, - PBC_flag=PBC_flag, - ) + orig_coord = np.array( + (0, orig_row["hdim_1"], orig_row["hdim_2"]) + ) + dist_from_orig_pt = dist_func(global_index, orig_coord) # The current point center is further away # than the original point center, so do nothing if dist_from_curr_pt > dist_from_orig_pt: continue else: # the current point center is closer. - marker_arr[global_index] = row["feature"] + marker_arr[*global_index] = row["feature"] # completely unseeded region so far. else: marker_arr[ diff --git a/tobac/tests/test_pbc_utils.py b/tobac/tests/test_pbc_utils.py index 95564048..2fd39d09 100644 --- a/tobac/tests/test_pbc_utils.py +++ b/tobac/tests/test_pbc_utils.py @@ -4,7 +4,7 @@ import tobac.testing as tb_test -def test_calc_distance_coords_pbc(): +def test_calc_distance_coords_pbc_2D(): """Tests ```tobac.utils.calc_distance_coords_pbc``` Currently tests: two points in normal space @@ -13,102 +13,226 @@ def test_calc_distance_coords_pbc(): import numpy as np # Test first two points in normal space with varying PBC conditions - for PBC_condition in ["none", "hdim_1", "hdim_2", "both"]: + for max_dims in [ + np.array([0, 0]), + np.array([10, 0]), + np.array([0, 10]), + np.array([10, 10]), + ]: assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 0)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0)), np.array((0, 0)), max_dims ) == pytest.approx(0) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 0)), np.array((0, 0, 1)), 0, 10, 0, 10, PBC_condition + np.array((0, 0)), np.array((0, 1)), max_dims ) == pytest.approx(1) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0)), np.array((0, 1)), 0, 10, 0, 10, PBC_condition + np.array((0, 0)), np.array((3, 1)), max_dims + ) == pytest.approx(10**0.5, rel=1e-3) + + # Now test two points that will be closer along the hdim_1 boundary for cases without PBCs + for max_dims in [np.array([10, 0]), np.array([10, 10])]: + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 0)), np.array((9, 0)), max_dims + ) == pytest.approx(1) + assert pbc_utils.calc_distance_coords_pbc( + np.array((9, 0)), np.array((0, 0)), max_dims + ) == pytest.approx(1) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 4)), np.array((7, 3)), max_dims + ) == pytest.approx(10**0.5) + + # Test the same points, except without PBCs + for max_dims in [np.array([0, 0]), np.array([0, 10])]: + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 0)), np.array((9, 0)), max_dims + ) == pytest.approx(9) + assert pbc_utils.calc_distance_coords_pbc( + np.array((9, 0)), np.array((0, 0)), max_dims + ) == pytest.approx(9) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 4)), np.array((7, 3)), max_dims + ) == pytest.approx(50**0.5) + + # Now test two points that will be closer along the hdim_2 boundary for cases without PBCs + for max_dims in [np.array([0, 10]), np.array([10, 10])]: + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 0)), np.array((0, 9)), max_dims + ) == pytest.approx(1) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 9)), np.array((0, 0)), max_dims + ) == pytest.approx(1) + + # Test the same points, except without PBCs + for max_dims in [np.array([0, 0]), np.array([10, 0])]: + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 0)), np.array((0, 9)), max_dims + ) == pytest.approx(9) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 9)), np.array((0, 0)), max_dims + ) == pytest.approx(9) + + # Test points that will be closer for the both + max_dims = np.array([10, 10]) + assert pbc_utils.calc_distance_coords_pbc( + np.array((9, 9)), np.array((0, 0)), max_dims + ) == pytest.approx(1.4142135, rel=1e-3) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 9)), np.array((9, 0)), max_dims + ) == pytest.approx(1.4142135, rel=1e-3) + + # Test the corner points for no PBCs + max_dims = np.array([0, 0]) + assert pbc_utils.calc_distance_coords_pbc( + np.array((9, 9)), np.array((0, 0)), max_dims + ) == pytest.approx(12.727922, rel=1e-3) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 9)), np.array((9, 0)), max_dims + ) == pytest.approx(12.727922, rel=1e-3) + + # Test the corner points for hdim_1 and hdim_2 + for max_dims in [np.array([10, 0]), np.array([0, 10])]: + assert pbc_utils.calc_distance_coords_pbc( + np.array((9, 9)), np.array((0, 0)), max_dims + ) == pytest.approx(9.055385) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 9)), np.array((9, 0)), max_dims + ) == pytest.approx(9.055385) + + +def test_calc_distance_coords_pbc_3D(): + """Tests ```tobac.utils.calc_distance_coords_pbc``` + Currently tests: + two points in normal space + Periodicity along hdim_1, hdim_2, and corners + """ + import numpy as np + + # Test first two points in normal space with varying PBC conditions + for max_dims in [ + np.array([0, 0, 0]), + np.array([0, 10, 0]), + np.array([0, 0, 10]), + np.array([0, 10, 10]), + ]: + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 0, 0)), np.array((0, 0, 0)), max_dims + ) == pytest.approx(0) + assert pbc_utils.calc_distance_coords_pbc( + np.array((0, 0, 0)), np.array((0, 0, 1)), max_dims ) == pytest.approx(1) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 0)), np.array((3, 3, 1)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 0)), np.array((3, 3, 1)), max_dims ) == pytest.approx(4.3588989, rel=1e-3) # Now test two points that will be closer along the hdim_1 boundary for cases without PBCs - for PBC_condition in ["hdim_1", "both"]: + for max_dims in [np.array([0, 10, 0]), np.array([0, 10, 10])]: assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 0)), np.array((0, 9, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 0)), np.array((0, 9, 0)), max_dims ) == pytest.approx(1) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 9, 0)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 9, 0)), np.array((0, 0, 0)), max_dims ) == pytest.approx(1) assert pbc_utils.calc_distance_coords_pbc( - np.array((8, 0)), np.array((0, 0)), 0, 10, 0, 10, PBC_condition - ) == pytest.approx(2) - assert pbc_utils.calc_distance_coords_pbc( - np.array((4, 0, 4)), np.array((3, 7, 3)), 0, 10, 0, 10, PBC_condition + np.array((4, 0, 4)), np.array((3, 7, 3)), max_dims ) == pytest.approx(3.3166247) assert pbc_utils.calc_distance_coords_pbc( - np.array((4, 0, 4)), np.array((3, 7, 3)), 0, 10, 0, 10, PBC_condition + np.array((4, 0, 4)), np.array((3, 7, 3)), max_dims ) == pytest.approx(3.3166247) # Test the same points, except without PBCs - for PBC_condition in ["none", "hdim_2"]: + for max_dims in [np.array([0, 0, 0]), np.array([0, 0, 10])]: assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 0)), np.array((0, 9, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 0)), np.array((0, 9, 0)), max_dims ) == pytest.approx(9) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 9, 0)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 9, 0)), np.array((0, 0, 0)), max_dims ) == pytest.approx(9) - assert pbc_utils.calc_distance_coords_pbc( - np.array((8, 0)), np.array((0, 0)), 0, 10, 0, 10, PBC_condition - ) == pytest.approx(8) # Now test two points that will be closer along the hdim_2 boundary for cases without PBCs - for PBC_condition in ["hdim_2", "both"]: + for max_dims in [np.array([0, 0, 10]), np.array([0, 10, 10])]: assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 0)), np.array((0, 0, 9)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 0)), np.array((0, 0, 9)), max_dims ) == pytest.approx(1) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 9)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 9)), np.array((0, 0, 0)), max_dims ) == pytest.approx(1) - assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 8)), np.array((0, 0)), 0, 10, 0, 10, PBC_condition - ) == pytest.approx(2) # Test the same points, except without PBCs - for PBC_condition in ["none", "hdim_1"]: + for max_dims in [np.array([0, 0, 0]), np.array([0, 10, 0])]: assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 0)), np.array((0, 0, 9)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 0)), np.array((0, 0, 9)), max_dims ) == pytest.approx(9) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 9)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 9)), np.array((0, 0, 0)), max_dims ) == pytest.approx(9) - assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 8)), np.array((0, 0)), 0, 10, 0, 10, PBC_condition - ) == pytest.approx(8) # Test points that will be closer for the both - PBC_condition = "both" + max_dims = np.array([0, 10, 10]) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 9, 9)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 9, 9)), np.array((0, 0, 0)), max_dims ) == pytest.approx(1.4142135, rel=1e-3) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 9)), np.array((0, 9, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 9)), np.array((0, 9, 0)), max_dims ) == pytest.approx(1.4142135, rel=1e-3) # Test the corner points for no PBCs - PBC_condition = "none" + max_dims = np.array([0, 0, 0]) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 9, 9)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 9, 9)), np.array((0, 0, 0)), max_dims ) == pytest.approx(12.727922, rel=1e-3) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 9)), np.array((0, 9, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 9)), np.array((0, 9, 0)), max_dims ) == pytest.approx(12.727922, rel=1e-3) # Test the corner points for hdim_1 and hdim_2 - for PBC_condition in ["hdim_1", "hdim_2"]: + for max_dims in [np.array([0, 10, 0]), np.array([0, 0, 10])]: assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 9, 9)), np.array((0, 0, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 9, 9)), np.array((0, 0, 0)), max_dims ) == pytest.approx(9.055385) assert pbc_utils.calc_distance_coords_pbc( - np.array((0, 0, 9)), np.array((0, 9, 0)), 0, 10, 0, 10, PBC_condition + np.array((0, 0, 9)), np.array((0, 9, 0)), max_dims ) == pytest.approx(9.055385) +def test_invalid_limit_names(): + """ + Test that invalid_limit_names will correctly return the names of keywords that equal None + """ + assert pbc_utils.invalid_limit_names() == [] + assert pbc_utils.invalid_limit_names(a=0, b=0) == [] + assert pbc_utils.invalid_limit_names(a=None, b=0) == ["a"] + assert pbc_utils.invalid_limit_names(a=0, b=None) == ["b"] + assert pbc_utils.invalid_limit_names(a=None, b=None) == ["a", "b"] + + +def test_validate_pbc_dims(): + """ + Test validate_pbc_dims works correctly and raise exceptions appropriately + """ + # Assert that size is only returned if PBCs are required in that dimension + assert pbc_utils.validate_pbc_dims(0, 10, 0, 15, "none") == (0, 0) + assert pbc_utils.validate_pbc_dims(0, 10, 0, 15, "hdim_1") == (10, 0) + assert pbc_utils.validate_pbc_dims(0, 10, 0, 15, "hdim_2") == (0, 15) + assert pbc_utils.validate_pbc_dims(0, 10, 0, 15, "both") == (10, 15) + + # Assert that giving None for limits of dimensions that are not PBCs is handled correctly + assert pbc_utils.validate_pbc_dims(None, None, None, None, "none") == (0, 0) + assert pbc_utils.validate_pbc_dims(0, 10, None, None, "hdim_1") == (10, 0) + assert pbc_utils.validate_pbc_dims(None, None, 0, 15, "hdim_2") == (0, 15) + + # Test that an invalid option for PBC_flag raises the correct error + with pytest.raises(pbc_utils.PBCflagError): + _ = pbc_utils.validate_pbc_dims(0, 10, 0, 15, "invalid_pbc_option") + + # Test that providing None for the limits of a PBC dimension raises the correct error + with pytest.raises(pbc_utils.PBCLimitError): + _ = pbc_utils.validate_pbc_dims(None, None, 0, 15, "hdim_1") + with pytest.raises(pbc_utils.PBCLimitError): + _ = pbc_utils.validate_pbc_dims(None, None, None, 15, "hdim_2") + with pytest.raises(pbc_utils.PBCLimitError): + _ = pbc_utils.validate_pbc_dims(0, None, 0, None, "both") + + @pytest.mark.parametrize( "loc_1, loc_2, bounds, PBC_flag, expected_dist", [ @@ -135,16 +259,23 @@ def test_calc_distance_coords_pbc_param(loc_1, loc_2, bounds, PBC_flag, expected expected_dist: float Expected distance between the two points """ - import numpy as np - - assert pbc_utils.calc_distance_coords_pbc( - np.array(loc_1), - np.array(loc_2), + h1_size, h2_size = pbc_utils.validate_pbc_dims( bounds[0], bounds[1], bounds[2], bounds[3], PBC_flag, + ) + + if len(loc_1) == 3: + max_dims = np.array([0, h1_size, h2_size]) + else: + max_dims = np.array([h1_size, h2_size]) + + assert pbc_utils.calc_distance_coords_pbc( + np.array(loc_1), + np.array(loc_2), + max_dims, ) == pytest.approx(expected_dist) @@ -292,18 +423,18 @@ def test_build_distance_function_3D(): that this produces an object that is suitable to call from trackpy """ - test_func = pbc_utils.build_distance_function(0, 10, 0, 10, "both") + test_func = pbc_utils.build_distance_function(0, 10, 0, 10, "both", True) assert test_func(np.array((0, 9, 9)), np.array((0, 0, 0))) == pytest.approx( 1.4142135 ) - test_func = pbc_utils.build_distance_function(0, 10, None, None, "hdim_1") + test_func = pbc_utils.build_distance_function(0, 10, None, None, "hdim_1", True) assert test_func(np.array((0, 9, 9)), np.array((0, 0, 9))) == pytest.approx(1) - test_func = pbc_utils.build_distance_function(None, None, 0, 10, "hdim_2") + test_func = pbc_utils.build_distance_function(None, None, 0, 10, "hdim_2", True) assert test_func(np.array((0, 9, 9)), np.array((0, 9, 0))) == pytest.approx(1) - test_func = pbc_utils.build_distance_function(None, None, None, None, "none") + test_func = pbc_utils.build_distance_function(None, None, None, None, "none", True) assert test_func(np.array((0, 9, 9)), np.array((0, 0, 0))) == pytest.approx( (2 * 81) ** 0.5 ) @@ -315,16 +446,16 @@ def test_build_distance_function_2D(): that this produces an object that is suitable to call from trackpy """ - test_func = pbc_utils.build_distance_function(0, 10, 0, 10, "both") + test_func = pbc_utils.build_distance_function(0, 10, 0, 10, "both", False) assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx(1.4142135) - test_func = pbc_utils.build_distance_function(0, 10, None, None, "hdim_1") + test_func = pbc_utils.build_distance_function(0, 10, None, None, "hdim_1", False) assert test_func(np.array((9, 9)), np.array((0, 9))) == pytest.approx(1) - test_func = pbc_utils.build_distance_function(None, None, 0, 10, "hdim_2") + test_func = pbc_utils.build_distance_function(None, None, 0, 10, "hdim_2", False) assert test_func(np.array((9, 9)), np.array((9, 0))) == pytest.approx(1) - test_func = pbc_utils.build_distance_function(None, None, None, None, "none") + test_func = pbc_utils.build_distance_function(None, None, None, None, "none", False) assert test_func(np.array((9, 9)), np.array((0, 0))) == pytest.approx( (2 * 81) ** 0.5 ) diff --git a/tobac/tracking.py b/tobac/tracking.py index 648bac3d..1fd5a03b 100644 --- a/tobac/tracking.py +++ b/tobac/tracking.py @@ -330,7 +330,7 @@ def linking_trackpy( # I think this shouldn't change results, but it will degrade performance. neighbor_strategy = "BTree" dist_func = pbc_utils.build_distance_function( - min_h1, max_h1, min_h2, max_h2, PBC_flag + min_h1, max_h1, min_h2, max_h2, PBC_flag, is_3D ) else: diff --git a/tobac/utils/periodic_boundaries.py b/tobac/utils/periodic_boundaries.py index 121dddc7..71ecb38e 100644 --- a/tobac/utils/periodic_boundaries.py +++ b/tobac/utils/periodic_boundaries.py @@ -1,5 +1,10 @@ +"""Utilities for handling indexing and distance calculation with periodic boundaries +""" from __future__ import annotations +import functools + import numpy as np + from tobac.utils.decorators import njit_if_available @@ -197,9 +202,9 @@ def get_pbc_coordinates( @njit_if_available def calc_distance_coords_pbc( - coords_1, coords_2, min_h1, max_h1, min_h2, max_h2, PBC_flag -): - """Function to calculate the distance between cartesian + coords_1: np.ndarray[float], coords_2: np.ndarray[float], max_dims: np.ndarray[int] +) -> float: + """Function to calculate the distance between 2D cartesian coordinate set 1 and coordinate set 2. Note that we assume both coordinates are within their min/max already. @@ -210,20 +215,10 @@ def calc_distance_coords_pbc( coordinates or (hdim_1, hdim_2) coordinates. coords_2: 2D or 3D array-like Similar to coords_1, but for the second pair of coordinates - min_h1: int - Minimum point in hdim_1 - max_h1: int - Maximum point in hdim_1, exclusive. max_h1-min_h1 should be the size. - min_h2: int - Minimum point in hdim_2 - max_h2: int - Maximum point in hdim_2, exclusive. max_h2-min_h2 should be the size. - PBC_flag : str('none', 'hdim_1', 'hdim_2', 'both') - 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 + max_dims: array-like + Array of same length as dimensionality of coords. Each item in max_dims + corresponds to a dimension of coords ([(vdim), hdim_1, hdim_2]) with + value equal to the size of that dimension if periodic, or 0 if not Returns ------- @@ -231,31 +226,12 @@ def calc_distance_coords_pbc( Distance between coords_1 and coords_2 in cartesian space. """ - - is_3D = len(coords_1) == 3 - - if not is_3D: - # Let's make the accounting easier. - coords_1 = np.array((0, coords_1[0], coords_1[1])) - coords_2 = np.array((0, coords_2[0], coords_2[1])) - - if PBC_flag in ["hdim_1", "both"]: - size_h1 = max_h1 - min_h1 - mod_h1 = size_h1 - else: - mod_h1 = 0 - if PBC_flag in ["hdim_2", "both"]: - size_h2 = max_h2 - min_h2 - mod_h2 = size_h2 - else: - mod_h2 = 0 - max_dims = np.array((0, mod_h1, mod_h2)) deltas = np.abs(coords_1 - coords_2) deltas = np.where(deltas > 0.5 * max_dims, deltas - max_dims, deltas) return np.sqrt(np.sum(deltas**2)) -def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag): +def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag, is_3D): """Function to build a partial ```calc_distance_coords_pbc``` function suitable for use with trackpy @@ -275,6 +251,8 @@ def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag): '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 + is_3D : bool + True if coordinates are to be provided in 3D, False if 2D Returns ------- @@ -283,17 +261,15 @@ def build_distance_function(min_h1, max_h1, min_h2, max_h2, PBC_flag): just f(coords_1, coords_2) """ - import functools - h1_size, h2_size = validate_pbc_dims(min_h1, max_h1, min_h2, max_h2, PBC_flag) + if is_3D: + max_dims = np.array([0, h1_size, h2_size]) + else: + max_dims = np.array([h1_size, h2_size]) return functools.partial( calc_distance_coords_pbc, - min_h1=0, - max_h1=h1_size, - min_h2=0, - max_h2=h2_size, - PBC_flag=PBC_flag, + max_dims=max_dims, ) @@ -360,7 +336,7 @@ def invalid_limit_names(**limits) -> list[str]: class PBCflagError(ValueError): def __init__(self): - super.init( + super().__init__( "PBC_flag keyword is not valid, must be one of ['none', 'hdim_1', 'hdim_2', 'both']" ) @@ -368,7 +344,7 @@ def __init__(self): class PBCLimitError(ValueError): def __init__(self, invalid_limits, PBC_flag): self.message = f"Keyword parameters {invalid_limits} must be provided for PBC_flag {PBC_flag}" - super.init(self.message) + super().__init__(self.message) def weighted_circmean( From 227a334a541f59a368203665627ae9359c6c528e Mon Sep 17 00:00:00 2001 From: William Jones Date: Thu, 7 Dec 2023 16:30:57 +0000 Subject: [PATCH 6/7] Avoid using argument expansion in indexing for backward compatibility --- tobac/segmentation.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/tobac/segmentation.py b/tobac/segmentation.py index a6b945db..987fa740 100644 --- a/tobac/segmentation.py +++ b/tobac/segmentation.py @@ -202,17 +202,16 @@ def add_markers( # Get its global index so that we can calculate # distance and set the array. local_index = it.multi_index - global_index = np.array( - ( - local_index[0] + z_seed_start, - local_index[1] + seed_box[0], - local_index[2] + seed_box[2], - ) + global_index = ( + local_index[0] + z_seed_start, + local_index[1] + seed_box[0], + local_index[2] + seed_box[2], ) + # If it's a background marker, we can just set it # with the feature we're working on. if curr_box_pt == bg_marker: - marker_arr[*global_index] = row["feature"] + marker_arr[global_index] = row["feature"] continue # it has another feature in it. Calculate the distance # from its current set feature and the new feature. @@ -223,7 +222,9 @@ def add_markers( else: curr_coord = np.array((0, row["hdim_1"], row["hdim_2"])) - dist_from_curr_pt = dist_func(global_index, curr_coord) + dist_from_curr_pt = dist_func( + np.array(global_index), curr_coord + ) # This is technically an O(N^2) operation, but # hopefully performance isn't too bad as this should @@ -243,14 +244,16 @@ def add_markers( orig_coord = np.array( (0, orig_row["hdim_1"], orig_row["hdim_2"]) ) - dist_from_orig_pt = dist_func(global_index, orig_coord) + dist_from_orig_pt = dist_func( + np.array(global_index), orig_coord + ) # The current point center is further away # than the original point center, so do nothing if dist_from_curr_pt > dist_from_orig_pt: continue else: # the current point center is closer. - marker_arr[*global_index] = row["feature"] + marker_arr[global_index] = row["feature"] # completely unseeded region so far. else: marker_arr[ From 20a6bea2653ee0c21d6000e6bc2f298c126971ac Mon Sep 17 00:00:00 2001 From: William Jones Date: Thu, 7 Dec 2023 17:07:47 +0000 Subject: [PATCH 7/7] Add tests for 2D data with test_add_markers_pbcs --- tobac/tests/test_segmentation.py | 101 +++++++++++++++++++++---------- 1 file changed, 70 insertions(+), 31 deletions(-) diff --git a/tobac/tests/test_segmentation.py b/tobac/tests/test_segmentation.py index 2f965a09..887d3e4b 100644 --- a/tobac/tests/test_segmentation.py +++ b/tobac/tests/test_segmentation.py @@ -815,6 +815,11 @@ def test_segmentation_timestep_3d_buddy_box( ((20, 30, 40), (8, 1, 1), (8, 28, 38), (0, 15, 15), None), ((20, 30, 40), (8, 0, 0), (8, 28, 38), (0, -8, -8), None), ((20, 30, 40), (8, 0, 0), (8, 28, 38), (0, -8, -8), (5, 5, 5)), + ((30, 40), (0, 0), (3, 3), (-8, -8), None), + ((30, 40), (0, 0), (3, 3), (-8, -8), None), + ((30, 40), (1, 1), (28, 38), (15, 15), None), + ((30, 40), (0, 0), (28, 38), (-8, -8), None), + ((30, 40), (0, 0), (28, 38), (-8, -8), (5, 5)), ], ) def test_add_markers_pbcs( @@ -852,20 +857,34 @@ def test_add_markers_pbcs( } # Generate dummy feature dataset only on the first feature. - test_feature_ds_1 = testing.generate_single_feature( - start_v=feat_1_loc[0], - start_h1=feat_1_loc[1], - start_h2=feat_1_loc[2], - feature_num=1, - **common_feat_opts, - ) - test_feature_ds_2 = testing.generate_single_feature( - start_v=feat_2_loc[0], - start_h1=feat_2_loc[1], - start_h2=feat_2_loc[2], - feature_num=2, - **common_feat_opts, - ) + if is_3D: + test_feature_ds_1 = testing.generate_single_feature( + start_v=feat_1_loc[0], + start_h1=feat_1_loc[1], + start_h2=feat_1_loc[2], + feature_num=1, + **common_feat_opts, + ) + test_feature_ds_2 = testing.generate_single_feature( + start_v=feat_2_loc[0], + start_h1=feat_2_loc[1], + start_h2=feat_2_loc[2], + feature_num=2, + **common_feat_opts, + ) + else: + test_feature_ds_1 = testing.generate_single_feature( + start_h1=feat_1_loc[0], + start_h2=feat_1_loc[1], + feature_num=1, + **common_feat_opts, + ) + test_feature_ds_2 = testing.generate_single_feature( + start_h1=feat_2_loc[0], + start_h2=feat_2_loc[1], + feature_num=2, + **common_feat_opts, + ) test_feature_ds = pd.concat([test_feature_ds_1, test_feature_ds_2]) common_marker_opts = dict() @@ -882,20 +901,35 @@ def test_add_markers_pbcs( ) # Now, shift the data over and re-run markers. - test_feature_ds_1 = testing.generate_single_feature( - start_v=feat_1_loc[0] + shift_domain[0], - start_h1=feat_1_loc[1] + shift_domain[1], - start_h2=feat_1_loc[2] + shift_domain[2], - feature_num=1, - **common_feat_opts, - ) - test_feature_ds_2 = testing.generate_single_feature( - start_v=feat_2_loc[0] + shift_domain[0], - start_h1=feat_2_loc[1] + shift_domain[1], - start_h2=feat_2_loc[2] + shift_domain[2], - feature_num=2, - **common_feat_opts, - ) + if is_3D: + test_feature_ds_1 = testing.generate_single_feature( + start_v=feat_1_loc[0] + shift_domain[0], + start_h1=feat_1_loc[1] + shift_domain[1], + start_h2=feat_1_loc[2] + shift_domain[2], + feature_num=1, + **common_feat_opts, + ) + test_feature_ds_2 = testing.generate_single_feature( + start_v=feat_2_loc[0] + shift_domain[0], + start_h1=feat_2_loc[1] + shift_domain[1], + start_h2=feat_2_loc[2] + shift_domain[2], + feature_num=2, + **common_feat_opts, + ) + else: + test_feature_ds_1 = testing.generate_single_feature( + start_h1=feat_1_loc[0] + shift_domain[0], + start_h2=feat_1_loc[1] + shift_domain[1], + feature_num=1, + **common_feat_opts, + ) + test_feature_ds_2 = testing.generate_single_feature( + start_h1=feat_2_loc[0] + shift_domain[0], + start_h2=feat_2_loc[1] + shift_domain[1], + feature_num=2, + **common_feat_opts, + ) + test_feature_ds_shifted = pd.concat([test_feature_ds_1, test_feature_ds_2]) marker_arr_shifted = seg.add_markers( @@ -903,9 +937,14 @@ def test_add_markers_pbcs( ) # Now, shift output back. - marker_arr_reshifted = np.roll( - marker_arr_shifted, tuple((-x for x in shift_domain)), axis=(0, 1, 2) - ) + if is_3D: + marker_arr_reshifted = np.roll( + marker_arr_shifted, tuple((-x for x in shift_domain)), axis=(0, 1, 2) + ) + else: + marker_arr_reshifted = np.roll( + marker_arr_shifted, tuple((-x for x in shift_domain)), axis=(0, 1) + ) assert np.all(marker_arr == marker_arr_reshifted)