Basic constraint programming modeling

basicModelingWeek3.py

from ortools.sat.python import cp_model

''' course: https://www.coursera.org/learn/basic-modeling#about '''


def luBuAssignmentProblem(heroList, spotList, damageList):
    model = cp_model.CpModel()
    numOfSpots = len(spotList)
    numOfRows = len(damageList)
    numOfCols = len(damageList[0])
    maxIndex = numOfRows * numOfCols
    damageListFlattened = sum(damageList, [])
    maxDamage = max(damageListFlattened)

    positionVarList = []
    for hero in heroList:
        positionVarList.append(model.NewIntVar(0, numOfSpots-1, name=hero))

    model.AddAllDifferent(positionVarList)

    indexVarList = []
    damageVarList = []
    for i in range(numOfRows):
        indexVarList.append(model.NewIntVar(0, maxIndex, name="intermediateVar"+str(i)))
        damageVarList.append(model.NewIntVar(0, maxDamage, name="damage"+str(i)))

    for i in range(numOfRows):
        model.Add(indexVarList[i] == (i * numOfCols) + (positionVarList[i]))
        model.AddElement(indexVarList[i], damageListFlattened, damageVarList[i])

    model.Maximize(sum(damageVarList))

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Damages: %i' % solver.ObjectiveValue(), "\n")

        for index, positionVar in enumerate(positionVarList):
            print(heroList[index], ":", spotList[solver.Value(positionVar)])

    print("----------------------------")




#----------------------------------------------------------------------------------------------------------------------#



def cellBlockProblem(costList, prisoners, dangerousPrisoners, femalePrisoners):

    model = cp_model.CpModel()
    malePrisoners = list(set(dangerousPrisoners) - set(femalePrisoners))
    numOfRows = len(costList)
    numOfCols = len(costList[0])
    costListFlattened = sum(costList, [])
    maxCost = max(costListFlattened)


    rowVarList = []
    colVarList = []
    for prisoner in prisoners:
        rowVarList.append(model.NewIntVar(0, numOfRows-1, name="rowForPrisoner"+str(prisoner)))
        colVarList.append(model.NewIntVar(0, numOfCols-1, name="colForPrisoner"+str(prisoner)))


    indexVarList = []
    costVarList = []
    for prisoner in prisoners:
        indexVar = model.NewIntVar(0, numOfRows*numOfCols, name="index"+str(prisoner))
        model.Add(indexVar == rowVarList[prisoner] * numOfCols + colVarList[prisoner])
        indexVarList.append(indexVar)
        costVarList.append(model.NewIntVar(0, maxCost, name="cost"+str(prisoner)))
    model.AddAllDifferent(indexVarList)


    for prisoner in prisoners:
        for dangerPrisoner in dangerousPrisoners:
            if prisoner != dangerPrisoner:
                absRowDiffVar = model.NewIntVar(0, numOfRows-1, name="absRowDiff"+str(prisoner)+str(dangerPrisoner))
                absColDiffVar = model.NewIntVar(0, numOfCols-1, name="absColDiff"+str(prisoner)+str(dangerPrisoner))
                rowDiffVar = model.NewIntVar(0, numOfRows-1, name="rowDiff"+str(prisoner)+str(dangerPrisoner))
                colDiffVar = model.NewIntVar(0, numOfCols-1, name="colDiff"+str(prisoner)+str(dangerPrisoner))

                model.Add(rowDiffVar == rowVarList[prisoner] - rowVarList[dangerPrisoner])
                model.Add(colDiffVar == colVarList[prisoner] - colVarList[dangerPrisoner])
                model.AddAbsEquality(absRowDiffVar, rowDiffVar)
                model.AddAbsEquality(absColDiffVar, colDiffVar)
                model.Add(absRowDiffVar + absColDiffVar > 1)


    for femalePrisoner in femalePrisoners:
        model.Add(rowVarList[femalePrisoner] <= (numOfRows-1)//2)

    for malePrisoner in malePrisoners:
        model.Add(rowVarList[malePrisoner] >= (numOfRows-1)//2 + 1)

    for prisoner in prisoners:
        model.AddElement(indexVarList[prisoner], costListFlattened, costVarList[prisoner])
    model.Maximize(sum(costVarList))


    '''rowSol = [1,3,4,4,4,2,1,3,4,2]
    colSol = [1,1,6,4,2,4,5,5,1,2]
    for i in range(len(rowSol)):
        model.Add(rowVarList[i] == rowSol[i]-1)
        model.Add(colVarList[i] == colSol[i]-1)'''

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Total cost: %i' % solver.ObjectiveValue(), "\n")

        for prisoner in prisoners:
            rowValue = solver.Value(rowVarList[prisoner])
            colValue = solver.Value(colVarList[prisoner])
            print("P"+str(prisoner), "in", "cell["+str(rowValue)+","+str(colValue)+"]")
    print("----------------------------")





#----------------------------------------------------------------------------------------------------------------------#


def getIndex(row, col, numOfCols):
    return (row * numOfCols) + col



def shiftPatrollingProblem(numOfSoldiers, numOfDays, numRequiredForNightShift, lbForEveningShift, ubForEveningShift):

    model = cp_model.CpModel()
    DAY_SHIFT = 0
    EVENING_SHIFT = 1
    NIGHT_SHIFT = 2

    rosterVarList = []
    workNightBoolVarList = []
    workEveBoolVarList = []
    for soldier in range(numOfSoldiers):
        for day in range(numOfDays):
            idx = getIndex(soldier, day, numOfDays)
            rosterVarList.append(model.NewIntVar(DAY_SHIFT, NIGHT_SHIFT, "soldier"+str(soldier)+"day"+str(day)))
            workNightBoolVarList.append(model.NewBoolVar("workingNightSoldier"+str(soldier)+"day"+str(day)))
            workEveBoolVarList.append(model.NewBoolVar("workingEveSoldier"+str(soldier)+"day"+str(day)))

            model.Add(rosterVarList[idx] == NIGHT_SHIFT).OnlyEnforceIf(workNightBoolVarList[idx])
            model.Add(rosterVarList[idx] != NIGHT_SHIFT).OnlyEnforceIf(workNightBoolVarList[idx].Not())
            model.Add(rosterVarList[idx] == EVENING_SHIFT).OnlyEnforceIf(workEveBoolVarList[idx])
            model.Add(rosterVarList[idx] != EVENING_SHIFT).OnlyEnforceIf(workEveBoolVarList[idx].Not())


    for soldier in range(numOfSoldiers):
        for day in range(numOfDays-2):
            idx = getIndex(soldier, day, numOfDays)
            idx1 = getIndex(soldier, day+1, numOfDays)
            idx2 = getIndex(soldier, day+2, numOfDays)
            working2NightsInARowBoolVar = model.NewBoolVar("working2NightsInARow" + str(soldier) + str(day))
            model.AddBoolOr([working2NightsInARowBoolVar, workNightBoolVarList[idx].Not(), workNightBoolVarList[idx1].Not()])
            model.AddImplication(working2NightsInARowBoolVar, workNightBoolVarList[idx])
            model.AddImplication(working2NightsInARowBoolVar, workNightBoolVarList[idx1])
            model.AddImplication(working2NightsInARowBoolVar, workNightBoolVarList[idx2].Not())
            # https://groups.google.com/forum/#!topic/or-tools-discuss/9trLOMSe_DA, https://activimetrics.com/blog/ortools/cp_sat/addboolor/


    for soldier in range(numOfSoldiers):
        for day in range(numOfDays-1):
            idx = getIndex(soldier, day, numOfDays)
            idx1 = getIndex(soldier, day + 1, numOfDays)
            model.AddImplication(workEveBoolVarList[idx], workNightBoolVarList[idx1].Not())


    for day in range(numOfDays):
        varCoeffListForNightWorkers = []
        varCoeffListForEveWorkers = []
        for soldier in range(numOfSoldiers):
            idx = getIndex(soldier, day, numOfDays)
            varCoeffListForNightWorkers.append(workNightBoolVarList[idx])
            varCoeffListForEveWorkers.append(workEveBoolVarList[idx])

        model.Add(sum(varCoeffListForNightWorkers) == numRequiredForNightShift)
        model.AddLinearConstraint(sum(varCoeffListForEveWorkers), lbForEveningShift, ubForEveningShift)


    model.Maximize(sum(workEveBoolVarList))

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Total cost: %i' % solver.ObjectiveValue(), "\n")

        for soldier in range(numOfSoldiers):
            for day in range(numOfDays):
                idx = getIndex(soldier, day, numOfDays)
                shift = solver.Value(rosterVarList[idx])
                if shift == DAY_SHIFT:
                    print("Soldier"+str(soldier)+" day"+str(day)+" -> DAY SHIFT")
                elif shift == EVENING_SHIFT:
                    print("Soldier"+str(soldier)+" day"+str(day)+" -> EVENING SHIFT")
                elif shift == NIGHT_SHIFT:
                    print("Soldier" + str(soldier) + " day" + str(day) + " -> NIGHT SHIFT")
            print("``````")

    print("----------------------------")





#----------------------------------------------------------------------------------------------------------------------#




def catapultClusteringProblem(spotList, distanceList, maxSeparation, lbUbForClusters):

    model = cp_model.CpModel()
    numOfSpots = len(spotList)
    numOfClusters = len(lbUbForClusters)
    maxDistance = max(sum(distanceList, []))

    shotVarList = []
    for spot in spotList:
        shotVarList.append(model.NewIntVar(0, numOfClusters-1, "Spot"+str(spot)))

    spotIndexVarDict = {}
    for cluster in range(numOfClusters):
        sameClusterVarList = []
        for i, shotVar in enumerate(shotVarList):
            spotIndexVarDict[(i, cluster)] = model.NewIntVar(0, numOfSpots-1, "Index"+str(i)+str(cluster))
            shotVarEqualsCluster = model.NewBoolVar(name="ShotVarEqualsCluster" + str(i) + str(cluster))
            model.Add(shotVar == cluster).OnlyEnforceIf(shotVarEqualsCluster)
            model.Add(shotVar != cluster).OnlyEnforceIf(shotVarEqualsCluster.Not())
            model.Add(spotIndexVarDict[(i, cluster)] == i).OnlyEnforceIf(shotVarEqualsCluster)
            model.Add(spotIndexVarDict[(i, cluster)] == numOfSpots - 1).OnlyEnforceIf(shotVarEqualsCluster.Not())
            sameClusterVarList.append(shotVarEqualsCluster)
        model.AddLinearConstraint(sum(sameClusterVarList), lbUbForClusters[cluster][0], lbUbForClusters[cluster][1])


    minSpotOfClusterVarList = []
    for cluster in range(numOfClusters):
        minSpotOfClusterVarList.append(model.NewIntVar(0, numOfSpots-1, "MinSpotOfClusterVar"+str(cluster)))

    for cluster in range(numOfClusters):
        varList = []
        for i, shotVar in enumerate(shotVarList):
            varList.append(spotIndexVarDict[(i, cluster)])
        model.AddMinEquality(minSpotOfClusterVarList[cluster], varList)

    # removing value symmetry
    for cluster in range(numOfClusters - 1):
        model.Add(minSpotOfClusterVarList[cluster] <= minSpotOfClusterVarList[cluster + 1])


    distBetweenClusterVarList = []
    for i in range(numOfSpots-1):
        for j in range(i+1, numOfSpots):
            distBetweenClusterVarList.append(model.NewIntVar(0, maxDistance, "Distance"+str(i)+str(j)))
            withInClusterVar = model.NewBoolVar(name="WithInClusterVar"+str(i)+str(j))
            model.Add(shotVarList[i] == shotVarList[j]).OnlyEnforceIf(withInClusterVar)
            model.Add(distBetweenClusterVarList[-1] == maxDistance).OnlyEnforceIf(withInClusterVar)
            model.Add(distanceList[i][j] <= maxSeparation).OnlyEnforceIf(withInClusterVar) # max distance within cluster <= maxSeparation
            model.Add(shotVarList[i] != shotVarList[j]).OnlyEnforceIf(withInClusterVar.Not())
            model.Add(distBetweenClusterVarList[-1] == distanceList[i][j]).OnlyEnforceIf(withInClusterVar.Not())


    objVar = model.NewIntVar(0, maxDistance, "ObjectiveVar")
    model.AddMinEquality(objVar, distBetweenClusterVarList)
    model.Maximize(objVar) # maximize min distance between clusters

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Minimum distance between clusters = %i\n' % solver.ObjectiveValue())
        print(["Spot"+str(i)+"|"+"Cluster"+str(solver.Value(shotVarList[i])) for i in range(numOfSpots)])
    print("----------------------------")




#----------------------------------------------------------------------------------------------------------------------#




if __name__ == '__main__':

    luBuAssignmentProblem(heroList=["Liu Bei", "Guan Yu", "Zhang Fei"],
                          spotList= ["Head", "Chest", "Abdomen", "Hand", "Feet"],
                          damageList= [[7, 1, 3, 4, 6], [8, 2, 5, 1, 4], [4, 3, 7, 2, 5]])

    cellBlockProblem(costList=[[2, 4, 3, 5, 1, 6], [6, 2, 3, 1, 5, 4], [3, 5, 6, 4, 2, 1], [1, 2, 6, 4, 5, 3]],
                     prisoners=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
                     dangerousPrisoners=[],
                     femalePrisoners=[0, 5, 6, 9])

    shiftPatrollingProblem(numOfSoldiers=6,
                           numOfDays=7,
                           numRequiredForNightShift=3,
                           lbForEveningShift=1,
                           ubForEveningShift=2)

    catapultClusteringProblem(spotList=[1,2,3,4,5,6,7,8],
                              distanceList=[[0,2,3,1,2,4,1,2], [1,0,3,2,4,2,3,1], [1,2,0,1,2,3,2,1],
                                            [3,1,2,0,3,1,2,1], [1,1,2,3,0,2,1,3], [1,2,3,2,2,0,4,3],
                                            [1,1,2,3,2,3,0,2], [1,2,3,3,4,4,2,0]],
                              maxSeparation=3,
                              lbUbForClusters=[(2, 8),(1, 8),(1, 8)])