diff --git a/requirements.txt b/requirements.txt index 8f863f3..53c81e4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ pyomo govuk_bank_holidays -pytest \ No newline at end of file +pytest +black \ No newline at end of file diff --git a/shifts.py b/shifts.py index 6ba6a24..c394329 100644 --- a/shifts.py +++ b/shifts.py @@ -1,6 +1,6 @@ import datetime import itertools -from typing import Iterable, List, Sequence, Tuple, Set +from typing import Dict, Iterable, List, Sequence, Tuple, Set import datetime @@ -82,7 +82,16 @@ class SingleShift(object): self.assign_as_block = assign_as_block self.force_as_block = force_as_block self.force_as_block_unless_nwd = force_as_block_unless_nwd - self.constraints = constraints + + self.constraint_options = {} + self.constraints = [] + for c in constraints: + match c: + case (constraint, options): + self.constraints.append(constraint) + self.constraint_options[constraint] = options + case constraint: + self.constraints.append(constraint) self.bank_holidays_only = bank_holidays_only @@ -170,7 +179,8 @@ class RotaBuilder(object): "balance_nights_across_sites": True, "balance_bank_holidays": True, "balance_blocks": True, - "balance_shifts": True, + "balance_shifts": True, # Does not use a quadratic function + "balance_shifts_quadratic": False, "balance_shifts_over_workers": True, "minimise_shift_diffs": False, # less sophisticated version of balance_shifts_over_workers "balance_weekends": True, @@ -436,7 +446,7 @@ class RotaBuilder(object): initialize=0, ) - if self.constraint_options["balance_shifts"]: + if self.constraint_options["balance_shifts"] or self.constraint_options["balance_shifts_quadratic"]: self.model.shift_count_t1 = Var( ( (worker.id, shift.name) @@ -457,6 +467,17 @@ class RotaBuilder(object): initialize=0, ) + if self.constraint_options["balance_shifts_quadratic"]: + self.model.shift_count_w = Var( + ( + (worker.id, shift.name) + for worker in self.workers + for shift in self.get_shifts() + ), + domain=NonNegativeReals, + initialize=0, + ) + if self.constraint_options["balance_nights"]: # We also try to even out the night shifts seperately self.model.night_shift_count = Var( @@ -722,6 +743,35 @@ class RotaBuilder(object): rule=nightShiftMaxSTRule, ) + # Limit grades + def limitGradesRule(model, week, shift, grade, limit): + workers = [w for w in self.workers if w.grade == grade] + if not workers: + return Constraint.Skip + return ( + sum( + model.shift_week_worker_assigned[shift.name, week, w.id] + for w in workers + ) + <= limit + ) + + for shift in self.get_shifts_with_constraint("limit_grade_number"): + if not shift.constraint_options["limit_grade_number"]: + raise ValueError("Constraint option must be defined for 'limit_grade_number'") + for grade in shift.constraint_options["limit_grade_number"]: + #self.model.limit_grades_constraint = Constraint( + setattr(self.model, f"limit_grades_constraint_{shift.name}_{grade}", Constraint( + [week for week in self.weeks], + #[shift for shift in self.get_shifts_with_constraint("limit_grade_number")], + [shift], + [grade], + [shift.constraint_options["limit_grade_number"][grade]], + rule=limitGradesRule, + #name=f"limit_grades_constraint_{shift.name}_{grade}" + )) + + def nightShiftMinST4Rule(model, week, shift): single_workers = [w for w in self.workers if w.grade >= 4] if not single_workers: @@ -785,13 +835,34 @@ class RotaBuilder(object): # # Count the number of workers from each site on each night shift # # As 1 or 0 is optimum we can simply subtract 1 from the number - if self.constraint_options["balance_nights_across_sites"]: - for site in self.sites: - for shift in self.get_shifts_with_constraint("night"): + #if self.constraint_options["balance_nights_across_sites"]: + # for site in self.sites: + # for shift in self.get_shifts_with_constraint("night"): + # block = shift.name + # for week in self.weeks: + # self.model.constraints.add( + # self.model.night_per_site[week, block, site] + # == sum( + # self.model.shift_week_worker_assigned[ + # block, week, worker.id + # ] + # for worker in self.workers + # if worker.site == site + # ) + # # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site) + # ) + # self.model.constraints.add( + # self.model.night_per_site_t1[week, block, site] + # - self.model.night_per_site_t2[week, block, site] + # == self.model.night_per_site[week, block, site] - 1 + # ) + for site in self.sites: + for shift in self.get_shifts_with_constraint("balance_across_groups"): + if site in shift.site: block = shift.name for week in self.weeks: self.model.constraints.add( - self.model.night_per_site[week, block, site] + self.model.shift_per_site[week, block, site] == sum( self.model.shift_week_worker_assigned[ block, week, worker.id @@ -802,30 +873,10 @@ class RotaBuilder(object): # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site) ) self.model.constraints.add( - self.model.night_per_site_t1[week, block, site] - - self.model.night_per_site_t2[week, block, site] - == self.model.night_per_site[week, block, site] - 1 + self.model.shift_per_site_t1[week, block, site] + - self.model.shift_per_site_t2[week, block, site] + == self.model.shift_per_site[week, block, site] - 1 ) - for site in self.sites: - for shift in self.get_shifts_with_constraint("balance_across_groups"): - block = shift.name - for week in self.weeks: - self.model.constraints.add( - self.model.shift_per_site[week, block, site] - == sum( - self.model.shift_week_worker_assigned[ - block, week, worker.id - ] - for worker in self.workers - if worker.site == site - ) - # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site) - ) - self.model.constraints.add( - self.model.shift_per_site_t1[week, block, site] - - self.model.shift_per_site_t2[week, block, site] - == self.model.shift_per_site[week, block, site] - 1 - ) weeks_days = self.get_week_day_combinations() @@ -1082,6 +1133,46 @@ class RotaBuilder(object): == self.model.shift_count_diff[worker.id, shift.name] ) + + if self.constraint_options["balance_shifts_quadratic"]: + # This may need to be updated + xU = 25 + xL = 1 + self.model.constraints.add( + inequality( + xL, + self.model.shift_count_t1[worker.id, shift.name] + + self.model.shift_count_t2[worker.id, shift.name] + + 1, + xU, + ) + ) + + self.model.constraints.add( + self.model.shift_count_w[worker.id, shift.name] + >= xL + * ( + self.model.shift_count_t1[worker.id, shift.name] + + self.model.shift_count_t2[worker.id, shift.name] + + 1 + ) + * 2 + - xL * xL + ) + + self.model.constraints.add( + self.model.shift_count_w[worker.id, shift.name] + >= xU + * ( + self.model.shift_count_t1[worker.id, shift.name] + + self.model.shift_count_t2[worker.id, shift.name] + + 1 + ) + * 2 + - xU * xU + ) + + # Define worker_shift_count_t1 and worker_shift_count_t2 constraints for the object # Thus bypassing the need for a quadratic solver # t1-t2 is the target @@ -1459,22 +1550,24 @@ class RotaBuilder(object): # self.model.shift_week_worker_assigned[s.name, week, worker.id] # for s in self.get_shifts_with_constraint("night"))) - for shift in self.get_shifts_with_constraint("night"): - # Force nights to be assigned in blocks - # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days) - # ) + #for shift in self.get_shifts_with_constraint("night"): + for shift in self.get_shifts(): + if shift.force_as_block: + # Force nights to be assigned in blocks + # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days) + # ) - # Force night shifts to be assigned in blocks - self.model.constraints.add( - self.model.shift_week_worker_assigned[ - shift.name, week, worker.id - ] - * len(shift.shift_days) - == sum( - self.model.works[worker.id, week, day, shift.name] - for day in shift.shift_days + # Force night shifts to be assigned in blocks + self.model.constraints.add( + self.model.shift_week_worker_assigned[ + shift.name, week, worker.id + ] + * len(shift.shift_days) + == sum( + self.model.works[worker.id, week, day, shift.name] + for day in shift.shift_days + ) ) - ) # @@ -1728,6 +1821,19 @@ class RotaBuilder(object): else: shift_balancing = 0 + if self.constraint_options["balance_shifts_quadratic"]: + shift_balancing = sum( + balance_modifier_constant + * ( + self.model.shift_count_t1[worker.id, shift.name] + + self.model.shift_count_t2[worker.id, shift.name] + ) + for worker in self.workers + for shift in self.get_shifts() + ) + else: + shift_balancing = 0 + if self.constraint_options["minimise_shift_diffs"]: shift_diff_modifier_constant = 10 shift_diff_balancing = sum( @@ -1808,6 +1914,22 @@ class RotaBuilder(object): else: nights_site_balancing = 0 + + if self.get_shifts_with_constraint("balance_across_groups"): + block_site_balancing = sum( + ( + self.model.shift_per_site_t1[week, shift.name, site] + + self.model.shift_per_site_t2[week, shift.name, site] + ) + * 2000 + # self.model.night_per_site2[week, block, site] + for week in self.weeks + for shift in self.get_shifts_with_constraint("balance_across_groups") + for site in self.sites + ) + else: + block_site_balancing = 0 + if self.constraint_options["balance_blocks"]: blocks_balancing = sum( 1 * self.model.blocks_assigned[week, shift] @@ -1833,6 +1955,7 @@ class RotaBuilder(object): + shift_diff_balancing + preferences + nights_site_balancing + + block_site_balancing + bank_holiday_balancing + blocks_balancing - work_requests @@ -2117,6 +2240,20 @@ class RotaBuilder(object): s.extend(self.shifts_to_force_as_blocks()) return s + def get_workers_by_group(self) -> Dict[str, Worker]: + group_workers = defaultdict(set) + for worker in self.workers: + group_workers[worker.site].add(worker) + + return group_workers + + def get_workers_by_grade(self) -> Dict[str, Worker]: + group_workers = defaultdict(set) + for worker in self.workers: + group_workers[worker.grade].add(worker) + + return group_workers + def get_workers_for_shift(self, shift: SingleShift) -> List[Worker]: return [worker for worker in self.workers if worker.site in shift.site] diff --git a/test_rota.py b/test_rota.py index 24ab50a..9bcbef2 100644 --- a/test_rota.py +++ b/test_rota.py @@ -1,3 +1,4 @@ +from copy import deepcopy from black import main import pytest from shifts import NoWorkers, RotaBuilder, SingleShift, days @@ -383,6 +384,11 @@ class TestDemoRotaShiftConstraints: ), ) + Rota3Workers = deepcopy(Rota) + worker3 = Worker(Rota3Workers, "worker3", "group1", 1) + + Rota3Workers.add_worker(worker3) + def test_max_shifts(self): self.Rota.constraint_options["max_shifts_per_month"] = 14 self.Rota.build_and_solve() @@ -397,23 +403,11 @@ class TestDemoRotaShiftConstraints: assert self.Rota.results.solver.termination_condition == "infeasible" def test_max_shifts_extra_worker(self): - self.Rota.constraint_options["max_shifts_per_month"] = 12 - worker3 = Worker(self.Rota, "worker3", "group1", 1) + self.Rota3Workers.constraint_options["max_shifts_per_month"] = 12 + self.Rota3Workers.build_and_solve() - self.Rota.add_worker(worker3) - self.Rota.build_and_solve() - - assert self.Rota.results.solver.status == "ok" - assert self.Rota.results.solver.termination_condition == "optimal" - - def test_max_shifts_extra_worker_wrong_group(self): - self.Rota.constraint_options["max_shifts_per_month"] = 12 - worker3 = Worker(self.Rota, "worker3", "group2", 1) - self.Rota.add_worker(worker3) - self.Rota.build_and_solve() - - assert self.Rota.results.solver.status == "warning" - assert self.Rota.results.solver.termination_condition == "infeasible" + assert self.Rota3Workers.results.solver.status == "ok" + assert self.Rota3Workers.results.solver.termination_condition == "optimal" def test_max_shifts_per_week_fail(self): self.Rota.constraint_options["max_shifts_per_week"] = 3 @@ -423,18 +417,19 @@ class TestDemoRotaShiftConstraints: assert self.Rota.results.solver.termination_condition == "infeasible" def test_max_shifts_per_week_pass(self): + self.Rota.constraint_options["max_shifts_per_month"] = 14 self.Rota.constraint_options["max_shifts_per_week"] = 4 self.Rota.build_and_solve() assert self.Rota.results.solver.status == "ok" def test_max_shifts_per_week_extra_worker_pass(self): - self.Rota.constraint_options["max_shifts_per_week"] = 3 - worker3 = Worker(self.Rota, "worker3", "group1", 1) - self.Rota.add_worker(worker3) - self.Rota.build_and_solve() + self.Rota3Workers.constraint_options["max_shifts_per_week"] = 3 + #worker3 = Worker(self.Rota, "worker3", "group1", 1) + #self.Rota.add_worker(worker3) + self.Rota3Workers.build_and_solve() - assert self.Rota.results.solver.status == "ok" + assert self.Rota3Workers.results.solver.status == "ok" class TestDemoRotaBalanceShiftSites: weeks_to_rota = 8 @@ -450,34 +445,238 @@ class TestDemoRotaBalanceShiftSites: worker2 = Worker(Rota, "worker2", "group1", 1) worker3 = Worker(Rota, "worker3", "group2", 1) worker4 = Worker(Rota, "worker4", "group2", 1) - worker5 = Worker(Rota, "worker3", "group2", 1) - worker6 = Worker(Rota, "worker4", "group2", 1) + worker5 = Worker(Rota, "worker5", "group3", 1) + worker6 = Worker(Rota, "worker6", "group3", 1) Rota.add_workers((worker1, worker2, worker3, worker4)) - #Rota.add_workers((worker5, worker6)) - Rota.add_shifts( - SingleShift(("group1", "group2"), "weekday_night", 12.5, days[:4], balance_offset=40, - workers_required=2, - force_as_block=True, - constraints=["balance_across_groups"] - ), - #SingleShift(("group1", "group2"), "weekend_night", 12.5, days[4:], balance_offset=40, - #workers_required=2, - #force_as_block=True, - #), - ) + Rota.add_workers((worker5, worker6)) - def test_balance_blocks_across_groups(self): - self.Rota.constraint_options["balance_nights_across_sites"] = True - #self.Rota.constraint_options["balance_bank_holidays"] = False - #self.Rota.constraint_options["balance_weekends"] = False - #self.Rota.constraint_options["balance_blocks"] = False - self.Rota.build_and_solve(options={"ratio": 0.0, "seconds": 1000, "threads": 10}) + def test_balance_blocks_across_2_groups(self): + self.Rota.add_shifts( + SingleShift(("group1", "group2"), "weekday_night", 12.5, days[:4], balance_offset=40, + workers_required=2, + force_as_block=True, + constraints=["balance_across_groups"] + ), + ) + self.Rota.constraint_options["balance_nights_across_sites"] = False + self.Rota.constraint_options["balance_nights"] = False + self.Rota.constraint_options["constrain_time_off_after_nights"] = False + self.Rota.build_and_solve(options={"ratio": 0.0}) self.Rota.export_rota_to_html("test5") + print(self.Rota.get_workers_by_group()) + + group_workers = self.Rota.get_workers_by_group() + for group in group_workers: + shift_patterns = [] + for w in group_workers[group]: + shift_patterns.append(self.Rota.get_worker_shift_list(w)) + + zipped_lists = list(zip(*shift_patterns)) + + for day_shifts in zipped_lists: + assert day_shifts.count("weekday_night") <= 1 + assert self.Rota.results.solver.status == "ok" assert self.Rota.results.solver.termination_condition == "optimal" + def test_balance_blocks_across_2_groups_unbalanced(self): + self.Rota.shifts = [] + self.Rota.add_shifts( + SingleShift(("group1", "group2"), "weekday_night", 12.5, days[:4], balance_offset=40, + workers_required=3, + force_as_block=True, + constraints=["balance_across_groups"] + ), + ) + self.Rota.constraint_options["balance_nights_across_sites"] = False + self.Rota.constraint_options["balance_nights"] = False + self.Rota.constraint_options["constrain_time_off_after_nights"] = False + self.Rota.build_and_solve(options={"ratio": 0.1}) + + group_workers = self.Rota.get_workers_by_group() + for group in group_workers: + shift_patterns = [] + for w in group_workers[group]: + shift_patterns.append(self.Rota.get_worker_shift_list(w)) + + zipped_lists = list(zip(*shift_patterns)) + + for day_shifts in zipped_lists: + assert day_shifts.count("weekday_night") <= 2 + + assert self.Rota.results.solver.status == "ok" + assert self.Rota.results.solver.termination_condition == "optimal" + + + def test_balance_blocks_across_3_groups(self): + self.Rota.shifts = [] + + self.Rota.add_shifts( + SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days[:4], balance_offset=40, + workers_required=3, + force_as_block=True, + constraints=["balance_across_groups"] + ), + ) + self.Rota.constraint_options["balance_nights_across_sites"] = False + self.Rota.constraint_options["balance_nights"] = False + self.Rota.constraint_options["constrain_time_off_after_nights"] = False + self.Rota.build_and_solve(options={"ratio": 0.1}) + + group_workers = self.Rota.get_workers_by_group() + for group in group_workers: + shift_patterns = [] + for w in group_workers[group]: + shift_patterns.append(self.Rota.get_worker_shift_list(w)) + + zipped_lists = list(zip(*shift_patterns)) + + for day_shifts in zipped_lists: + assert day_shifts.count("weekday_night") <= 1 + + assert self.Rota.results.solver.status == "ok" + assert self.Rota.results.solver.termination_condition == "optimal" + +class TestLimitConstraints: + weeks_to_rota = 8 + start_date = datetime.date(2022, 3, 7) + + Rota = RotaBuilder( + start_date, + weeks_to_rota=weeks_to_rota, + ) + #Rota.constraint_options["max_shifts_per_week"] = 5 + #Rota.constraint_options["max_shifts_per_month"] = 20 + worker1 = Worker(Rota, "worker1", "group1", 2) + worker2 = Worker(Rota, "worker2", "group1", 2) + worker3 = Worker(Rota, "worker3", "group1", 2) + worker4 = Worker(Rota, "worker4", "group3", 3) + worker5 = Worker(Rota, "worker5", "group3", 3) + worker6 = Worker(Rota, "worker6", "group3", 3) + + Rota.add_workers((worker1, worker2, worker3, worker4)) + Rota.add_workers((worker5, worker6)) + + Rota.shifts = [] + + + def test_constraint_limit_grades(self): + self.Rota.shifts = [] + self.Rota.add_shifts( + SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40, + workers_required=4, + force_as_block=True, + constraints=[("limit_grade_number", {2: 1})] + ), + ) + #self.Rota.constraint_options["balance_nights_across_sites"] = False + #self.Rota.constraint_options["balance_shifts_over_workers"] = False + self.Rota.constraint_options["balance_shifts_quadratic"] = True + self.Rota.build_and_solve(options={"ratio": 0.00}) + + grade_workers = self.Rota.get_workers_by_grade() + for grade in grade_workers: + shift_patterns = [] + for w in grade_workers[grade]: + shift_patterns.append(self.Rota.get_worker_shift_list(w)) + + zipped_lists = list(zip(*shift_patterns)) + + limit = 1 + if grade > 2: + limit = 3 + for day_shifts in zipped_lists: + assert day_shifts.count("weekday_night") == limit + + assert self.Rota.results.solver.status == "ok" + assert self.Rota.results.solver.termination_condition == "optimal" + + def test_constraint_limit_grades2(self): + self.Rota.shifts = [] + self.Rota.add_shifts( + SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40, + workers_required=4, + force_as_block=True, + constraints=[("limit_grade_number", {3:2})] + ), + ) + #self.Rota.constraint_options["balance_nights_across_sites"] = False + #self.Rota.constraint_options["balance_shifts_over_workers"] = False + self.Rota.constraint_options["balance_shifts_quadratic"] = True + self.Rota.build_and_solve(options={"ratio": 0.1}) + + self.Rota.export_rota_to_html("test9") + + grade_workers = self.Rota.get_workers_by_grade() + for grade in grade_workers: + shift_patterns = [] + for w in grade_workers[grade]: + shift_patterns.append(self.Rota.get_worker_shift_list(w)) + + zipped_lists = list(zip(*shift_patterns)) + + limit = 2 + if grade > 2: + limit = 2 + for day_shifts in zipped_lists: + assert day_shifts.count("weekday_night") == limit + + assert self.Rota.results.solver.status == "ok" + assert self.Rota.results.solver.termination_condition == "optimal" + + def test_constraint_limit_grades3(self): + self.Rota.shifts = [] + self.Rota.add_shifts( + SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40, + workers_required=4, + force_as_block=True, + constraints=[("limit_grade_number", {2: 3, 3:1})] + ), + ) + #self.Rota.constraint_options["balance_nights_across_sites"] = False + #self.Rota.constraint_options["balance_shifts_over_workers"] = False + self.Rota.constraint_options["balance_shifts_quadratic"] = True + self.Rota.build_and_solve(options={"ratio": 0.1}) + + self.Rota.export_rota_to_html("test9") + + grade_workers = self.Rota.get_workers_by_grade() + for grade in grade_workers: + shift_patterns = [] + for w in grade_workers[grade]: + shift_patterns.append(self.Rota.get_worker_shift_list(w)) + + zipped_lists = list(zip(*shift_patterns)) + + limit = 3 + if grade > 2: + limit = 1 + for day_shifts in zipped_lists: + assert day_shifts.count("weekday_night") == limit + + assert self.Rota.results.solver.status == "ok" + assert self.Rota.results.solver.termination_condition == "optimal" + + def test_constraint_limit_grades4(self): + self.Rota.shifts = [] + self.Rota.add_shifts( + SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40, + workers_required=4, + force_as_block=True, + constraints=[("limit_grade_number", {2: 4, 3:0})] + ), + ) + #self.Rota.constraint_options["balance_nights_across_sites"] = False + #self.Rota.constraint_options["balance_shifts_over_workers"] = False + self.Rota.constraint_options["balance_shifts_quadratic"] = True + self.Rota.build_and_solve(options={"ratio": 0.1}) + + self.Rota.export_rota_to_html("test9") + + assert self.Rota.results.solver.termination_condition == "infeasible" + #class TestNoWorkerRota: # weeks_to_rota = 10 # start_date = datetime.date(2022, 3, 7) @@ -493,5 +692,5 @@ class TestDemoRotaBalanceShiftSites: if __name__ == "__main__": - t = TestDemoRotaShiftConstraints() - t.test_balance_blocks_across_groups() + t = TestLimitConstraints() + t.test_constraint_limit_grades4()