Basic constraint programming modeling

basicModelingWeek4.py

from ortools.sat.python import cp_model

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


def cookingWineProblem():

    foodList = ["CHILIFISHHEAD", "MAPOTOFU", "SNAKESOUP", "GONGBAOFROG"]
    foodDict = {"CHILIFISHHEAD": 0, "MAPOTOFU": 1, "SNAKESOUP": 2, "GONGBAOFROG": 3}
    tasteList = [5, 2, 4, 2]
    wineList = ["HUADIAO", "GRAPE", "RICE", "GAOLIANG"]
    wineDict = {"HUADIAO": 0, "GRAPE": 1, "RICE": 2, "GAOLIANG": 3}
    alcoholList = [5, 8, 2, 4]
    joyList = [[6, 1, 3, 4], [8, 2, 5, 1], [4, 3, 7, 2], [3, 1, 6, 3]]

    numOfFoods = len(foodDict)
    numOfWines = len(wineDict)
    model = cp_model.CpModel()

    maxIndex = numOfFoods * numOfWines
    joyListFlattened = sum(joyList,[])
    maxJoy = max(joyListFlattened)
    maxTaste = max(tasteList)
    maxAlcohol = max(alcoholList)

    drinkVarList = []
    for food in foodList:
        drinkVarList.append(model.NewIntVar(0, numOfWines-1, str(food)))
    model.AddAllDifferent(drinkVarList)

    eatVarList = []
    for wine in wineList:
        eatVarList.append(model.NewIntVar(0, numOfFoods-1, str(wine)))

    model.AddInverse(drinkVarList, eatVarList)

    tasteVarForGrape = model.NewIntVar(0, maxTaste, "tasteForGrape")
    tasteVarForGaoLiang = model.NewIntVar(0, maxTaste, "tasteForGaoLiang")
    model.AddElement(eatVarList[wineDict["GRAPE"]], tasteList, tasteVarForGrape)
    model.AddElement(eatVarList[wineDict["GAOLIANG"]], tasteList, tasteVarForGaoLiang)
    model.Add(tasteVarForGrape < tasteVarForGaoLiang)

    alcoholVarForSnakeSoup = model.NewIntVar(0, maxAlcohol, "alcoholForSnakeSoup")
    alcoholVarForGongbaOFrog = model.NewIntVar(0, maxAlcohol, "alcoholForGongbaOFrog")
    model.AddElement(drinkVarList[foodDict["SNAKESOUP"]], alcoholList, alcoholVarForSnakeSoup)
    model.AddElement(drinkVarList[foodDict["GONGBAOFROG"]], alcoholList, alcoholVarForGongbaOFrog)
    model.Add(alcoholVarForSnakeSoup < alcoholVarForGongbaOFrog)

    isRiceWithMapotofuVar = model.NewBoolVar("isRiceWithMapotofu")
    model.Add(eatVarList[wineDict["RICE"]] == foodDict["MAPOTOFU"]).OnlyEnforceIf(isRiceWithMapotofuVar)
    model.Add(eatVarList[wineDict["HUADIAO"]] == foodDict["SNAKESOUP"]).OnlyEnforceIf(isRiceWithMapotofuVar)
    model.Add(eatVarList[wineDict["RICE"]] != foodDict["MAPOTOFU"]).OnlyEnforceIf(isRiceWithMapotofuVar.Not())
    model.Add(eatVarList[wineDict["HUADIAO"]] != foodDict["SNAKESOUP"]).OnlyEnforceIf(isRiceWithMapotofuVar.Not())

    indexVarList = []
    joyVarList = []
    for i in range(numOfFoods):
        indexVarList.append(model.NewIntVar(0, maxIndex, name="index"+str(i)))
        joyVarList.append(model.NewIntVar(0, maxJoy, name="joy"+str(i)))

    for i in range(numOfFoods):
        model.Add(indexVarList[i] == (i * numOfWines) + (drinkVarList[i]))
        model.AddElement(indexVarList[i], joyListFlattened, joyVarList[i])
    model.Maximize(sum(joyVarList))

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

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

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

        for index, food in enumerate(foodList):
            print("Pair Food "+str(food)+" with Wine "+str(wineList[solver.Value(drinkVarList[index])]))

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




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




def theBeltProblem(numOfDigits, numOfCopies):

    model = cp_model.CpModel()
    sizeOfBelt = numOfDigits * numOfCopies

    positionVarList = []
    positionDict = {}
    digitCopyIntegerDict = {} # digit copy to a single integer: 1'1, 1'2, 2'1, 2'2 => 1, 2, 3, 4
    counter = 0
    for digit in range(1, numOfDigits+1):
        for copy in range(1, numOfCopies+1):
            positionVarList.append(model.NewIntVar(0, sizeOfBelt-1, "PositionFor"+str(digit)+str(copy)))

            positionDict[(digit, copy)] = counter
            digitCopyIntegerDict[counter] = str(digit)+"'"+str(copy)
            counter += 1


    digitCopyVarList = []
    for position in range(sizeOfBelt):
        digitCopyVarList.append(model.NewIntVar(0, sizeOfBelt-1, "DigitCopyFor"+str(position)))


    model.AddInverse(positionVarList, digitCopyVarList)


    for digit in range(1, numOfDigits+1):
        for copy in range(1, numOfCopies):
            model.Add(positionVarList[positionDict[(digit, copy+1)]] == positionVarList[positionDict[(digit, copy)]] + digit + 1)


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

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

    if status == cp_model.OPTIMAL:
        print("Belt: ")
        for position in range(sizeOfBelt):
            digitCopyInteger = solver.Value(digitCopyVarList[position])
            print(digitCopyIntegerDict[digitCopyInteger], end=" | ")

    print("\n----------------------------")




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




def theTunnelProblem(startPoint, pivotList, coordinateList, pickupList, deliverList):

    model = cp_model.CpModel()
    numOfPrecendes = len(pickupList)
    numOfPivots = len(pivotList)
    maxCoordinate = max(coordinateList)

    orderVarList = []
    routeVarList = []
    for pivot in pivotList:
        orderVarList.append(model.NewIntVar(0, numOfPivots-1, "positionFor"+pivot)) # for each pivot, which position it has
        routeVarList.append(model.NewIntVar(0, numOfPivots-1, "routeFor" + pivot))  # for each position, which pivot it has


    model.Add(routeVarList[0] == startPoint)  # model.Add(orderVarList[4] == 0
    model.AddInverse(orderVarList, routeVarList)

    for i in range(numOfPrecendes):
        model.Add(orderVarList[pickupList[i]] < orderVarList[deliverList[i]])


    coordinateVarList = []
    for i in range(numOfPivots):
        coordinateVarList.append(model.NewIntVar(-maxCoordinate, maxCoordinate, "CoordinatePivot"+str(i)))
        model.AddElement(routeVarList[i], coordinateList, coordinateVarList[i])


    absDistanceVarList = []
    for i in range(numOfPivots-1):
        absDistanceVarList.append(model.NewIntVar(0, maxCoordinate, "AbsDistanceRoute"+str(i)+"Route"+str(i+1)))
        distanceVar =  model.NewIntVar(-maxCoordinate, maxCoordinate, "DistanceRoute"+str(i)+"Route"+str(i+1))
        model.Add(distanceVar == coordinateVarList[i] - coordinateVarList[i+1])
        model.AddAbsEquality(absDistanceVarList[i], distanceVar)

    model.Minimize(sum(absDistanceVarList))


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

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

    if status == cp_model.OPTIMAL:
        print('Total Distance: %i' % solver.ObjectiveValue(), "\n")
        print("Route:", [pivotList[solver.Value(routeVarList[i])] for i in range(numOfPivots)])

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




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




def theAcupunctureProblem():

    model = cp_model.CpModel()

    spotList = ["DAZHUI", "TAODAO", "SHENZHU", "LINGTAI", "ZHIYANG", "ZHONGSHU", "JIZHONG",
                       "XUANSHU", "YAOYANGGUAN", "YAOYU"]
    ZHIYANG = 4
    ZHONGSHU = 5
    numOfJabs = 3
    numOfStages = 4
    numOfSpots = len(spotList)

    pointVarList = []
    pointIndexDict = {}
    counter = 0
    for stage in range(numOfStages):
        for jab in range(numOfJabs):
            pointVarList.append(model.NewIntVar(0, numOfSpots-1, "Stage"+str(stage)+"Jab"+str(jab)))
            pointIndexDict[(stage, jab)] = counter
            counter += 1


    pointEqualsSpotVarList = []
    pointIndexWithoutSpotDict = {}
    counter = 0
    for stage in range(numOfStages):
        for jab in range(numOfJabs):
            pointIndexWithoutSpotDict[(stage, jab)] = counter
            for spot in range(numOfSpots):
                pointEqualsSpotVarList.append(model.NewBoolVar("stage"+str(stage)+"jab"+str(jab)+"spot"+str(spot)))
                counter += 1


    for stage in range(numOfStages):
        for jab in range(numOfJabs):
            indexVar = model.NewIntVar(0, numOfStages*numOfJabs*numOfSpots, "Index"+str(stage)+str(jab))
            model.Add(indexVar == pointIndexWithoutSpotDict[(stage, jab)] + pointVarList[pointIndexDict[(stage, jab)]])
            model.AddElement(indexVar, pointEqualsSpotVarList, 1)


    # remove symmetry
    for stage in range(numOfStages):
        for jab in range(numOfJabs-1):
            model.Add(pointVarList[pointIndexDict[(stage, jab)]] < pointVarList[pointIndexDict[(stage, jab+1)]])


    for stage in range(numOfStages-1):
        for jab in range(numOfJabs):
            model.Add(pointVarList[pointIndexDict[(stage, jab)]] < pointVarList[pointIndexDict[(stage+1, jab)]])


    for spot in range(numOfSpots):
        varListForSettingBounds = []
        for stage in range(numOfStages):
            for jab in range(numOfJabs):
                index = spot + pointIndexWithoutSpotDict[(stage, jab)]
                varListForSettingBounds.append(pointEqualsSpotVarList[index])

        if spot == ZHIYANG or spot == ZHONGSHU:
            model.AddLinearConstraint(sum(varListForSettingBounds), 0, 1)
        else:
            model.AddLinearConstraint(sum(varListForSettingBounds), 0, 2)


    # for stage in range(numOfStages-1):
    #     varList = []
    #     for jab in range(numOfJabs):
    #         for spot in range(numOfSpots):
    #             varList.append(pointEqualsSpotVarList[spot + pointIndexWithoutSpotDict[(stage, jab)]])
    #             varList.append(pointEqualsSpotVarList[spot + pointIndexWithoutSpotDict[(stage+1, jab)]])
    #             varList.append(pointEqualsSpotVarList[spot + pointIndexWithoutSpotDict[(stage+2, jab)]])
    #     model.AddLinearConstraint(sum(varList), 0, 1)

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

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

    if status == cp_model.OPTIMAL:
        print()
        for stage in range(numOfStages):
            for jab in range(numOfJabs):
                spotIndex = solver.Value(pointVarList[pointIndexDict[(stage, jab)]])
                print("Stage: "+str(stage)+", Jab: "+str(jab)+", Spot: "+str(spotList[spotIndex]+", SpotIndex: "+str(spotIndex)))

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




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




if __name__ == "__main__":

    cookingWineProblem()

    theBeltProblem(numOfDigits=9,
                   numOfCopies=3)

    theTunnelProblem(startPoint= 4,
                     pivotList=["P1", "P2", "P3", "P4", "P5", "P6", "P7", "P8", "P9"],
                     coordinateList=[0, 3, 7, 10, 16, 18, 21, 25, 28],
                     pickupList=[2, 3, 5, 7],
                     deliverList=[7, 6, 1, 5])

    theAcupunctureProblem()