Refactor main algorithm

task_queue/grid_optimizer.py

@@ -103,6 +103,11 @@ def optimize_grid(grid_opt_json):
 
 
 class GridOptimizer:
+    _ROAD_POLE_CLUSTER_OFFSET = 100_000  # road-pole cluster labels; k-means uses 0..n
+    _ENDPOINT_SNAP_TOL_M = 1             # drop road-path poles this close to an endpoint
+    _COST_EPS = 1e-9                     # division guard in cost calculations
+    _MAX_SHS_ITER = 100                  # iteration cap for _cut_leaf_poles_on_condition
+
     def __init__(self, grid_opt_json):
         print("Initiating grid optimizer...")
         # TODO go through the helper functions and figure out what they do / document
@@ -178,23 +183,44 @@ def __init__(self, grid_opt_json):
 
     def optimize(self):
         print("Optimizing distribution grid...")
+        n_grid_consumers = self._setup_coordinates()
+        print("Determining number of poles...")
+        self._place_poles(n_grid_consumers)
+        self._connect_nodes()
+        print("Determining distribution links...")
+        self._insert_long_link_poles()
+        print("Determining power house location...")
+        self._run_cost_and_shs_loop()
+        return self._process_results()
+
+    def _setup_coordinates(self) -> int:
+        """Convert to projected coords, build road graph, clear stale poles.
+
+        Returns n_grid_consumers for use by _place_poles.
+        """
         self.convert_lonlat_xy()
         if self.roads is not None:
             self._build_road_graph()
         # Drop ALL poles from previous runs -> _clear_poles() preserves road-sampled poles
         self.nodes = self.nodes[self.nodes["node_type"] != "pole"]
-        n_total_consumers = len(self.nodes)
         n_shs_consumers = len(self.nodes[self.nodes["is_connected"] == False])  # noqa: E712
-        n_grid_consumers = n_total_consumers - n_shs_consumers
+        n_grid_consumers = len(self.nodes) - n_shs_consumers
         self.nodes = self.nodes.sort_index(key=lambda x: x.astype("int64"))
+        return n_grid_consumers
+
+    def _place_poles(self, n_grid_consumers: int):
+        """Place k-means or road-sampled poles. Handles power-house consumer placeholder.
+
+        Stores _power_house_idx on self for use by _connect_nodes.
+        """
         if self.power_house is not None:
             power_house_consumers = self._connect_power_house_consumer_manually(
                 self.connection_cable_max_length,
             )
             self._placeholder_consumers_for_power_house()
         else:
             power_house_consumers = None
-        print("Determining number of poles...")
+
         if self.roads is not None:
             self._place_poles_with_roads()
         else:
@@ -204,52 +230,44 @@ def optimize(self):
 
         if self.power_house is not None:
             cluster_label = self.nodes.loc["100000", "cluster_label"]
-            power_house_idx = self.nodes[
+            self._power_house_idx = self.nodes[
                 (self.nodes["node_type"] == "pole")
                 & (self.nodes["cluster_label"] == cluster_label)
             ].index
             power_house_consumers["cluster_label"] = cluster_label
             power_house_consumers["consumer_type"] = np.nan
             self.nodes = pd.concat([self.nodes, power_house_consumers])
             self._placeholder_consumers_for_power_house(remove=True)
-            # Drop old power house index to avoid duplicates
             self.nodes = self.nodes.drop(index=self.power_house.index)
+        else:
+            self._power_house_idx = None
 
+    def _connect_nodes(self):
+        """Build MST distribution backbone and connect consumers to their poles."""
         self.create_minimum_spanning_tree()
         self.connect_grid_consumers()
+        if self.power_house is not None and self._power_house_idx is not None:
+            ph = self._power_house_idx[0]
+            self.nodes.loc[self.nodes.index == ph, "node_type"] = "power-house"
+            self.nodes.loc[self.nodes.index == ph, "how_added"] = "manual"
 
-        if self.power_house is not None:
-            self.nodes.loc[self.nodes.index == power_house_idx[0], "node_type"] = "power-house"
-            self.nodes.loc[self.nodes.index == power_house_idx[0], "how_added"] = "manual"
-
-        # Populate self.links with distribution links so find_index_longest_distribution_link
-        # can detect links that exceed max_length. A second call follows after intermediate
-        # poles have been inserted to rebuild distribution links with long links broken.
+    def _insert_long_link_poles(self):
+        """Find over-length distribution links, insert intermediate poles, and reconnect."""
         self.connect_grid_poles()
-        # Find the connection links_df in the network with lengths greater than the
-        # maximum allowed length for `connection` cables, specified by the user.
         long_links = self.find_index_longest_distribution_link()
-        # Add poles to the identified long `distribution` links_df, so that the
-        # distance between all poles remains below the maximum allowed distance.
         self._add_fixed_poles_on_long_links(long_links=long_links)
-        # Update the (lon,lat) coordinates based on the newly inserted poles
-        # which only have (x,y) coordinates.
         self.convert_lonlat_xy(inverse=True)
-        # Connect all poles together using the minimum spanning tree algorithm.
-        print("Determining distribution links...")
         self.connect_grid_poles(long_links=long_links)
-        # Calculate distances of all poles from the load centroid.
-        # Find the location of the power house.
         self.add_number_of_distribution_and_connection_cables()
+
+    def _run_cost_and_shs_loop(self):
+        """Iteratively place power house, compute costs, and determine SHS consumers."""
         n_iter = 2 if self.power_house is None else 1
-        print("Determining power house location...")
         for i in range(n_iter):
             if self.power_house is None and i == 0:
                 self._select_location_of_power_house()
             self._set_direction_of_links()
-            self.allocate_poles_to_branches()
-            self.allocate_subbranches_to_branches()
-            self.label_branch_of_consumers()
+            self._build_branch_hierarchy()
             self.determine_cost_per_pole()
             self._connection_cost_per_consumer()
             self.determine_costs_per_branch()
@@ -271,8 +289,6 @@ def optimize(self):
             else:
                 break
 
-        return self._process_results()
-
     def _process_results(self):
         """
         Returns a json object with processed nodes and links
@@ -547,8 +563,8 @@ def _place_pole(label, x, y):
                 # Drop positions that coincide with an endpoint (1m tolerance)
                 road_positions = [
                     (px, py) for px, py in raw_positions
-                    if math.sqrt((px - x_from) ** 2 + (py - y_from) ** 2) > 1
-                    and math.sqrt((px - x_to) ** 2 + (py - y_to) ** 2) > 1
+                    if math.sqrt((px - x_from) ** 2 + (py - y_from) ** 2) > self._ENDPOINT_SNAP_TOL_M
+                    and math.sqrt((px - x_to) ** 2 + (py - y_to) ** 2) > self._ENDPOINT_SNAP_TOL_M
                 ]
 
             if road_positions:
@@ -906,7 +922,14 @@ def add_number_of_distribution_and_connection_cables(self):
             self.nodes["n_connection_links"].fillna(0).astype(int)
         )
 
-    def allocate_poles_to_branches(self):
+    def _build_branch_hierarchy(self):
+        """Assign branch and parent_branch labels to all poles and consumers.
+
+        Runs the three allocation steps in the required order:
+        1. assign branch labels to poles (allocate_poles_to_branches)
+        2. assign parent_branch labels to poles (allocate_subbranches_to_branches)
+        3. propagate branch labels to consumers (label_branch_of_consumers)
+        """
         poles = self._poles().copy()
         leaf_poles = pd.Series(
             poles[poles["n_distribution_links"] == 1].index
@@ -952,7 +975,27 @@ def allocate_poles_to_branches(self):
             "branch",
         ] = power_house
 
-    def label_branch_of_consumers(self):
+        # --- parent_branch assignment ---
+        poles = self._poles().copy()
+        self.nodes["parent_branch"] = None
+
+        if len(poles["branch"].unique()) > 1:
+            leaf_poles_idx = poles[poles["n_distribution_links"] == 1].index
+            self._determine_parent_branches(leaf_poles_idx)
+            poles_excl_ph = poles[poles["node_type"] != "power-house"]
+            split_poles_idx = poles_excl_ph[
+                poles_excl_ph["n_distribution_links"] > 2  # noqa: PLR2004
+            ].index
+            if len(split_poles_idx) > 0:
+                self._determine_parent_branches(split_poles_idx)
+
+        self.nodes.loc[
+            (self.nodes["parent_branch"].isna())
+            & (self.nodes["node_type"].isin(["pole", "power-house"])),
+            "parent_branch",
+        ] = power_house
+
+        # --- propagate branch to consumers ---
         branch_map = self.nodes.loc[
             self.nodes.node_type.isin(["pole", "power-house"]),
             "branch",
@@ -962,7 +1005,7 @@ def label_branch_of_consumers(self):
             "parent",
         ].map(branch_map)
 
-    def determine_parent_branches(self, start_poles):
+    def _determine_parent_branches(self, start_poles):
         poles = self._poles().copy()
         for pole in start_poles:
             branch_start_pole = poles[poles.index == pole]["branch"].iloc[0]
@@ -975,27 +1018,6 @@ def determine_parent_branches(self, start_poles):
                 "parent_branch",
             ] = parent_branch
 
-    def allocate_subbranches_to_branches(self):
-        poles = self._poles().copy()
-        self.nodes["parent_branch"] = None
-        power_house = poles[poles["node_type"] == "power-house"].index[0]
-
-        if len(poles["branch"].unique()) > 1:
-            leaf_poles = poles[poles["n_distribution_links"] == 1].index
-            self.determine_parent_branches(leaf_poles)
-            poles_expect_power_house = poles[poles["node_type"] != "power-house"]
-            split_poles = poles_expect_power_house[
-                poles_expect_power_house["n_distribution_links"] > 2  # noqa: PLR2004
-            ].index
-            if len(split_poles) > 0:
-                self.determine_parent_branches(split_poles)
-
-        self.nodes.loc[
-            (self.nodes["parent_branch"].isna())
-            & (self.nodes["node_type"].isin(["pole", "power-house"])),
-            "parent_branch",
-        ] = power_house
-
     def determine_cost_per_pole(self):
         poles = self._poles().copy()
         links = self.get_links().copy()
@@ -1188,7 +1210,7 @@ def _cut_leaf_poles_on_condition(self):
         exclude_lst.extend(
             self.nodes[self.nodes["shs_options"] == 1]["parent"].unique()
         )
-        for _ in range(100):
+        for _ in range(self._MAX_SHS_ITER):
             counter = 0
             leaf_poles = self.nodes[self.nodes["n_distribution_links"] == 1].index
             for pole in leaf_poles:
@@ -1212,7 +1234,7 @@ def _cut_leaf_poles_on_condition(self):
                     "cost_per_branch",
                 ].iloc[0] / (
                     self.nodes.loc[consumer_of_branch, "yearly_consumption"].sum()
-                    + 1e-9
+                    + self._COST_EPS
                 )
                 if average_marginal_cost_of_pole > self.max_levelized_grid_cost or (
                     average_total_cost_of_pole > self.max_levelized_grid_cost
@@ -1735,7 +1757,7 @@ def _add_pole(x, y):
         poles["custom_specification"] = ""
         poles["shs_options"] = 0
         # High-offset cluster labels avoid collision with k-means labels (0..n)
-        poles["cluster_label"] = range(100000, 100000 + len(poles))
+        poles["cluster_label"] = range(self._ROAD_POLE_CLUSTER_OFFSET, self._ROAD_POLE_CLUSTER_OFFSET + len(poles))
 
         self.nodes = pd.concat(
             [self.nodes, poles],

tests/test_grid_optimizer.py

@@ -1260,9 +1260,7 @@ def test_allocate_branches_assigns_branch_and_propagates_to_consumers(
     optimizer._set_direction_of_links()
     optimizer.distribution_links = optimizer.links[optimizer.links["link_type"] == "distribution"]
 
-    optimizer.allocate_poles_to_branches()
-    optimizer.allocate_subbranches_to_branches()
-    optimizer.label_branch_of_consumers()
+    optimizer._build_branch_hierarchy()
 
     assert optimizer.nodes.at["p-0", "branch"] == "p-0"
     assert optimizer.nodes.at["p-1", "branch"] == "p-0"