Use the Linear Programming Package PuLP to solve an Integer Programming Problem, subject to certain constraints

schedule_shifts_in_pulp.py

"""[summary]
My Question Described: 
https://stackoverflow.com/questions/67004787/pulp-linear-programming-scheduling-optimisation-constraint-on-consecutive-nigh
https://or.stackexchange.com/questions/6070/pulp-scheduling-optimisation-constraint-on-consecutive-shifts-consecutive-ni

Minimise:
The maximum difference in the numbers of shifts per person...?
The cost function defined by hand-coded weights in a data table...?

Objective Function
Decision Variables
Constraints
Solver
"""
import itertools
import pulp
from pulp import LpProblem
from pulp import LpVariable, LpBinary, LpMinimize, lpSum
import numpy as np
from collections import defaultdict
import pandas as pd
from pandas.api.types import CategoricalDtype
from typing import Tuple, List, Union


def assign_weights_to_avoid_consecutive_night_shifts(
    data: pd.DataFrame,
    weight_for_consecutive_night: int = 1000,
    weight_for_2nd_night_off: int = 50,
    weight_for_consecutive_morning_after_night: int = 1000,
) -> pd.DataFrame:
    """Encode costs associated with certain shift patterns

    Args:
        data (pd.DataFrame): [description]
        weight_for_consecutive_night (int, optional): [description]. Defaults to 1000.
        weight_for_2nd_night_off (int, optional): [description]. Defaults to 50.
        weight_for_consecutive_morning_after_night (int, optional): [description]. Defaults to 10_000.

    Returns:
        pd.DataFrame: DataFrame with `weight` column (to be minimized)
            1) weight on consecutive nights
            2) weight on consecutive morning after night
            3) weight on less than 2 nights off
    """
    # initialise weights column 
    data["weight"] = np.ones(shape=data.shape[0]).astype(int)

    # create groups of three that cycle through ...
    all_people = data["person"].unique()
    first_choice = np.random.choice(all_people, 3, replace=False)
    second_choice = np.random.choice(
        [p for p in all_people if p not in first_choice], 3, replace=False
    )
    third_choice = [
        p for p in all_people if (p not in first_choice) & (p not in second_choice)
    ]

    for day in data["day"].unique():
        if (day + 1) % 3 == 0:
            choice = third_choice
            second_day_off_choice = first_choice
        elif (day + 1) % 2 == 0:
            choice = second_choice
            second_day_off_choice = third_choice
        else:
            choice = first_choice
            second_day_off_choice = third_choice
        
        # one day of rest = required
        data.loc[
            (
                (data["shift"] == "night") 
                & (data["day"] == day) 
                & (~np.isin(data["person"], choice))
            ), "weight"
        ] = weight_for_consecutive_night
        
        # two days of rest = where possible
        data.loc[
            (
                (data["shift"] == "night") 
                & (data["day"] == day) 
                & (np.isin(data["person"], second_day_off_choice))
            ), "weight"
        ] = weight_for_2nd_night_off

    # penalize night-morning consecutive shifts
    on_nights = data.loc[
        (
            (data["weight"] == 1) 
            & (data["shift"] == "night") 
        )
    ]
    next_day = (on_nights.loc[:, "day"].astype(int) + 1)
    on_nights.loc[:, "day"] = next_day
    on_nights.loc[:, "shift"] = "morning"
    # drop the final morning shift outside of our day range
    on_nights = on_nights.loc[on_nights["day"] <= data["day"].max()]
    
    # get indexes matching the on_nights object 
    right = data.reset_index().set_index(["day", "shift", "person"])
    left = on_nights.set_index(["day", "shift", "person"])
    mornings_to_weight = left.join(right, how="left", lsuffix="1")
    
    # assign the weight to those mornings
    data.loc[mornings_to_weight["index"], "weight"] = weight_for_consecutive_morning_after_night

    return data, first_choice

# data
N_DAYS = 5
all_people = [
    "clinton",
    "tommy",
    "callum",
    "bonface",
    "rose",
    "dennis",
    "geoffrey",
    "sebastian",
]
all_shifts = ["morning", "afternoon", "night"]
all_days = np.arange(N_DAYS)
index = np.array(list(itertools.product(all_days, all_shifts, all_people)))

# initialise the data in DataFrame
cat_type = CategoricalDtype(categories=["morning", "afternoon", "night"], ordered=True)
data = pd.DataFrame(index, columns=["day", "shift", "person"])
data = data.astype({"shift": cat_type, "day": int})
data = data.sort_values(["day", "shift"])

## assign weights to:
# 1. try and avoid consecutive night shifts
# 2. Avoid consecutive night-morning shifts
data, first_choice = assign_weights_to_avoid_consecutive_night_shifts(data)
data.loc[data["shift"] == "night"]

# Decision Variables
decision_variables = []
for rownum, row in data.iterrows():
    variable = f"{row['day']}_{row['shift']}_{row['person']}"
    variable = pulp.LpVariable(variable, cat=LpBinary)
    decision_variables.append(variable)

# decision_variable_dict = pulp.LpVariable.dicts('VAR', (all_days, all_shifts, all_people), 0, 1, 'Binary')
decision_variable_dict = {
    (day, shift, person): LpVariable(f"{person} on day {day} {shift}", cat=LpBinary)
    for day in all_days
    for shift in all_shifts
    for person in all_people
}

#  initialise problem
# use the += operator to assign obectives to the solver/model
allocation_model = LpProblem("Allocation", LpMinimize)

# Objective Function defined by the weight column
def get_cost(key: Tuple[Union[int, str], str, str], data: pd.DataFrame) -> int:
    """Extract the weight value from the pd.DataFrame for that person-shift"""
    day = key[0]
    shift = key[1]
    person = key[2]
    return int(data.loc[
        (
            (data["day"] == day) 
            & (data["shift"] == shift) 
            & (data["person"] == person)
        ), "weight"
    ])

allocation_model += lpSum(
    [var * get_cost(k, data) for k, var in decision_variable_dict.items()]
)

# add people constraints
#  assign each person to max 2 shifts per day
for day in all_days:
    for person in all_people:
        allocation_model.addConstraint(
            sum(decision_variable_dict[(day, shift, person)] for shift in all_shifts)
            <= 2
        )
#  assign each person to minimum 1 shift per day
for day in all_days:
    for person in all_people:
        allocation_model.addConstraint(
            sum(decision_variable_dict[(day, shift, person)] for shift in all_shifts)
            == 1
        )

#  add shift constraints
# assign all shifts to three people
for day in all_days:
    for shift in all_shifts:
        allocation_model.addConstraint(
            sum(decision_variable_dict[(day, shift, person)] for person in all_people)
            == 3
        )

## TODO: HOW TO ENCODE THE NUMBER OF RESTS FROM NIGHT SHIFT >= 1
## TODO: HOW TO ENCODE NO CONSECUTIVE SHIFTS

# Solve and get results:
solver = None
allocation_model.solve()

if allocation_model.status == 1:
    print("Solution Found.")


results = defaultdict(list)
for (day, shift, person), assigned in sorted(decision_variable_dict.items()):
    if pulp.value(assigned):
        #  print(f"Day {day}: Assign {person} to {shift}")
        results["day"].append(day)
        results["person"].append(person)
        results["shift"].append(shift)

df = pd.DataFrame(results)
df = df.astype({"shift": cat_type})
df = df.sort_values(["day", "shift"])

# print solution
print("Calendar Assignment (nights, mornings)")
print("-" * 50)
print(df.loc[np.isin(df["shift"], ["morning", "night"])])
print()
print(f"Choice for first shift: {first_choice}")


# test equity of solution
print()
print(
    f"How many shifts for each individual? \n{'-'*50}\n",
    df.groupby("person").count()["shift"],
)
print()
print(
    f"How many days does each individual work? \n{'-'*50}\n",
    df.groupby("person").nunique()["day"],
)
print()
print(
    f"How many night shifts per individual? \n{'-'*50}\n",
    df.loc[df["shift"] == "night"].groupby("person").count()["shift"],
)