diff --git a/task_queue/grid_optimizer.py b/task_queue/grid_optimizer.py index 962d649..327f4e3 100644 --- a/task_queue/grid_optimizer.py +++ b/task_queue/grid_optimizer.py @@ -452,7 +452,8 @@ def get_shs_consumers(self): return df.copy() def find_index_longest_distribution_link(self): - # Find the links longer than two times of the allowed distance + # Find the links longer than the allowed distance + self.distribution_links = self.links[self.links["link_type"] == "distribution"] critical_link = self.distribution_links[ self.distribution_links["length"] > self.distribution_cable_max_length ] @@ -1183,6 +1184,8 @@ def _cut_leaf_poles_on_condition(self): for pole in leaf_poles: if pole in exclude_lst: continue + if pole not in self.nodes.index: + continue consumer_of_pole = self.nodes[self.nodes["parent"] == pole] branch = self.nodes[self.nodes.index == pole]["branch"].iloc[0] consumer_of_branch = self.nodes[self.nodes["branch"] == branch].index @@ -1249,6 +1252,8 @@ def _remove_poles_and_links(self, pole): def _correct_n_distribution_links_of_parent_poles(self, pole): parent_pole = self.nodes[self.nodes.index == pole]["parent"].iloc[0] + if parent_pole == "unknown" or pd.isna(parent_pole): + return self.nodes.loc[parent_pole, "n_distribution_links"] -= 1 def _determine_shs_consumers(self, max_iter=20): @@ -1580,19 +1585,20 @@ def _break_long_link(self, mst_pole_from, mst_pole_to, added_poles): label_node_from=index_added_pole if to_from else mst_pole_from, label_node_to=mst_pole_to if to_from else index_added_pole, ) - elif counter == n_added_poles - 1: - # The last `added poles` should be connected to + else: + # Connect each subsequent pole to the previous added pole. + self._add_links( + label_node_from=added_pole_indices[counter - 1], + label_node_to=index_added_pole, + ) + if counter == n_added_poles - 1: + # The last `added poles` should also be connected to # the end or to the beginning of the long link, # depending on the `to_from` flag. self._add_links( label_node_from=mst_pole_from if to_from else index_added_pole, label_node_to=index_added_pole if to_from else mst_pole_to, ) - else: - self._add_links( - label_node_from=added_pole_indices[counter - 1], - label_node_to=index_added_pole, - ) self.nodes.loc[index_added_pole, "how_added"] = "long-distance" def create_minimum_spanning_tree(self): diff --git a/tests/test_grid_optimizer.py b/tests/test_grid_optimizer.py index 047f197..4558470 100644 --- a/tests/test_grid_optimizer.py +++ b/tests/test_grid_optimizer.py @@ -334,14 +334,16 @@ def test_clear_links_and_clear_all_links(optimizer: GridOptimizer) -> None: def test_find_index_longest_distribution_link(optimizer: GridOptimizer) -> None: - optimizer.distribution_links = pd.DataFrame( - { - "length": [99.0, 100.0, 101.0], - }, - index=["short", "equal", "long"], - ) + # Considering that max_length_distribution_link is 100m + optimizer._add_node("p-0", node_type="pole", x=0.0, y=0.0) + optimizer._add_node("p-1", node_type="pole", x=0.0, y=100.0) + optimizer._add_node("p-2", node_type="pole", x=0.0, y=200.0) + optimizer._add_node("p-3", node_type="pole", x=0.0, y=350.0) + optimizer._add_links("p-0", "p-1") + optimizer._add_links("p-1", "p-2") + optimizer._add_links("p-2", "p-3") - assert optimizer.find_index_longest_distribution_link() == ["long"] + assert optimizer.find_index_longest_distribution_link() == ['(p-2, p-3)'] def test_add_number_of_distribution_and_connection_cables(