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
-                    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))