diff --git a/gen_cons.py b/gen_cons.py index 0a32e7f..e00beb1 100644 --- a/gen_cons.py +++ b/gen_cons.py @@ -29,7 +29,7 @@ def main( time_to_run: int = 60 * 60, ratio: float = 0.001, start_date: datetime.datetime = "2025-09-01", - weeks: int = 16, + weeks: int = 12, bom: int = 1, ): rota_start_date = start_date.date() @@ -42,61 +42,42 @@ def main( name="cons rota", ) # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1) - Rota.constraint_options["balance_shifts"] = False - Rota.constraint_options["balance_shifts_quadratic"] = True Rota.constraint_options["balance_weekends"] = True - Rota.constraint_options["max_night_frequency"] = 3 # 3 - Rota.constraint_options["max_night_frequency_week_exclusions"] = [] Rota.constraint_options["max_weekend_frequency"] = 4 - Rota.constraint_options["hard_constrain_pair_separation"] = True - Rota.constraint_options["max_shifts_per_week"] = 3 + Rota.constraint_options["max_shifts_per_week"] = 20 # Rota.constraint_options["avoid_st2_first_month"] = True Rota.add_shifts( SingleShift( sites=("rota a",), - name="oncall weekday", + name="oncall", length=12.5, - days=days[:4], - balance_offset=2, + days=days, + balance_offset=5, + assign_as_block=False, constraint=[ - { - "name": "max_shifts_per_week", - "options": 3, - } + #{ + # "name": "max_shifts_per_week", + # "options": 3, + #} ], ), SingleShift( sites=("rota a",), - name="Oncall weekend", - length=12.5, - days=days[4:], - balance_offset=2, - assign_as_block=True, - constraint=[ - { - "name": "max_shifts_per_week", - "options": 3, - }, - {"name": "pre", "options": 2}, - {"name": "post", "options": 2}, - ], - ), - SingleShift( - sites=("rota a",), - name="weekend a", + name="weekend", workers_required=1, length=12.5, days=days[5:], balance_offset=2, constraint=[ - {"name": "pre", "options": 2}, - {"name": "post", "options": 2}, + #{"name": "pre", "options": 2}, + #{"name": "post", "options": 1}, ], display_char="a", + #force_assign_with=["oncall"] ), SingleShift( - sites=("rota b", "rota c"), + sites=("rota b",), name="weekend b", workers_required=1, length=12.5, @@ -113,22 +94,25 @@ def main( "rota a", "rota b", ), - name="evening", + name="twilight", workers_required=1, length=12.5, days=days[:5], - balance_offset=2, - constraint=[], + balance_offset=5, + #constraint=[ + # {"name": "pre", "options": 1}, + # {"name": "post", "options": 1}, + #], ), ) - Rota.allow_shifts_together( - - "oncall weekday", - "Oncall weekend", - "weekend a", - "evening", - - ) + #Rota.allow_shifts_together_for_all_workers( + # "oncall", + # "twilight", + #) + #Rota.allow_shifts_together_for_all_workers( + # "oncall", + # "twilight", + #) # Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"]) @@ -136,8 +120,24 @@ def main( Rota.build_shifts() # Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names()) - for w in range(1, 12): - w = Worker(name=f"A{w:02d}", site="rota a", grade=1) + for n in range(1, 12): + w = Worker(name=f"A{n:02d}", site="rota a", grade=1) + #w.prefer_multi_shift_together = 10 + #w.allow_shifts_together("oncall", "twilight") + #w.allow_shifts_together("oncall", "weekend") + + #w.add_hard_day_dependency("Fri", "Sat", "Sun") + + w.add_force_assign_with("twilight", "oncall") + w.add_force_assign_with("weekend", "oncall") + w.set_max_shifts_per_week("twilight", 1) + + if n <= 2: + pass + w.add_hard_day_dependency("Fri", "Sat", "Sun") + else: + w.add_hard_day_dependency("Fri", "Sun") + ## #w.add_hard_day_dependency("Sat", "Sun") print(w) @@ -154,6 +154,9 @@ def main( ) ], ) + #w.add_hard_day_exclusion("Sat", "Sun") + w.set_max_shifts_per_week("twilight", 1) + w.set_max_shifts_per_week("weekend b", 1) print(w) diff --git a/rota/shifts.py b/rota/shifts.py index 3ada902..cc87b2b 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -1385,6 +1385,7 @@ class RotaBuilder(object): ) # Most of our constraints apply per worker + # Worker constraint loop worker: Worker for worker in self.workers: console.rule( @@ -1392,6 +1393,22 @@ class RotaBuilder(object): ) logging.debug(f"Generate worker constraints: {worker.name}") + + if hasattr(worker, "max_shifts_per_week_by_shift_name"): + for shift_type, max_per_week in worker.max_shifts_per_week_by_shift_name.items(): + if shift_type not in self.get_shift_names(): + raise InvalidShift( + f"Worker '{worker.name}' has max_shifts_per_week_by_shift_name referencing non-existent shift: {shift_type}" + ) + for week in self.weeks: + self.model.constraints.add( + sum( + self.model.works[worker.id, week, day, shift_type] + for day in days + if shift_type in self.get_shift_names_by_week_day(week, day) + ) <= max_per_week + ) + for week, day, shift in self.get_all_shiftclass_combinations(): # Only apply if this shift has force_assign_with set if hasattr(shift, "force_assign_with") and shift.force_assign_with: diff --git a/rota/workers.py b/rota/workers.py index 47b4bb7..79855ae 100644 --- a/rota/workers.py +++ b/rota/workers.py @@ -117,6 +117,7 @@ class Worker(BaseModel): hard_day_dependencies: set[frozenset[str]] = set() # e.g. {frozenset({"Mon", "Tue"}), frozenset({"Fri", "Sun"})} hard_day_exclusions: list[tuple[str, str]] = [] # e.g. [("Fri", "Sun")] force_assign_with: dict[str, list[str]] = {} # e.g. {"oncall": ["weekend"]} + max_shifts_per_week_by_shift_name: dict[str, int] = {} # e.g. {"night": 2, "a": 3} assign_as_block_preferences: dict[str, float] = {} # shift_name -> weight (positive = prefer block) @@ -125,6 +126,8 @@ class Worker(BaseModel): weekend_shift_target_number: int = 0 + + model_config = ConfigDict( extra="allow", ) @@ -421,4 +424,10 @@ class Worker(BaseModel): self.force_assign_with = {} if isinstance(with_shifts, str): with_shifts = [with_shifts] - self.force_assign_with[shift] = with_shifts \ No newline at end of file + self.force_assign_with[shift] = with_shifts + + def set_max_shifts_per_week(self, shift_name: str, max_per_week: int): + """Helper to set the maximum number of shifts of a given type per week for this worker.""" + if not hasattr(self, "max_shifts_per_week_by_shift_name"): + self.max_shifts_per_week_by_shift_name = {} + self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week \ No newline at end of file diff --git a/test/test_workers.py b/test/test_workers.py index 87055e9..17540d1 100644 --- a/test/test_workers.py +++ b/test/test_workers.py @@ -1380,4 +1380,79 @@ def test_workers_double_shifts_on_weekends(): Rota.build_and_solve(options={"ratio": 0.0}) Rota.export_rota_to_html("test_worker_double_shifts", folder="tests") - assert Rota.results.solver.status == "ok" \ No newline at end of file + assert Rota.results.solver.status == "ok" + + +def test_max_shifts_per_week_by_shift_name_basic(): + Rota = RotaBuilder(weeks_to_rota=2, start_date=datetime.date(2025, 6, 16)) + w = Worker(name="A01", site="site1", grade=1) + w.max_shifts_per_week_by_shift_name = {"a": 1} + Rota.add_worker(w) + Rota.add_shifts( + SingleShift( + sites=("site1",), + name="a", + length=8, + days=days[:5], + workers_required=1, + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + # Worker should not be assigned more than 1 "a" shift per week + for week in Rota.weeks: + count = sum( + Rota.model.works[w.id, week, day, "a"].value > 0.5 + for day in days[:5] + ) + assert count <= 1 + +def test_max_shifts_per_week_by_shift_name_multiple_weeks(): + Rota =generate_basic_rota(workers=4, weeks_to_rota=2) + for worker in Rota.get_workers(): + worker.set_max_shifts_per_week("a", 1) + Rota.add_shifts( + SingleShift( + sites=("group1",), + name="a", + length=8, + days=days[:5], + workers_required=1, + ), + ) + Rota.build_and_solve(options={"ratio": 0.0}) + assert Rota.results.solver.status == "error" + + Rota.add_worker( + Worker(name="extraworker", site="group1", grade=1, max_shifts_per_week_by_shift_name={"a": 1}) + ) + Rota.build_and_solve(options={"ratio": 0.0}) + assert Rota.results.solver.status == "ok" + + Rota.export_rota_to_html("test_worker_double_shifts", folder="tests") + # Worker should not be assigned more than 2 "a" shifts per week + for worker in Rota.get_workers(): + for week in Rota.weeks: + count = sum( + Rota.model.works[worker.id, week, day, "a"].value > 0.5 + for day in days[:5] + ) + assert count <= 2 + + +def test_max_shifts_per_week_by_shift_name_invalid_shift(): + Rota = RotaBuilder(weeks_to_rota=1, start_date=datetime.date(2025, 6, 16)) + w = Worker(name="A03", site="site1", grade=1) + # "b" does not exist + w.max_shifts_per_week_by_shift_name = {"b": 1} + Rota.add_worker(w) + Rota.add_shifts( + SingleShift( + sites=("site1",), + name="a", + length=8, + days=days[:5], + workers_required=1, + ), + ) + with pytest.raises(InvalidShift): + Rota.build_and_solve(options={"ratio": 0.0}) \ No newline at end of file