kCommonFlowDecomp.py

import networkx as nx
import gurobipy as gb
import utils

class KCommonFlowDecomp:
    def __init__(self, G: nx.DiGraph, num_flows: int, k: int, flow_attr: str = "flow", subpath_constr: list = [], console = False):
        if not nx.is_directed_acyclic_graph(G):
            print("uh oh")
            raise ValueError('Input graph is not a directed acyclic graph')
        if not utils.check_st_graph(G):
            print("uh oh")
            raise ValueError('Input graph is not an st graph')
        if not utils.check_correct_num_flows(G, num_flows, flow_attr):
            print("uh oh")
            raise ValueError('Number of flows does not match')
        if not utils.check_valid_flow_format(G, num_flows, flow_attr):
            print("uh oh")
            raise ValueError('Flow value must be int or float')
        if not utils.check_multi_flow_conservation(G, num_flows, flow_attr):
            print("uh oh")
            raise ValueError('Input graph does not conserve flow')
        """
        if subpath_constr and not utils.check_subpath_constr(G, subpath_constr):
            print("uh oh")
        """
        
        self.model = gb.Model()
        self.model.setParam('OutputFlag', 0)
        self.G = G
        self.num_flows = num_flows
        self.k = k
        self.flow_attr = flow_attr
        self.w_max = utils.get_max_flow(self.G, self.num_flows, self.flow_attr)
        self.console = console
        
        self.subpath_constr = subpath_constr
                
        self.path_indexes = [(i, j) for i in range(self.k) for j in range(self.num_flows)]
        self.edge_indexes = [(u, v, i) for u, v in self.G.edges() for i in range(self.k)]
        self.edge_flows = {(u, v, j): data[self.flow_attr][j] for u, v, data in self.G.edges(data=True) for j in
                           range(self.num_flows)}
        self.pi_indexes = [(u, v, i, j) for u, v in self.G.edges() for i in range(self.k) for j in
                           range(self.num_flows)]
        self.subpath_indexes = [(i, p) for i in range(self.k) for p in range(len(self.subpath_constr))]
  
    def fix_zero_edges(self, safe_sequences, reachable_nodes_from, nodes_reaching):
        count = 0
        for i in range(min(len(safe_sequences), self.k)):
            path = safe_sequences[i]
            if not path:
                continue

            first_node = path[0][0]
            last_node  = path[-1][1]

            protected_edges = set((u, v) for (u, v) in path if self.G.has_edge(u, v))

            # protect edges reachable from last node or that can reach first node
            for (u, v) in self.G.edges():
                if (u in reachable_nodes_from[last_node]) or (v in nodes_reaching[first_node]):
                    protected_edges.add((u, v))

            # protect edges that bridge gaps in the sequence
            for idx in range(len(path) - 1):
                end_prev   = path[idx][1]
                start_next = path[idx + 1][0]
                if end_prev != start_next:
                    for (u, v) in self.G.edges():
                        if (u in reachable_nodes_from[end_prev]) and (v in nodes_reaching[start_next]):
                            protected_edges.add((u, v))

            # fix everything else to 0
            
            for (u, v) in self.G.edges():
                if (u, v) not in protected_edges:
                    count += 1
                    self.model.addConstr(self.edge_vars[u, v, i] == 0)
        return count
                      
    def add_safe_sequences(self, safe_sequences):
        for i, safe_seq in enumerate(safe_sequences):
            for u, v in safe_seq:
                self.model.addConstr(self.edge_vars[u, v, i] == 1)

    def build_model(self):

        self.variable_name_prefixes = []

        self.path_vars = self.add_variables(indexes=self.path_indexes, name_prefix='w', ub=self.w_max)
        self.edge_vars = self.add_variables(indexes=self.edge_indexes, name_prefix='x', var_type="binary")
        self.pi_vars = self.add_variables(indexes=self.pi_indexes, name_prefix='pi', ub=self.w_max)
        self.subpath_vars = self.add_variables(indexes=self.subpath_indexes, name_prefix='r', var_type="binary")

        for v in self.G.nodes():
            predecessors = list(self.G.predecessors(v))
            successors = list(self.G.neighbors(v))
            if len(predecessors) == 0:
                for i in range(self.k):
                    self.model.addConstr(gb.quicksum(self.edge_vars[v, w, i] for w in successors) == 1,
                                         name=f"single_path_i={i}")
            elif len(successors) != 0:
                for i in range(self.k):
                    self.model.addConstr(gb.quicksum(self.edge_vars[u, v, i] for u in predecessors) ==
                                         gb.quicksum(self.edge_vars[v, w, i] for w in successors),
                                         name=f"flow_cons_v={v}_i={i}")
        for u, v in self.G.edges():
            for j in range(self.num_flows):
                self.model.addConstr(gb.quicksum(self.pi_vars[u, v, i, j] for i in range(self.k)) ==
                                     self.edge_flows[u, v, j], name=f"correct_flow_u={u}_v={v}_j={j}")
                for i in range(self.k):
                    self.add_binary_continuous_product_constraint(binary_var=self.edge_vars[u, v, i],
                                                                  continuous_var=self.path_vars[i, j],
                                                                  product_var=self.pi_vars[u, v, i, j], lb=0,
                                                                  ub=self.w_max, name=f"pi_u={u}_v={v}_i={i}_j={j}")
        
        ###PRIMARY FORMULATION -- EACH SUBPATH CONSTRAINT SATISFIED BY A SINGLE FLOW
        if self.subpath_constr:
            for p in range(len(self.subpath_constr)):
                self.model.addConstr(gb.quicksum(self.subpath_vars[i,p] for i in range(self.k)) >= 1,
                                     name=f"subpath_claim_p={p}")
                for i in range(self.k):
                    self.model.addConstr(gb.quicksum(self.edge_vars[u,v,i] for u, v in self.subpath_constr[p]) >=
                                         len(self.subpath_constr[p]) * self.subpath_vars[i,p],
                                         name=f"subpath_proof_i={i}_p={p}")

            for i in range(self.k):
                self.model.addConstr(gb.quicksum(self.path_vars[i,j] for j in range(self.num_flows)) >= 1,
                                     name=f"path_used_i={i}")
        
        ###ALTERNATIVE FORMULATION -- EACH SUBPATH CONSTRAINT SATISFIED BY ALL FLOWS
        # if self.subpath_constr:
        #     for j in range(self.num_flows):
        #         for p in range(len(self.subpath_constr)):
        #             self.model.addConstr(gb.quicksum(self.subpath_vars[i,j,p] for i in range(self.k)) >= 1,
        #                                  name=f"subpath_flow_claim_j={j}_p={p}")
        #             for i in range(self.k):
        #                 self.model.addConstr(gb.quicksum(self.edge_vars[u,v,i] for u, v in self.subpath_constr[p]) >=
        #                                      (len(self.subpath_constr[p]) - 1) * self.subpath_vars[i,j,p],
        #                                      name=f"subpath_proof_i={i}_j={j}_p={p}")
        #
        #         for i in range(self.k):
        #             self.model.addConstr(self.path_vars[i,j] >= self.path_vars[i,j,p],
        #                                  name=f"path_flow_used_i={i}_j={j}_p={p}")

    def solve_model(self):
        self.model.optimize()
        if self.model.status == gb.GRB.Status.OPTIMAL:
            return True
        else:
            return False

    def get_model_solution(self):
        solution = ""
        for i in range(self.k):
            solution = solution + f"Path {i+1} (carries weight"
            for j in range(self.num_flows):
                if j == self.num_flows - 1:
                    solution = solution + " and"
                solution = solution + f" {self.path_vars[(i, j)].X} for flow {j+1}"
                if j < self.num_flows - 1 and self.num_flows > 2:
                    solution = solution + ","
            solution = solution + "):\n"
            for u,v in self.G.edges():
                if self.edge_vars[u, v, i].X != 0:
                    solution = solution + f"{u}, "
            solution = solution + "t\n"
        for i in range(self.k):
            for p in range(len(self.subpath_constr)):
                if self.subpath_vars[(i,p)].X != 0:
                    solution = solution + f"Path {i+1} satisfies constraint {j}\n"
        return solution

    def get_model_weights(self):
        weights = []
        for i in range(self.k):
            weight_vector = [self.path_vars[i, j].X for j in range(self.num_flows)]
            weights.append(weight_vector)
        return weights

    def get_model_paths(self):
        paths = []
        for i in range(self.k):
            path = []
            for u in nx.topological_sort(self.G):
                for v in self.G.successors(u):
                    if self.edge_vars[u, v, i].X != 0:
                        path.append(u)
            path.append(list(nx.topological_sort(self.G))[-1])
            paths.append(path)
        return paths    
            
    def add_variables(self, indexes, name_prefix: str, lb=0, ub=1, var_type="continuous"):
        for prefix in self.variable_name_prefixes:
            if prefix.startswith(name_prefix) or name_prefix.startswith(prefix):
                print("uh oh")
                raise ValueError(
                    f"Variable name prefix {name_prefix} conflicts with existing variable name prefix {prefix}. "
                    f"Use a different name prefix."
                )

        self.variable_name_prefixes.append(name_prefix)

        var_type_map = {
            "integer": gb.GRB.INTEGER,
            "continuous": gb.GRB.CONTINUOUS,
            "binary": gb.GRB.BINARY,
        }
        vars = {}
        for index in indexes:
            vars[index] = self.model.addVar(
                lb=lb,
                ub=ub,
                vtype=var_type_map[var_type],
                name=f"{name_prefix}{index}",
            )
        self.model.update()
        return vars

    def add_binary_continuous_product_constraint(self, binary_var, continuous_var, product_var, lb, ub, name: str):
        self.model.addConstr(product_var <= ub * binary_var, name=name + "_a")
        self.model.addConstr(product_var >= lb * binary_var, name=name + "_b")
        self.model.addConstr(product_var <= continuous_var - lb * (1 - binary_var), name=name + "_c")
        self.model.addConstr(product_var >= continuous_var - ub * (1 - binary_var), name=name + "_d")