Advanced constraint programming modeling - 2

advancedModeling2Week1.py

from ortools.sat.python import cp_model

''' course: https://www.coursera.org/learn/solving-algorithms-discrete-optimization#about 
About how cp works and intro to solver strategies'''



def patchingThe9HeavensProblem(): #sudoku problem
    model = cp_model.CpModel()

    size = 9
    availableSolution = {(1,0):6, (1,1):7, (1,4):0, (2,4):1, (2,7):2, (3,3):2, (3,4):3, (4,1):5, (4,4):4, (4,7):0,
                         (5,4):5, (6,4):6, (7,0):4, (7,1):3, (7,4):7, (7,5):5, (7,6):8, (7,7):6, (8,4):8}

    cellVarDict = {}
    for i in range(size):
        for j in range(size):
            cellVarDict[(i, j)] = model.NewIntVar(0, size-1, "cellValueAt"+str(i)+str(j))

    for ijTuple, sol in availableSolution.items():
        model.Add(cellVarDict[ijTuple] == sol)

    for i in range(size):
        varList = []
        for j in range(size):
            varList.append(cellVarDict[(i, j)])
        model.AddAllDifferent(varList)

    for j in range(size):
        varList = []
        for i in range(size):
            varList.append(cellVarDict[(i, j)])
        model.AddAllDifferent(varList)

    for i in range(3):
        for j in range(3):
            varList = []
            for s in range(3):
                for t in range(3):
                    varList.append(cellVarDict[(3*i + s, 3*j + t)])
            model.AddAllDifferent(varList)

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

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

    if status == cp_model.OPTIMAL:
        for i in range(size):
            for j in range(size):
                print(solver.Value(cellVarDict[i,j]), end=" | ")
            print()
    print("----------------------------")



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




def buildingHerbGardenProblem():
    model = cp_model.CpModel()

    numOfRows = 10
    numOfCols = 10
    elementList = ["METAL", "WOOD", "WATER", "FIRE", "EARTH"]
    compatibilityList = [(2, 1), (2, 3), (1, 3), (1, 4), (4, 0), (3, 4), (3, 0), (0, 2), (0, 1), (4, 2)]
    numOfElements = len(elementList)

    herbVarDict = {}
    for row in range(numOfRows):
        for col in range(numOfCols):
            herbVarDict[(row, col)] = model.NewIntVar(0, numOfElements-1, "herbFor"+str(row)+str(col))

    for row in range(numOfRows):
        for col in range(numOfCols-1):
            model.AddAllowedAssignments([herbVarDict[(row, col)], herbVarDict[(row, col+1)]], compatibilityList)

    model.Add(herbVarDict[0,0] != 0)

    for col in range(numOfCols):
        varList1 = []
        varList2 = []
        for row in range(numOfRows):
            isMetalVar = model.NewBoolVar("isMetal"+str(row)+str(col))
            isEarthVar = model.NewBoolVar("isEarth"+str(row)+str(col))
            model.Add(herbVarDict[(row, col)] == 0).OnlyEnforceIf(isMetalVar)
            model.Add(herbVarDict[(row, col)] != 0).OnlyEnforceIf(isMetalVar.Not())
            model.Add(herbVarDict[(row, col)] == 4).OnlyEnforceIf(isEarthVar)
            model.Add(herbVarDict[(row, col)] != 4).OnlyEnforceIf(isEarthVar.Not())
            varList1.append(isMetalVar)
            varList2.append(isEarthVar)
        model.AddLinearConstraint(sum(varList1), 1, 2)
        model.AddLinearConstraint(sum(varList2), 1, 2)


    model.AddDecisionStrategy(list(herbVarDict.values()), cp_model.CHOOSE_FIRST, cp_model.SELECT_MIN_VALUE)
    solver = cp_model.CpSolver()
    solver.parameters.search_branching = cp_model.FIXED_SEARCH
    status = solver.Solve(model)

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

    if status == cp_model.OPTIMAL:
        for row in range(numOfRows):
            for col in range(numOfCols):
                print(elementList[solver.Value(herbVarDict[(row,col)])], end=" | ")
            print()
    print("----------------------------")




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




if __name__ == '__main__':

    patchingThe9HeavensProblem()
    buildingHerbGardenProblem()