Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Refactor of PBC distance functions #386

Merged
merged 8 commits into from
Jun 4, 2024
Merged

Refactor of PBC distance functions #386

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
17c55ee
Update build_distance_function to always provide integer values to ca…
w-k-jones Dec 7, 2023
17e12d2
Add additional tests for build_distance_function with different PBC o…
w-k-jones Dec 7, 2023
2249508
Improved validation of inputs in build_distance_function
w-k-jones Dec 7, 2023
f6d3555
Move build_distance_func to pbc_utils
w-k-jones Dec 7, 2023
0099958
Refactor and reorganise PBC distance functions
w-k-jones Dec 7, 2023
227a334
Avoid using argument expansion in indexing for backward compatibility
w-k-jones Dec 7, 2023
20a6bea
Add tests for 2D data with test_add_markers_pbcs
w-k-jones Dec 7, 2023
319d7a3
Merge branch 'RC_v1.5.x' of https://github.com/climate-processes/toba…
w-k-jones Dec 7, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions tobac/feature_detection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -1541,8 +1540,8 @@ 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(
min_h1 * dxy, max_h1 * dxy, min_h2 * dxy, max_h2 * dxy, PBC_flag
dist_func = pbc_utils.build_distance_function(
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)
Expand Down
47 changes: 23 additions & 24 deletions tobac/segmentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -203,6 +207,7 @@ def add_markers(
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:
Expand All @@ -211,18 +216,14 @@ def add_markers(
# 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(
np.array(global_index), curr_coord
)

# This is technically an O(N^2) operation, but
Expand All @@ -232,21 +233,19 @@ 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(
np.array(global_index), orig_coord
)
# The current point center is further away
# than the original point center, so do nothing
Expand Down
Loading
Loading