Advanced constraint programming modeling - 2

advancedModeling2Week2.py

from ortools.sat.python import cp_model
import random

''' course: https://www.coursera.org/learn/solving-algorithms-discrete-optimization#about 
About solver strategies and how global constraints work'''



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

    sizeOfWindow = 2
    numOfLeavesInEachWindow = 8
    capacityOfNutrients = 19
    nutrientList = [[4,2,3,7], [3,8,20,20], [5,4,6,20], [3,6,2,20]] # 20 (>capacity) for non-existent segments
    leavesToGrowList = [[7,6,4,8], [5,8,0,0], [5,4,7,0], [6,9,3,0]] # 0 for non-existent segments
    numOfPotions = len(nutrientList)
    maxNumOfGrownSegments = len(nutrientList[0])

    nutrientListFlattened = sum(nutrientList, [])
    leavesToGrowListFlattened = sum(leavesToGrowList, [])
    maxLeaves = max(leavesToGrowListFlattened)

    segmentChoiceVarList = []
    for potion in range(numOfPotions):
        segmentChoiceVarList.append(model.NewIntVar(0, maxNumOfGrownSegments-1, "segmentForPotion"+str(potion)))

    nutrientVarList = []
    leaveVarList = []
    for potion in range(numOfPotions):
        nutrientVarList.append(model.NewIntVar(0, 20, "nutrientFor"+str(potion)))
        leaveVarList.append(model.NewIntVar(0, maxLeaves, "leavesFor"+str(potion)))
        indexVar = model.NewIntVar(0, (numOfPotions*maxNumOfGrownSegments)-1, "indexFor"+str(potion))
        model.Add(indexVar == (potion * numOfPotions) + segmentChoiceVarList[potion])
        model.AddElement(indexVar, nutrientListFlattened, nutrientVarList[-1])
        model.AddElement(indexVar, leavesToGrowListFlattened, leaveVarList[-1])

    model.Add(sum(nutrientVarList) <= capacityOfNutrients)

    for tail in range(sizeOfWindow, numOfPotions):
        varList = []
        for i in range(tail-sizeOfWindow, tail):
            varList.append(leaveVarList[i])
        model.Add(sum(varList) >= numOfLeavesInEachWindow)

    model.Maximize(sum(leaveVarList))

    model.AddDecisionStrategy(leaveVarList, cp_model.CHOOSE_HIGHEST_MAX, cp_model.SELECT_MAX_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:
        print('Max number of leaves: %i' % solver.ObjectiveValue())
        print("Segments selected for each potion:", [solver.Value(var) for var in segmentChoiceVarList])
    print("----------------------------")



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



def theArcheryProblem():  # Quasigroup completion problem. Exhibits a heavy tailed phenomenon.
    model = cp_model.CpModel()

    numOfColors = 11
    sizeOfMatrix = 11
    availableSol = {}
    for i in range(int(sizeOfMatrix * 0.3)):  # 30% random pre-assignment
        randRow = random.randint(0, sizeOfMatrix-1)
        randCol = random.randint(0, sizeOfMatrix-1)
        randColor = random.randint(0, numOfColors-1)
        availableSol[(randRow, randCol)] = randColor

    cellVarDict = {}
    for row in range(sizeOfMatrix):
        for col in range(sizeOfMatrix):
            cellVarDict[(row, col)] = model.NewIntVar(0, numOfColors-1, "colorFor"+str(row)+str(col))

    for row in range(sizeOfMatrix):
        for col in range(sizeOfMatrix):
            if (row,col) in availableSol:
                model.Add(cellVarDict[(row, col)] == availableSol[(row,col)])

    for row in range(sizeOfMatrix):
        varList = []
        for col in range(sizeOfMatrix):
            varList.append(cellVarDict[(row, col)])
        model.AddAllDifferent(varList)

    for col in range(sizeOfMatrix):
        varList = []
        for row in range(sizeOfMatrix):
            varList.append(cellVarDict[(row, col)])
        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 row in range(sizeOfMatrix):
            for col in range(sizeOfMatrix):
                print(solver.Value(cellVarDict[row,col]), end=" | ")
            print()
    print("----------------------------")



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




if __name__ == '__main__':

    theYaoCaoModel()
    theArcheryProblem()