Add visit_water_buckets.py

modules/visit_water_buckets.py

@@ -0,0 +1,98 @@
+"""
+Find the optimal itinerary if we wish to distribute water from a water source to
+buckets.
+"""
+
+import itertools
+
+from .waypoint import Waypoint
+from .waypoint import waypoint_distance
+
+
+def _calculate_travel_distance(
+    origin: Waypoint, buckets: "list[Waypoint]", buckets_at_once
+) -> "tuple[bool, float]":
+    """Calculate the distance (in meters) by visting `buckets` in order, if you
+    had to visit the origin at the start and end, and each time after visiting
+    `buckets_at_once` buckets.
+
+    Args:
+        origin (Waypoint): The start (and end) waypoint
+        buckets (list[Waypoint]): The buckets to visit, in order
+        buckets_at_once (_type_): Max number of buckets to visit before
+            returning to origin
+
+    Returns:
+        tuple[bool, float]: Returns (False, 0) in the case of failure, otherwise
+            (True, distance) where distance is the total travel distance.
+    """
+    total = 0
+    for i in range(0, len(buckets), buckets_at_once):
+        success, dist = waypoint_distance(origin, buckets[i])
+        if not success:
+            return False, 0
+        total += dist
+
+        for j in range(i + 1, min(len(buckets), i + buckets_at_once)):
+            success, dist = waypoint_distance(buckets[j - 1], buckets[j])
+            if not success:
+                return False, 0
+            total += dist
+
+        success, dist = waypoint_distance(
+            origin, buckets[min(len(buckets), i + buckets_at_once) - 1]
+        )
+        if not success:
+            return False, 0
+        total += dist
+
+    return True, total
+
+
+def find_optimal_path(
+    origin: Waypoint, buckets: "list[Waypoint]", buckets_at_once: int = 2
+) -> "tuple[bool, list[Waypoint]]":
+    """Find an optimal itinerary of waypoints, given that we must
+        * Visit every waypoint in `buckets`
+        * Visit `origin` at the start and end of our itinerary
+        * Visit `origin` each time we visit `buckets_at_once` buckets. \\
+
+    Args:
+        origin (Waypoint): The waypoint that we start (and end) at.
+        buckets (list[Waypoint]): The waypoints of the buckets.
+        buckets_at_once (int, optional): Maximum number of buckets that we can
+            visit before returning to `origin`. Defaults to 2.
+
+    Returns:
+        tuple[bool, list[Waypoint]]: Returns (False, None) upon failure,
+            otherwise (True, path) where `path` is the list of waypoints that we
+            visit in order, including `origin` at the start, end, and middle.
+            If multiple optimal paths exist, return the lexicographically
+            smallest path based on (latitude, longitude).
+    """
+    sorted_buckets = sorted(
+        buckets,
+        key=lambda bucket: (bucket.location_ground.latitude, bucket.location_ground.longitude),
+    )
+
+    optimal_permutation = sorted_buckets
+    success, shortest_distance = _calculate_travel_distance(origin, buckets, buckets_at_once)
+    if not success:
+        return False, None
+
+    for permutation in itertools.permutations(sorted_buckets):
+        success, distance = _calculate_travel_distance(origin, permutation, buckets_at_once)
+        if not success:
+            return False, None
+
+        if distance < shortest_distance:
+            shortest_distance = distance
+            optimal_permutation = permutation
+
+    # the optimal permutation consists of only the buckets. We need to add the
+    # origin
+    path = [origin]
+    for i in range(0, len(optimal_permutation), buckets_at_once):
+        path.extend(optimal_permutation[i : i + buckets_at_once])
+        path.append(origin)
+    return True, path

tests/unit/test_visit_water_buckets.py

@@ -0,0 +1,61 @@
+"""
+Test finding optimal itinerary of waypoints to visit.
+"""
+
+from modules.visit_water_buckets import find_optimal_path
+from modules.waypoint import Waypoint
+
+
+def test_2_buckets_at_once() -> None:
+    """
+    Test visiting at most 2 buckets at once.
+    """
+    origin = Waypoint("Origin", 43.47073179293396, -80.53501978862127, 1)  # UWP Beck Hall
+    buckets = [
+        Waypoint("Bucket 0", 43.46925477790072, -80.54034107786745, 1),  # SCH
+        Waypoint("Bucket 1", 43.46834804571286, -80.54341048064786, 1),  # E3
+        Waypoint("Bucket 2", 43.46983001550084, -80.54225704095209, 1),  # DP
+        Waypoint("Bucket 3", 43.47153126326250, -80.54211697668684, 1),  # EIT
+        Waypoint("Bucket 4", 43.47076658531946, -80.54311788821606, 1),  # STC
+    ]
+
+    success, path = find_optimal_path(origin, buckets, 2)
+    assert success
+    assert path == [
+        origin,
+        buckets[4],
+        buckets[3],
+        origin,
+        buckets[2],
+        buckets[1],
+        origin,
+        buckets[0],
+        origin,
+    ]
+
+
+def test_3_buckets_at_once() -> None:
+    """
+    Test visiting at most 3 buckets at once.
+    """
+    origin = Waypoint("Origin", 43.47073179293396, -80.53501978862127, 1)  # UWP Beck Hall
+    buckets = [
+        Waypoint("Bucket 0", 43.46925477790072, -80.54034107786745, 1),  # SCH
+        Waypoint("Bucket 1", 43.46834804571286, -80.54341048064786, 1),  # E3
+        Waypoint("Bucket 2", 43.46983001550084, -80.54225704095209, 1),  # DP
+        Waypoint("Bucket 3", 43.47153126326250, -80.54211697668684, 1),  # EIT
+        Waypoint("Bucket 4", 43.47076658531946, -80.54311788821606, 1),  # STC
+    ]
+
+    success, path = find_optimal_path(origin, buckets, 3)
+    assert success
+    assert path == [
+        origin,
+        buckets[0],
+        buckets[1],
+        buckets[2],
+        origin,
+        buckets[4],
+        buckets[3],
+        origin,
+    ]