diff --git a/tobac/feature_detection.py b/tobac/feature_detection.py
index a21d8db1..b937a0c0 100644
--- a/tobac/feature_detection.py
+++ b/tobac/feature_detection.py
@@ -31,7 +31,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 decorators
 
@@ -1545,8 +1544,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)
diff --git a/tobac/segmentation.py b/tobac/segmentation.py
index 486100df..fe2eda2e 100644
--- a/tobac/segmentation.py
+++ b/tobac/segmentation.py
@@ -181,6 +181,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.
@@ -205,6 +209,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:
@@ -213,18 +218,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
@@ -234,21 +235,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
diff --git a/tobac/tests/test_pbc_utils.py b/tobac/tests/test_pbc_utils.py
index c88ac2e6..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)
 
 
@@ -286,6 +417,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", 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", 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", 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", True)
+    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", 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", 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", 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", False)
+    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_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)
 
diff --git a/tobac/tests/test_tracking.py b/tobac/tests/test_tracking.py
index 94f410c3..52fab643 100644
--- a/tobac/tests/test_tracking.py
+++ b/tobac/tests/test_tracking.py
@@ -3,16 +3,17 @@
 """
 
 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
@@ -275,29 +276,6 @@ def test_linking_trackpy():
     )
 
 
-def test_build_distance_function():
-    """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
-    )
-
-
 @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 558dd85e..7b21d1db 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, is_3D
+        )
 
     else:
         neighbor_strategy = "KDTree"
@@ -609,43 +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
-
-    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,
-        PBC_flag=PBC_flag,
-    )
diff --git a/tobac/utils/periodic_boundaries.py b/tobac/utils/periodic_boundaries.py
index 52c93c4b..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,49 +215,136 @@ 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
+    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
+    -------
+    float
+        Distance between coords_1 and coords_2 in cartesian space.
+
+    """
+    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, is_3D):
+    """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, exclusive. max_h1-min_h1 should be the size.
+        Maximum point in hdim_1
     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.
+        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
+    is_3D : bool
+        True if coordinates are to be provided in 3D, False if 2D
 
     Returns
     -------
-    float
-        Distance between coords_1 and coords_2 in cartesian space.
+    function object
+        A version of calc_distance_coords_pbc suitable to be called by
+        just f(coords_1, coords_2)
 
     """
+    h1_size, h2_size = validate_pbc_dims(min_h1, max_h1, min_h2, max_h2, PBC_flag)
 
-    is_3D = len(coords_1) == 3
+    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,
+        max_dims=max_dims,
+    )
 
-    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 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(