feat: add exact shifts functionality to Worker class and update related logic in RotaBuilder

djangorota/rota/forms.py

@@ -45,6 +45,7 @@ class WorkerForm(forms.ModelForm):
     # Complex structured options edited as JSON
     assign_as_block_preferences = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON mapping shift_name->weight", label="Block preferences (JSON)")
     shift_fte_overrides = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON mapping shift_name->fte", label="Shift FTE overrides (JSON)")
+    exact_shifts = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON mapping shift_name->exact_count", label="Exact shifts (JSON)")
     previous_shifts = forms.CharField(required=False, widget=forms.Textarea, help_text="Free JSON structure for previous shifts", label="Previous shifts (JSON)")
     shift_balance_extra = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON structure for extra shift-balance penalties", label="Shift balance extra (JSON)")
     allowed_multi_shift_sets = forms.CharField(required=False, widget=forms.Textarea, help_text="JSON list of sets (e.g. [[\"a\",\"b\"], ...])", label="Allowed multi-shift sets (JSON)")
@@ -94,6 +95,7 @@ class WorkerForm(forms.ModelForm):
         for json_fld in [
             "assign_as_block_preferences",
             "shift_fte_overrides",
+            "exact_shifts",
             "previous_shifts",
             "shift_balance_extra",
             "allowed_multi_shift_sets",
@@ -146,6 +148,7 @@ class WorkerForm(forms.ModelForm):
         json_fields = [
             "assign_as_block_preferences",
             "shift_fte_overrides",
+            "exact_shifts",
             "previous_shifts",
             "shift_balance_extra",
             "allowed_multi_shift_sets",

djangorota/rota/models.py

@@ -229,6 +229,7 @@ class Worker(models.Model):
 					"max_shifts_per_week_by_shift_name",
 					"assign_as_block_preferences",
 					"shift_fte_overrides",
+					"exact_shifts",
 				]
 				list_keys = [
 					"oop",

gen_proc.py

@@ -30,6 +30,10 @@ sites = (
     "plymouth",
     "taunton",
     "proc",
+    "plymouth ir",
+    "truro ir",
+    "exeter ir",
+    "torbay ir",
 )
 
 from rota_generator.workers import (
@@ -49,7 +53,7 @@ def main(
     time_to_run: int = 60 * 60 * 10,
     ratio: float = 0.1,
     start_date: datetime.datetime = "2026-09-07",
-    weeks: int = 26,
+    weeks: int = 12,
     bom: float = 1,
 ):
     rota_start_date = start_date.date()
@@ -89,6 +93,7 @@ def main(
                 "exeter twilights and weekends",
                 "exeter no nights",
                 "exeter twilights",
+                "exeter ir",
             ),
             name="exeter_twilight",
             length=12.5,
@@ -104,6 +109,7 @@ def main(
                 "truro twilights and weekends",
                 "truro twilights and weekend nights",
                 "truro twilights, weekends and weekend nights",
+                "truro ir",
             ),
             name="truro_twilight",
             length=12.5,
@@ -112,7 +118,12 @@ def main(
             constraints=[MaxShiftsPerWeekConstraint(max_shifts=1)],
         ),
         SingleShift(
-            sites=("torbay", "torbay twilights", "torbay twilights and weekends"),
+            sites=(
+                "torbay",
+                "torbay twilights",
+                "torbay twilights and weekends",
+                "torbay ir",
+            ),
             name="torbay_twilight",
             length=12.5,
             days=days[:4],
@@ -125,6 +136,7 @@ def main(
                 "plymouth twilights",
                 "plymouth twilights and weekends",
                 "plymouth twilights, weekends and weekend nights",
+                "plymouth ir",
             ),
             name="plymouth_twilight",
             length=12.5,
@@ -216,6 +228,7 @@ def main(
                 "plymouth twilights",
                 "plymouth twilights and weekends",
                 "plymouth twilights, weekends and weekend nights",
+                "plymouth ir",
             ),
             name="plymouth_bank_holidays",
             length=8,
@@ -227,7 +240,9 @@ def main(
             bank_holidays_only=True,
         ),
         SingleShift(
-            sites=sites,
+            sites=[
+                *sites,
+            ],
             name="night_weekday",
             length=12.25,
             days=days[:4],
@@ -279,8 +294,12 @@ def main(
     # Prevent shifts for ST2s for 2 weeks
     # Rota.add_grade_constraint_by_week([2], [1, 2], [shift.name for shift in Rota.shifts])
 
-    Rota.add_min_summed_grade_by_shifts_per_day_constraint(["weekend_plymouth1", "weekend_plymouth2"], 6)
+    Rota.add_min_summed_grade_by_shifts_per_day_constraint(
-    Rota.add_min_summed_grade_by_shifts_per_day_constraint(["plymouth_twilight", "plymouth_bank_holidays"], 6)
+        ["weekend_plymouth1", "weekend_plymouth2"], 6
+    )
+    Rota.add_min_summed_grade_by_shifts_per_day_constraint(
+        ["plymouth_twilight", "plymouth_bank_holidays"], 6
+    )
 
     Rota.add_min_summed_grade_by_shifts_per_day_constraint(["night_weekday"], 12)
     Rota.add_min_summed_grade_by_shifts_per_day_constraint(["night_weekend"], 12)
@@ -406,7 +425,9 @@ def main(
                     elif "to" in raw:
                         s, e = raw.split("to", 1)
                     else:
-                        raise ValueError(f"Cannot parse OOP dates: '{oop}' for {worker}")
+                        raise ValueError(
+                            f"Cannot parse OOP dates: '{oop}' for {worker}"
+                        )
 
                     s = s.strip().strip("()[]\"' ")
                     e = e.strip().strip("()[]\"' ")
@@ -416,7 +437,9 @@ def main(
                         try:
                             start_dt = datetime.datetime.strptime(s, fmt).date()
                             end_dt = datetime.datetime.strptime(e, fmt).date()
-                            formatted_oops.append({"start_date": start_dt, "end_date": end_dt})
+                            formatted_oops.append(
+                                {"start_date": start_dt, "end_date": end_dt}
+                            )
                             parsed = True
                             break
                         except ValueError:
@@ -425,7 +448,9 @@ def main(
                     if not parsed:
                         print("WORKER", worker)
                         print("DATES", s, e)
-                        raise ValueError(f"Cannot parse OOP dates: '{oop}' for {worker}")
+                        raise ValueError(
+                            f"Cannot parse OOP dates: '{oop}' for {worker}"
+                        )
 
                 oop = formatted_oops
             else:
@@ -465,6 +490,19 @@ def main(
 
             shift_fte_overrides = {}
 
+            exact_shifts = {}
+
+            if worker in ["hannah.lewis29@nhs.net", "shruti.bodapati@nhs.net", "a.abouelatta@nnhs.net"]:
+                exact_shifts = {
+                    "night_weekday": 4,
+                    "night_weekend": 3,
+                }
+            elif worker in ["g.abdelhalim@nhs.net", ""]:
+                exact_shifts = {
+                    "night_weekend": 3,
+                    "night_weekday": 0,
+                }
+
             # if worker_name == "Hadi Mohamed":
             #    shift_fte_overrides = {
             #        "weekend_exeter": 50,
@@ -497,6 +535,7 @@ def main(
                 shift_balance_extra=w["shift_balance_extra"],
                 bank_holiday_extra=w["bank_holiday_extra"],
                 shift_fte_overrides=shift_fte_overrides,
+                exact_shifts=exact_shifts,
             )
 
             # print(w)
@@ -507,7 +546,6 @@ def main(
     # Rota.build_workers()
     # Rota.build_model()
 
-
     solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
     # solver_options = {"seconds": time_to_run, "threads": 10}
 

leave.py

@@ -219,6 +219,21 @@ def load_leave(Rota):
 
                 #elif row_title == "Academy":
                 #    pass
+                #elif "night specific" in row_title:
+                #    if lower_item:
+                #        for shift in lower_item.split(","):
+                #            shift = shift.strip()
+                #            if shift == "nights":
+                #                shift = "night_weekday"
+                #            elif shift == "nights only":
+                #                if date.weekday() < 5:
+                #                    shift = "night_weekday"
+                #                else:
+                #                    shift = "night_weekend"
+                #            worker["single_nights"].append(
+                #                WorkRequests(date=date, shift=shift)
+                #            )
+
 
                 elif re.match(date_re, row_date) is not None:
                     if lower_item != "":

rota_generator/shifts.py

@@ -2648,39 +2648,48 @@ class RotaBuilder(object):
             ):
                 if (
                     worker.site in shift.sites
-                ):  # Each site specfies which sites self.workers can fullfill it
+                ):  # Each site specfies which sites self.workers can fullfill
-                    # calaculate the total number of shifts that need assigning over the rota
-                    # total_shifts = (len(self.weeks) * len(shift.days) *
-                    # TODO: check if shift_counts is still required?
-                    # total_shifts = (
-                    #    self.shift_counts[shift.name] * shift.workers_required
-                    # )
                     total_shifts = self.shift_worker_counts[shift.name]
 
+                    # Find workers who have exact shifts for this shift type
+                    exact_workers = [
+                        w for w in self.workers 
+                        if w.site in shift.sites and shift.name in getattr(w, "exact_shifts", {})
+                    ]
+                    exact_shifts_sum = sum(w.exact_shifts[shift.name] for w in exact_workers)
+
+                    if exact_shifts_sum > total_shifts:
+                        warning_msg = f"For shift {shift.name}, total exact shifts ({exact_shifts_sum}) exceeds total required shifts ({total_shifts})."
+                        if not any(warn[1] == warning_msg for warn in self.warnings):
+                            self.add_warning("Exact shifts exceed total shifts", warning_msg)
+
+                    if worker in exact_workers:
+                        target_shifts = worker.exact_shifts[shift.name]
+                    else:
+                        remaining_shifts = total_shifts - exact_shifts_sum
+                        # Sum FTE only for workers who do not have exact shifts for this shift type
                         full_time_equivalent_joined = sum(
-                        self.full_time_equivalent_sites[i][shift.name]
+                            w.get_fte(shift=shift.name)
-                        for i in shift.sites
+                            for w in self.workers
+                            if w.site in shift.sites and w not in exact_workers
                         )
 
-                    # Guard against division by zero: if no full-time
-                    # equivalent is available for the shift's sites the
-                    # denominator may be zero (for example when all site
-                    # FTEs are 0). In that case emit a warning and set the
-                    # target to 0 rather than raising an exception.
                         if not full_time_equivalent_joined:
                             # avoid ZeroDivisionError
+                            if remaining_shifts > 0:
                                 self.add_warning(
-                            "Zero FTE for sites",
+                                    "Zero FTE for sites with remaining shifts",
-                            f"full_time_equivalent sum is zero for shift {shift.name}; setting target_shifts=0 for worker {worker.name}",
+                                    f"full_time_equivalent sum is zero for shift {shift.name} but remaining_shifts = {remaining_shifts}; setting target_shifts=0 for worker {worker.name}",
                                 )
                             target_shifts = 0
                         else:
                             target_shifts = (
-                            total_shifts
+                                max(0.0, float(remaining_shifts))
                                 / full_time_equivalent_joined
                                 * worker.get_fte(shift=shift.name)
                             )
 
+                    if worker not in exact_workers:
                         if self.use_previous_shifts:
                             if shift.name in worker.previous_shifts:
                                 entry = worker.previous_shifts[shift.name]
@@ -2730,7 +2739,11 @@ class RotaBuilder(object):
                     # print(worker.name, shift.name, target_shifts)
                     worker.shift_target_number[shift.name] = target_shifts
 
-                    if shift.hard_constrain_shift:
+                    if worker in exact_workers:
+                        self.model.constraints.add(
+                            self.model.shift_count[worker.id, shift.name] == target_shifts
+                        )
+                    elif shift.hard_constrain_shift:
                         adjusted_balance_offset = shift.balance_offset
                         if worker.get_fte(shift=shift.name) < 100:
                             adjusted_balance_offset = (
@@ -4244,6 +4257,9 @@ class RotaBuilder(object):
 
         Must be called prior to attempting to solve
         """
+        if not hasattr(self, "shifts_by_name"):
+            self.build_shifts()
+
         if not self.workers:
             raise NoWorkers("Workers must be added prior to calling build_workers")
 
@@ -4253,6 +4269,19 @@ class RotaBuilder(object):
         for worker in track(self.workers, description="Building workers"):
             print(worker.name)
             worker.load_rota(self)
+            for s_name in getattr(worker, "exact_shifts", {}):
+                try:
+                    s = self.get_shift_by_name(s_name)
+                    if worker.site not in s.sites:
+                        self.add_warning(
+                            "Invalid exact shift site",
+                            f"Worker {worker.name} requested exact shifts for shift {s_name} but their site {worker.site} is not in the shift sites {s.sites}",
+                        )
+                except KeyError:
+                    self.add_warning(
+                        "Invalid exact shift",
+                        f"Worker {worker.name} requested exact shifts for non-existent shift {s_name}",
+                    )
             wid = worker.id
             if wid in self.workers_id_map:
                 message = f"Worker with id '{wid}' has been added twice"

rota_generator/workers.py

@@ -213,6 +213,8 @@ class Worker(BaseModel):
 
     shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist
 
+    exact_shifts: dict[str, int] = {} # Map shift_name to exact number of shifts
+
     weekend_shift_target_number: float = 0.0
     avoid_shifts_on_dates: list[AvoidShiftOnDates] = []
 

test/test_workers.py

@@ -1625,3 +1625,211 @@ def test_worker_hard_day_exclusion2():
             sat_idx = week * 7 + days.index("Sat")
             sun_idx = week * 7 + days.index("Sun")
             assert not (shift_list[sat_idx] == "a" and shift_list[sun_idx] == "a"), f"Worker {worker.name} has 'a' on both Mon and Tue in week {week}, which should not happen due to hard day exclusion"
+
+
+def test_exact_shifts_distribution():
+    weeks_to_rota = 10
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    workers = [
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"a": 4}),
+        Worker(name="worker2", site="group1", grade=1, fte=100),
+        Worker(name="worker3", site="group1", grade=1, fte=50),
+    ]
+    Rota.add_workers(workers)
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            hard_constrain_shift=True,
+            balance_offset=0,
+        ),
+    )
+
+    Rota.build_and_solve()
+
+    assert Rota.workers_name_map["worker1"].shift_target_number["a"] == 4
+    assert Rota.workers_name_map["worker2"].shift_target_number["a"] == 4
+    assert Rota.workers_name_map["worker3"].shift_target_number["a"] == 2
+
+    # Verify actual assigned counts are exactly correct
+    worker_shift_counts = {w.name: 0 for w in Rota.get_workers()}
+    for week, day, shiftname in Rota.get_all_shiftname_combinations():
+        for worker in Rota.get_workers():
+            val = Rota.model.works[worker.id, week, day, shiftname].value
+            if val is not None and val > 0.5:
+                if shiftname == "a":
+                    worker_shift_counts[worker.name] += 1
+
+    assert worker_shift_counts["worker1"] == 4
+    assert worker_shift_counts["worker2"] == 4
+    assert worker_shift_counts["worker3"] == 2
+
+
+def test_exact_shifts_distribution_unbalanced():
+    weeks_to_rota = 16
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+
+    workers = [
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"a": 4}),
+        Worker(name="worker2", site="group1", grade=1, fte=100),
+        Worker(name="worker3", site="group1", grade=1, fte=50),
+    ]
+    Rota.add_workers(workers)
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+            hard_constrain_shift=True,
+            balance_offset=0,
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=12.5,
+            days=("Tue",),
+            workers_required=1,
+            hard_constrain_shift=True,
+            balance_offset=0.8,
+        ),
+    )
+
+    Rota.build_and_solve()
+    Rota.export_rota_to_html("test_exact_shifts_distribution_unbalanced", folder="tests")
+
+    assert Rota.results.solver.status == "ok"
+
+    assert Rota.workers_name_map["worker1"].shift_target_number["a"] == 4
+    assert Rota.workers_name_map["worker2"].shift_target_number["a"] == 8
+    assert Rota.workers_name_map["worker3"].shift_target_number["a"] == 4
+
+    assert Rota.workers_name_map["worker1"].shift_target_number["b"] == 6.4
+    assert Rota.workers_name_map["worker2"].shift_target_number["b"] == 6.4
+    assert Rota.workers_name_map["worker3"].shift_target_number["b"] == 3.2
+
+    # Verify actual assigned counts are exactly correct
+    worker_shift_counts_a = {w.name: 0 for w in Rota.get_workers()}
+    worker_shift_counts_b = {w.name: 0 for w in Rota.get_workers()}
+    for week, day, shiftname in Rota.get_all_shiftname_combinations():
+        for worker in Rota.get_workers():
+            val = Rota.model.works[worker.id, week, day, shiftname].value
+            if val is not None and val > 0.5:
+                if shiftname == "a":
+                    worker_shift_counts_a[worker.name] += 1
+                elif shiftname == "b":
+                    worker_shift_counts_b[worker.name] += 1
+
+    assert worker_shift_counts_a["worker1"] == 4
+    assert worker_shift_counts_a["worker2"] == 8
+    assert worker_shift_counts_a["worker3"] == 4
+
+    assert worker_shift_counts_b["worker1"] in (6, 7)  # Allowing for rounding differences
+    assert worker_shift_counts_b["worker2"] in (6, 7)  # Allowing for rounding differences
+    assert worker_shift_counts_b["worker3"] in (3, 4)  # Allowing for rounding differences
+
+
+def test_exact_shifts_invalid_site_warning():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group2", grade=1, fte=100, exact_shifts={"a": 4}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+        SingleShift(
+            sites=("group2",),
+            name="b",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+    )
+    Rota.build_workers()
+    
+    warnings = Rota.get_warnings("Invalid exact shift site")
+    assert len(warnings) > 0
+    assert "group2 is not in the shift sites" in warnings[0][1]
+
+
+def test_exact_shifts_invalid_shift_warning():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"nonexistent": 4}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+    )
+    Rota.build_workers()
+    
+    warnings = Rota.get_warnings("Invalid exact shift")
+    assert len(warnings) > 0
+    assert "non-existent shift nonexistent" in warnings[0][1]
+
+
+def test_exact_shifts_exceed_total_warning():
+    weeks_to_rota = 2
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    Rota.add_workers([
+        Worker(name="worker1", site="group1", grade=1, fte=100, exact_shifts={"a": 4}),
+    ])
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=("Mon",),
+            workers_required=1,
+        ),
+    )
+    Rota.build_workers()
+    Rota.build_model()
+    
+    warnings = Rota.get_warnings("Exact shifts exceed total shifts")
+    assert len(warnings) > 0
+    assert "total exact shifts (4) exceeds total required shifts" in warnings[0][1]