<filename>sample/assignment_sat.py
# This sample is a simple rewrite of one of the OR-Tools samples.
# The original file can be found at here.
# https://github.com/google/or-tools/blob/stable/ortools/sat/samples/assignment_sat.py

from ortools.sat.python import cp_model
from orseries.array import CpArray
import pandas as pd

def data_initialize():
    # initialize data (wide)
    # -------------------------------------------
    #           task01  task02  task03  task04
    # worker_1      90      80      75      70
    # worker_2      35      85      55      65
    # worker_3     125      95      90      95
    # worker_4      45     110      95     115
    # worker_5      50     100      90     100
    # -------------------------------------------
    df_pre = pd.DataFrame([ 
            [90, 80, 75, 70], 
            [35, 85, 55, 65], 
            [125, 95, 90, 95], 
            [45, 110, 95, 115], 
            [50, 100, 90, 100], 
        ], 
        index=["worker_1","worker_2","worker_3","worker_4","worker_5"], 
        columns=["task01", "task02", "task03", "task04"], 
    )

    # converts to a common data format (vertical).
    # -------------------------------------------
    #   worker    task  cost
    # worker_1  task01    90
    # worker_1  task02    80
    # worker_1  task03    75
    # worker_1  task04    70
    # worker_2  task01    35
    # worker_2  task02    85
    # ...
    # -------------------------------------------
    df = df_pre.stack().reset_index().rename(
        {"level_0":"worker", "level_1":"task", 0: "cost"}, axis=1
    )

    return df

def main():
    # in a real case, pd.read_csv would be used here.
    df = data_initialize()
    # set index to define a variable and constraints.
    df = df.set_index(["worker", "task"])

    # create empty cp model.
    model = cp_model.CpModel()
    # define the variables for which task each worker will execute.
    df["var_is_assigned"] = CpArray.BooleanVariables(model, df, name="is_assigned")

    # each worker will be assigned at most one task (and possibly zero).
    cons_worker = df["var_is_assigned"].groupby("worker").sum() <= 1
    # each task is always assigned to one of the workers.
    cons_task = df["var_is_assigned"].groupby("task").sum() == 1
    # add these constraint expressions to the model.
    CpArray.AddToModel(model, cons_worker)
    CpArray.AddToModel(model, cons_task)

    # define the objective variable.
    # the objective variable is the sum of the costs for each worker,
    # and the goal is to minimize it.
    objective_cost = df["var_is_assigned"] * df["cost"]
    model.Minimize(objective_cost.sum())

    # and solve it!
    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    # print solution
    if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
        print(f'Total cost = {solver.ObjectiveValue()}\n')
        # Retrieve and filter the tasks that have actually been assigned.
        df["assigned_flag"] = CpArray.ToBooleanValues(solver, df["var_is_assigned"])
        rst_df = df[df["assigned_flag"] == True].reset_index()
        
        for i,row in rst_df.iterrows():
            print('{worker} assigned to {task} Cost = {cost}'.format_map(row))

        return True
    else:
        print('No solution found.')
        return False


if __name__ == '__main__':
    main()
<filename>sample/assignment_sat.py
# This sample is a simple rewrite of one of the OR-Tools samples.
# The original file can be found at here.
# https://github.com/google/or-tools/blob/stable/ortools/sat/samples/assignment_sat.py

from ortools.sat.python import cp_model
from orseries.array import CpArray
import pandas as pd

def data_initialize():
    # initialize data (wide)
    # -------------------------------------------
    #           task01  task02  task03  task04
    # worker_1      90      80      75      70
    # worker_2      35      85      55      65
    # worker_3     125      95      90      95
    # worker_4      45     110      95     115
    # worker_5      50     100      90     100
    # -------------------------------------------
    df_pre = pd.DataFrame([ 
            [90, 80, 75, 70], 
            [35, 85, 55, 65], 
            [125, 95, 90, 95], 
            [45, 110, 95, 115], 
            [50, 100, 90, 100], 
        ], 
        index=["worker_1","worker_2","worker_3","worker_4","worker_5"], 
        columns=["task01", "task02", "task03", "task04"], 
    )

    # converts to a common data format (vertical).
    # -------------------------------------------
    #   worker    task  cost
    # worker_1  task01    90
    # worker_1  task02    80
    # worker_1  task03    75
    # worker_1  task04    70
    # worker_2  task01    35
    # worker_2  task02    85
    # ...
    # -------------------------------------------
    df = df_pre.stack().reset_index().rename(
        {"level_0":"worker", "level_1":"task", 0: "cost"}, axis=1
    )

    return df

def main():
    # in a real case, pd.read_csv would be used here.
    df = data_initialize()
    # set index to define a variable and constraints.
    df = df.set_index(["worker", "task"])

    # create empty cp model.
    model = cp_model.CpModel()
    # define the variables for which task each worker will execute.
    df["var_is_assigned"] = CpArray.BooleanVariables(model, df, name="is_assigned")

    # each worker will be assigned at most one task (and possibly zero).
    cons_worker = df["var_is_assigned"].groupby("worker").sum() <= 1
    # each task is always assigned to one of the workers.
    cons_task = df["var_is_assigned"].groupby("task").sum() == 1
    # add these constraint expressions to the model.
    CpArray.AddToModel(model, cons_worker)
    CpArray.AddToModel(model, cons_task)

    # define the objective variable.
    # the objective variable is the sum of the costs for each worker,
    # and the goal is to minimize it.
    objective_cost = df["var_is_assigned"] * df["cost"]
    model.Minimize(objective_cost.sum())

    # and solve it!
    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    # print solution
    if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
        print(f'Total cost = {solver.ObjectiveValue()}\n')
        # Retrieve and filter the tasks that have actually been assigned.
        df["assigned_flag"] = CpArray.ToBooleanValues(solver, df["var_is_assigned"])
        rst_df = df[df["assigned_flag"] == True].reset_index()
        
        for i,row in rst_df.iterrows():
            print('{worker} assigned to {task} Cost = {cost}'.format_map(row))

        return True
    else:
        print('No solution found.')
        return False


if __name__ == '__main__':
    main()
# This sample is a simple rewrite of one of the OR-Tools samples. # The original file can be found at here. # https://github.com/google/or-tools/blob/stable/ortools/sat/samples/assignment_sat.py # initialize data (wide) # ------------------------------------------- #           task01  task02  task03  task04 # worker_1      90      80      75      70 # worker_2      35      85      55      65 # worker_3     125      95      90      95 # worker_4      45     110      95     115 # worker_5      50     100      90     100 # ------------------------------------------- # converts to a common data format (vertical). # ------------------------------------------- #   worker    task  cost # worker_1  task01    90 # worker_1  task02    80 # worker_1  task03    75 # worker_1  task04    70 # worker_2  task01    35 # worker_2  task02    85 # ... # ------------------------------------------- # in a real case, pd.read_csv would be used here. # set index to define a variable and constraints. # create empty cp model. # define the variables for which task each worker will execute. # each worker will be assigned at most one task (and possibly zero). # each task is always assigned to one of the workers. # add these constraint expressions to the model. # define the objective variable. # the objective variable is the sum of the costs for each worker, # and the goal is to minimize it. # and solve it! # print solution # Retrieve and filter the tasks that have actually been assigned.