diff --git a/kadmos/graph/graph_data.py b/kadmos/graph/graph_data.py
index 00c6a2b6a3407f2156ec5d360b576c7badaa4a1e..a2c8e5a1857b407482a7576719839321bddd5fb8 100644
--- a/kadmos/graph/graph_data.py
+++ b/kadmos/graph/graph_data.py
@@ -542,7 +542,8 @@ class DataGraph(KadmosGraph):
 
         return coupling_dict
 
-    def get_possible_function_order(self, method, multi_start=None, check_graph=False, node_selection=None):
+    def get_possible_function_order(self, method, multi_start=None, check_graph=False, coupling_dict=None,
+                                    node_selection=None, rcb=1.0, use_runtime_info=False):
         """ Method to find a possible function order, in the order: pre-coupled, coupled, post-coupled functions
 
         :param method: algorithm which will be used to minimize the feedback loops
@@ -553,6 +554,9 @@ class DataGraph(KadmosGraph):
         :type check_graph: bool
         :param node_selection: possibility to get the order of only a selection of nodes instead of the entire graph
         :type node_selection: list
+        :param rcb: runtime-coupling balance, relative importance between feedback and runtime while optimizing
+                    function order. 1: min feedback, 0: min runtime
+        :type rcb: float
         :return Possible function order
         :rtype list
         """
@@ -561,6 +565,11 @@ class DataGraph(KadmosGraph):
         if check_graph:
             assert not self.find_all_nodes(subcategory='all problematic variables'), \
                 'Graph still has problematic variables.'
+        assert 0 <= rcb <= 1, 'Runtime-coupling balance should be between zero and one.'
+
+        # Get coupling dictionary
+        if not coupling_dict:
+            coupling_dict = self.get_coupling_dictionary()
 
         # Check for partitions
         if 'problem_formulation' in self.graph and 'coupled_functions_groups' in self.graph['problem_formulation'] and \
@@ -581,12 +590,16 @@ class DataGraph(KadmosGraph):
         function_graph.add_node('super_node', category='function')
         coupled_functions = []
 
-        # If partitions check if all coupled nodes are assigned to a partition
         if partitions:
+            # Merge the nodes in the partitions into the super node
             for partition in partitions:
                 for function_id in partition:
                     function_graph = nx.contracted_nodes(function_graph, 'super_node', function_id)
+
+            # Remove selfloops that arise due to the merging
             function_graph.remove_edges_from(function_graph.selfloop_edges())
+
+            # Check if all coupled nodes are assigned to a partition
             if not nx.is_directed_acyclic_graph(function_graph):
                 cycle = nx.find_cycle(function_graph)
                 functions_in_cycle = set()
@@ -623,12 +636,15 @@ class DataGraph(KadmosGraph):
         if partitions:
             coupled_functions_order = []
             for partition, nodes in enumerate(partitions):
-                nodes = self.minimize_feedback(list(nodes), method, multi_start=multi_start)
+                nodes = self.minimize_feedback(list(nodes), method, multi_start=multi_start, rcb=rcb,
+                                               use_runtime_info=use_runtime_info, coupling_dict=coupling_dict)
                 partitions[partition] = nodes
                 coupled_functions_order.extend(nodes)
             self.graph['problem_formulation']['coupled_functions_groups'] = partitions
         else:
-            coupled_functions_order = self.minimize_feedback(coupled_functions, method, multi_start=multi_start)
+            coupled_functions_order = self.minimize_feedback(coupled_functions, method, multi_start=multi_start,
+                                                             rcb=rcb, use_runtime_info=use_runtime_info,
+                                                             coupling_dict=coupling_dict)
 
         # Get post-coupling functions and sort
         post_coupling_functions = function_order[function_order.index('super_node') + 1:]
@@ -702,7 +718,38 @@ class DataGraph(KadmosGraph):
 
         return function_order
 
-    def minimize_feedback(self, nodes, method, multi_start=None, get_evaluations=False):
+    def get_runtime_sequence(self, sequence, use_runtime_info=False, coupling_dict=None):
+        """Function to calculate the runtime of a sequence of nodes"""
+
+        # Input assertion
+        if use_runtime_info:
+            for node in sequence:
+                assert 'performance_info' in self.nodes[node]
+                assert 'run_time' in self.nodes[node]['performance_info'], 'Run time missing for node {}'.format(node)
+
+        # Get coupling dictionary
+        if not coupling_dict:
+            coupling_dict = self.get_coupling_dictionary()
+
+        # Calculate runtime
+        nodes = copy.deepcopy(sequence)
+        run_time = 0
+        while nodes:
+            parallel_nodes, run_times = [], []
+            for idx, node in enumerate(nodes):
+                if not set(nodes[:idx]).intersection(coupling_dict[node]):
+                    parallel_nodes.append(node)
+                    run_times.append(self.nodes[node]['performance_info']['run_time'] if use_runtime_info else 1)
+                else:
+                    break
+            run_time += max(run_times)
+            for node in parallel_nodes:
+                nodes.pop(nodes.index(node))
+
+        return run_time
+
+    def minimize_feedback(self, nodes, method, multi_start=None, get_evaluations=False, coupling_dict=None, rcb=1,
+                          use_runtime_info=False):
         """Function to find the function order with minimum feedback
 
         :param nodes: nodes for which the feedback needs to be minimized
@@ -713,16 +760,33 @@ class DataGraph(KadmosGraph):
         :type multi_start: int
         :param get_evaluations: option to give the number of evaluations as output
         :type get_evaluations: bool
+        :param coupling_dict:
+        :type coupling_dict: dict
+        :param rcb: runtime-coupling balance, relative importance between feedback and runtime while optimizing
+                    function order. 1: min feedback, 0: min runtime
+        :type rcb: float
         :return function order
         :rtype list
         :return number of evaluations
         :rtype int
         """
 
-        # TODO: check if input nodes > 1
+        # Input assertions
+        assert 0 <= rcb <= 1, 'Runtime-coupling balance should be between zero and one.'
+        if use_runtime_info:
+            for node in nodes:
+                assert 'performance_info' in self.nodes[node]
+                assert 'run_time' in self.nodes[node]['performance_info'], 'Run time missing for node {}'.format(node)
 
         # Get coupling dictionary
-        coupling_dict = self.get_coupling_dictionary()
+        if not coupling_dict:
+            coupling_dict = self.get_coupling_dictionary()
+
+        # Get total number of couplings and total runtime
+        total_couplings = sum([coupling_dict[function1][function2] for function1 in nodes for function2 in nodes if
+                               function2 in coupling_dict[function1]])
+        total_time = sum([self.nodes[node]['performance_info']['run_time'] for node in nodes]) if use_runtime_info else\
+            len(nodes)
 
         # Get random starting points for a multi-start
         if isinstance(multi_start, (int, long)):
@@ -734,8 +798,7 @@ class DataGraph(KadmosGraph):
 
         if multi_start:
             best_order = list(nodes)
-            min_feedback = float("inf")
-            min_time = float("inf")
+            min_f, min_feedback, min_time = float("inf"), float("inf"), float("inf")
 
             # Start algorithm for each starting point
             n_eval = 0
@@ -743,12 +806,16 @@ class DataGraph(KadmosGraph):
                 if method == 'brute-force' or method == 'branch-and-bound':
                     raise IOError('No multi start possible for an exact algorithm')
                 elif method == 'single-swap':
-                    function_order, n_eval_iter = self._single_swap(multi_start[start_point], coupling_dict)
+                    function_order, n_eval_iter = self._single_swap(multi_start[start_point], coupling_dict, rcb,
+                                                                    use_runtime_info)
                 elif method == 'two-swap':
-                    function_order, n_eval_iter = self._two_swap(multi_start[start_point], coupling_dict)
+                    function_order, n_eval_iter = self._two_swap(multi_start[start_point], coupling_dict, rcb,
+                                                                 use_runtime_info)
                 elif method == 'hybrid-swap':
-                    function_order, n_eval_iter1 = self._two_swap(multi_start[start_point], coupling_dict)
-                    function_order, n_eval_iter2 = self._single_swap(function_order, coupling_dict)
+                    function_order, n_eval_iter1 = self._two_swap(multi_start[start_point], coupling_dict, rcb,
+                                                                  use_runtime_info)
+                    function_order, n_eval_iter2 = self._single_swap(function_order, coupling_dict, rcb,
+                                                                     use_runtime_info)
                     n_eval_iter = n_eval_iter1 + n_eval_iter2
                 else:
                     raise IOError('Selected method (' + method + ') is not a valid method for sequencing, supported ' +
@@ -757,32 +824,31 @@ class DataGraph(KadmosGraph):
                 n_eval += n_eval_iter
 
                 # Get feedback info
-                feedback, time = self.get_feedback_info(function_order, coupling_dict)
+                feedback, time = self.get_feedback_info(function_order, coupling_dict, use_runtime_info)
 
                 # Remember best order found
-                if feedback < min_feedback or (feedback == min_feedback and time < min_time):
-                    best_order = list(function_order)
-                    min_feedback = feedback
-                    min_time = time
+                f = rcb*(feedback/float(total_couplings)) + (1-rcb)*(time/float(total_time))
+                if (feedback == min_feedback and time < min_time) or (time == min_time and feedback < min_feedback) or \
+                        (f < min_f):
+                    best_order, min_f, min_feedback, min_time = list(function_order), f, feedback, time
 
             function_order = list(best_order)
 
         else:
             if method == 'brute-force':
-                function_order, n_eval = self._brute_force(nodes, coupling_dict)
+                function_order, n_eval = self._brute_force(nodes, coupling_dict, rcb, use_runtime_info)
             elif method == 'branch-and-bound':
-                function_order, n_eval = self._branch_and_bound(nodes, coupling_dict)
+                function_order, n_eval = self._branch_and_bound(nodes, coupling_dict, rcb, use_runtime_info)
             elif method == 'single-swap':
-                function_order, n_eval = self._single_swap(nodes, coupling_dict)
+                function_order, n_eval = self._single_swap(nodes, coupling_dict, rcb, use_runtime_info)
             elif method == 'two-swap':
-                function_order, n_eval = self._two_swap(nodes, coupling_dict)
+                function_order, n_eval = self._two_swap(nodes, coupling_dict, rcb, use_runtime_info)
             elif method == 'hybrid-swap':
-                function_order, n_eval1 = self._two_swap(nodes, coupling_dict)
-                function_order, n_eval2 = self._single_swap(function_order, coupling_dict)
+                function_order, n_eval1 = self._two_swap(nodes, coupling_dict, rcb, use_runtime_info)
+                function_order, n_eval2 = self._single_swap(function_order, coupling_dict, rcb, use_runtime_info)
                 n_eval = n_eval1 + n_eval2
             elif method == 'genetic-algorithm':
-                function_order = self._genetic_algorithm(nodes, coupling_dict)  # TODO: add n_eval to genetic algorithm
-                n_eval = 1
+                function_order, n_eval = self._genetic_algorithm(nodes, coupling_dict, rcb, use_runtime_info)
             else:
                 raise IOError('Selected method (' + method + ') is not a valid method for sequencing, supported ' +
                               'methods are: brute-force, single-swap, two-swap, hybrid-swap, branch-and-bound')
@@ -792,7 +858,7 @@ class DataGraph(KadmosGraph):
         else:
             return function_order
 
-    def _brute_force(self, nodes, coupling_dict):
+    def _brute_force(self, nodes, coupling_dict, rcb, use_runtime_info):
         """Function to find the minimum number of feedback loops using the brute-force method: try all possible
         combinations and select the best one
 
@@ -803,14 +869,19 @@ class DataGraph(KadmosGraph):
         """
 
         # Calculate the number of runs that are needed and give a warning when it exceeds a threshold
-        n_runs = np.math.factorial(len(nodes))
-        if n_runs > 3e5:
-            logger.warning(str(n_runs) + ' tool combinations need to be evaluated for the brute-force method. Be ' +
-                           'aware that this can take up a considerable amount of time and resources')
+        if len(nodes) > 9:
+            logger.warning(str(np.math.factorial(len(nodes))) + ' tool combinations need to be evaluated for the '
+                           'brute-force method. Be aware that this can take up a considerable amount of time and '
+                           'resources')
+
+        # Get total number of couplings and total runtime
+        total_couplings = sum([coupling_dict[function1][function2] for function1 in nodes for function2 in nodes if
+                               function2 in coupling_dict[function1]])
+        total_time = sum([self.nodes[node]['performance_info']['run_time'] for node in nodes]) if use_runtime_info \
+            else len(nodes)
 
-        function_order = list(nodes)
-        min_feedback = float("inf")
-        min_time = float("inf")
+        best_order = list(nodes)
+        min_f, min_feedback, min_time = float("inf"), float("inf"), float("inf")
 
         # Keep track of number of evaluated function orders
         n_eval = 0
@@ -818,17 +889,19 @@ class DataGraph(KadmosGraph):
         # Get all possible combinations
         for current_order in itertools.permutations(nodes):
             n_eval += 1
-            feedback, time = self.get_feedback_info(current_order, coupling_dict)
 
-            # Evaluate whether current solution is better than the one found so far
-            if feedback < min_feedback or (feedback == min_feedback and time < min_time):
-                min_feedback = feedback
-                min_time = time
-                function_order = list(current_order)
+            # Get feedback and runtime information of current solution
+            feedback, time = self.get_feedback_info(list(current_order), coupling_dict, use_runtime_info)
 
-        return function_order, n_eval
+            # Evaluate whether current solution is better than the best one found so far
+            f = rcb * (feedback / float(total_couplings)) + (1 - rcb) * (time / float(total_time))
+            if (feedback == min_feedback and time < min_time) or (time == min_time and feedback < min_feedback) or \
+                    (f < min_f):
+                best_order, min_f, min_feedback, min_time = list(current_order), f, feedback, time
 
-    def _single_swap(self, nodes, coupling_dict):
+        return best_order, n_eval
+
+    def _single_swap(self, nodes, coupling_dict, rcb, use_runtime_info):
         """Function to find the minimum number of feedback loops using the single-swap method: improve the solution
          by searching for a better position for each node
 
@@ -840,9 +913,16 @@ class DataGraph(KadmosGraph):
 
         converged = False
 
+        # Get total number of couplings and total runtime
+        total_couplings = sum([coupling_dict[function1][function2] for function1 in nodes for function2 in nodes if
+                               function2 in coupling_dict[function1]])
+        total_time = sum([self.nodes[node]['performance_info']['run_time'] for node in nodes]) if use_runtime_info \
+            else len(nodes)
+
         # Take the input order as start point
         best_order = list(nodes)
-        min_feedback, min_time = self.get_feedback_info(nodes, coupling_dict)
+        min_feedback, min_time = self.get_feedback_info(nodes, coupling_dict, use_runtime_info)
+        min_f = rcb*(min_feedback/float(total_couplings)) + (1-rcb)*(min_time/float(total_time))
 
         # Keep track of number of evaluated function orders
         n_eval = 0
@@ -867,12 +947,13 @@ class DataGraph(KadmosGraph):
                     # Insert node at new position (starting from the back)
                     new_order.insert(len(best_order) - position - 1, node)
                     n_eval += 1
-                    feedback, time = self.get_feedback_info(new_order, coupling_dict)
-                    # Check whether new order is an improvement
-                    if feedback < min_feedback or (feedback == min_feedback and time < min_time):
-                        best_order = list(new_order)
-                        min_feedback = feedback
-                        min_time = time
+                    # Get feedback and runtime information of current solution
+                    feedback, time = self.get_feedback_info(new_order, coupling_dict, use_runtime_info)
+                    # Evaluate whether current solution is better than the best one found so far
+                    f = rcb * (feedback / float(total_couplings)) + (1 - rcb) * (time / float(total_time))
+                    if (feedback == min_feedback and time < min_time) or (time == min_time and feedback < min_feedback)\
+                            or (f < min_f):
+                        best_order, min_f, min_feedback, min_time = list(new_order), f, feedback, time
                         new_iteration = True
 
                     # When a better solution is found, the current iteration is stopped and a new iteration is
@@ -886,11 +967,9 @@ class DataGraph(KadmosGraph):
             if not new_iteration:
                 converged = True
 
-        function_order = list(best_order)
+        return best_order, n_eval
 
-        return function_order, n_eval
-
-    def _two_swap(self, nodes, coupling_dict):
+    def _two_swap(self, nodes, coupling_dict, rcb, use_runtime_info):
         """Function to find the minimum number of feedback loops using the two-swap method: improve the solution
         by swapping two nodes
 
@@ -905,9 +984,16 @@ class DataGraph(KadmosGraph):
         # Keep track of number of evaluated function orders
         n_eval = 0
 
+        # Get total number of couplings and total runtime
+        total_couplings = sum([coupling_dict[function1][function2] for function1 in nodes for function2 in nodes if
+                               function2 in coupling_dict[function1]])
+        total_time = sum([self.nodes[node]['performance_info']['run_time'] for node in nodes]) if use_runtime_info \
+            else len(nodes)
+
         # Take the input order as start point
         best_order = list(nodes)
-        min_feedback, min_time = self.get_feedback_info(best_order, coupling_dict)
+        min_feedback, min_time = self.get_feedback_info(best_order, coupling_dict, use_runtime_info)
+        min_f = rcb * (min_feedback / float(total_couplings)) + (1 - rcb) * (min_time / float(total_time))
 
         while not converged:
             new_iteration = False
@@ -922,12 +1008,13 @@ class DataGraph(KadmosGraph):
                     neighbor_solution[i], neighbor_solution[-j - 1] = best_order[-j - 1], best_order[i]
                     # Get feedback information of the neighbor solution
                     n_eval += 1
-                    feedback, time = self.get_feedback_info(neighbor_solution, coupling_dict)
-                    # Check whether the neighbor solution is a better than the current solution
-                    if feedback < min_feedback or (feedback == min_feedback and time < min_time):
-                        best_order = list(neighbor_solution)
-                        min_feedback = feedback
-                        min_time = time
+                    # Get feedback and runtime information of current solution
+                    feedback, time = self.get_feedback_info(neighbor_solution, coupling_dict, use_runtime_info)
+                    # Evaluate whether current solution is better than the best one found so far
+                    f = rcb * (feedback / float(total_couplings)) + (1 - rcb) * (time / float(total_time))
+                    if (feedback == min_feedback and time < min_time) or (time == min_time and feedback < min_feedback)\
+                            or (f < min_f):
+                        best_order, min_f, min_feedback, min_time = list(neighbor_solution), f, feedback, time
                         new_iteration = True
 
                     # When a better solution is found, the current iteration is stopped and a new iteration is
@@ -941,11 +1028,9 @@ class DataGraph(KadmosGraph):
             if not new_iteration:
                 converged = True
 
-        function_order = best_order
+        return best_order, n_eval
 
-        return function_order, n_eval
-
-    def _branch_and_bound(self, nodes, coupling_dict):
+    def _branch_and_bound(self, nodes, coupling_dict, rcb, use_runtime_info):
         """Function to find the minimum number of feedback loops using the branch-and-bound method: search the solution
         space in a systematic way to find the exact solution
 
@@ -960,17 +1045,35 @@ class DataGraph(KadmosGraph):
         # Keep track of number of evaluated function orders
         n_eval = 0
 
+        # Get total number of couplings and maximum runtime of graph
+        total_couplings = sum([coupling_dict[function1][function2] for function1 in nodes for function2 in nodes if
+                               function2 in coupling_dict[function1]])
+        total_time = sum([self.nodes[node]['performance_info']['run_time'] for node in nodes]) if use_runtime_info \
+            else len(nodes)
+
         # Calculate lower bound for each initial branch
-        for node in nodes:
-            node = [node]
-            n_eval += 1
-            lower_bound_feedback = self._get_lower_bound_branch_and_bound(node, nodes, coupling_dict)
-            active_branches.append([node, lower_bound_feedback])
+        if len(nodes) != 1:
+            for node in nodes:
+                node = [node]
+                n_eval += 1
+                feedback, time = self._get_lower_bound_branch_and_bound(node, nodes, coupling_dict, use_runtime_info)
+                f = rcb * (feedback / float(total_couplings)) + (1 - rcb) * (time / float(total_time))
+                active_branches.append([node, f, feedback, time])
 
         while True:
+            if len(nodes) == 1:
+                selected_branch = [nodes]
+                break
+
+            # Get the branches with the lowest objective
+            min_f = min(branch[1] for branch in active_branches)
+            best_branches = [branch for branch in active_branches if branch[1] == min_f]
             # Get the branches with the lowest feedback
-            min_feedback = min(branch[1] for branch in active_branches)
-            best_branches = [branch for branch in active_branches if branch[1] == min_feedback]
+            min_feedback = min(branch[2] for branch in best_branches)
+            best_branches = [branch for branch in best_branches if branch[2] == min_feedback]
+            # Get the branches with the lowest runtime
+            min_runtime = min(branch[3] for branch in best_branches)
+            best_branches = [branch for branch in best_branches if branch[3] == min_runtime]
 
             # Select the longest branch as best branch
             selected_branch = max(best_branches, key=lambda x: len(x[0]))
@@ -986,17 +1089,17 @@ class DataGraph(KadmosGraph):
                     node = [node]
                     current_order = selected_branch[0] + node
                     n_eval += 1
-                    lower_bound_feedback = self._get_lower_bound_branch_and_bound(current_order, nodes, coupling_dict)
-                    active_branches.append([current_order, lower_bound_feedback])
+                    feedback, time = self._get_lower_bound_branch_and_bound(current_order, nodes, coupling_dict,
+                                                                            use_runtime_info)
+                    f = rcb * (feedback / float(total_couplings)) + (1 - rcb) * (time / float(total_time))
+                    active_branches.append([current_order, f, feedback, time])
 
             # Delete explored branch
             del active_branches[active_branches.index(selected_branch)]
 
-        function_order = selected_branch[0]
-
-        return function_order, n_eval
+        return selected_branch[0], n_eval
 
-    def _get_lower_bound_branch_and_bound(self, branch, nodes, coupling_dict):
+    def _get_lower_bound_branch_and_bound(self, branch, nodes, coupling_dict, use_runtime_info):
         """Function to calculate the lower bound of a branch in the branch and bound algorithm.
         The lower bound is defined as the amount of feedback loops that are guaranteed to occur if the
         selected nodes are placed at the beginning of the order
@@ -1009,20 +1112,18 @@ class DataGraph(KadmosGraph):
         :rtype int
         """
 
-        # Get a random function order with the nodes of the branch at the start
-        function_order = list(branch)
-        for node in nodes:
-            if node not in function_order:
-                function_order.append(node)
-
-        # Calculate lower bound for the number of feedback loops
-        feedback = 0
+        # Calculate the lower bound for the number of feedback loops
+        n_feedback = 0
         for idx1, function1 in enumerate(branch):
-            for idx2, function2 in enumerate(function_order[idx1 + 1:]):
-                if function2 in coupling_dict[function1]:
-                    feedback += coupling_dict[function1][function2]
+            for idx2, function2 in enumerate(nodes):
+                if function2 in coupling_dict[function1] and function2 not in branch[:idx1]:
+                    n_feedback += coupling_dict[function1][function2]
 
-        return feedback
+        # Calculate the lower bound for the runtime by calculating the runtime of the branch
+        run_time_branch = self.get_runtime_sequence(branch, coupling_dict=coupling_dict,
+                                                    use_runtime_info=use_runtime_info)
+
+        return n_feedback, run_time_branch
 
     def _genetic_algorithm(self, nodes, coupling_dict):
 
@@ -1080,7 +1181,7 @@ class DataGraph(KadmosGraph):
         function_order = []
         for idx in individual:
             function_order.append(mapping[idx])
-        feedback, size = self.get_feedback_info(function_order, coupling_dict)
+        feedback, time = self.get_feedback_info(function_order, coupling_dict)
 
         return feedback,
 
@@ -1115,26 +1216,10 @@ class DataGraph(KadmosGraph):
                     n_feedback_loops += coupling_dict[function1][function2]
 
         # Calculate runtime of the sequence
-        nodes = copy.deepcopy(function_order)
-        run_time_partition = 0
-        while nodes:
-            parallel_nodes = []
-            run_times = []
-            for idx, node in enumerate(nodes):
-                if not set(nodes[:idx]).intersection(coupling_dict[node]):
-                    parallel_nodes.append(node)
-                    if use_runtime_info:
-                        run_times.append(self.nodes[node]['performance_info']['run_time'])
-                    else:
-                        run_times.append(1)
-                else:
-                    break
-            run_time_partition += max(run_times)
-            for node in parallel_nodes:
-                nodes.pop(nodes.index(node))
-
-        return n_feedback_loops, run_time_partition
+        run_time = self.get_runtime_sequence(function_order, coupling_dict=coupling_dict,
+                                             use_runtime_info=use_runtime_info)
 
+        return n_feedback_loops, run_time
 
 
 class RepositoryConnectivityGraph(DataGraph):
@@ -2832,9 +2917,20 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
                 if not nx.is_directed_acyclic_graph(subgraph):
                     convergers.append(part_nr)
 
-        return best_solution['partitions'], convergers
+        # Update the function order
+
+
+        function_order = []
 
-    def get_partition_info(self, partitions, include_run_time=False):
+        self.graph['problem_formulation']['coupled_functions_groups'] = best_solution['partitions']
+        self.graph['problem_formulation']['local_convergers'] = convergers
+        self.graph['problem_formulation']['partition_convergence'] = []
+        self.graph['problem_formulation']['sequence_partitions'] = []
+        self.graph['problem_formulation']['function_order'] = function_order
+
+        return
+
+    def get_partition_info(self, partitions, include_run_time=False, coupling_dict=None):
         """ Function to get the number of feedback variables in the partitions and the number system variables (feedback
         and feedforward variables between partitions)
 
@@ -2858,7 +2954,8 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
                                                                            '{}'.format(node)
 
         # Get coupling dictionary
-        coupling_dict = self.get_coupling_dictionary()
+        if not coupling_dict:
+            coupling_dict = self.get_coupling_dictionary()
 
         # For each node in the partitions check whether its input comes from the same partition, another partition or
         # neither
@@ -3201,7 +3298,7 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
                             if nx.has_path(subgraph, func, post_func):
                                 linked_groups_post_func.append(idx)
                                 break
-                    if len(linked_groups_post_func) == 1:
+                    if len(linked_groups_post_func) >= 1:
                         local_groups.update(linked_groups_post_func)
                 # If the variable can only be associated to one group it is a local design variable
                 if len(local_groups) == 1:
@@ -3209,8 +3306,8 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
                 else:
                     global_des_vars.append(des_var)
                 if not post_funcs:
-                    raise AssertionError("Design variable {} could not be associated with a coupled or post-coupling "
-                                         "function.".format(des_var))
+                    logger.warning("Design variable {} could not be associated with a coupled or post-coupling "
+                                   "function.".format(des_var))
             elif len(linked_groups) == 1:
                 local_des_vars.append(des_var)
             else:
@@ -3291,8 +3388,6 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
                                 break
                     if len(linked_groups_pre_func) >= 1:
                         connected_groups.update(linked_groups_pre_func)
-                    if len(connected_groups) > 1:
-                        break
                 # If the constraint can only be associated to one group it is a local design variable
                 if len(connected_groups) == 1:
                     local_cnstrnt_funcs.append(cnstrnt_func)
@@ -3301,8 +3396,8 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
                     global_cnstrnt_funcs.append(cnstrnt_func)
                     global_cnstrnt_vars.extend(cnstrnt_funcs[cnstrnt_func])
                 if not pre_funcs:
-                    raise AssertionError("Constraint function {} could not be associated with a pre-coupling or "
-                                         "coupled function.".format(cnstrnt_func))
+                    logger.warning("Constraint function {} could not be associated with a pre-coupling or "
+                                   "coupled function.".format(cnstrnt_func))
             elif len(linked_groups) == 1:
                 local_cnstrnt_funcs.append(cnstrnt_func)
                 local_cnstrnt_vars.extend(cnstrnt_funcs[cnstrnt_func])
@@ -3393,8 +3488,8 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
 
         return sys_post_couplings, sys_post_couplings_groups_idx
 
-    def get_system_level_functions(self, global_objective_function, global_cnstrnt_functions, subsys_functions_dicts,
-                                   mg_function_ordering=None):
+    def get_system_level_functions(self, global_objective_function, global_cnstrnt_functions, mg_function_ordering=None,
+                                   add_system_functions_to_subsystems=True):
         # TODO: Add docstring
 
         # Initiate dictionary
@@ -3403,7 +3498,7 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
         system_level_function_dict[self.FUNCTION_ROLES[1]] = []
         system_level_function_dict[self.FUNCTION_ROLES[3]] = []
         system_level_function_dict[self.FUNCTION_ROLES[4]] = []
-        system_level_function_dict[self.FUNCTION_ROLES[2]] = global_functions
+        system_level_function_dict[self.FUNCTION_ROLES[2]] = []
 
         # Get and check function groups
         if mg_function_ordering is None:
@@ -3412,124 +3507,106 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
         post_desvars = mg_function_ordering[self.FUNCTION_ROLES[4]]
         post_couplings = mg_function_ordering[self.FUNCTION_ROLES[2]]
 
-        # Get sublevel nodes
-        sublevel_nodes = [node for subsys_functions_dict in subsys_functions_dicts for type in
-                          subsys_functions_dict for node in subsys_functions_dict[type]]
-
         # Add pre-design variables functions to the dictionary
         system_level_function_dict[self.FUNCTION_ROLES[3]] = pre_desvars
 
-        # Get system level post design functions and post-coupling functions
-        system_post_desvars = [post_desvar for post_desvar in post_desvars if post_desvar not in sublevel_nodes]
-        system_post_coupling = [post_coupling for post_coupling in post_couplings if post_coupling not in sublevel_nodes]
-
-        # This operation is done to keep the right order of the functions
-        system_level_function_dict[self.FUNCTION_ROLES[4]] = system_post_desvars
-        system_level_function_dict[self.FUNCTION_ROLES[2]] = [post_coupling for post_coupling in post_couplings if
-                                                              post_coupling in global_functions + system_post_coupling]
+        if add_system_functions_to_subsystems:
+            # Add post-design variables and post-coupling functions to the dictionary if they have a dependency to one
+            # of the global functions
+            # Create a subgraph
+            subgraph = self.get_subgraph_by_function_nodes(post_desvars + post_couplings)
+            # Check each function
+            for post_desvar in post_desvars:
+                for global_function in global_functions:
+                    if nx.has_path(subgraph, post_desvar, global_function):
+                        system_level_function_dict[self.FUNCTION_ROLES[4]].append(post_desvar)
+            additional_post_couplings = []
+            for post_coupling in post_couplings:
+                if post_coupling not in global_functions:
+                    for global_function in global_functions:
+                        if nx.has_path(subgraph, post_coupling, global_function):
+                            additional_post_couplings.append(post_coupling)
+            # This operation is done to keep the right order of the functions
+            system_level_function_dict[self.FUNCTION_ROLES[2]] = [fun for fun in post_couplings if fun in
+                                                                  global_functions + additional_post_couplings]
+        else:
+            # Add functions to the dictionary
+            system_level_function_dict[self.FUNCTION_ROLES[4]] = post_desvars
+            system_level_function_dict[self.FUNCTION_ROLES[2]] = post_couplings
 
         return system_level_function_dict
 
-    def get_sub_level_functions(self, coupled_functions_groups, mg_function_ordering=None, add_system_functions_to_subsystems=True):
+    def get_sub_level_functions(self, local_objective_function, local_cnstrnt_funcs, coupled_functions_group,
+                                mg_function_ordering=None, add_system_functions_to_subsystems=True):
         # TODO: Add docstring
 
-        # Get and check function groups
-        if mg_function_ordering is None:
-            mg_function_ordering = self.get_mg_function_ordering()
-        post_desvars = copy.deepcopy(mg_function_ordering[self.FUNCTION_ROLES[4]])
-        post_couplings = copy.deepcopy(mg_function_ordering[self.FUNCTION_ROLES[2]])
-        nodes_in_partitions = copy.deepcopy(coupled_functions_groups)
+        # Initiate dictionary
+        local_objective_function_list = [] if local_objective_function is None else [local_objective_function]
+        local_functions = local_objective_function_list + local_cnstrnt_funcs
+        sub_level_function_dict = dict()
+        sub_level_function_dict[self.FUNCTION_ROLES[3]] = []
+        sub_level_function_dict[self.FUNCTION_ROLES[4]] = []
+        sub_level_function_dict[self.FUNCTION_ROLES[1]] = []
+        sub_level_function_dict[self.FUNCTION_ROLES[2]] = []
+
+        # Evaluate coupled functions group
+        # Get pre-coupled functions on subsystem level
         coupling_dict = self.get_coupling_dictionary()
-
-        not_converged = True
-        nr_nodes_previous_iteration = len([node for partition in nodes_in_partitions for node in partition])
-        if add_system_functions_to_subsystems:
-            while not_converged:
-                # Check which pre-desvars can be added to the partitions
-                linked_nodes = []
-                for post_desvar in post_desvars:
-                    linked_groups = []
-                    for idx, nodes_partition in enumerate(nodes_in_partitions):
-                        subgraph = self.get_subgraph_by_function_nodes(post_desvars+nodes_partition)
-                        for func in nodes_partition:
-                            if nx.has_path(subgraph, post_desvar, func) or nx.has_path(subgraph, func, post_desvar):
-                                linked_groups.append(idx)
-                                break
-                    if len(linked_groups) == 1:
-                        nodes_in_partitions[linked_groups[0]].append(post_desvar)
-                        linked_nodes.append(post_desvar)
-                # Delete linked nodes from the post_desvars list
-                for post_desvar in linked_nodes:
-                    post_desvars.pop(post_desvars.index(post_desvar))
-
-                # Check which post-coupling can be added to the partitions
-                linked_nodes = []
-                for post_coupling in post_couplings:
-                    linked_groups = []
-                    for idx, nodes_partition in enumerate(nodes_in_partitions):
-                        subgraph = self.get_subgraph_by_function_nodes(nodes_partition+post_couplings)
-                        for func in nodes_partition:
-                            if nx.has_path(subgraph, post_coupling, func) or nx.has_path(subgraph, func, post_coupling):
-                                linked_groups.append(idx)
-                                break
-                    if len(linked_groups) == 1:
-                        nodes_in_partitions[linked_groups[0]].append(post_coupling)
-                        linked_nodes.append(post_coupling)
-                # Delete linked nodes from post_couplings list
-                for post_coupling in linked_nodes:
-                    post_couplings.pop(post_couplings.index(post_coupling))
-
-                # Check whether all nodes are added
-                nr_nodes = len([node for partition in nodes_in_partitions for node in partition])
-                if nr_nodes != nr_nodes_previous_iteration:
-                    nr_nodes_previous_iteration = nr_nodes
+        pre_coupling_idx = -1
+        for idx, node in enumerate(coupled_functions_group):
+            if not set(coupled_functions_group[idx + 1:]).intersection(coupling_dict[node]):
+                pre_coupling_idx = idx
+            else:
+                break
+        # Get post-coupled functions on subsystem level
+        post_coupling_idx = len(coupled_functions_group)
+        if pre_coupling_idx + 1 != len(coupled_functions_group):
+            input_functions = []
+            for idx, node in enumerate(coupled_functions_group):
+                input_functions.extend([node for node in coupling_dict[node] if node in coupled_functions_group[idx:]])
+            for idx, node in reversed(list(enumerate(coupled_functions_group))):
+                if node not in input_functions:
+                    post_coupling_idx = idx
                 else:
                     break
+        sublevel_post_couplings = coupled_functions_group[post_coupling_idx:]
+        sublevel_pre_couplings = coupled_functions_group[:pre_coupling_idx + 1]
 
-        # Make sublevel dictionaries
-        sub_level_function_dicts = []
-        for index, coupled_functions_group in enumerate(coupled_functions_groups):
-
-            # Initialize function dict
-            sub_level_function_dict = dict()
-            sub_level_function_dict[self.FUNCTION_ROLES[3]] = []
-            sub_level_function_dict[self.FUNCTION_ROLES[4]] = []
-            sub_level_function_dict[self.FUNCTION_ROLES[1]] = []
-            sub_level_function_dict[self.FUNCTION_ROLES[2]] = []
+        # Get and check function groups
+        if mg_function_ordering is None:
+            mg_function_ordering = self.get_mg_function_ordering()
+        post_desvars = mg_function_ordering[self.FUNCTION_ROLES[4]]
+        post_couplings = mg_function_ordering[self.FUNCTION_ROLES[2]]
+        additional_post_couplings = []
+        additional_post_desvars = []
 
-            # Evaluate coupled functions group
-            # Get pre-coupled functions
-            pre_coupling_idx = -1
-            for idx, node in enumerate(coupled_functions_group):
-                if not set(coupled_functions_group[idx + 1:]).intersection(coupling_dict[node]):
-                    pre_coupling_idx = idx
-                else:
-                    break
-            # Get post-coupled functions
-            post_coupling_idx = len(coupled_functions_group)
-            if pre_coupling_idx + 1 != len(coupled_functions_group):
-                input_functions = []
-                for idx, node in enumerate(coupled_functions_group):
-                    input_functions.extend(
-                        [node for node in coupling_dict[node] if node in coupled_functions_group[idx:]])
-                for idx, node in reversed(list(enumerate(coupled_functions_group))):
-                    if node not in input_functions:
-                        post_coupling_idx = idx
-                    else:
-                        break
-            sublevel_post_couplings = coupled_functions_group[post_coupling_idx:]
-            sublevel_pre_couplings = coupled_functions_group[:pre_coupling_idx + 1]
-            sub_level_function_dict[self.FUNCTION_ROLES[1]].extend(coupled_functions_group[pre_coupling_idx + 1:post_coupling_idx])
+        if add_system_functions_to_subsystems:
+            # Add local functions to the dictionary
+            additional_post_couplings = list(local_functions)
+            # Add post-design variables and post-coupling functions to the dictionary if they have a dependency to one
+            # of the global functions
+            # Create a subgraph
+            subgraph = self.get_subgraph_by_function_nodes(post_desvars + coupled_functions_group + post_couplings)
+            # Check each function
+            for post_desvar in post_desvars:
+                for local_function in local_functions:
+                    if nx.has_path(subgraph, post_desvar, local_function):
+                        additional_post_desvars.append(post_desvar)
+            for post_coupling in post_couplings:
+                if post_coupling not in local_functions:
+                    for local_function in local_functions:
+                        if nx.has_path(subgraph, post_coupling, local_function):
+                            additional_post_couplings.append(post_coupling)
 
-            sub_level_function_dict[self.FUNCTION_ROLES[4]] = [func for func in sublevel_pre_couplings +
-                                                               mg_function_ordering[self.FUNCTION_ROLES[4]] if func in
-                                                               nodes_in_partitions[index]]
-            sub_level_function_dict[self.FUNCTION_ROLES[2]] = [func for func in sublevel_post_couplings +
-                                                               mg_function_ordering[self.FUNCTION_ROLES[2]] if func in
-                                                               nodes_in_partitions[index]]
-            sub_level_function_dicts.append(sub_level_function_dict)
+        # This operation is done to keep the right order of the functions
+        sub_level_function_dict[self.FUNCTION_ROLES[4]] = [func for func in post_desvars + coupled_functions_group if
+                                                           func in sublevel_pre_couplings + additional_post_desvars]
+        sub_level_function_dict[self.FUNCTION_ROLES[1]] = coupled_functions_group[
+                                                          pre_coupling_idx + 1:post_coupling_idx]
+        sub_level_function_dict[self.FUNCTION_ROLES[2]] = [func for func in coupled_functions_group + post_couplings if
+                                                           func in sublevel_post_couplings + additional_post_couplings]
 
-        return sub_level_function_dicts
+        return sub_level_function_dict
 
     def check_and_get_pre_coupling_functions(self, pre_coupling_functions=None):
         if not pre_coupling_functions:
@@ -3632,35 +3709,38 @@ class FundamentalProblemGraph(DataGraph, KeChainMixin):
         cnstrnt_vars_group_idxs, \
         cnstrnt_funcs_group_idxs = self.determine_scope_constraint_functions(cnstrnt_vars=constraint_nodes)
 
-        # Create dictionaries of post-desvar, coupled, and post-coupling functions per each subgroup
-        local_cnstrnt_funcs_groups = []
-        for idx, coupled_functions_group in enumerate(coupled_functions_groups):
-            # Get the local constraint functions of the current group
-            local_cnstrnt_funcs_group = []
-            for cnstrnt_func, groups in cnstrnt_funcs_group_idxs.iteritems():
-                if idx in groups and len(groups) == 1:
-                    local_cnstrnt_funcs_group.append(cnstrnt_func)
-            local_cnstrnt_funcs_groups.append(local_cnstrnt_funcs_group)
-        subsys_functions_dicts = self.get_sub_level_functions(coupled_functions_groups,
-                                                              mg_function_ordering=mg_function_ordering,
-                                                              add_system_functions_to_subsystems=
-                                                              add_system_functions_to_subsystems)
-
         # Create dictionary of pre-desvar, post-desvar, and post-coupling functions for the system optimizer
         sys_functions_dict = self.get_system_level_functions(global_objective_function, global_cnstrnt_funcs,
-                                                             subsys_functions_dicts,
-                                                             mg_function_ordering=mg_function_ordering)
-
-        # Assert that each node is part of only one function dict
-        sublevel_functions = [func for subsys_functions_dict in subsys_functions_dicts for type in
-                              subsys_functions_dict for func in subsys_functions_dict[type]]
-        system_functions = [func for type in sys_functions_dict for func in sys_functions_dict[type]]
-        assert len(sublevel_functions+system_functions) == len(set(sublevel_functions+system_functions))
+                                                             mg_function_ordering=mg_function_ordering,
+                                                             add_system_functions_to_subsystems=
+                                                             add_system_functions_to_subsystems)
 
         # Determine couplings between coupled groups and system-level post-coupling functions
         add_group_couplings, \
         add_group_couplings_groups_idx = self.get_sys_post_couplings(sys_functions_dict[self.FUNCTION_ROLES[2]])
 
+        # Add additional couplings to the group_couplings
+        for add_group_coupling in add_group_couplings:
+            if add_group_coupling not in group_couplings:
+                group_couplings.append(add_group_coupling)
+                group_couplings_groups_idx[add_group_coupling] = add_group_couplings_groups_idx[add_group_coupling]
+
+        # Create dictionaries of post-desvar, coupled, and post-coupling functions per each subgroup
+        subsys_functions_dicts = []
+        for idx, coupled_functions_group in enumerate(coupled_functions_groups):
+            # Get the local constraint functions of the current group
+            local_cnstrnt_funcs_group = []
+            for cnstrnt_func, groups in cnstrnt_funcs_group_idxs.iteritems():
+                if idx in groups:
+                    local_cnstrnt_funcs_group.append(cnstrnt_func)
+            subsys_functions_dict = self.get_sub_level_functions(None, local_cnstrnt_funcs_group,
+                                                                 coupled_functions_group,
+                                                                 mg_function_ordering=mg_function_ordering,
+                                                                 add_system_functions_to_subsystems=
+                                                                 add_system_functions_to_subsystems)
+            # Create dict collecting the subsystem functions dictionary
+            subsys_functions_dicts.append(subsys_functions_dict)
+
         # Add additional couplings to the group_couplings
         for add_group_coupling in add_group_couplings:
             if add_group_coupling not in group_couplings:
@@ -5065,7 +5145,13 @@ class MdaoDataGraph(DataGraph, MdaoMixin):
                           label=label,
                           instance=1,
                           settings={'convergence_tolerance_absolute': 1e-6, 'convergence_tolerance_relative': 1e-6,
-                                    'last_iterations_to_consider': 1, 'maximum_iterations': 100})
+                                    'last_iterations_to_consider': 1, 'maximum_iterations': 100} if not
+                                    converger_is_optimizer else {'algorithm': 'Dakota Quasi-Newton method',
+                                      'maximum_iterations': 1000,
+                                      'maximum_function_evaluations': 1000,
+                                      'constraint_tolerance': 1e-4,
+                                      'convergence_tolerance': 1e-4,
+                                      'apply_scaling': True})
         assert conv_type in self.OPTIONS_CONVERGERS + [self.OPTIONS_ARCHITECTURES[2]], \
             'Invalid converger type %s specified.' % conv_type
         assert isinstance(coupling_functions, list)
@@ -5481,20 +5567,26 @@ class MdaoDataGraph(DataGraph, MdaoMixin):
         # Manipulate the coupling nodes accordingly
         idx = 0
         for coupling in couplings:
-            if system_converger and 'part_id' in self.nodes[coupling[0]] and 'part_id' in self.nodes[coupling[1]]:
-                # Do not manipulate nodes if they are in the same partition
-                if self.nodes[coupling[0]]['part_id'] == self.nodes[coupling[1]]['part_id']:
-                    continue
-                # Do not manipulate nodes if the partitions are solved in sequence
-                if 'sequence_partitions' in self.graph['problem_formulation']:
-                    skip_coupling = False
-                    for sequence in self.graph['problem_formulation']['sequence_partitions']:
-                        if self.nodes[coupling[0]]['part_id'] in sequence:
-                            index = sequence.index(self.nodes[coupling[0]]['part_id'])
-                            if self.nodes[coupling[1]]['part_id'] in sequence[index:]:
-                                skip_coupling = True
-                    if skip_coupling:
+            if system_converger:
+                assert 'coupled_functions_groups' in self.graph['problem_formulation'], 'Graph is not partitioned'
+                partitions = self.graph['problem_formulation']['coupled_functions_groups']
+                partition_nodes = [node for nodes in partitions for node in nodes]
+                if coupling[0] in partition_nodes and coupling[1] in partition_nodes:
+                    part_id_0 = [i for i in range(len(partitions)) if coupling[0] in partitions[i]][0]
+                    part_id_1 = [i for i in range(len(partitions)) if coupling[1] in partitions[i]][0]
+                    # Do not manipulate nodes if they are in the same partition
+                    if part_id_0 == part_id_1:
                         continue
+                    # Do not manipulate nodes if the partitions are solved in sequence
+                    if 'sequence_partitions' in self.graph['problem_formulation']:
+                        skip_coupling = False
+                        for sequence in self.graph['problem_formulation']['sequence_partitions']:
+                            if part_id_0 in sequence:
+                                index = sequence.index(part_id_0)
+                                if part_id_1 in sequence[index:]:
+                                    skip_coupling = True
+                        if skip_coupling:
+                            continue
             # Create initial guess coupling variable node
             ini_guess_node = self.copy_node_as(coupling[2], self.ARCHITECTURE_ROLES_VARS[0])
             # If there is no converger node, then just add an initial guess of the coupled node
diff --git a/kadmos/graph/graph_kadmos.py b/kadmos/graph/graph_kadmos.py
index ccaedbb0693602f272c7e08629b9ac8a111bde0e..64a9917f95f27e081998ad7dd9fe4d556aa01ed8 100644
--- a/kadmos/graph/graph_kadmos.py
+++ b/kadmos/graph/graph_kadmos.py
@@ -664,8 +664,6 @@ class KadmosGraph(nx.DiGraph, EquationMixin, VistomsMixin):
         :type keep_tex_file: bool
         :param abbreviate_keywords: optional argument to keep make keywords shorter (input -> inp., output -> outp.)
         :type abbreviate_keywords: bool
-        :param compile_pdf: optional argument to compile the PDF
-        :type compile_pdf: bool
         :param colors_based_on: option to base the colors either on the problem role or the partitions
         :type colors_based_on: str
         """
@@ -680,6 +678,8 @@ class KadmosGraph(nx.DiGraph, EquationMixin, VistomsMixin):
         assert isinstance(keep_tex_file, bool)
         assert keep_tex_file or compile_pdf, 'The settings do not make sense, set either keep_tex_file or compile_pdf' \
                                              ' to True, or both!'
+        if colors_based_on == 'partitions':
+            assert 'coupled_functions_groups' in self.graph['problem_formulation'], 'Graph is not partitioned'
 
         # Check if MPG is applicable
         if mpg is not None:
@@ -763,8 +763,11 @@ class KadmosGraph(nx.DiGraph, EquationMixin, VistomsMixin):
             for idx, item in enumerate(diagonal_nodes):
                 node = diagonal_nodes[idx]
                 if isinstance(graph, FundamentalProblemGraph):
-                    if colors_based_on == 'partitions' and 'part_id' in graph.nodes[node]:
-                        node_style = 'EvenPartitions' if int(graph.nodes[node]['part_id']) % 2 == 0 else 'OddPartitions'
+                    if colors_based_on == 'partitions' and node in [node1 for nodes in self.graph[
+                            'problem_formulation']['coupled_functions_groups'] for node1 in nodes]:
+                        partitions = self.graph['problem_formulation']['coupled_functions_groups']
+                        part_id = [i for i in range(len(partitions)) if node in partitions[i]]
+                        node_style = 'EvenPartitions' if part_id[0] % 2 == 0 else 'OddPartitions'
                     elif 'problem_role' in graph.node[node]:
                         if graph.node[node]['problem_role'] == graph.FUNCTION_ROLES[0]:
                             node_style = 'PreAnalysis'
@@ -782,8 +785,11 @@ class KadmosGraph(nx.DiGraph, EquationMixin, VistomsMixin):
                         except ValueError:
                             raise AssertionError('Architecture role %s is not supported for creation of XDSMs.'
                                                  % graph.node[node]['architecture_role'])
-                        if colors_based_on == 'partitions' and 'part_id' in graph.nodes[node]:
-                            node_style = 'EvenPartitions' if int(graph.nodes[node]['part_id']) % 2 == 0 else 'OddPartitions'
+                        if colors_based_on == 'partitions' and node in [node1 for nodes in self.graph[
+                                'problem_formulation']['coupled_functions_groups'] for node1 in nodes]:
+                            partitions = self.graph['problem_formulation']['coupled_functions_groups']
+                            part_id = [i for i in range(len(partitions)) if node in partitions[i]]
+                            node_style = 'EvenPartitions' if part_id[0] % 2 == 0 else 'OddPartitions'
                         else:
                             node_style = self.ARCHITECTURE_ROLES_NODESTYLES[role_index]
                     else:
@@ -809,8 +815,11 @@ class KadmosGraph(nx.DiGraph, EquationMixin, VistomsMixin):
                 except ValueError:
                     raise AssertionError('Architecture role %s is not supported for creation of XDSMs.'
                                          % graph_mpg.node[node]['architecture_role'])
-                if colors_based_on == 'partitions' and 'part_id' in graph.nodes[node]:
-                    node_style = 'EvenPartitions' if int(graph.nodes[node]['part_id']) % 2 == 0 else 'OddPartitions'
+                if colors_based_on == 'partitions' and node in [node1 for nodes in self.graph['problem_formulation'][
+                        'coupled_functions_groups'] for node1 in nodes]:
+                    partitions = self.graph['problem_formulation']['coupled_functions_groups']
+                    part_id = [i for i in range(len(partitions)) if node in partitions[i]]
+                    node_style = 'EvenPartitions' if part_id[0] % 2 == 0 else 'OddPartitions'
                 else:
                     node_style = self.ARCHITECTURE_ROLES_NODESTYLES[role_index]
                 # Determine node text