From 2443ac285bf91b69aededdd9145c5a898943dcd6 Mon Sep 17 00:00:00 2001 From: Ross Date: Sun, 22 Jun 2025 22:20:52 +0100 Subject: [PATCH] fix a few more things --- gen.py | 41 +++--- leave.py | 30 +++-- rota/shifts.py | 117 +++++++++++++++--- rota/workers.py | 2 +- .../test_summed_shift_grade_constraints.py | 94 ++++++++++++++ test/test_limit_constraints.py | 6 +- test/test_weekend_padding.py | 2 +- 7 files changed, 248 insertions(+), 44 deletions(-) create mode 100644 test/rota_constraints/test_summed_shift_grade_constraints.py diff --git a/gen.py b/gen.py index cb1faa0..89a5451 100644 --- a/gen.py +++ b/gen.py @@ -37,7 +37,7 @@ def main( solve: bool = True, time_to_run: int = 60 * 60, ratio: float = 0.001, - start_date: datetime.datetime = "2025-03-03", + start_date: datetime.datetime = "2025-09-01", weeks: int = 26, bom: int = 1, ): @@ -60,16 +60,16 @@ def main( Rota.constraint_options["hard_constrain_pair_separation"] = True # Rota.constraint_options["avoid_st2_first_month"] = True - wr = [ - WorkerRequirement( - end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(), - number=3, - ), - WorkerRequirement( - start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(), - number=4, - ), - ] + #wr = [ + # WorkerRequirement( + # end_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(), + # number=3, + # ), + # WorkerRequirement( + # start_date=datetime.datetime.strptime("2025-06-02", "%Y-%m-%d").date(), + # number=4, + # ), + #] Rota.add_shifts( SingleShift( @@ -226,7 +226,7 @@ def main( balance_offset=3.9, balance_weighting=1, # hard_constrain_shift=False, - workers_required=wr, + workers_required=4, force_as_block=True, rota_on_nwds=True, constraint=[ @@ -256,7 +256,7 @@ def main( balance_offset=2.9, balance_weighting=1, # hard_constrain_shift=False, - workers_required=wr, + workers_required=4, force_as_block=True, rota_on_nwds=True, constraint=[ @@ -330,6 +330,8 @@ def main( # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"]) + Rota.add_min_summed_grade_by_shifts_per_day_constraint(["weekend_plymouth1", "weekend_plymouth2"], 6) + load_leave = True Rota.build_shifts() # Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names()) @@ -341,10 +343,17 @@ def main( n = 0 for worker in workers: - n = n + 1 worker_name = worker + + if worker_name == "Name": + continue + n = n + 1 w = workers[worker] - site = w["site"] + try: + site = w["site"] + except KeyError: + print(f"Worker {worker} has no site") + raise KeyError grade = w["grade"] try: fte = float(w["fte"]) * 100 @@ -481,7 +490,7 @@ def main( w = Worker( name=worker_name, site=site.lower(), - grade=int(grade[2]), + grade=int(grade), id=n, fte=int(fte), nwds=nwds, diff --git a/leave.py b/leave.py index 6f0eddd..e6c28c9 100644 --- a/leave.py +++ b/leave.py @@ -27,7 +27,8 @@ def load_leave(Rota): } if live_rota: - download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSRx9VWXSlRubPyA0RhiI-Oqf5eHNYYEc6rFzlraDbR5_8qqr5g13-4uV-gn4u-TjZxiSMv1fBUaESq/pub?gid=814517272&single=true&output=csv", headers=headers) + download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSfcQi5Qs__A8hE2CVvUA_6ULNJALrtZOBHPUZ3xyrxPif9obtQF2IqhioT_4nebjHV1Ac5iqhNtuq4/pub?gid=2024511103&single=true&output=csv", headers=headers) + #download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSRx9VWXSlRubPyA0RhiI-Oqf5eHNYYEc6rFzlraDbR5_8qqr5g13-4uV-gn4u-TjZxiSMv1fBUaESq/pub?gid=814517272&single=true&output=csv", headers=headers) decoded_content = download.content.decode('utf-8') reader = csv.reader(decoded_content.splitlines(), delimiter=',') @@ -56,6 +57,7 @@ def load_leave(Rota): names = r order = names for name in names: + name = name.strip() workers[name] = {} workers[name]["leave"] = [] @@ -67,13 +69,16 @@ def load_leave(Rota): workers[name]["start_date"] = None workers[name]["pair"] = None + n = n + 1 + continue + for i in range(len(r)): try: - a = order[i] + a = order[i].strip() except IndexError: continue - worker = workers[order[i]] + worker = workers[a] lower_item = r[i].lower() if lower_item == "derriford": @@ -82,16 +87,18 @@ def load_leave(Rota): if lower_item == "derriford twilights": lower_item = "plymouth_twilights" - if "PROC site" in row_title or row_title == "PROC nights from": + if "PROC site" in row_title or row_title == "PROC nights from" or "proc nights" in row_title.lower(): worker["site_pref"] = lower_item elif row_title in ("Rotation", "Site", "Placement", "Placement location"): worker["site"] = lower_item - elif row_title in ("Grade", "Grade (ST)"): - worker["grade"] = lower_item + elif row_title in ("Grade", "Grade (ST)") or "Year of training" in row_title: + worker["grade"] = lower_item.lstrip("ST") - elif row_title in ("%FTE", "%FTE on-call", "FTE", "FTE on-call"): + elif row_title in ("%FTE", "%FTE on-call", "FTE", "FTE on-call") or "on-call commitment" in row_title: + if lower_item.endswith("%"): + lower_item = float(lower_item.rstrip("%")) / 100 worker["fte"] = lower_item elif "NWD" in row_title: @@ -102,7 +109,10 @@ def load_leave(Rota): elif row_title in ("End Date", "CCT date"): if lower_item: - date = datetime.strptime(lower_item, "%d/%m/%y").date() + try: + date = datetime.strptime(lower_item, "%d/%m/%y").date() + except ValueError: + date = datetime.strptime(lower_item, "%d/%m/%Y").date() worker["end_date"] = date else: worker["end_date"] = None @@ -121,7 +131,7 @@ def load_leave(Rota): else: worker["start_date"] = None - elif "OOP" in row_title: + elif "OOP" in row_title or "out of programme" in row_title.lower(): worker["oop"] = lower_item elif row_title == "Group": @@ -178,6 +188,8 @@ def load_leave(Rota): n = n + 1 + + return workers diff --git a/rota/shifts.py b/rota/shifts.py index cc87b2b..de1cfd2 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -1,10 +1,10 @@ from calendar import week import datetime import itertools -from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal +from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal, Optional import time -from pydantic import BaseModel, ConfigDict, Extra, constr, model_validator, StrictInt +from pydantic import BaseModel, ConfigDict, Extra, constr, model_validator, StrictInt, Field from pathlib import Path @@ -72,11 +72,24 @@ VALID_SHIFT_CONSTRAINTS = Literal[ "require_remote_site_presence_week", ] +def _pydantic_to_dict(obj): + if isinstance(obj, list): + return [_pydantic_to_dict(x) for x in obj] + if hasattr(obj, "model_dump"): + return obj.model_dump() + return obj + class ShiftConstraint(BaseModel): name: VALID_SHIFT_CONSTRAINTS options: bool | int | Dict | tuple = False +class MinSummedGradeByShiftsPerDayConstraint(BaseModel): + shifts: List[str] = Field(..., description="List of shift names to sum grades over") + min_grade_sum: int = Field(..., description="Minimum total grade required") + weeks: Optional[List[int]] = None # If None, applies to all weeks + days: Optional[List[str]] = None # If None, applies to all days + class WorkerRequirement(BaseModel): """ @@ -310,6 +323,7 @@ class RotaBuilder(object): "maximum_allowed_locum_shifts_per_worker": None, "minimum_allowed_locum_shifts_per_worker": None, } + self.constraint_options["min_summed_grade_by_shifts_per_day"]: List[MinSummedGradeByShiftsPerDayConstraint] = [] self.terminate_on_warning = [ "Worker/duplicate id", @@ -1007,6 +1021,24 @@ class RotaBuilder(object): def build_model_constraints(self): self.model.constraints = ConstraintList() # Create a set of constraints + + for constraint in self.constraint_options.get("min_summed_grade_by_shifts_per_day", []): + weeks = constraint.weeks if constraint.weeks is not None else self.weeks + days_list = constraint.days if constraint.days is not None else self.days + for week in weeks: + for day in days_list: + # Only consider if all shifts are present that day + if all(shift in self.get_shift_names_by_week_day(week, day) for shift in constraint.shifts): + self.model.constraints.add( + sum( + worker.grade * self.model.works[worker.id, week, day, shift] + for shift in constraint.shifts + for worker in self.get_workers_for_shift(self.get_shift_by_name(shift)) + ) >= constraint.min_grade_sum + ) + + + # Ensure each shift has the requisit number of workers assigned for ( week, @@ -2269,17 +2301,17 @@ class RotaBuilder(object): # Try disabling tihs to allow more than 2 shifts ot be assigned on a weekend - #self.model.constraints.add( - # self.model.works_weekend[worker.id, week] - # >= sum( - # self.model.works[worker.id, week, day, shiftname] - # for w, day, shiftname in self.get_all_shiftname_combinations( - # week=week - # ) - # if day in self.days[5:] - # ) - # / 2 - #) + self.model.constraints.add( + self.model.works_weekend[worker.id, week] + >= sum( + self.model.works[worker.id, week, day, shiftname] + for w, day, shiftname in self.get_all_shiftname_combinations( + week=week + ) + if day in self.days[5:] + ) + / 2 + ) self.model.constraints.add( self.model.works_weekend[worker.id, week] <= sum( @@ -3090,6 +3122,24 @@ class RotaBuilder(object): {"grades": grades, "weeks": weeks, "shifts": shifts} ) + def add_min_summed_grade_by_shifts_per_day_constraint( + self, + shifts: list[str], + min_grade_sum: int, + weeks: list[int] = None, + days: list[str] = None, + ): + """ + Helper to add a minimum summed grade constraint for a set of shifts per day. + """ + constraint = MinSummedGradeByShiftsPerDayConstraint( + shifts=shifts, + min_grade_sum=min_grade_sum, + weeks=weeks, + days=days, + ) + self.constraint_options["min_summed_grade_by_shifts_per_day"].append(constraint) + def add_shift(self, shift: SingleShift) -> None: """Add a shift to the collection @@ -3766,6 +3816,23 @@ class RotaBuilder(object): week_table[week][day][shift].append(worker.get_details()) return week_table + def get_shift_timetable_by_week(self): + """Build a timetable of the week as a dictionary from the model's optimal solution.""" + works = self.model.works + shift_timetable = { + shift.name: {week: {day: [] for day in days} for week in self.weeks} + for shift in self.shifts + } + for shift in self.shifts: + for week in self.weeks: + for day in days: + assigned_workers = [] + for worker in self.workers: + if works[worker.id, week, day, shift.name].value > 0.5: + assigned_workers.append(worker) + shift_timetable[shift.name][week][day] = assigned_workers + return shift_timetable + def get_worker_timetable(self): works = self.model.works timetable = { @@ -4118,7 +4185,7 @@ class RotaBuilder(object):

Rota settings

 
-        {json.dumps(self.constraint_options, indent=4)}
+        {json.dumps({k: _pydantic_to_dict(v) for k, v in self.constraint_options.items()}, indent=4)}
@@ -4222,6 +4289,28 @@ class RotaBuilder(object): return shifts + def get_shift_worker_list(self, shift_name: str) -> List[set[Worker]]: + """ + Returns a list of workers assigned to a specific shift. + + Args: + shift_name (str): The name of the shift. + + Returns: + List[Worker]: List of workers assigned to the specified shift. + """ + shift_list = [] + for week, day in self.get_week_day_combinations(): + workers: set = set() + for worker in self.workers: + if self.model.works[worker.id, week, day, shift_name].value > 0.5: + workers.add(worker) + + shift_list.append(workers) + + return shift_list + + def get_worker_shift_list(self, worker: Worker, include_locums=False, search_multiple_assignments=False) -> List: """ This function returns a list of shifts assigned to a worker for each week and day. diff --git a/rota/workers.py b/rota/workers.py index 79855ae..766cd5b 100644 --- a/rota/workers.py +++ b/rota/workers.py @@ -237,7 +237,7 @@ class Worker(BaseModel): end_oop_date = datetime.datetime.strptime(end_oop, "%d/%m/%y").date() if start_oop_date >= end_oop_date: - raise ValueError("End OOP date must be after start date") + raise ValueError(f"End OOP date must be after start date [{self.name} - Start date: {self.start_date} / End date: {self.end_date}]") # ignore oops if they finish before the rota (or worker) start date if end_oop_date > self.calculated_start_date: diff --git a/test/rota_constraints/test_summed_shift_grade_constraints.py b/test/rota_constraints/test_summed_shift_grade_constraints.py new file mode 100644 index 0000000..40754f2 --- /dev/null +++ b/test/rota_constraints/test_summed_shift_grade_constraints.py @@ -0,0 +1,94 @@ +import datetime +import pytest +from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days, MinSummedGradeByShiftsPerDayConstraint + +from rota.workers import Worker + +import itertools + + +def generate_basic_rota(weeks_to_rota=9): + start_date = datetime.date(2022, 3, 7) + + Rota = RotaBuilder( + start_date, + weeks_to_rota=weeks_to_rota, + ) + + # Add a few workers + Rota.add_workers( + [ + Worker(name="worker1", site="group1", grade=1), + Worker(name="worker2", site="group1", grade=2), + Worker(name="worker3", site="group1", grade=3), + Worker(name="worker4", site="group1", grade=1), + Worker(name="worker5", site="group1", grade=2), + Worker(name="worker6", site="group1", grade=3), + ] + ) + + return Rota + + +def test_summed_shift_grade_constraints(): + Rota = generate_basic_rota(5) + + Rota.add_shifts( + SingleShift( + sites=("group1",), name="a", length=12.5, days=days, + workers_required=1, + force_as_block=False, + ), + SingleShift( + sites=("group1",), name="b", length=12.5, days=days, + workers_required=1, + force_as_block=False, + ), + ) + + # Add summed shift grade constraints + Rota.constraint_options["min_summed_grade_by_shifts_per_day"].append( + MinSummedGradeByShiftsPerDayConstraint( + shifts=["a", "b"], + min_grade_sum=4, + ), + ) + + Rota.build_and_solve(options={"ratio": 0.000}) + + assert Rota.results.solver.status == "ok" + + Rota.export_rota_to_html("summed_shift_grade_constraints", folder="tests") + + a_grades = [] + b_grades = [] + + for shift, shift_timetable in Rota.get_shift_timetable_by_week().items(): + for week, week_timetable in shift_timetable.items(): + for day, workers in week_timetable.items(): + if shift == "a": + for worker in workers: + a_grades.append(worker.grade) + elif shift == "b": + for worker in workers: + b_grades.append(worker.grade) + + # Check that the minimum summed grade is respected + summed_grades = [a + b for a, b in zip(a_grades, b_grades)] + for summed_grade in summed_grades: + assert summed_grade >= 4, f"Summed grade {summed_grade} is less than 4" + + Rota.constraint_options["min_summed_grade_by_shifts_per_day"] = [ + MinSummedGradeByShiftsPerDayConstraint( + shifts=["a", "b"], + min_grade_sum=5, + ), + ] + + Rota.build_and_solve(options={"ratio": 0.000}) + assert Rota.results.solver.status == "ok" + + for worker in Rota.get_workers(): + if worker.grade == 1: + assert Rota.get_worker_shift_list(worker).count("a") == 0, f"Worker {worker.name} with grade 1 should not have shift 'a'" + assert Rota.get_worker_shift_list(worker).count("b") == 0, f"Worker {worker.name} with grade 1 should not have shift 'b'" diff --git a/test/test_limit_constraints.py b/test/test_limit_constraints.py index 7736f49..43245be 100644 --- a/test/test_limit_constraints.py +++ b/test/test_limit_constraints.py @@ -60,17 +60,17 @@ def test_constraint_limit_grades2(limit_constraint_rota): balance_offset=40, workers_required=4, force_as_block=True, - constraint=[{"name": "limit_grade_number", "options": {3: 2}}], + constraint=[{"name": "limit_grade_number", "options": {3: 1}}], ), ) Rota.constraint_options["balance_shifts_quadratic"] = True Rota.build_and_solve(options={"ratio": 0.1}) - Rota.export_rota_to_html("test9") + Rota.export_rota_to_html("test9", folder="tests") grade_workers = Rota.get_workers_by_grade() for grade in grade_workers: shift_patterns = [Rota.get_worker_shift_list(w) for w in grade_workers[grade]] zipped_lists = list(zip(*shift_patterns)) - limit = 2 if grade == 3 else 2 + limit = 1 if grade == 3 else 3 for day_shifts in zipped_lists: assert day_shifts.count("weekday_night") == limit assert Rota.results.solver.status == "ok" diff --git a/test/test_weekend_padding.py b/test/test_weekend_padding.py index 3b5d816..701e386 100644 --- a/test/test_weekend_padding.py +++ b/test/test_weekend_padding.py @@ -46,7 +46,7 @@ class TestDemoRota: self.Rota.build_and_solve(options={"ratio": 0.00}) - self.Rota.export_rota_to_html("balance_weekend") + self.Rota.export_rota_to_html("balance_weekend", folder="tests") assert self.Rota.results.solver.status == "ok" assert self.Rota.results.solver.termination_condition == "optimal"