make multi-task amendment

CDDLeiden · Jun 4, 2024 · f3547f0 · f3547f0
1 parent 1ee41e7
commit f3547f0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/gbmtsplits/split.py b/src/gbmtsplits/split.py
@@ -502,8 +502,8 @@ def _merge_clusters_with_balancing_mapping(
             for m in range(S):
                 for t in range(M):
                     cs = [c for c in range(N) if A[t,c] != 0]
-                    prob += LpAffineExpression([(x[c+m*N],A[t,c]) for c in cs]) - X[m] <= fractional_sizes[m]
-                    prob += LpAffineExpression([(x[c+m*N],A[t,c]) for c in cs]) + X[m] >= fractional_sizes[m]
+                    prob += LpAffineExpression([(x[c+m*N],A[t,c]) for c in cs]) - X[m*M+t] <= fractional_sizes[m]
+                    prob += LpAffineExpression([(x[c+m*N],A[t,c]) for c in cs]) + X[m*M+t] >= fractional_sizes[m]
 
             # Solve the model
             prob.solve(PULP_CBC_CMD(gapAbs = absolute_gap, timeLimit = time_limit_seconds, threads = max_N_threads, msg=False))