many improvements

run.py

@@ -47,7 +47,7 @@ Rota.constraint_options["constrain_time_off_after_nights"] = True
 Rota.constraint_options["balance_nights_across_sites"] = True
 Rota.constraint_options["balance_blocks"] = True
 Rota.constraint_options["balance_weekends"] = True
-Rota.constraint_options["balance_shifts"] = True
+Rota.constraint_options["balance_shifts_over_workers"] = True
 Rota.constraint_options["balance_bank_holidays"] = True
 Rota.constraint_options["avoid_st2_first_month"] = False
 Rota.constraint_options["hard_constrain_pair_separation"] = True

runfixnights.py

@@ -47,7 +47,7 @@ Rota.constraint_options["constrain_time_off_after_nights"] = True
 Rota.constraint_options["balance_nights_across_sites"] = True
 Rota.constraint_options["balance_blocks"] = True
 Rota.constraint_options["balance_weekends"] = True
-Rota.constraint_options["balance_shifts"] = True
+Rota.constraint_options["balance_shifts_over_workers"] = True
 Rota.constraint_options["balance_bank_holidays"] = True
 Rota.constraint_options["avoid_st2_first_month"] = False
 Rota.constraint_options["hard_constrain_pair_separation"] = True

test_rota.py

@@ -1,6 +1,6 @@
 from black import main
 import pytest
-from shifts import NoWorkers, RotaBuilder, RotaResults, SingleShift, days
+from shifts import NoWorkers, RotaBuilder, SingleShift, days
 
 import datetime
 
@@ -27,8 +27,10 @@ class TestDemoRota:
     Rota.constraint_options["constrain_time_off_after_nights"] = False
     Rota.constraint_options["balance_nights_across_site"] = False
     Rota.constraint_options["balance_shifts"] = True
+    Rota.constraint_options["balance_shifts_over_workers"] = True
     Rota.constraint_options["balance_nights"] = False
-    Rota.constraint_options["minimise_shift_diffs"] = True
+    Rota.constraint_options["minimise_shift_diffs"] = False
+    Rota.constraint_options["balance_blocks"] = False
     #Rota.constraint_options["balance_weekends"] = False
 
     # Add a few workers
@@ -52,7 +54,7 @@ class TestDemoRota:
             workers_required=1,
             # We use a different balance weighting for each shift as otherwise
             # the shifts can be treated equally (and no target diff)
-            balance_weighting=0.5, 
+            #balance_weighting=0.6, 
         ),
         SingleShift(("group1",), "weekend", 12.5, days[5:], balance_offset=40,
         #balance_weighting=0.5
@@ -80,15 +82,17 @@ class TestDemoRota:
     # result.Solver.Status = SolverStatus.warning
     # prob.solutions.load_from(result)
 
-    ResultsHolder = RotaResults(Rota)
+    #worker_timetable_brief = Rota.get_worker_timetable_brief(
+    #    show_prefs=False, show_unavailable=False
+    #)
 
+    #print(worker_timetable_brief)
+    #print(Rota.get_shift_summary_dict())
+    Rota.export_rota_to_html("test2")
 
-    worker_timetable_brief = ResultsHolder.get_worker_timetable_brief(
+    def test_optimal_solution(self):
-        show_prefs=False, show_unavailable=False
+        assert self.Rota.results.solver.status == "ok"
-    )
+        assert self.Rota.results.solver.termination_condition == "optimal"
-
-    print(worker_timetable_brief)
-    ResultsHolder.export_rota_to_html("test")
 
     def test_start_date(self):
         assert self.Rota.start_date == self.start_date
@@ -145,23 +149,349 @@ class TestDemoRota:
                 case _:
                     assert 1 == 0
 
+    def test_worker_shift_targets(self):
+        shift_number = self.Rota.get_shift_summary_dict()
+
+        for worker in self.Rota.workers:
+            match worker.site:
+                case "group1":
+                    assert shift_number[worker.name]["weekday"] == 10
+                    assert shift_number[worker.name]["weekend"] == 5
+                case "group2":
+                    assert shift_number[worker.name]["weekend"] == 0
+                    match worker.fte:
+                        case 60:
+                            assert shift_number[worker.name]["weekday"] == 6
+                        case 40:
+                            assert shift_number[worker.name]["weekday"] == 4
+                        case _:
+                            assert 1 == 0
+                case _:
+                    assert 1 == 0
+
+
     def test_rota_sites(self):
         assert self.Rota.sites == {"group1", "group2"}
 
+class TestDemoRotaNights:
+    # Set up rota
+    weeks_to_rota = 12
+    start_date = datetime.date(2022, 3, 14)
 
-class TestNoWorkerRota:
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+        max_weekend_frequency=1,
+    )
+
+    Rota.constraint_options["max_shifts_per_week"] = 5
+    Rota.constraint_options["max_shifts_per_month"] = 31
+    Rota.constraint_options["ensure_derriford_reg_for_nights"] = False
+    Rota.constraint_options["limit_to_1_st2_on_nights"] = False
+    Rota.constraint_options["ensure_1_st4_plus_on_nights"] = False
+    Rota.constraint_options["constrain_time_off_after_nights"] = False
+    Rota.constraint_options["balance_nights_across_site"] = False
+    Rota.constraint_options["balance_shifts"] = True
+    Rota.constraint_options["balance_shifts_over_workers"] = True
+    Rota.constraint_options["balance_nights"] = False
+    Rota.constraint_options["minimise_shift_diffs"] = False
+    Rota.constraint_options["balance_blocks"] = False
+    #Rota.constraint_options["balance_weekends"] = False
+
+    # Add a few workers
+    worker1 = Worker(Rota, "worker1", "group1", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+    worker3 = Worker(Rota, "worker3", "group1", 1)
+    worker4 = Worker(Rota, "worker4", "group1", 1)
+    worker5 = Worker(Rota, "worker5", "group2", 1, fte=60)
+    worker6 = Worker(Rota, "worker6", "group2", 1, fte=40)
+    worker7 = Worker(Rota, "worker7", "group2", 1)
+    worker8 = Worker(Rota, "worker8", "group2", 1)
+
+    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6, worker7, worker8))
+
+    # Add a weekday and weekend shift
+    Rota.add_shifts(
+        SingleShift(
+            ("group1", "group2"),
+            "night_weekday",
+            12.5,
+            days[:4],
+            balance_offset=10,
+            workers_required=1,
+            constraints=["preclear2", "postclear2"],
+            #constraints=["night"],
+            force_as_block=True
+        
+            # We use a different balance weighting for each shift as otherwise
+            # the shifts can be treated equally (and no target diff)
+            #balance_weighting=0.6, 
+        ),
+        SingleShift(("group1","group2"), "night_weekend", 12.5, days[4:], balance_offset=40,
+        #balance_weighting=0.5
+            constraints=["preclear2", "postclear2"],
+            #constraints=["night"],
+            force_as_block=True
+        ),
+        SingleShift(("group1","group2"), "twilight", 12.5, days[:5], balance_offset=40,
+        #balance_weighting=0.5
+        workers_required=2,
+            #constraints=["preclear2", "postclear2"],
+        ),
+        #SingleShift(("group1","group2"), "atwilight", 12.5, days[:5], balance_offset=40,
+        ##balance_weighting=0.5
+        #workers_required=2,
+        #    #constraints=["postclear2"],
+        #),
+
+    )
+
+    Rota.build_shifts()
+    Rota.build_workers()
+    Rota.build_model()
+
+    print(Rota.get_worker_details())
+
+    for w in Rota.get_workers():
+        print(w)
+        print(w.get_shift_targets())
+
+    solver_options = {"ratio": 0.1, "seconds": 1000, "threads": 10}
+
+    Rota.solve_model(options=solver_options)
+
+    # print(Rota.get_worker_details())
+
+    # optimizer = SolverFactory('cbc')
+    # result = optimizer.solve(prob,tee=True)
+    # result.Solver.Status = SolverStatus.warning
+    # prob.solutions.load_from(result)
+
+    worker_timetable_brief = Rota.get_worker_timetable_brief(
+        show_prefs=False, show_unavailable=False
+    )
+
+    print(worker_timetable_brief)
+    print(Rota.get_shift_summary_dict())
+    Rota.export_rota_to_html("test3")
+
+    print(Rota.get_worker_shift_list(worker1))
+
+    def test_start_date(self):
+        assert self.Rota.start_date == self.start_date
+
+class TestDemoRotaClear:
     weeks_to_rota = 10
     start_date = datetime.date(2022, 3, 7)
 
     Rota = RotaBuilder(
         start_date,
         weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+        max_weekend_frequency=1,
+    )
+    #Rota.constraint_options["max_shifts_per_week"] = 5
+    Rota.constraint_options["max_shifts_per_month"] = 20
+    #Rota.constraint_options["ensure_derriford_reg_for_nights"] = False
+    #Rota.constraint_options["limit_to_1_st2_on_nights"] = False
+    #Rota.constraint_options["ensure_1_st4_plus_on_nights"] = False
+    #Rota.constraint_options["constrain_time_off_after_nights"] = False
+    #Rota.constraint_options["balance_nights_across_site"] = False
+    #Rota.constraint_options["balance_shifts"] = True
+    #Rota.constraint_options["balance_shifts_over_workers"] = True
+    #Rota.constraint_options["balance_nights"] = False
+    #Rota.constraint_options["minimise_shift_diffs"] = False
+    #Rota.constraint_options["balance_blocks"] = False
+
+    worker1 = Worker(Rota, "worker1", "group1", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+    worker3 = Worker(Rota, "worker3", "group1", 1)
+    worker4 = Worker(Rota, "worker4", "group1", 1)
+    worker5 = Worker(Rota, "worker5", "group2", 1, fte=60)
+    worker6 = Worker(Rota, "worker6", "group2", 1, fte=40)
+    worker7 = Worker(Rota, "worker7", "group2", 1)
+    worker8 = Worker(Rota, "worker8", "group2", 1)
+
+    Rota.add_workers((worker1, worker2, worker3, worker4, worker5, worker6, worker7, worker8))
+
+    def test_preclear(self):
+        self.Rota.add_shifts(
+            SingleShift(("group1","group2"), "a", 12.5, days[:5], balance_offset=40,
+            workers_required=2,
+                constraints=["preclear"],
+            ),
+            SingleShift(("group1","group2"), "b", 12.5, days[:5], balance_offset=40,
+            workers_required=2,
+            ),
+            SingleShift(("group1","group2"), "c", 12.5, days[3], balance_offset=40,
+            workers_required=1,
+                constraints=["preclear2", "postclear2"],
+            ),
+        )
+
+        self.Rota.build_shifts()
+        self.Rota.build_workers()
+        self.Rota.build_model()
+
+
+        solver_options = {"ratio": 0.1, "seconds": 1000, "threads": 10}
+
+        self.Rota.solve_model(options=solver_options)
+        self.Rota.export_rota_to_html("test4")
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+        for worker in self.Rota.workers:
+            shifts = self.Rota.get_worker_shift_list(worker)
+
+            # Convert shift to a string representation
+            shifts_string = "".join([i if i != "" else "-" for i in shifts])
+
+            # check for patterns that should not occur
+            assert "ba" not in shifts_string
+            assert "ca" not in shifts_string
+            assert "ac" not in shifts_string
+            assert "ca" not in shifts_string
+            assert "a-c" not in shifts_string
+            assert "c-a" not in shifts_string
+            assert "bc" not in shifts_string
+            assert "cb" not in shifts_string
+            assert "b-c" not in shifts_string
+            assert "c-b" not in shifts_string
+
+class TestDemoRotaShiftConstraints:
+    weeks_to_rota = 10
+    start_date = datetime.date(2022, 3, 7)
+
+    Rota = RotaBuilder(
+        start_date,
+        weeks_to_rota=weeks_to_rota,
+        balance_offset_modifier=1,
+        max_weekend_frequency=1,
+    )
+    #Rota.constraint_options["max_shifts_per_week"] = 5
+    #Rota.constraint_options["max_shifts_per_month"] = 20
+    worker1 = Worker(Rota, "worker1", "group1", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+
+    Rota.add_workers((worker1, worker2))
+    Rota.add_shifts(
+        SingleShift(("group1",), "a", 12.5, days, balance_offset=40,
+        workers_required=1,
+        ),
     )
 
-    def test_no_workers(self):
+    def test_max_shifts(self):
-        with pytest.raises(NoWorkers):
+        self.Rota.constraint_options["max_shifts_per_month"] = 14
-            self.Rota.build_workers()
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+    def test_max_shifts_fail(self):
+        self.Rota.constraint_options["max_shifts_per_month"] = 13
+        self.Rota.build_and_solve()
+        assert self.Rota.results.solver.status == "warning"
+        assert self.Rota.results.solver.termination_condition == "infeasible"
+
+    def test_max_shifts_extra_worker(self):
+        self.Rota.constraint_options["max_shifts_per_month"] = 12
+        worker3 = Worker(self.Rota, "worker3", "group1", 1)
+
+        self.Rota.add_worker(worker3)
+        self.Rota.build_and_solve()
+
+        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
+        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_pass(self):
+        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
+        worker3 = Worker(self.Rota, "worker3", "group1", 1)
+        self.Rota.add_worker(worker3)
+        self.Rota.build_and_solve()
+
+        assert self.Rota.results.solver.status == "ok"
+
+class TestDemoRotaBalanceShiftSites:
+    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", 1)
+    worker2 = Worker(Rota, "worker2", "group1", 1)
+    worker3 = Worker(Rota, "worker3", "group2", 1)
+    worker4 = Worker(Rota, "worker4", "group2", 1)
+    worker5 = Worker(Rota, "worker3", "group2", 1)
+    worker6 = Worker(Rota, "worker4", "group2", 1)
+
+    Rota.add_workers((worker1, worker2, worker3, worker4))
+    #Rota.add_workers((worker5, worker6))
+    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,
+        #),
+    )
+
+    def test_balance_blocks_across_groups(self):
+        self.Rota.constraint_options["balance_nights_across_sites"] = True
+        #self.Rota.constraint_options["balance_bank_holidays"] = False
+        #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")
+
+        assert self.Rota.results.solver.status == "ok"
+        assert self.Rota.results.solver.termination_condition == "optimal"
+
+#class TestNoWorkerRota:
+#    weeks_to_rota = 10
+#    start_date = datetime.date(2022, 3, 7)
+#
+#    Rota = RotaBuilder(
+#        start_date,
+#        weeks_to_rota=weeks_to_rota,
+#    )
+#
+#    def test_no_workers(self):
+#        with pytest.raises(NoWorkers):
+#            self.Rota.build_workers()
 
 
 if __name__ == "__main__":
-    t = TestDemoRota()
+    t = TestDemoRotaShiftConstraints()
+    t.test_balance_blocks_across_groups()