def find_ranking(comparisons, equal_width=0.2, max_rank=-1, verbose=False): """ Find the least changes to a set of comparisons so that they are consistent (transitive), it returns a topological ranking. comparisons A dictionary with tuple keys in the form of (i, j), values are scalars indicating the probability of i > j. It is assumed that comparisons are symmetric. Use 0 for i < j, 0.5 for i == j, and 1 for i > j (and any value in between). equal_width 0..0.5-equal_width/2 is considered '<=' and 0.5..0.5+equal_width/2 is considered '>='. In between it is considered to be '=='. max_rank Maximal rank, a low value forces the model to choose more equal cases. verbose Whether to print gurobi's progress. Returns: A tuple of size two: 0) Ranking derived from topological sort (list of ranks in order of nodes); 1) Sum of absolute changes to the comparisons. """ # remove unnecessary variables comparisons = {(i, j) if i < j else (j, i): value if i < j else 1 - value for (i, j), value in comparisons.items()} nodes = np.unique([i for ij in comparisons.keys() for i in ij]) # define variables model = Model('comparison') model.setParam('OutputFlag', verbose) values = np.fromiter(comparisons.values(), dtype=float) assert values.max() <= 1 and values.min() >= 0 # variables to encode the error of comparisons E_ij = model.addVars(comparisons.keys(), name='e_ij', vtype=GRB.CONTINUOUS, ub=1.0 - values, lb=-values) # variables to encode hard choice of >=, <=, == Ge_ij = model.addVars(comparisons.keys(), name='ge_ij', vtype=GRB.BINARY) Le_ij = model.addVars(comparisons.keys(), name='le_ij', vtype=GRB.BINARY) Eq_ij = model.addVars(comparisons.keys(), name='eq_ij', vtype=GRB.BINARY) # variables to help with transitivity in non-fully connected graphs if max_rank < 1: max_rank = len(nodes) R_i = model.addVars(nodes, name='r_i', vtype=GRB.CONTINUOUS, lb=0, ub=max_rank) # variables to emulate abs T_ij_pos = {} T_ij_neg = {} index = (values != 1) & (values != 0) T_ij_pos = model.addVars( (ij for ij, value in comparisons.items() if value not in [0.0, 1.0]), vtype=GRB.CONTINUOUS, name='T_ij_pos', lb=0, ub=1 - values[index]) T_ij_neg = model.addVars( (ij for ij, value in comparisons.items() if value not in [0.0, 1.0]), vtype=GRB.CONTINUOUS, name='T_ij_neg', lb=0, ub=values[index]) model.update() # emulate abs for non-binary comparisons: E_ij = T_ij_pos - T_ij_neg model.addConstrs((E_ij[ij] == T_ij_pos[ij] - T_ij_neg[ij] for ij in T_ij_pos), 'E_ij = T_ij_pos - T_ij_neg') # hard decision of >=, <=, and == lower_bound = 0.5 - equal_width / 2.0 upper_bound = 0.5 + equal_width / 2.0 # <= model.addConstrs((E_ij[ij] + comparisons[ij] - upper_bound <= ge_ij for ij, ge_ij in Ge_ij.items()), 'ge_ij_lower_bound') model.addConstrs((E_ij[ij] + comparisons[ij] - upper_bound >= -1 + ge_ij for ij, ge_ij in Ge_ij.items()), 'ge_ij_upper_bound') # >= model.addConstrs((E_ij[ij] + comparisons[ij] - lower_bound >= -le_ij for ij, le_ij in Le_ij.items()), 'le_ij_lower_bound') model.addConstrs((E_ij[ij] + comparisons[ij] - lower_bound <= 1 - le_ij for ij, le_ij in Le_ij.items()), 'le_ij_upper_bound') # == model.addConstrs(( le + eq + ge == 1 for le, eq, ge in zip(Le_ij.values(), Eq_ij.values(), Ge_ij.values())), 'eq_ij') # transitivity for (i, j), eq_a in Eq_ij.items(): le_a = Le_ij[i, j] ge_a = Ge_ij[i, j] for k in nodes: j_, k_ = j, k if j > k: j_, k_ = k, j eq_b = Eq_ij.get((j_, k_), None) if eq_b is None: continue else: le_b = Le_ij[j_, k_] ge_b = Ge_ij[j_, k_] if j_ != j: le_b, ge_b = ge_b, le_b i_, k_ = i, k if i > k: i_, k_ = k, i eq_c = Eq_ij.get((i_, k_), None) if eq_c is None: continue else: le_c = Le_ij[i_, k_] ge_c = Ge_ij[i_, k_] if i_ != i: le_c, ge_c = ge_c, le_c # a <= b and b <= c -> a <= c model.addLConstr(ge_a + ge_b, GRB.LESS_EQUAL, 1 + ge_c, f'transitivity_ge_{i},{j},{k}') # a >= b and b >= c -> a >= c model.addLConstr(le_a + le_b, GRB.LESS_EQUAL, 1 + le_c, f'transitivity_le_{i},{j},{k}') # a <= b and b == c -> a <= c model.addLConstr(le_a + eq_b, GRB.LESS_EQUAL, 1 + le_c, f'transitivity_leeq_{i},{j},{k}') # a == b and b <= c -> a <= c model.addLConstr(eq_a + le_b, GRB.LESS_EQUAL, 1 + le_c, f'transitivity_eqle_{i},{j},{k}') # a >= b and b == c --> a >= c model.addLConstr(ge_a + eq_b, GRB.LESS_EQUAL, 1 + ge_c, f'transitivity_geeq_{i},{j},{k}') # a == b and b >= c --> a >= c model.addLConstr(eq_a + ge_b, GRB.LESS_EQUAL, 1 + ge_c, f'transitivity_eqge_{i},{j},{k}') # a == b and b == c --> a == c model.addLConstr(eq_a + eq_b, GRB.LESS_EQUAL, 1 + eq_c, f'transitivity_eq_{i},{j},{k}') # transitivity helper (for not-fully connected graphs) # also provides a latent rank big_m = max_rank model.addConstrs(((1 - ge_ij) * big_m + R_i[i] >= R_i[j] + 1 for (i, j), ge_ij in Ge_ij.items()), 'rank_transitivity_larger') model.addConstrs(((1 - le_ij) * big_m + R_i[j] >= R_i[i] + 1 for (i, j), le_ij in Le_ij.items()), 'rank_transitivity_smaller') model.addConstrs(((1 - eq_ij) * big_m + R_i[j] >= R_i[i] for (i, j), eq_ij in Eq_ij.items()), 'rank_transitivity_equal1') model.addConstrs(((1 - eq_ij) * big_m + R_i[i] >= R_i[j] for (i, j), eq_ij in Eq_ij.items()), 'rank_transitivity_equal2') # objective function objective = LinExpr() for ij, value in comparisons.items(): if value == 1.0: objective += -E_ij[ij] elif value == 0.0: objective += E_ij[ij] else: objective += T_ij_pos[ij] + T_ij_neg[ij] model.setObjective(objective, GRB.MINIMIZE) # solve model.optimize() # verify abs emulation: one T_ij has to be 0 for ij, value in T_ij_pos.items(): assert value.X == 0 or T_ij_neg[ij] == 0, \ f'T_{ij} pos {value.X} neg {T_ij_neg[ij]}' # find minimal Rs model_ = Model('comparison') model_.setParam('OutputFlag', verbose) R_i = model_.addVars(nodes, name='r_i', vtype=GRB.CONTINUOUS, lb=0, ub=len(nodes)) for ((i, j), ge_ij), le_ij in zip(Ge_ij.items(), Le_ij.values()): if ge_ij.x == 1: model_.addConstr(R_i[i] >= R_i[j] + 1) elif le_ij.x == 1: model_.addConstr(R_i[j] >= R_i[i] + 1) else: model_.addConstr(R_i[j] == R_i[i]) model_.setObjective(R_i.sum(), GRB.MINIMIZE) model_.optimize() return [model_.getVarByName(f'r_i[{i}]').X for i in range(len(nodes))], \ model.objVal
def find_ranking(comparisons, verbose=False): """ Find the least changes to a set of comparisons so that they are consistent (transitive), it returns a topological ranking. comparisons A dictionary with tuple keys in the form of (i, j), values are scalars indicating the probability of i >= j. It is assumed that comparisons are symmetric. The only way this function it can handle '=' is by setting the probability to 0.5. In this case the MILP can treat it as > or < (but not as =). verbose Whether to print gurobi's progress. Returns: A tuple of size two: 0) Ranking derived from topological sort (list of ranks in order of nodes); 1) Sum of absolute changes to the comparisons. """ # remove unnecessary variables comparisons = {(i, j) if i < j else (j, i): value if i < j else 1 - value for (i, j), value in comparisons.items()} nodes = np.unique([i for ij in comparisons.keys() for i in ij]) # define variables model = Model('comparison') model.setParam('OutputFlag', verbose) values = np.fromiter(comparisons.values(), dtype=float) assert values.max() <= 1 and values.min() >= 0 # variables to encode the error of comparisons E_ij = model.addVars(comparisons.keys(), name='e_ij', vtype=GRB.CONTINUOUS, ub=1.0 - values, lb=-values) # variable to encode hard choice of >= and <= Z_ij = model.addVars(comparisons.keys(), name='z_ij', vtype=GRB.BINARY) # variables to help with transitivity in non-fully connected graphs R_i = model.addVars(nodes, name='r_i', vtype=GRB.CONTINUOUS, lb=0, ub=len(nodes)) # variables to emulate abs T_ij_pos = {} T_ij_neg = {} index = (values != 1) & (values != 0) T_ij_pos = model.addVars( (ij for ij, value in comparisons.items() if value not in [0.0, 1.0]), vtype=GRB.CONTINUOUS, name='T_ij_pos', lb=0, ub=1 - values[index]) T_ij_neg = model.addVars( (ij for ij, value in comparisons.items() if value not in [0.0, 1.0]), vtype=GRB.CONTINUOUS, name='T_ij_neg', lb=0, ub=values[index]) model.update() # emulate abs for non-binary comparisons: E_ij = T_ij_pos - T_ij_neg model.addConstrs((E_ij[ij] == T_ij_pos[ij] - T_ij_neg[ij] for ij in T_ij_pos), 'E_ij = T_ij_pos - T_ij_neg') # hard decision of >= and <=: z_ij == 1 <-> i > j model.addConstrs((E_ij[ij] + comparisons[ij] - 0.5 >= -1 + z_ij for ij, z_ij in Z_ij.items()), 'z_ij_upper_bound') model.addConstrs((E_ij[ij] + comparisons[ij] - 0.5 <= z_ij for ij, z_ij in Z_ij.items()), 'z_ij_lower_bound') # transitivity for (i, j), a in Z_ij.items(): for k in nodes: j_, k_ = j, k if j > k: j_, k_ = k, j b = Z_ij.get((j_, k_), None) if b is None: continue elif j_ != j: b = 1 - b i_, k_ = i, k if i > k: i_, k_ = k, i c = Z_ij.get((i_, k_), None) if c is None: continue elif i_ != i: c = 1 - c # a <= b and b <= c -> a <= c model.addLConstr(a + b, GRB.LESS_EQUAL, 1 + c, f'transitivity_ge_{i},{j},{k}') # a >= b and b >= c -> a >= c model.addLConstr(a + b, GRB.GREATER_EQUAL, c, f'transitivity_le_{i},{j},{k}') # transitivity helper (for not-fully connected graphs) # also provides a latent rank big_m = len(nodes) model.addConstrs(((1 - z_ij) * big_m + R_i[i] >= R_i[j] + 1 for (i, j), z_ij in Z_ij.items()), 'rank_transitivity_larger') model.addConstrs( (z_ij * big_m + R_i[j] >= R_i[i] + 1 for (i, j), z_ij in Z_ij.items()), 'rank_transitivity_smaller') # objective function objective = LinExpr() for ij, value in comparisons.items(): if value == 1.0: objective += -E_ij[ij] elif value == 0.0: objective += E_ij[ij] else: objective += T_ij_pos[ij] + T_ij_neg[ij] model.setObjective(objective, GRB.MINIMIZE) # solve model.optimize() # verify abs emulation: one T_ij has to be 0 for ij, value in T_ij_pos.items(): assert value.X == 0 or T_ij_neg[ij] == 0, \ f'T_{ij} pos {value.X} neg {T_ij_neg[ij]}' # find minimal Rs model_ = Model('comparison') model_.setParam('OutputFlag', verbose) R_i = model_.addVars(nodes, name='r_i', vtype=GRB.CONTINUOUS, lb=0, ub=len(nodes)) for (i, j), z_ij in Z_ij.items(): if z_ij.x == 1: model_.addConstr(R_i[i] >= R_i[j] + 1) else: model_.addConstr(R_i[j] >= R_i[i] + 1) model_.setObjective(R_i.sum(), GRB.MINIMIZE) model_.optimize() return [model_.getVarByName(f'r_i[{i}]').X for i in range(len(nodes))], \ model.objVal
class ILPSolver: def __init__( self, g: DFGraph, budget: int, eps_noise=None, seed_s=None, integral=True, imposed_schedule: ImposedSchedule = ImposedSchedule.FULL_SCHEDULE, solve_r=True, write_model_file: Optional[PathLike] = None, gurobi_params: Dict[str, Any] = None): self.GRB_CONSTRAINED_PRESOLVE_TIME_LIMIT = 300 # todo (paras): read this from gurobi_params self.gurobi_params = gurobi_params self.num_threads = self.gurobi_params.get('Threads', 1) self.model_file = write_model_file self.seed_s = seed_s self.integral = integral self.imposed_schedule = imposed_schedule self.solve_r = solve_r self.eps_noise = eps_noise self.budget = budget self.g: DFGraph = g self.solve_time = None if not self.integral: assert not self.solve_r, "Can't solve for R if producing a fractional solution" self.init_constraints = [] # used for seeding the model self.m = Model("checkpointmip_gc_{}_{}".format(self.g.size, self.budget)) if gurobi_params is not None: for k, v in gurobi_params.items(): setattr(self.m.Params, k, v) T = self.g.size self.ram_gcd = self.g.ram_gcd(self.budget) if self.integral: self.R = self.m.addVars(T, T, name="R", vtype=GRB.BINARY) self.S = self.m.addVars(T, T, name="S", vtype=GRB.BINARY) self.Free_E = self.m.addVars(T, len(self.g.edge_list), name="FREE_E", vtype=GRB.BINARY) else: self.R = self.m.addVars(T, T, name="R", vtype=GRB.CONTINUOUS, lb=0.0, ub=1.0) self.S = self.m.addVars(T, T, name="S", vtype=GRB.CONTINUOUS, lb=0.0, ub=1.0) self.Free_E = self.m.addVars(T, len(self.g.edge_list), name="FREE_E", vtype=GRB.CONTINUOUS, lb=0.0, ub=1.0) self.U = self.m.addVars(T, T, name="U", lb=0.0, ub=float(budget) / self.ram_gcd) for x in range(T): for y in range(T): self.m.addLConstr(self.U[x, y], GRB.GREATER_EQUAL, 0) self.m.addLConstr(self.U[x, y], GRB.LESS_EQUAL, float(budget) / self.ram_gcd) def build_model(self): T = self.g.size dict_val_div = lambda cost_dict, divisor: { k: v / divisor for k, v in cost_dict.items() } permute_ram = dict_val_div(self.g.cost_ram, self.ram_gcd) budget = self.budget / self.ram_gcd permute_eps = lambda cost_dict, eps: { k: v * (1. + eps * np.random.randn()) for k, v in cost_dict.items() } permute_cpu = dict_val_div(self.g.cost_cpu, self.g.cpu_gcd()) if self.eps_noise: permute_cpu = permute_eps(permute_cpu, self.eps_noise) with Timer("Gurobi model construction", extra_data={ 'T': str(T), 'budget': str(budget) }): with Timer("Objective construction", extra_data={ 'T': str(T), 'budget': str(budget) }): # seed solver with a baseline strategy if self.seed_s is not None: for x in range(T): for y in range(T): if self.seed_s[x, y] < 1: self.init_constraints.append( self.m.addLConstr(self.S[x, y], GRB.EQUAL, 0)) self.m.update() # define objective function self.m.setObjective( quicksum(self.R[t, i] * permute_cpu[i] for t in range(T) for i in range(T)), GRB.MINIMIZE) with Timer("Variable initialization", extra_data={ 'T': str(T), 'budget': str(budget) }): if self.imposed_schedule == ImposedSchedule.FULL_SCHEDULE: self.m.addLConstr( quicksum(self.R[t, i] for t in range(T) for i in range(t + 1, T)), GRB.EQUAL, 0) self.m.addLConstr( quicksum(self.S[t, i] for t in range(T) for i in range(t, T)), GRB.EQUAL, 0) self.m.addLConstr(quicksum(self.R[t, t] for t in range(T)), GRB.EQUAL, T) elif self.imposed_schedule == ImposedSchedule.COVER_ALL_NODES: self.m.addLConstr(quicksum(self.S[0, i] for i in range(T)), GRB.EQUAL, 0) for i in range(T): self.m.addLConstr( quicksum(self.R[t, i] for t in range(T)), GRB.GREATER_EQUAL, 1) elif self.imposed_schedule == ImposedSchedule.COVER_LAST_NODE: self.m.addLConstr(quicksum(self.S[0, i] for i in range(T)), GRB.EQUAL, 0) # note: the integrality gap is very large as this constraint # is only applied to the last node (last column of self.R). self.m.addLConstr( quicksum(self.R[t, T - 1] for t in range(T)), GRB.GREATER_EQUAL, 1) with Timer("Correctness constraints", extra_data={ 'T': str(T), 'budget': str(budget) }): # ensure all checkpoints are in memory for t in range(T - 1): for i in range(T): self.m.addLConstr(self.S[t + 1, i], GRB.LESS_EQUAL, self.S[t, i] + self.R[t, i]) # ensure all computations are possible for (u, v) in self.g.edge_list: for t in range(T): self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL, self.R[t, u] + self.S[t, u]) # define memory constraints def _num_hazards(t, i, k): if t + 1 < T: return 1 - self.R[t, k] + self.S[t + 1, i] + quicksum( self.R[t, j] for j in self.g.successors(i) if j > k) return 1 - self.R[t, k] + quicksum( self.R[t, j] for j in self.g.successors(i) if j > k) def _max_num_hazards(t, i, k): num_uses_after_k = sum(1 for j in self.g.successors(i) if j > k) if t + 1 < T: return 2 + num_uses_after_k return 1 + num_uses_after_k with Timer("Constraint: upper bound for 1 - Free_E", extra_data={ 'T': str(T), 'budget': str(budget) }): for t in range(T): for eidx, (i, k) in enumerate(self.g.edge_list): self.m.addLConstr(1 - self.Free_E[t, eidx], GRB.LESS_EQUAL, _num_hazards(t, i, k)) with Timer("Constraint: lower bound for 1 - Free_E", extra_data={ 'T': str(T), 'budget': str(budget) }): for t in range(T): for eidx, (i, k) in enumerate(self.g.edge_list): self.m.addLConstr( _max_num_hazards(t, i, k) * (1 - self.Free_E[t, eidx]), GRB.GREATER_EQUAL, _num_hazards(t, i, k)) with Timer( "Constraint: initialize memory usage (includes spurious checkpoints)", extra_data={ 'T': str(T), 'budget': str(budget) }): for t in range(T): self.m.addLConstr( self.U[t, 0], GRB.EQUAL, self.R[t, 0] * permute_ram[0] + quicksum(self.S[t, i] * permute_ram[i] for i in range(T))) with Timer("Constraint: memory recurrence", extra_data={ 'T': str(T), 'budget': str(budget) }): for t in range(T): for k in range(T - 1): mem_freed = quicksum( permute_ram[i] * self.Free_E[t, eidx] for (eidx, i) in self.g.predecessors_indexed(k)) self.m.addLConstr( self.U[t, k + 1], GRB.EQUAL, self.U[t, k] + self.R[t, k + 1] * permute_ram[k + 1] - mem_freed) if self.model_file is not None and self.g.size < 200: # skip for big models to save runtime with Timer("Saving model", extra_data={ 'T': str(T), 'budget': str(budget) }): self.m.write(self.model_file) return None # return value ensures ray remote call can be chained def solve(self): T = self.g.size with Timer('Gurobi model optimization', extra_data={ 'T': str(T), 'budget': str(self.budget) }): if self.seed_s is not None: self.m.Params.TimeLimit = self.GRB_CONSTRAINED_PRESOLVE_TIME_LIMIT self.m.optimize() if self.m.status == GRB.INFEASIBLE: print(f"Infeasible ILP seed at budget {self.budget:.2E}") self.m.remove(self.init_constraints) self.m.Params.TimeLimit = self.gurobi_params.get('TimeLimit', 0) self.m.message("\n\nRestarting solve\n\n") with Timer("ILPSolve") as solve_ilp: self.m.optimize() self.solve_time = solve_ilp.elapsed infeasible = (self.m.status == GRB.INFEASIBLE) if infeasible: raise ValueError( "Infeasible model, check constraints carefully. Insufficient memory?" ) if self.m.solCount < 1: raise ValueError( f"Model status is {self.m.status} (not infeasible), but solCount is {self.m.solCount}" ) Rout = np.zeros((T, T), dtype=remat.core.utils.solver_common.SOLVER_DTYPE if self.integral else np.float) Sout = np.zeros((T, T), dtype=remat.core.utils.solver_common.SOLVER_DTYPE if self.integral else np.float) Uout = np.zeros((T, T), dtype=remat.core.utils.solver_common.SOLVER_DTYPE if self.integral else np.float) Free_Eout = np.zeros((T, len(self.g.edge_list)), dtype=remat.core.utils.solver_common.SOLVER_DTYPE) solver_dtype_cast = int if self.integral else float try: for t in range(T): for i in range(T): try: Rout[t][i] = solver_dtype_cast(self.R[t, i].X) except (AttributeError, TypeError) as e: Rout[t][i] = solver_dtype_cast(self.R[t, i]) try: Sout[t][i] = solver_dtype_cast(self.S[t, i]) except (AttributeError, TypeError) as e: Sout[t][i] = solver_dtype_cast(self.S[t, i].X) try: Uout[t][i] = self.U[t, i].X * self.ram_gcd except (AttributeError, TypeError) as e: Uout[t][i] = self.U[t, i] * self.ram_gcd for e in range(len(self.g.edge_list)): try: Free_Eout[t][e] = solver_dtype_cast(self.Free_E[t, e].X) except (AttributeError, TypeError) as e: Free_Eout[t][e] = solver_dtype_cast(self.Free_E[t, e]) except AttributeError as e: logging.exception(e) return None, None, None, None # prune R using closed-form solver if self.solve_r and self.integral: Rout = solve_r_opt(self.g, Sout) return Rout, Sout, Uout, Free_Eout
class MaxBatchILPSolver: def __init__( self, g: DFGraph, budget: int, eps_noise=None, model_file=None, remote=False, gurobi_params: Dict[str, Any] = None, cpu_fwd_factor: int = 2, ): self.cpu_fwd_factor = cpu_fwd_factor self.logger = logging.getLogger(__name__) self.remote = remote self.profiler = functools.partial(Timer, print_results=True) self.gurobi_params = gurobi_params self.num_threads = self.gurobi_params.get("Threads", 1) self.model_file = model_file self.eps_noise = eps_noise self.budget = budget self.g = g self.solve_time = None self.init_constraints = [] # used for seeding the model self.m = Model("checkpointmip_gc_maxbs_{}_{}".format(self.g.size, self.budget)) if gurobi_params is not None: for k, v in gurobi_params.items(): setattr(self.m.Params, k, v) T = self.g.size CPU_VALS = list(self.g.cost_cpu.values()) RAM_VALS = list(self.g.cost_ram.values()) self.logger.info( "RAM: [{:.2E}, {:.2E}], {:.2E} +- {:.2E}".format( np.min(RAM_VALS), np.max(RAM_VALS), np.mean(RAM_VALS), np.std(RAM_VALS) ) ) self.logger.info( "CPU: [{:.2E}, {:.2E}], {:.2E} +- {:.2E}".format( np.min(CPU_VALS), np.max(CPU_VALS), np.mean(CPU_VALS), np.std(CPU_VALS) ) ) self.ram_gcd = int(max(self.g.ram_gcd(self.budget), 1)) self.cpu_gcd = int(max(self.g.cpu_gcd(), 1)) self.logger.info("ram_gcd = {} cpu_gcd = {}".format(self.ram_gcd, self.cpu_gcd)) self.batch_size = self.m.addVar(lb=1, ub=1024 * 16, name="batch_size") self.R = self.m.addVars(T, T, name="R", vtype=GRB.BINARY) self.S = self.m.addVars(T, T, name="S", vtype=GRB.BINARY) self.Free_E = self.m.addVars(T, len(self.g.edge_list), name="FREE_E", vtype=GRB.BINARY) self.U = self.m.addVars(T, T, name="U", lb=0.0, ub=self.budget) for x in range(T): for y in range(T): self.m.addLConstr(self.U[x, y], GRB.GREATER_EQUAL, 0) self.m.addLConstr(self.U[x, y], GRB.LESS_EQUAL, float(budget) / self.ram_gcd) def build_model(self): T = self.g.size dict_val_div = lambda cost_dict, divisor: {k: np.ceil(v / divisor) for k, v in cost_dict.items()} permute_ram = dict_val_div(self.g.cost_ram, self.ram_gcd) budget = self.budget / self.ram_gcd permute_eps = lambda cost_dict, eps: {k: v * (1.0 + eps * np.random.randn()) for k, v in cost_dict.items()} permute_cpu = dict_val_div(self.g.cost_cpu, self.cpu_gcd) if self.eps_noise: permute_cpu = permute_eps(permute_cpu, self.eps_noise) with self.profiler("Gurobi model construction", extra_data={"T": str(T), "budget": str(budget)}): with self.profiler("Objective construction", extra_data={"T": str(T), "budget": str(budget)}): # define objective function self.m.setObjective(self.batch_size, GRB.MAXIMIZE) with self.profiler("Variable initialization", extra_data={"T": str(T), "budget": str(budget)}): self.m.addLConstr(quicksum(self.R[t, i] for t in range(T) for i in range(t + 1, T)), GRB.EQUAL, 0) self.m.addLConstr(quicksum(self.S[t, i] for t in range(T) for i in range(t, T)), GRB.EQUAL, 0) self.m.addLConstr(quicksum(self.R[t, t] for t in range(T)), GRB.EQUAL, T) with self.profiler("Correctness constraints", extra_data={"T": str(T), "budget": str(budget)}): # ensure all checkpoints are in memory for t in range(T - 1): for i in range(T): self.m.addLConstr(self.S[t + 1, i], GRB.LESS_EQUAL, self.S[t, i] + self.R[t, i]) # ensure all computations are possible for (u, v) in self.g.edge_list: for t in range(T): self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL, self.R[t, u] + self.S[t, u]) # define memory constraints def _num_hazards(t, i, k): if t + 1 < T: return 1 - self.R[t, k] + self.S[t + 1, i] + quicksum(self.R[t, j] for j in self.g.successors(i) if j > k) return 1 - self.R[t, k] + quicksum(self.R[t, j] for j in self.g.successors(i) if j > k) def _max_num_hazards(t, i, k): num_uses_after_k = sum(1 for j in self.g.successors(i) if j > k) if t + 1 < T: return 2 + num_uses_after_k return 1 + num_uses_after_k with self.profiler("Constraint: upper bound for 1 - Free_E", extra_data={"T": str(T), "budget": str(budget)}): for t in range(T): for eidx, (i, k) in enumerate(self.g.edge_list): self.m.addLConstr(1 - self.Free_E[t, eidx], GRB.LESS_EQUAL, _num_hazards(t, i, k)) with self.profiler("Constraint: lower bound for 1 - Free_E", extra_data={"T": str(T), "budget": str(budget)}): for t in range(T): for eidx, (i, k) in enumerate(self.g.edge_list): self.m.addLConstr( _max_num_hazards(t, i, k) * (1 - self.Free_E[t, eidx]), GRB.GREATER_EQUAL, _num_hazards(t, i, k) ) with self.profiler( "Constraint: initialize memory usage (includes spurious checkpoints)", extra_data={"T": str(T), "budget": str(budget)}, ): for t in range(T): self.m.addConstr( self.U[t, 0] == self.batch_size * (self.R[t, 0] * permute_ram[0] + quicksum(self.S[t, i] * permute_ram[i] for i in range(T))), name="init_mem", ) with self.profiler("Constraint: memory recurrence", extra_data={"T": str(T), "budget": str(budget)}): for t in range(T): for k in range(T - 1): mem_freed = quicksum( permute_ram[i] * self.Free_E[t, eidx] for (eidx, i) in self.g.predecessors_indexed(k) ) self.m.addConstr( self.U[t, k + 1] == self.U[t, k] + self.batch_size * (self.R[t, k + 1] * permute_ram[k + 1] - mem_freed), name="update_mem", ) if self.cpu_fwd_factor: with self.profiler("Constraint: recomputation overhead"): compute_fwd = sum([permute_cpu[i] for i in self.g.vfwd]) bwd_compute = sum([permute_cpu[i] for i in self.g.v if i not in self.g.vfwd]) max_mem = self.cpu_fwd_factor * compute_fwd + bwd_compute self.logger.info("Solver using compute overhead ceiling of {}".format(max_mem)) self.m.addConstr( quicksum(self.R[t, i] * permute_cpu[i] for t in range(T) for i in range(T)) <= max_mem, name="limit_cpu", ) else: self.logger.info("No compute limit") with self.profiler("Model update"): self.m.update() if self.model_file is not None and self.g.size < 200: # skip for big models to save runtime with self.profiler("Saving model", extra_data={"T": str(T), "budget": str(budget)}): self.m.write(self.model_file) return None # return value ensures ray remote call can be chained def solve(self): T = self.g.size with self.profiler("Gurobi model optimization", extra_data={"T": str(T), "budget": str(self.budget)}): with Timer("ILPSolve") as solve_ilp: self.m.optimize() self.solve_time = solve_ilp.elapsed infeasible = self.m.status == GRB.INFEASIBLE try: _ = self.R[0, 0].X _ = self.S[0, 0].X _ = self.U[0, 0].X _ = self.batch_size.X except AttributeError as e: infeasible = True if infeasible: raise ValueError("Infeasible model, check constraints carefully. Insufficient memory?") Rout = np.zeros((T, T), dtype=SOLVER_DTYPE) Sout = np.zeros((T, T), dtype=SOLVER_DTYPE) Uout = np.zeros((T, T), dtype=SOLVER_DTYPE) Free_Eout = np.zeros((T, len(self.g.edge_list)), dtype=SOLVER_DTYPE) batch_size = self.batch_size.X try: for t in range(T): for i in range(T): Rout[t][i] = int(self.R[t, i].X) Sout[t][i] = int(self.S[t, i].X) Uout[t][i] = self.U[t, i].X * self.ram_gcd for e in range(len(self.g.edge_list)): Free_Eout[t][e] = int(self.Free_E[t, e].X) except AttributeError as e: logging.exception(e) return None, None, None, None Rout = solve_r_opt(self.g, Sout) # prune R using optimal recomputation solver ilp_aux_data = ILPAuxData( U=Uout, Free_E=Free_Eout, ilp_approx=False, ilp_time_limit=0, ilp_eps_noise=0, ilp_num_constraints=self.m.numConstrs, ilp_num_variables=self.m.numVars, ) schedule, aux_data = schedule_from_rs(self.g, Rout, Sout) return ( ScheduledResult( solve_strategy=SolveStrategy.OPTIMAL_ILP_GC, solver_budget=self.budget, feasible=True, schedule=schedule, schedule_aux_data=aux_data, solve_time_s=self.solve_time, ilp_aux_data=ilp_aux_data, ), batch_size, )
class ILPSolver: def __init__(self, si: SolverInfo, budget, gurobi_params: Dict[str, Any] = None, ablation=False, overhead=False): self.gurobi_params = gurobi_params self.num_threads = self.gurobi_params.get('Threads', 1) self.budget = int(budget * MEM_GCD_MULTIPLIER * GB_TO_KB) self.si: SolverInfo = si self.solve_time = None self.ablation = ablation self.overhead = overhead V = self.si.loss + 1 T = len(self.si.nodes) - self.si.loss Y = 3 budget = self.budget self.m = Model("monet{}".format(self.budget)) if gurobi_params is not None: for k, v in gurobi_params.items(): setattr(self.m.Params, k, v) self.ram = np.array([math.ceil(self.si.nodes[i].mem*MEM_GCD_MULTIPLIER/1024) for i in self.si.nodes]) # Convert to KB self.cpu = dict(( i, [math.ceil(val*CPU_GCD_MULTIPLIER) for val in self.si.nodes[i].workspace_compute] ) for i in self.si.nodes) self.cpu_recompute = dict(( i, [math.ceil(val*CPU_GCD_MULTIPLIER) for val in self.si.nodes[i].recompute_workspace_compute] ) for i in self.si.nodes) self.cpu_inplace = dict(( i, [math.ceil(val*CPU_GCD_MULTIPLIER) for val in self.si.nodes[i].inplace_workspace_compute] ) for i in self.si.nodes) self.R = self.m.addVars(T, V, name="R", vtype=GRB.BINARY) # Recomputation self.P = self.m.addVars(T, V, name="P", vtype=GRB.BINARY) # In-memory self.S = self.m.addVars(T+1, V, name="S", vtype=GRB.BINARY) # Stored self.M = self.m.addVars(T, Y, name="M", vtype=GRB.BINARY) # Backward operator implementation self.SM = self.m.addVars(T, V, Y, name="SM", vtype=GRB.BINARY) # Linearization of S * M if self.si.select_conv_algo: self.RF = self.m.addVars(T, len(self.si.conv_list), self.si.num_conv_algos, name="RF", vtype=GRB.BINARY) # Conv operator implementations if self.si.do_inplace: self.IP = self.m.addVars(T, V, name="IP", vtype=GRB.BINARY) # In-place def cmem(self, j, recompute=False, inplace=False): if inplace: return [math.ceil(val * MEM_GCD_MULTIPLIER * 1024) for val in self.si.nodes[j].inplace_workspace_mem] else: if recompute: return [math.ceil(val * MEM_GCD_MULTIPLIER * 1024) for val in self.si.nodes[j].recompute_workspace_mem] else: return [math.ceil(val * MEM_GCD_MULTIPLIER * 1024) for val in self.si.nodes[j].workspace_mem] def local_memory(self, j): return math.ceil(self.si.nodes[j].local_memory * MEM_GCD_MULTIPLIER / 1024) def fixed_ram(self, j): return math.ceil(self.si.nodes[j].fixed_mem * MEM_GCD_MULTIPLIER / 1024) def build_model(self): V = self.si.loss + 1 T = len(self.si.nodes) - self.si.loss Y = 3 budget = self.budget # define objective function if self.si.select_conv_algo: fwd_compute = quicksum(self.R[0, i] * self.cpu[i][0] for i in range(V) if i not in self.si.conv_list) fwd_recompute = quicksum(quicksum(self.R[t, i]*self.cpu_recompute[i][0] for t in range(1,T)) for i in range(V) if i not in self.si.conv_list) conv_fwd_compute = quicksum(quicksum(self.RF[t,self.si.conv_list[i],c] * self.cpu[i][c] for t in range(T)) for i in self.si.conv_list if i<=self.si.loss for c in range(len(self.cpu[i]))) # conv's compute and recompute same if self.si.do_inplace: inplace_fwd_compute = quicksum(quicksum(self.IP[t,i] * (self.cpu_inplace[i][0]-self.cpu[i][0]) for t in range(T)) for i in self.si.inplace_list) # relu's compute and recompute same conv_bwd_compute = quicksum(quicksum(self.RF[t,self.si.conv_list[i],c] * self.cpu[i][c] for t in range(T)) for i in self.si.conv_list if i>self.si.loss for c in range(len(self.cpu[i]))) # conv's compute and recompute same bwd_compute = quicksum(self.M[t, p] * self.cpu[t+V-1][p] for t in range(T) for p in range(Y) if (p<len(self.cpu[t+V-1]) and (t+V-1) not in self.si.conv_list)) if self.si.do_inplace: compute_fwd = fwd_compute + fwd_recompute + conv_fwd_compute + inplace_fwd_compute else: compute_fwd = fwd_compute + fwd_recompute + conv_fwd_compute compute_bwd = conv_bwd_compute + bwd_compute compute_cost = compute_fwd + compute_bwd else: fwd_compute = quicksum(self.R[0, i] * self.cpu[i][0] for i in range(V)) fwd_recompute = quicksum(self.R[t, i] * self.cpu_recompute[i][0] for t in range(1,T) for i in range(V)) if self.si.do_inplace: inplace_fwd_compute = quicksum(quicksum(self.IP[t,i] * (self.cpu_inplace[i][0]-self.cpu[i][0]) for t in range(T)) for i in self.si.inplace_list) # relu's compute and recompute same bwd_compute = quicksum(self.M[t, p] * self.cpu[t+V-1][p] for t in range(T) for p in range(Y) if p<len(self.cpu[t+V-1])) if self.si.do_inplace: compute_fwd = fwd_compute + fwd_recompute + inplace_fwd_compute else: compute_fwd = fwd_compute + fwd_recompute compute_bwd = bwd_compute compute_cost = compute_fwd + compute_bwd self.m.setObjective(compute_cost, GRB.MINIMIZE) self.compute_cost = compute_cost self.compute_fwd = compute_fwd self.compute_bwd = compute_bwd # Add in-place constraints for t in range(T): for v in range(V): if self.si.do_inplace: self.m.addLConstr(self.P[t,v], GRB.LESS_EQUAL, self.R[t,v]) self.m.addLConstr(self.IP[t,v], GRB.LESS_EQUAL, self.R[t,v]) # print(list(self.si.inplace_list.keys())) if v not in self.si.inplace_list: self.m.addLConstr(self.IP[t,v], GRB.EQUAL, 0) elif not isinstance(self.si.nodes[v], IntNode): # If IP[v], then P[u] = 0, else P[u] = R[u] self.m.addLConstr(self.P[t,self.si.inplace_list[v]], GRB.GREATER_EQUAL, self.R[t,self.si.inplace_list[v]] - 2*self.IP[t,v]) self.m.addLConstr(self.P[t,self.si.inplace_list[v]], GRB.LESS_EQUAL, 2 - 2*self.IP[t,v]) self.m.addLConstr(self.S[t+1,self.si.inplace_list[v]], GRB.LESS_EQUAL, 2 - 2*self.IP[t,v]) else: self.m.addLConstr(self.P[t,v], GRB.EQUAL, self.R[t,v]) # store nothing in the beginning for i in range(V): self.m.addLConstr(self.S[0,i] , GRB.EQUAL, 0) # Recompute full forward for i in range(V): if not isinstance(self.si.nodes[i], IntNode): self.m.addLConstr(self.R[0,i] , GRB.EQUAL, 1) # All M which have no paths are set as 0 self.m.addLConstr(self.M[0,0], GRB.EQUAL, 1) for p in range(1,Y): self.m.addLConstr(self.M[0,p], GRB.EQUAL, 0) for t in range(1, T): bkwd_t = V + t - 1 paths = len(self.si.nodes[bkwd_t].args) for p in range(paths, Y): self.m.addLConstr(self.M[t,p], GRB.EQUAL, 0) # Set failed conv's RF to be 0 if self.si.select_conv_algo: for i in self.si.conv_list: for c in range(self.si.num_conv_algos): if c >= len(self.cpu[i]) or self.cpu[i][c] < 0: for t in range(T): self.m.addLConstr(self.RF[t,self.si.conv_list[i],c], GRB.EQUAL, 0) # create constraints for boolean multiplication linearization for p in range(Y): for i in range(V): for t in range(T): self.m.addLConstr(self.SM[t,i, p], GRB.GREATER_EQUAL, self.M[t,p] + self.S[t+1,i] - 1) self.m.addLConstr(self.SM[t,i, p], GRB.LESS_EQUAL, self.M[t,p]) self.m.addLConstr(self.SM[t,i, p], GRB.LESS_EQUAL, self.S[t+1,i]) if self.ablation: print("Doing ablation") # Disable all recomputation for t in range(1,T): for v in range(V): self.m.addLConstr(self.R[t,v], GRB.EQUAL, 0) # Fix all operations to be inplace for t in range(T): for v in range(V): if self.si.do_inplace and v in self.si.inplace_list: self.m.addLConstr(self.IP[t,v], GRB.EQUAL, self.R[t,v]) # Correctness constraints # Ensure all checkpoints are in memory for t in range(T): for i in range(V): self.m.addLConstr(self.S[t + 1, i], GRB.LESS_EQUAL, self.S[t, i] + self.P[t, i]) # At least one path should be used for t in range(T): self.m.addLConstr(quicksum(self.M[t,p] for p in range(Y)), GRB.GREATER_EQUAL, 1) # Ensure all computations are possible for (u, v, p) in self.si.edge_list: if u < V and v <V: for t in range(T): self.m.addLConstr(self.R[t, v], GRB.LESS_EQUAL, self.R[t, u] + self.S[t, u]) if u < V and v >= V: t = v + 1 - V self.m.addLConstr(self.M[t,p], GRB.LESS_EQUAL, self.P[t, u] + self.S[t+1, u]) # Ensure that new nodes only computed if main node computed for t in range(T): for i in range(V): if self.si.nodes[i].has_intermediates: for (_,intid) in self.si.nodes[i].intermediates: self.m.addLConstr(self.R[t, intid], GRB.LESS_EQUAL, self.R[t,i]) if self.si.do_inplace and i in self.si.inplace_list: self.m.addLConstr(self.IP[:,intid], GRB.GREATER_EQUAL, self.IP[:,i] + self.R[:,intid] -1) self.m.addLConstr(self.IP[:,intid], GRB.LESS_EQUAL, self.IP[:,i]) # Constraints for conv selection if self.si.select_conv_algo: for i in self.si.conv_list: if i < self.si.loss: for t in range(T): self.m.addLConstr(quicksum(self.RF[t,self.si.conv_list[i],c] for c in range(self.si.num_conv_algos)), GRB.GREATER_EQUAL, self.R[t,i]) else: t_bwdi = i + 1 - V self.m.addLConstr(quicksum(self.RF[t_bwdi,self.si.conv_list[i],c] for c in range(self.si.num_conv_algos)), GRB.GREATER_EQUAL, 1) for c in range(self.si.num_conv_algos): for t in range(t_bwdi): self.m.addLConstr(self.RF[t,self.si.conv_list[i],c], GRB.EQUAL, 0) for t in range(t_bwdi+1, T): self.m.addLConstr(self.RF[t,self.si.conv_list[i],c], GRB.EQUAL, 0) # Memory constraints for t in range(T): bkwd_t = V + t - 1 bwd_node = self.si.nodes[bkwd_t] for t in range(T): bkwd_t = V + t - 1 bwd_node = self.si.nodes[bkwd_t] bwd_local_memory = self.local_memory(bkwd_t) if t!=0 else 0 gradm = bwd_local_memory + self.fixed_ram(bkwd_t-1) for j in range(V): keep_tensors = [[fwdin for fwdin in bwd_node.args[p] if fwdin <= self.si.loss] for p in range(len(bwd_node.args))] # Forward checkpoint constraint self.m.addLConstr( gradm + quicksum(self.S[t,i]*self.ram[i] for i in range(j, V)) + quicksum(self.S[t+1,i]*self.ram[i] for i in range(j) if t<T-1) + quicksum( quicksum( self.M[t,p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ) - quicksum( quicksum( self.SM[t,i, p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ), GRB.LESS_EQUAL, budget) workspace_mem = self.cmem(j) if t==0 else self.cmem(j, recompute=True, inplace=False) if self.si.select_conv_algo and j in self.si.conv_list: wm = quicksum(self.RF[t,self.si.conv_list[j],c] * workspace_mem[c] for c in range(len(workspace_mem))) elif self.si.do_inplace and j in self.si.inplace_list: inplace_workspace_mem = self.cmem(j, recompute=False, inplace=True) wm = self.IP[t,j]*(inplace_workspace_mem[0] - workspace_mem[0]) + self.R[t,j] * workspace_mem[0] else: wm = self.R[t,j] * workspace_mem[0] # Forward recomputation constraint self.m.addLConstr( gradm + wm + self.R[t,j]*self.ram[j] + self.R[t,j]*self.local_memory(j) + quicksum(self.S[t,i]*self.ram[i] for i in range(j+1, V)) + quicksum(self.S[t+ 1,i]*self.ram[i] for i in range(j) if i not in self.si.nodes[j].local_tensors) + quicksum( quicksum( self.M[t,p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ) - quicksum( quicksum( self.SM[t,i, p]*self.ram[i] for i in range(j) if i in keep_tensors[p]) for p in range(len(bwd_node.args)) ), GRB.LESS_EQUAL, budget) if t>0: if self.si.select_conv_algo and bkwd_t in self.si.conv_list: wmem = self.cmem(bkwd_t) bwd_wm = quicksum(self.RF[t,self.si.conv_list[bkwd_t],c] * wmem[c] for c in range(len(wmem))) else: wmem = self.cmem(bkwd_t) bwd_wm = quicksum(self.M[t,p]*wmem[p] for p in range(len(bwd_node.args))) rammem = [sum([self.ram[i] for i in bwd_node.args[p] if i<V]) for p in range(len(bwd_node.args))] # Backward constraint self.m.addLConstr( bwd_wm + bwd_local_memory + self.fixed_ram(bkwd_t) + self.ram[bkwd_t] + quicksum(self.M[t,p]*rammem[p] for p in range(len(bwd_node.args))) + quicksum(quicksum(self.SM[t,i,p]*self.ram[i] for i in range(V) if i not in bwd_node.args[p]) for p in range(len(bwd_node.args))), GRB.LESS_EQUAL, budget) return None def solve(self): from time import time V = self.si.loss + 1 T = len(self.si.nodes) - self.si.loss Y = 3 budget = self.budget self.m.Params.TimeLimit = self.time_limit t0 = time() self.m.optimize() self.solve_time = time() - t0 infeasible = (self.m.status == GRB.INFEASIBLE) if infeasible: self.m.computeIIS() self.m.write("model.ilp") raise ValueError("Infeasible model, check constraints carefully. Insufficient memory?") if self.m.solCount < 1: raise ValueError(f"Model status is {self.m.status} (not infeasible), but solCount is {self.m.solCount}") Rout = np.zeros((T, V), dtype=bool) Pout = np.zeros((T, V), dtype=bool) Sout = np.zeros((T+1, V), dtype=bool) Mout = np.zeros((T, Y), dtype=bool) SMout = np.zeros((T, V, Y), dtype=bool) RFout = np.zeros((T, len(self.si.conv_list), self.si.num_conv_algos), dtype=bool) IPout = np.zeros((T, V), dtype=bool) try: for t in range(T): for i in range(V): Rout[t][i] = round(self.R[t, i].X) Pout[t][i] = round(self.P[t, i].X) Sout[t][i] = round(self.S[t, i].X) if self.si.do_inplace: IPout[t][i] = round(self.IP[t, i].X) for p in range(Y): SMout[t][i][p] = round(self.SM[t, i, p].X) for p in range(Y): Mout[t][p] = round(self.M[t, p].X) if self.si.select_conv_algo: for i in range(len(self.si.conv_list)): for c in range(self.si.num_conv_algos): RFout[t][i][c] = round(self.RF[t,i,c].X) for i in range(V): Sout[T][i] = round(self.S[T, i].X) except AttributeError as e: logging.exception(e) return None, None, None, None, None, None, None solution = Solution(Rout, Sout, Mout, RFout, IPout, Pout, self.solve_time, -1, -1, -1) return solution
def run_ilp(tint, remaining_rids, incomp_rids, ilp_settings, log_prefix): # Variables directly based on the input ------------------------------------ # I[i,j] = 1 if reads[i]['data'][j]==1 and 0 if reads[i]['data'][j]==0 or 2 # C[i,j] = 1 if exon j is between the first and last exons (inclusively) # covered by read i and is not in read i but can be turned into a 1 ISOFORM_INDEX_START = 1 M = len(tint['segs']) MAX_ISOFORM_LG = sum(seg[2] for seg in tint['segs']) I = tint['ilp_data']['I'] C = tint['ilp_data']['C'] INCOMP_READ_PAIRS = incomp_rids GARBAGE_COST = tint['ilp_data']['garbage_cost'] informative = informative_segs(tint, remaining_rids) # ILP model ------------------------------------------------------ ILP_ISOFORMS = Model('isoforms_v8_20210209') ILP_ISOFORMS.setParam('OutputFlag', 0) ILP_ISOFORMS.setParam(GRB.Param.Threads, ilp_settings['threads']) # Decision variables # R2I[i,k] = 1 if read i assigned to isoform k R2I = {} R2I_C1 = {} # Constraint enforcing that each read is assigned to exactly one isoform for i in remaining_rids: R2I[i] = {} for k in range(ilp_settings['K']): R2I[i][k] = ILP_ISOFORMS.addVar( vtype=GRB.BINARY, name='R2I[{i}][{k}]'.format(i=i, k=k) ) R2I_C1[i] = ILP_ISOFORMS.addLConstr( lhs=quicksum(R2I[i][k] for k in range(0, ilp_settings['K'])), sense=GRB.EQUAL, rhs=1, name='R2I_C1[{i}]'.format(i=i) ) # Implied variable: canonical exons presence in isoforms # E2I[j,k] = 1 if canonical exon j is in isoform k # E2I_min[j,k] = 1 if canonical exon j is in isoform k and is shared by all reads of that isoform # E2IR[j,k,i] = 1 if read i assigned to isoform k AND exon j covered by read i # Auxiliary variable # E2IR[j,k,i] = R2I[i,k] AND I[i,j] # E2I[j,k] = max over all reads i of E2IR[j,k,i] # E2I_min[j,k] = min over all reads i of E2IR[j,k,i] E2I = {} E2I_C1 = {} E2I_min = {} E2I_min_C1 = {} E2IR = {} E2IR_C1 = {} for j in range(0, M): if not informative[j]: continue E2I[j] = {} E2I_C1[j] = {} E2I_min[j] = {} E2I_min_C1[j] = {} E2IR[j] = {} E2IR_C1[j] = {} # No exon is assignd to the garbage isoform E2I[j][0] = ILP_ISOFORMS.addVar( vtype=GRB.BINARY, name='E2I[{j}][{k}]'.format(j=j, k=0) ) E2I_C1[j][0] = ILP_ISOFORMS.addLConstr( lhs=E2I[j][0], sense=GRB.EQUAL, rhs=0, name='E2I_C1[{j}][{k}]'.format(j=j, k=0) ) # We start assigning exons from the first isoform for k in range(ISOFORM_INDEX_START, ilp_settings['K']): E2I[j][k] = ILP_ISOFORMS.addVar( vtype=GRB.BINARY, name='E2I[{j}][{k}]'.format(j=j, k=k) ) E2I_min[j][k] = ILP_ISOFORMS.addVar( vtype=GRB.BINARY, name='E2I_min[{j}][{k}]'.format(j=j, k=k) ) E2IR[j][k] = {} E2IR_C1[j][k] = {} for i in remaining_rids: E2IR[j][k][i] = ILP_ISOFORMS.addVar( vtype=GRB.BINARY, name='E2IR[{j}][{k}][{i}]'.format(j=j, k=k, i=i) ) E2IR_C1[j][k][i] = ILP_ISOFORMS.addLConstr( lhs=E2IR[j][k][i], sense=GRB.EQUAL, rhs=R2I[i][k]*I[i][j], name='E2IR_C1[{j}][{k}][{i}]'.format(j=j, k=k, i=i) ) E2I_C1[j][k] = ILP_ISOFORMS.addGenConstrMax( resvar=E2I[j][k], vars=[E2IR[j][k][i] for i in remaining_rids], constant=0.0, name='E2I_C1[{j}][{k}]'.format(j=j, k=k) ) E2I_min_C1[j][k] = ILP_ISOFORMS.addGenConstrMin( resvar=E2I_min[j][k], vars=[E2IR[j][k][i] for i in remaining_rids], constant=0.0, name='E2I_min_C1[{j}][{k}]'.format(j=j, k=k) ) # Adding constraints for unaligned gaps # If read i is assigned to isoform k, and reads[i]['gaps'] contains ((j1,j2),l), and # the sum of the lengths of exons in isoform k between exons j1 and j2 is L # then (1-EPSILON)L <= l <= (1+EPSILON)L # GAPI[(j1,j2,k)] = sum of the length of the exons between exons j1 and j2 (inclusively) in isoform k GAPI = {} GAPI_C1 = {} # Constraint fixing the value of GAPI GAPR_C1 = {} # Constraint ensuring that the unaligned gap is not too short for every isoform and gap GAPR_C2 = {} # Constraint ensuring that the unaligned gap is not too long for every isoform and gap for i in remaining_rids: for ((j1, j2), l) in tint['reads'][tint['read_reps'][i][0]]['gaps'].items(): # No such constraint on the garbage isoform if any for k in range(ISOFORM_INDEX_START, ilp_settings['K']): if not (j1, j2, k) in GAPI: assert informative[j1 % M] assert informative[j2 % M] assert not any(informative[j+1:j2]) GAPI[(j1, j2, k)] = ILP_ISOFORMS.addVar( vtype=GRB.INTEGER, name='GAPI[({j1},{j2},{k})]'.format(j1=j1, j2=j2, k=k) ) GAPI_C1[(j1, j2, k)] = ILP_ISOFORMS.addLConstr( lhs=GAPI[(j1, j2, k)], sense=GRB.EQUAL, rhs=quicksum(E2I[j][k]*tint['segs'][j][2] for j in range(j1+1, j2) if informative[j]), name='GAPI_C1[({j1},{j2},{k})]'.format( j1=j1, j2=j2, k=k) ) GAPR_C1[(i, j1, j2, k)] = ILP_ISOFORMS.addLConstr( lhs=(1.0-ilp_settings['epsilon'])*GAPI[(j1, j2, k)] - ilp_settings['offset']-((1-R2I[i][k])*MAX_ISOFORM_LG), sense=GRB.LESS_EQUAL, rhs=l, name='GAPR_C1[({i},{j1},{j2},{k})]'.format( i=i, j1=j1, j2=j2, k=k) ) GAPR_C2[(i, j1, j2, k)] = ILP_ISOFORMS.addLConstr( lhs=(1.0+ilp_settings['epsilon'])*GAPI[(j1, j2, k)] + ilp_settings['offset']+((1-R2I[i][k])*MAX_ISOFORM_LG), sense=GRB.GREATER_EQUAL, rhs=l, name='GAPR_C2[({i},{j1},{j2},{k})]'.format( i=i, j1=j1, j2=j2, k=k) ) # Adding constraints for incompatible read pairs INCOMP_READ_PAIRS_C1 = {} for (i1, i2) in INCOMP_READ_PAIRS: if not (i1 in remaining_rids and i2 in remaining_rids): continue # Again, no such constraint on the garbage isoform if any for k in range(ISOFORM_INDEX_START, ilp_settings['K']): INCOMP_READ_PAIRS_C1[(i1, i2, k)] = ILP_ISOFORMS.addLConstr( lhs=R2I[i1][k]+R2I[i2][k], sense=GRB.LESS_EQUAL, rhs=1, name='INCOMP_READ_PAIRS_C1[({i1},{i2},{k})]'.format( i1=i1, i2=i2, k=k) ) # [OPTIONAL] Labeling non-garbage isoforms by their exon content and forcing them to occur in increasing label order # LABEL_I = {} # LABEL_I_C1 = {} # LABEL_I_C2 = {} # for k in range(ISOFORM_INDEX_START,ilp_settings['K']): # LABEL_I[k] = ILP_ISOFORMS.addVar( # vtype = GRB.INTEGER, # name = 'LABEL_I[{k}]'.format(k=k) # ) # LABEL_I_C1[k] = ILP_ISOFORMS.addLConstr( # lhs = LABEL_I[k], # sense = GRB.EQUAL, # rhs = quicksum(E2I[j][k]*(2**j) for j in range(0,M)), # name = 'LABEL_I_C1[{k}]'.format(k =k) # ) # if k > ISOFORM_INDEX_START: # LABEL_I_C2[k] = ILP_ISOFORMS.addLConstr( # lhs = LABEL_I[k], # sense = GRB.LESS_EQUAL, # rhs = LABEL_I[k-1]-0.1, # name = 'LABEL_I_C2[{k}]'.format(k=k) # ) # Objective function # For i,j,k such that i ∈ remaining_rids, C[i,j]=1 (read has a zero that can be # corrected), and E2I[j,k]=1 (isoform k has exon j), OBJ[i][j][k] = 1 OBJ = {} OBJ_C1 = {} OBJ_SUM = LinExpr(0.0) for i in remaining_rids: OBJ[i] = {} OBJ_C1[i] = {} for j in range(0, M): if not informative[j]: continue if C[i][j] > 0: # 1 if exon j not in read i but can be added to it OBJ[i][j] = {} OBJ_C1[i][j] = {} for k in range(ISOFORM_INDEX_START, ilp_settings['K']): OBJ[i][j][k] = ILP_ISOFORMS.addVar( vtype=GRB.BINARY, name='OBJ[{i}][{j}][{k}]'.format(i=i, j=j, k=k) ) OBJ_C1[i][j][k] = ILP_ISOFORMS.addGenConstrAnd( resvar=OBJ[i][j][k], vars=[R2I[i][k], E2I[j][k]], name='OBJ_C1[{i}][{j}][{k}]'.format(i=i, j=j, k=k) ) OBJ_SUM.addTerms(1.0*C[i][j], OBJ[i][j][k]) # coeffs = 1.0, # vars = OBJ[i][j][k] # ) # We add the chosen cost for each isoform assigned to the garbage isoform if any GAR_OBJ = {} GAR_OBJ_C = {} for i in remaining_rids: if ilp_settings['recycle_model'] in ['constant', 'exons', 'introns']: OBJ_SUM.addTerms(1.0*GARBAGE_COST[i], R2I[i][0]) elif ilp_settings['recycle_model'] == 'relative': GAR_OBJ[i] = {} GAR_OBJ_C[i] = {} for j in range(0, M): if not informative[j]: continue GAR_OBJ[i][j] = {} GAR_OBJ_C[i][j] = {} for k in range(ISOFORM_INDEX_START, ilp_settings['K']): if I[i][j] == 1: GAR_OBJ[i][j][k] = ILP_ISOFORMS.addVar( vtype=GRB.BINARY, name='GAR_OBJ[{i}][{j}][{k}]'.format(i=i, j=j, k=k) ) GAR_OBJ_C[i][j][k] = ILP_ISOFORMS.addGenConstrAnd( resvar=GAR_OBJ[i][j][k], vars=[R2I[i][0], E2I_min[j][k]], name='GAR_OBJ_C[{i}][{j}][{k}]'.format( i=i, j=j, k=k) ) OBJ_SUM.addTerms(1.0, GAR_OBJ[i][j][k]) elif I[i][j] == 0 and C[i][j] == 1: pass ILP_ISOFORMS.setObjective( expr=OBJ_SUM, sense=GRB.MINIMIZE ) # Optimization # ILP_ISOFORMS.Params.PoolSearchMode=2 # ILP_ISOFORMS.Params.PoolSolutions=5 ILP_ISOFORMS.setParam('TuneOutput', 1) if not log_prefix == None: ILP_ISOFORMS.setParam('LogFile', '{}.glog'.format(log_prefix)) ILP_ISOFORMS.write('{}.lp'.format(log_prefix)) ILP_ISOFORMS.setParam('TimeLimit', ilp_settings['timeout']*60) ILP_ISOFORMS.optimize() ILP_ISOFORMS_STATUS = ILP_ISOFORMS.Status isoforms = {k: dict() for k in range(ISOFORM_INDEX_START, ilp_settings['K'])} # print('STATUS: {}'.format(ILP_ISOFORMS_STATUS)) # if ILP_ISOFORMS_STATUS == GRB.Status.TIME_LIMIT: # status = 'TIME_LIMIT' if ILP_ISOFORMS_STATUS != GRB.Status.OPTIMAL: status = 'NO_SOLUTION' else: status = 'OPTIMAL' # Writing the optimal solution to disk if not log_prefix == None: solution_file = open('{}.sol'.format(log_prefix), 'w+') for v in ILP_ISOFORMS.getVars(): solution_file.write('{}\t{}\n'.format(v.VarName, v.X)) solution_file.close() # Isoform id to isoform structure for k in range(ISOFORM_INDEX_START, ilp_settings['K']): isoforms[k]['exons'] = list() for j in range(0, M): if informative[j]: isoforms[k]['exons'].append( int(E2I[j][k].getAttr(GRB.Attr.X) > 0.9)) else: isoforms[k]['exons'].append( I[next(iter(remaining_rids))][j]) isoforms[k]['rid_to_corrections'] = dict() # Isoform id to read ids set for i in remaining_rids: isoform_id = -1 for k in range(0, ilp_settings['K']): if R2I[i][k].getAttr(GRB.Attr.X) > 0.9: assert isoform_id == - \ 1, 'Read {} has been assigned to multiple isoforms!'.format( i) isoform_id = k assert isoform_id != - \ 1, 'Read {} has not been assigned to any isoform!'.format(i) if isoform_id == 0: continue isoforms[isoform_id]['rid_to_corrections'][i] = -1 # Read id to its exon corrections for k in range(ISOFORM_INDEX_START, ilp_settings['K']): for i in isoforms[k]['rid_to_corrections'].keys(): isoforms[k]['rid_to_corrections'][i] = [ str(tint['reads'][tint['read_reps'][i][0]]['data'][j]) for j in range(M)] for j in range(0, M): if not informative[j]: isoforms[k]['rid_to_corrections'][i][j] = '-' elif C[i][j] == 1 and OBJ[i][j][k].getAttr(GRB.Attr.X) > 0.9: isoforms[k]['rid_to_corrections'][i][j] = 'X' return ILP_ISOFORMS_STATUS, status, isoforms