basic multishift per day assignement

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
+                    # single shift per day (unless multi-shift allowed)
-                    self.model.constraints.add(
+                    # This is signifantly slower so only enable if required
-                        1
+                    if self.allowed_multi_shift_sets:
-                        >= sum(
+                        shifts_today = self.get_shift_names_by_week_day(week, day)
-                            self.model.works[worker.id, week, day, shift]
+                        if shifts_today:
-                            for shift in self.get_shift_names_by_week_day(week, day)
+                            # 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))

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