From e7701cb02509f55580d71cb5b3fc3b0754c5877d Mon Sep 17 00:00:00 2001 From: paulapreuss Date: Mon, 8 Jun 2026 17:19:24 +0200 Subject: [PATCH 1/3] Replace magic numbers with variables --- task_queue/grid_optimizer.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/task_queue/grid_optimizer.py b/task_queue/grid_optimizer.py index 675351d..876d1f0 100644 --- a/task_queue/grid_optimizer.py +++ b/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 @@ -547,8 +552,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: @@ -1188,7 +1193,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 +1217,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 +1740,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], From 2cc6eceb859dc5e9bd200d79d286dced0df46f7b Mon Sep 17 00:00:00 2001 From: paulapreuss Date: Mon, 8 Jun 2026 17:30:01 +0200 Subject: [PATCH 2/3] Refactor main algorithm into submodules --- task_queue/grid_optimizer.py | 121 ++++++++++++++++++++--------------- 1 file changed, 69 insertions(+), 52 deletions(-) diff --git a/task_queue/grid_optimizer.py b/task_queue/grid_optimizer.py index 876d1f0..6f51dc7 100644 --- a/task_queue/grid_optimizer.py +++ b/task_queue/grid_optimizer.py @@ -183,15 +183,36 @@ 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, @@ -199,7 +220,7 @@ def optimize(self): 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: @@ -209,7 +230,7 @@ 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 @@ -217,44 +238,36 @@ def optimize(self): 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() @@ -276,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 @@ -911,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 @@ -957,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", @@ -967,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] @@ -980,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() From c7eb81f0827fc6252c33e067b5aef5d30dfa5eaf Mon Sep 17 00:00:00 2001 From: paulapreuss Date: Tue, 9 Jun 2026 11:26:52 +0200 Subject: [PATCH 3/3] Update branches test --- tests/test_grid_optimizer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_grid_optimizer.py b/tests/test_grid_optimizer.py index 64e8fa9..49ecf64 100644 --- a/tests/test_grid_optimizer.py +++ b/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"