add limit grades rule

requirements.txt

@@ -1,3 +1,4 @@
 pyomo
 govuk_bank_holidays
 pytest
+black

shifts.py

@@ -1,6 +1,6 @@
 import datetime
 import itertools
-from typing import Iterable, List, Sequence, Tuple, Set
+from typing import Dict, Iterable, List, Sequence, Tuple, Set
 
 import datetime
 
@@ -82,7 +82,16 @@ class SingleShift(object):
         self.assign_as_block = assign_as_block
         self.force_as_block = force_as_block
         self.force_as_block_unless_nwd = force_as_block_unless_nwd
-        self.constraints = constraints
+
+        self.constraint_options = {}
+        self.constraints = []
+        for c in constraints:
+            match c:
+                case (constraint, options):
+                    self.constraints.append(constraint)
+                    self.constraint_options[constraint] = options
+                case constraint:
+                    self.constraints.append(constraint)
 
         self.bank_holidays_only = bank_holidays_only
 
@@ -170,7 +179,8 @@ class RotaBuilder(object):
             "balance_nights_across_sites": True,
             "balance_bank_holidays": True,
             "balance_blocks": True,
-            "balance_shifts": True,
+            "balance_shifts": True, # Does not use a quadratic function
+            "balance_shifts_quadratic": False, 
             "balance_shifts_over_workers": True,
             "minimise_shift_diffs": False,  # less sophisticated version of balance_shifts_over_workers
             "balance_weekends": True,
@@ -436,7 +446,7 @@ class RotaBuilder(object):
                 initialize=0,
             )
 
-        if self.constraint_options["balance_shifts"]:
+        if self.constraint_options["balance_shifts"] or self.constraint_options["balance_shifts_quadratic"]:
             self.model.shift_count_t1 = Var(
                 (
                     (worker.id, shift.name)
@@ -457,6 +467,17 @@ class RotaBuilder(object):
                 initialize=0,
             )
 
+        if self.constraint_options["balance_shifts_quadratic"]:
+            self.model.shift_count_w = Var(
+                (
+                    (worker.id, shift.name)
+                    for worker in self.workers
+                    for shift in self.get_shifts()
+                ),
+                domain=NonNegativeReals,
+                initialize=0,
+            )
+
         if self.constraint_options["balance_nights"]:
             # We also try to even out the night shifts seperately
             self.model.night_shift_count = Var(
@@ -722,6 +743,35 @@ class RotaBuilder(object):
                 rule=nightShiftMaxSTRule,
             )
 
+        # Limit grades
+        def limitGradesRule(model, week, shift, grade, limit):
+            workers = [w for w in self.workers if w.grade == grade]
+            if not workers:
+                return Constraint.Skip
+            return (
+                sum(
+                    model.shift_week_worker_assigned[shift.name, week, w.id]
+                    for w in workers
+                )
+                <= limit
+            )
+
+        for shift in self.get_shifts_with_constraint("limit_grade_number"):
+            if not shift.constraint_options["limit_grade_number"]:
+                raise ValueError("Constraint option must be defined for 'limit_grade_number'")
+            for grade in shift.constraint_options["limit_grade_number"]:
+                #self.model.limit_grades_constraint = Constraint(
+                setattr(self.model, f"limit_grades_constraint_{shift.name}_{grade}", Constraint(
+                    [week for week in self.weeks],
+                    #[shift for shift in self.get_shifts_with_constraint("limit_grade_number")],
+                    [shift],
+                    [grade],
+                    [shift.constraint_options["limit_grade_number"][grade]],
+                    rule=limitGradesRule,
+                    #name=f"limit_grades_constraint_{shift.name}_{grade}"
+                ))
+
+
         def nightShiftMinST4Rule(model, week, shift):
             single_workers = [w for w in self.workers if w.grade >= 4]
             if not single_workers:
@@ -785,29 +835,30 @@ class RotaBuilder(object):
 
         # # Count the number of workers from each site on each night shift
         # # As 1 or 0 is optimum we can simply subtract 1 from the number
-        if self.constraint_options["balance_nights_across_sites"]:
+        #if self.constraint_options["balance_nights_across_sites"]:
-            for site in self.sites:
+        #    for site in self.sites:
-                for shift in self.get_shifts_with_constraint("night"):
+        #        for shift in self.get_shifts_with_constraint("night"):
-                    block = shift.name
+        #            block = shift.name
-                    for week in self.weeks:
+        #            for week in self.weeks:
-                        self.model.constraints.add(
+        #                self.model.constraints.add(
-                            self.model.night_per_site[week, block, site]
+        #                    self.model.night_per_site[week, block, site]
-                            == sum(
+        #                    == sum(
-                                self.model.shift_week_worker_assigned[
+        #                        self.model.shift_week_worker_assigned[
-                                    block, week, worker.id
+        #                            block, week, worker.id
-                                ]
+        #                        ]
-                                for worker in self.workers
+        #                        for worker in self.workers
-                                if worker.site == site
+        #                        if worker.site == site
-                            )
+        #                    )
-                            # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site)
+        #                    # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site)
-                        )
+        #                )
-                        self.model.constraints.add(
+        #                self.model.constraints.add(
-                            self.model.night_per_site_t1[week, block, site]
+        #                    self.model.night_per_site_t1[week, block, site]
-                            - self.model.night_per_site_t2[week, block, site]
+        #                    - self.model.night_per_site_t2[week, block, site]
-                            == self.model.night_per_site[week, block, site] - 1
+        #                    == self.model.night_per_site[week, block, site] - 1
-                        )
+        #                )
         for site in self.sites:
             for shift in self.get_shifts_with_constraint("balance_across_groups"):
+                if site in shift.site:
                     block = shift.name
                     for week in self.weeks:
                         self.model.constraints.add(
@@ -1082,6 +1133,46 @@ class RotaBuilder(object):
                         == self.model.shift_count_diff[worker.id, shift.name]
                     )
 
+
+                if self.constraint_options["balance_shifts_quadratic"]:
+                    # This may need to be updated
+                    xU = 25
+                    xL = 1
+                    self.model.constraints.add(
+                        inequality(
+                            xL,
+                            self.model.shift_count_t1[worker.id, shift.name]
+                            + self.model.shift_count_t2[worker.id, shift.name]
+                            + 1,
+                            xU,
+                        )
+                    )
+
+                    self.model.constraints.add(
+                        self.model.shift_count_w[worker.id, shift.name]
+                        >= xL
+                        * (
+                            self.model.shift_count_t1[worker.id, shift.name]
+                            + self.model.shift_count_t2[worker.id, shift.name]
+                            + 1
+                        )
+                        * 2
+                        - xL * xL
+                    )
+
+                    self.model.constraints.add(
+                        self.model.shift_count_w[worker.id, shift.name]
+                        >= xU
+                        * (
+                            self.model.shift_count_t1[worker.id, shift.name]
+                            + self.model.shift_count_t2[worker.id, shift.name]
+                            + 1
+                        )
+                        * 2
+                        - xU * xU
+                    )
+
+
             # Define worker_shift_count_t1 and worker_shift_count_t2 constraints for the object
             # Thus bypassing the need for a quadratic solver
             # t1-t2 is the target
@@ -1459,7 +1550,9 @@ class RotaBuilder(object):
                 #     self.model.shift_week_worker_assigned[s.name, week, worker.id]
                 #     for s in self.get_shifts_with_constraint("night")))
 
-                for shift in self.get_shifts_with_constraint("night"):
+                #for shift in self.get_shifts_with_constraint("night"):
+                for shift in self.get_shifts():
+                    if shift.force_as_block:
                         # Force nights to be assigned in blocks
                         # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days)
                         #     )
@@ -1728,6 +1821,19 @@ class RotaBuilder(object):
             else:
                 shift_balancing = 0
 
+            if self.constraint_options["balance_shifts_quadratic"]:
+                shift_balancing = sum(
+                    balance_modifier_constant
+                    * (
+                        self.model.shift_count_t1[worker.id, shift.name]
+                        + self.model.shift_count_t2[worker.id, shift.name]
+                    )
+                    for worker in self.workers
+                    for shift in self.get_shifts()
+                )
+            else:
+                shift_balancing = 0
+
             if self.constraint_options["minimise_shift_diffs"]:
                 shift_diff_modifier_constant = 10
                 shift_diff_balancing = sum(
@@ -1808,6 +1914,22 @@ class RotaBuilder(object):
             else:
                 nights_site_balancing = 0
 
+
+            if self.get_shifts_with_constraint("balance_across_groups"):
+                block_site_balancing = sum(
+                    (
+                        self.model.shift_per_site_t1[week, shift.name, site]
+                        + self.model.shift_per_site_t2[week, shift.name, site]
+                    )
+                    * 2000
+                    # self.model.night_per_site2[week, block, site]
+                    for week in self.weeks
+                    for shift in self.get_shifts_with_constraint("balance_across_groups")
+                    for site in self.sites
+                )
+            else:
+                block_site_balancing = 0
+
             if self.constraint_options["balance_blocks"]:
                 blocks_balancing = sum(
                     1 * self.model.blocks_assigned[week, shift]
@@ -1833,6 +1955,7 @@ class RotaBuilder(object):
                 + shift_diff_balancing
                 + preferences
                 + nights_site_balancing
+                + block_site_balancing
                 + bank_holiday_balancing
                 + blocks_balancing
                 - work_requests
@@ -2117,6 +2240,20 @@ class RotaBuilder(object):
         s.extend(self.shifts_to_force_as_blocks())
         return s
 
+    def get_workers_by_group(self) -> Dict[str, Worker]:
+        group_workers = defaultdict(set)
+        for worker in self.workers:
+            group_workers[worker.site].add(worker)
+
+        return group_workers
+
+    def get_workers_by_grade(self) -> Dict[str, Worker]:
+        group_workers = defaultdict(set)
+        for worker in self.workers:
+            group_workers[worker.grade].add(worker)
+
+        return group_workers
+
     def get_workers_for_shift(self, shift: SingleShift) -> List[Worker]:
         return [worker for worker in self.workers if worker.site in shift.site]
 

test_rota.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from black import main
 import pytest
 from shifts import NoWorkers, RotaBuilder, SingleShift, days
@@ -383,6 +384,11 @@ class TestDemoRotaShiftConstraints:
         ),
     )
 
+    Rota3Workers = deepcopy(Rota)
+    worker3 = Worker(Rota3Workers, "worker3", "group1", 1)
+
+    Rota3Workers.add_worker(worker3)
+
     def test_max_shifts(self):
         self.Rota.constraint_options["max_shifts_per_month"] = 14
         self.Rota.build_and_solve()
@@ -397,23 +403,11 @@ class TestDemoRotaShiftConstraints:
         assert self.Rota.results.solver.termination_condition == "infeasible"
 
     def test_max_shifts_extra_worker(self):
-        self.Rota.constraint_options["max_shifts_per_month"] = 12
+        self.Rota3Workers.constraint_options["max_shifts_per_month"] = 12
-        worker3 = Worker(self.Rota, "worker3", "group1", 1)
+        self.Rota3Workers.build_and_solve()
 
-        self.Rota.add_worker(worker3)
+        assert self.Rota3Workers.results.solver.status == "ok"
-        self.Rota.build_and_solve()
+        assert self.Rota3Workers.results.solver.termination_condition == "optimal"
-
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
-
-    def test_max_shifts_extra_worker_wrong_group(self):
-        self.Rota.constraint_options["max_shifts_per_month"] = 12
-        worker3 = Worker(self.Rota, "worker3", "group2", 1)
-        self.Rota.add_worker(worker3)
-        self.Rota.build_and_solve()
-
-        assert self.Rota.results.solver.status == "warning"
-        assert self.Rota.results.solver.termination_condition == "infeasible"
 
     def test_max_shifts_per_week_fail(self):
         self.Rota.constraint_options["max_shifts_per_week"] = 3
@@ -423,18 +417,19 @@ class TestDemoRotaShiftConstraints:
         assert self.Rota.results.solver.termination_condition == "infeasible"
 
     def test_max_shifts_per_week_pass(self):
+        self.Rota.constraint_options["max_shifts_per_month"] = 14
         self.Rota.constraint_options["max_shifts_per_week"] = 4
         self.Rota.build_and_solve()
 
         assert self.Rota.results.solver.status == "ok"
         
     def test_max_shifts_per_week_extra_worker_pass(self):
-        self.Rota.constraint_options["max_shifts_per_week"] = 3
+        self.Rota3Workers.constraint_options["max_shifts_per_week"] = 3
-        worker3 = Worker(self.Rota, "worker3", "group1", 1)
+        #worker3 = Worker(self.Rota, "worker3", "group1", 1)
-        self.Rota.add_worker(worker3)
+        #self.Rota.add_worker(worker3)
-        self.Rota.build_and_solve()
+        self.Rota3Workers.build_and_solve()
 
-        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota3Workers.results.solver.status == "ok"
 
 class TestDemoRotaBalanceShiftSites:
     weeks_to_rota = 8
@@ -450,34 +445,238 @@ class TestDemoRotaBalanceShiftSites:
     worker2 = Worker(Rota, "worker2", "group1", 1)
     worker3 = Worker(Rota, "worker3", "group2", 1)
     worker4 = Worker(Rota, "worker4", "group2", 1)
-    worker5 = Worker(Rota, "worker3", "group2", 1)
+    worker5 = Worker(Rota, "worker5", "group3", 1)
-    worker6 = Worker(Rota, "worker4", "group2", 1)
+    worker6 = Worker(Rota, "worker6", "group3", 1)
 
     Rota.add_workers((worker1, worker2, worker3, worker4))
-    #Rota.add_workers((worker5, worker6))
+    Rota.add_workers((worker5, worker6))
-    Rota.add_shifts(
+
+    def test_balance_blocks_across_2_groups(self):
+        self.Rota.add_shifts(
             SingleShift(("group1", "group2"), "weekday_night", 12.5, days[:4], balance_offset=40,
             workers_required=2,
             force_as_block=True,
             constraints=["balance_across_groups"]
             ),
-        #SingleShift(("group1", "group2"), "weekend_night", 12.5, days[4:], balance_offset=40,
-        #workers_required=2,
-        #force_as_block=True,
-        #),
         )
-
+        self.Rota.constraint_options["balance_nights_across_sites"] = False
-    def test_balance_blocks_across_groups(self):
+        self.Rota.constraint_options["balance_nights"] = False
-        self.Rota.constraint_options["balance_nights_across_sites"] = True
+        self.Rota.constraint_options["constrain_time_off_after_nights"] = False
-        #self.Rota.constraint_options["balance_bank_holidays"] = False
+        self.Rota.build_and_solve(options={"ratio": 0.0})
-        #self.Rota.constraint_options["balance_weekends"] = False
-        #self.Rota.constraint_options["balance_blocks"] = False
-        self.Rota.build_and_solve(options={"ratio": 0.0, "seconds": 1000, "threads": 10})
         self.Rota.export_rota_to_html("test5")
 
+        print(self.Rota.get_workers_by_group())
+
+        group_workers = self.Rota.get_workers_by_group()
+        for group in group_workers:
+            shift_patterns = []
+            for w in group_workers[group]:
+                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+
+            zipped_lists = list(zip(*shift_patterns))
+            
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("weekday_night") <= 1
+
         assert self.Rota.results.solver.status == "ok"
         assert self.Rota.results.solver.termination_condition == "optimal"
 
+    def test_balance_blocks_across_2_groups_unbalanced(self):
+        self.Rota.shifts = []
+        self.Rota.add_shifts(
+            SingleShift(("group1", "group2"), "weekday_night", 12.5, days[:4], balance_offset=40,
+            workers_required=3,
+            force_as_block=True,
+            constraints=["balance_across_groups"]
+            ),
+        )
+        self.Rota.constraint_options["balance_nights_across_sites"] = False
+        self.Rota.constraint_options["balance_nights"] = False
+        self.Rota.constraint_options["constrain_time_off_after_nights"] = False
+        self.Rota.build_and_solve(options={"ratio": 0.1})
+
+        group_workers = self.Rota.get_workers_by_group()
+        for group in group_workers:
+            shift_patterns = []
+            for w in group_workers[group]:
+                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+
+            zipped_lists = list(zip(*shift_patterns))
+            
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("weekday_night") <= 2
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+
+    def test_balance_blocks_across_3_groups(self):
+        self.Rota.shifts = []
+
+        self.Rota.add_shifts(
+            SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days[:4], balance_offset=40,
+            workers_required=3,
+            force_as_block=True,
+            constraints=["balance_across_groups"]
+            ),
+        )
+        self.Rota.constraint_options["balance_nights_across_sites"] = False
+        self.Rota.constraint_options["balance_nights"] = False
+        self.Rota.constraint_options["constrain_time_off_after_nights"] = False
+        self.Rota.build_and_solve(options={"ratio": 0.1})
+
+        group_workers = self.Rota.get_workers_by_group()
+        for group in group_workers:
+            shift_patterns = []
+            for w in group_workers[group]:
+                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+
+            zipped_lists = list(zip(*shift_patterns))
+            
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("weekday_night") <= 1
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+class TestLimitConstraints:
+    weeks_to_rota = 8
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+    )
+    #Rota.constraint_options["max_shifts_per_week"] = 5
+    #Rota.constraint_options["max_shifts_per_month"] = 20
+    worker1 = Worker(Rota, "worker1", "group1", 2)
+    worker2 = Worker(Rota, "worker2", "group1", 2)
+    worker3 = Worker(Rota, "worker3", "group1", 2)
+    worker4 = Worker(Rota, "worker4", "group3", 3)
+    worker5 = Worker(Rota, "worker5", "group3", 3)
+    worker6 = Worker(Rota, "worker6", "group3", 3)
+
+    Rota.add_workers((worker1, worker2, worker3, worker4))
+    Rota.add_workers((worker5, worker6))
+
+    Rota.shifts = []
+
+
+    def test_constraint_limit_grades(self):
+        self.Rota.shifts = []
+        self.Rota.add_shifts(
+            SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraints=[("limit_grade_number", {2: 1})]
+            ),
+        )
+        #self.Rota.constraint_options["balance_nights_across_sites"] = False
+        #self.Rota.constraint_options["balance_shifts_over_workers"] = False
+        self.Rota.constraint_options["balance_shifts_quadratic"] = True
+        self.Rota.build_and_solve(options={"ratio": 0.00})
+
+        grade_workers = self.Rota.get_workers_by_grade()
+        for grade in grade_workers:
+            shift_patterns = []
+            for w in grade_workers[grade]:
+                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+
+            zipped_lists = list(zip(*shift_patterns))
+            
+            limit = 1
+            if grade > 2:
+                limit = 3
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("weekday_night") == limit
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+    def test_constraint_limit_grades2(self):
+        self.Rota.shifts = []
+        self.Rota.add_shifts(
+            SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraints=[("limit_grade_number", {3:2})]
+            ),
+        )
+        #self.Rota.constraint_options["balance_nights_across_sites"] = False
+        #self.Rota.constraint_options["balance_shifts_over_workers"] = False
+        self.Rota.constraint_options["balance_shifts_quadratic"] = True
+        self.Rota.build_and_solve(options={"ratio": 0.1})
+
+        self.Rota.export_rota_to_html("test9")
+
+        grade_workers = self.Rota.get_workers_by_grade()
+        for grade in grade_workers:
+            shift_patterns = []
+            for w in grade_workers[grade]:
+                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+
+            zipped_lists = list(zip(*shift_patterns))
+            
+            limit = 2
+            if grade > 2:
+                limit = 2
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("weekday_night") == limit
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+    def test_constraint_limit_grades3(self):
+        self.Rota.shifts = []
+        self.Rota.add_shifts(
+            SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraints=[("limit_grade_number", {2: 3, 3:1})]
+            ),
+        )
+        #self.Rota.constraint_options["balance_nights_across_sites"] = False
+        #self.Rota.constraint_options["balance_shifts_over_workers"] = False
+        self.Rota.constraint_options["balance_shifts_quadratic"] = True
+        self.Rota.build_and_solve(options={"ratio": 0.1})
+
+        self.Rota.export_rota_to_html("test9")
+
+        grade_workers = self.Rota.get_workers_by_grade()
+        for grade in grade_workers:
+            shift_patterns = []
+            for w in grade_workers[grade]:
+                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+
+            zipped_lists = list(zip(*shift_patterns))
+            
+            limit = 3
+            if grade > 2:
+                limit = 1
+            for day_shifts in zipped_lists:
+                assert day_shifts.count("weekday_night") == limit
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+    def test_constraint_limit_grades4(self):
+        self.Rota.shifts = []
+        self.Rota.add_shifts(
+            SingleShift(("group1", "group2", "group3"), "weekday_night", 12.5, days, balance_offset=40,
+            workers_required=4,
+            force_as_block=True,
+            constraints=[("limit_grade_number", {2: 4, 3:0})]
+            ),
+        )
+        #self.Rota.constraint_options["balance_nights_across_sites"] = False
+        #self.Rota.constraint_options["balance_shifts_over_workers"] = False
+        self.Rota.constraint_options["balance_shifts_quadratic"] = True
+        self.Rota.build_and_solve(options={"ratio": 0.1})
+
+        self.Rota.export_rota_to_html("test9")
+
+        assert self.Rota.results.solver.termination_condition == "infeasible"
+
 #class TestNoWorkerRota:
 #    weeks_to_rota = 10
 #    start_date = datetime.date(2022, 3, 7)
@@ -493,5 +692,5 @@ class TestDemoRotaBalanceShiftSites:
 
 
 if __name__ == "__main__":
-    t = TestDemoRotaShiftConstraints()
+    t = TestLimitConstraints()
-    t.test_balance_blocks_across_groups()
+    t.test_constraint_limit_grades4()