From b81570f4579e99f21892328ddf52efb81d32a510 Mon Sep 17 00:00:00 2001 From: Ross Date: Fri, 6 Jun 2025 22:36:58 +0100 Subject: [PATCH] basic multishift per day assignement --- rota/shifts.py | 50 +++++++++++++--- test/test_multiple_shifts_per_day.py | 88 ++++++++++++++++++++++++++++ 2 files changed, 130 insertions(+), 8 deletions(-) create mode 100644 test/test_multiple_shifts_per_day.py diff --git a/rota/shifts.py b/rota/shifts.py index 06d0f38..c6bb27e 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -257,6 +257,7 @@ class RotaBuilder(object): self.SHIFT_BOUNDS = SHIFT_BOUNDS self.shifts: List[SingleShift] = [] # type List[SingleShift] + self.allowed_multi_shift_sets: list[set[str]] = [] self.use_previous_shifts = use_previous_shifts self.use_shift_balance_extra = use_shift_balance_extra @@ -856,10 +857,11 @@ class RotaBuilder(object): # within=Binary, initialize=0) def availability_init(model, wid, week, day): + # This needs to be higher than the max number of possible shifts per day if (wid, week, day) in self.unavailable_to_work: # print((wid, week, day)) return 0 - return 1 + return 3 self.model.available = Param( ( @@ -2412,14 +2414,40 @@ class RotaBuilder(object): ) ) - # single shift per day - self.model.constraints.add( - 1 - >= sum( - self.model.works[worker.id, week, day, shift] - for shift in self.get_shift_names_by_week_day(week, day) + # single shift per day (unless multi-shift allowed) + # This is signifantly slower so only enable if required + if self.allowed_multi_shift_sets: + shifts_today = self.get_shift_names_by_week_day(week, day) + if shifts_today: + # For each allowed set, allow those specific shifts to be assigned together + for allowed_set in self.allowed_multi_shift_sets: + if allowed_set.issubset(shifts_today): + # Allow up to len(allowed_set) shifts, but only for those shifts + self.model.constraints.add( + sum( + self.model.works[worker.id, week, day, shift] + for shift in allowed_set + ) <= len(allowed_set) + ) + # Prevent assigning any other shift together with the allowed set + for shift in shifts_today: + if shift not in allowed_set: + self.model.constraints.add( + sum( + self.model.works[worker.id, week, day, s] + for s in allowed_set + ) + self.model.works[worker.id, week, day, shift] + <= len(allowed_set) + ) + # Otherwise only one shift per day + else: + self.model.constraints.add( + 1 + >= sum( + self.model.works[worker.id, week, day, shift] + for shift in self.get_shift_names_by_week_day(week, day) + ) ) - ) # This applies to locums as well if self.get_workers_who_require_locums(): self.model.constraints.add( @@ -2800,6 +2828,12 @@ class RotaBuilder(object): # add objective function to the model. rule (pass function) or expr (pass expression directly) self.model.obj = Objective(rule=obj_rule, sense=minimize) + def allow_shifts_together(self, *shift_names: str): + """ + Allow the given set of shift names to be assigned together on a single day. + """ + self.allowed_multi_shift_sets.append(set(shift_names)) + def add_warning(self, warning_type: str, message: str): print(f"[bold red]WARNING:[/bold red] {warning_type} - {message}") self.warnings.append((warning_type, message)) diff --git a/test/test_multiple_shifts_per_day.py b/test/test_multiple_shifts_per_day.py new file mode 100644 index 0000000..bb90624 --- /dev/null +++ b/test/test_multiple_shifts_per_day.py @@ -0,0 +1,88 @@ + +import pytest +import datetime +from rota.shifts import RotaBuilder, SingleShift, days +from rota.workers import Worker + +@pytest.fixture +def basic_rota(): + start_date = datetime.date(2022, 4, 4) # must be a Monday! + Rota = RotaBuilder(start_date=start_date, weeks_to_rota=2) + worker = Worker(name="worker1", site="site1", grade=1) + Rota.add_worker(worker) + return Rota + +def test_single_shift_per_day_default(basic_rota): + Rota = basic_rota + Rota.add_shifts( + SingleShift(sites=["site1"], name="a", length=8, days=days[:5], workers_required=1), + SingleShift(sites=["site1"], name="b", length=8, days=days[:5], workers_required=1), + ) + Rota.build_and_solve() + assert Rota.results.solver.status == "error" + # Only one shift per day should be assigned + for week, day in Rota.get_week_day_combinations(): + assigned = [ + shift for shift in Rota.get_shift_names_by_week_day(week, day) + if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5 + ] + assert len(assigned) <= 1 + +def test_allow_two_shifts_per_day(basic_rota): + Rota = basic_rota + Rota.add_shifts( + SingleShift(sites=["site1"], name="a", length=8, days=days[:1], workers_required=1), + SingleShift(sites=["site1"], name="b", length=8, days=days[:1], workers_required=1), + ) + Rota.allow_shifts_together("a", "b") + Rota.build_and_solve() + Rota.export_rota_to_html("test_allow_two_shifts", folder="tests") + assert Rota.results.solver.status == "ok" + + ## Now both shifts can be assigned together + #for week, day in Rota.get_week_day_combinations(): + # assigned = [ + # shift for shift in Rota.get_shift_names_by_week_day(week, day) + # if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5 + # ] + # assert len(assigned) <= 2 + +def test_only_specified_shifts_can_be_double_assigned(basic_rota): + Rota = basic_rota + Rota.add_shifts( + SingleShift(sites=["site1"], name="a", length=8, days=days[:5], workers_required=1), + SingleShift(sites=["site1"], name="b", length=8, days=days[:5], workers_required=1), + SingleShift(sites=["site1"], name="c", length=8, days=days[:5], workers_required=1), + ) + Rota.allow_shifts_together("a", "b") + Rota.build_and_solve() + # Only "a" and "b" can be assigned together, never "a"+"c" or "b"+"c" + for week, day in Rota.get_week_day_combinations(): + assigned = [ + shift for shift in Rota.get_shift_names_by_week_day(week, day) + if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5 + ] + if set(["a", "b"]).issubset(set(Rota.get_shift_names_by_week_day(week, day))): + assert len(assigned) <= 2 + else: + assert len(assigned) <= 1 + +def test_triple_shifts_allowed(basic_rota): + Rota = basic_rota + Rota.add_shifts( + SingleShift(sites=["site1"], name="a", length=8, days=days[:5], workers_required=1), + SingleShift(sites=["site1"], name="b", length=8, days=days[:5], workers_required=1), + SingleShift(sites=["site1"], name="c", length=8, days=days[:5], workers_required=1), + ) + Rota.allow_shifts_together("a", "b", "c") + Rota.build_and_solve() + # All three can be assigned together if allowed + for week, day in Rota.get_week_day_combinations(): + assigned = [ + shift for shift in Rota.get_shift_names_by_week_day(week, day) + if Rota.model.works[Rota.workers[0].id, week, day, shift].value > 0.5 + ] + if set(["a", "b", "c"]).issubset(set(Rota.get_shift_names_by_week_day(week, day))): + assert len(assigned) <= 3 + else: + assert len(assigned) <= 1 \ No newline at end of file