diff --git a/kadmos/graph/graph_data.py b/kadmos/graph/graph_data.py index f7dce1f8f4270a878e58c369b2dde856c55e4f2a..67980dd70878108600926709117ff193c917ed7a 100644 --- a/kadmos/graph/graph_data.py +++ b/kadmos/graph/graph_data.py @@ -5363,11 +5363,11 @@ class MdaoDataGraph(DataGraph, MdaoMixin): mapping[global_des_var] = local_des_var_copy for target in local_targets: + # Connect the local copy to the targets + self.copy_edge((global_des_var, target), (local_des_var_copy, target)) + # Remove the connection between the global design variable and the target self.remove_edge(global_des_var, target) - - # Connect the local copy to the targets - self.add_edge(local_des_var_copy, target) # The local design variables get copies at the global level and are connected accordingly for local_des_var in local_des_vars: # Find the functions outside the local level for which the design variable is input @@ -5383,12 +5383,11 @@ class MdaoDataGraph(DataGraph, MdaoMixin): mapping[local_des_var] = global_des_var_copy for target in external_targets: - # Remove the connection between the local design variable and the external target - self.remove_edge(local_des_var, target) - # Connect the global copy to the targets - self.add_edge(global_des_var_copy, target) + self.copy_edge((local_des_var, target), (global_des_var_copy, target)) + # Remove the connection between the local design variable and the external target + self.remove_edge(local_des_var, target) return local_des_var_copies, global_des_var_copies, mapping def localize_group_couplings(self, group_functions, external_couplings, local_couplings, @@ -5430,11 +5429,12 @@ class MdaoDataGraph(DataGraph, MdaoMixin): external_couplings_copies.append(external_coupling_copy) for target in local_targets: + # Connect the local copy to the targets + self.copy_edge((external_coupling, target), (external_coupling_copy, target)) + # Remove the connection between the global coupling variable and the target self.remove_edge(external_coupling, target) - # Connect the local copy to the targets - self.add_edge(external_coupling_copy, target) # Local couplings should only be handled by the functions inside the group, outside, they are handled by copies for local_coupling in local_couplings: # Find the external functions for which the coupling variable is input @@ -5459,12 +5459,11 @@ class MdaoDataGraph(DataGraph, MdaoMixin): mapping_locals[local_coupling] = local_coupling_copy for target in external_targets: - # Remove the connection between the local coupling variable and the target - self.remove_edge(local_coupling, target) - # Connect the local copy to the targets - self.add_edge(local_coupling_copy, target) + self.copy_edge((local_coupling, target), (local_coupling_copy, target)) + # Remove the connection between the local coupling variable and the target + self.remove_edge(local_coupling, target) return external_couplings_copies, local_couplings_copies, mapping_locals def connect_nodes_as_output(self, nodes, function): @@ -5479,7 +5478,6 @@ class MdaoDataGraph(DataGraph, MdaoMixin): for node in nodes: assert self.has_node(node) self.add_edge(function, node) - return def connect_coordinator(self, additional_inputs=[], additional_outputs=[]):