a few more fixes and improvements

This commit is contained in:
Ross
2025-06-22 11:48:33 +01:00
parent 8402cc92fd
commit 82bc624a33
4 changed files with 153 additions and 49 deletions
+50 -47
View File
@@ -29,7 +29,7 @@ def main(
time_to_run: int = 60 * 60,
ratio: float = 0.001,
start_date: datetime.datetime = "2025-09-01",
weeks: int = 16,
weeks: int = 12,
bom: int = 1,
):
rota_start_date = start_date.date()
@@ -42,61 +42,42 @@ def main(
name="cons rota",
)
# Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
Rota.constraint_options["balance_shifts"] = False
Rota.constraint_options["balance_shifts_quadratic"] = True
Rota.constraint_options["balance_weekends"] = True
Rota.constraint_options["max_night_frequency"] = 3 # 3
Rota.constraint_options["max_night_frequency_week_exclusions"] = []
Rota.constraint_options["max_weekend_frequency"] = 4
Rota.constraint_options["hard_constrain_pair_separation"] = True
Rota.constraint_options["max_shifts_per_week"] = 3
Rota.constraint_options["max_shifts_per_week"] = 20
# Rota.constraint_options["avoid_st2_first_month"] = True
Rota.add_shifts(
SingleShift(
sites=("rota a",),
name="oncall weekday",
name="oncall",
length=12.5,
days=days[:4],
balance_offset=2,
days=days,
balance_offset=5,
assign_as_block=False,
constraint=[
{
"name": "max_shifts_per_week",
"options": 3,
}
#{
# "name": "max_shifts_per_week",
# "options": 3,
#}
],
),
SingleShift(
sites=("rota a",),
name="Oncall weekend",
length=12.5,
days=days[4:],
balance_offset=2,
assign_as_block=True,
constraint=[
{
"name": "max_shifts_per_week",
"options": 3,
},
{"name": "pre", "options": 2},
{"name": "post", "options": 2},
],
),
SingleShift(
sites=("rota a",),
name="weekend a",
name="weekend",
workers_required=1,
length=12.5,
days=days[5:],
balance_offset=2,
constraint=[
{"name": "pre", "options": 2},
{"name": "post", "options": 2},
#{"name": "pre", "options": 2},
#{"name": "post", "options": 1},
],
display_char="a",
#force_assign_with=["oncall"]
),
SingleShift(
sites=("rota b", "rota c"),
sites=("rota b",),
name="weekend b",
workers_required=1,
length=12.5,
@@ -113,22 +94,25 @@ def main(
"rota a",
"rota b",
),
name="evening",
name="twilight",
workers_required=1,
length=12.5,
days=days[:5],
balance_offset=2,
constraint=[],
balance_offset=5,
#constraint=[
# {"name": "pre", "options": 1},
# {"name": "post", "options": 1},
#],
),
)
Rota.allow_shifts_together(
"oncall weekday",
"Oncall weekend",
"weekend a",
"evening",
)
#Rota.allow_shifts_together_for_all_workers(
# "oncall",
# "twilight",
#)
#Rota.allow_shifts_together_for_all_workers(
# "oncall",
# "twilight",
#)
# Rota.add_grade_constraint_by_week([2], [1, 2], ["night_weekday", "night_weekend"])
@@ -136,8 +120,24 @@ def main(
Rota.build_shifts()
# Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names())
for w in range(1, 12):
w = Worker(name=f"A{w:02d}", site="rota a", grade=1)
for n in range(1, 12):
w = Worker(name=f"A{n:02d}", site="rota a", grade=1)
#w.prefer_multi_shift_together = 10
#w.allow_shifts_together("oncall", "twilight")
#w.allow_shifts_together("oncall", "weekend")
#w.add_hard_day_dependency("Fri", "Sat", "Sun")
w.add_force_assign_with("twilight", "oncall")
w.add_force_assign_with("weekend", "oncall")
w.set_max_shifts_per_week("twilight", 1)
if n <= 2:
pass
w.add_hard_day_dependency("Fri", "Sat", "Sun")
else:
w.add_hard_day_dependency("Fri", "Sun")
## #w.add_hard_day_dependency("Sat", "Sun")
print(w)
@@ -154,6 +154,9 @@ def main(
)
],
)
#w.add_hard_day_exclusion("Sat", "Sun")
w.set_max_shifts_per_week("twilight", 1)
w.set_max_shifts_per_week("weekend b", 1)
print(w)
+17
View File
@@ -1385,6 +1385,7 @@ class RotaBuilder(object):
)
# Most of our constraints apply per worker
# Worker constraint loop
worker: Worker
for worker in self.workers:
console.rule(
@@ -1392,6 +1393,22 @@ class RotaBuilder(object):
)
logging.debug(f"Generate worker constraints: {worker.name}")
if hasattr(worker, "max_shifts_per_week_by_shift_name"):
for shift_type, max_per_week in worker.max_shifts_per_week_by_shift_name.items():
if shift_type not in self.get_shift_names():
raise InvalidShift(
f"Worker '{worker.name}' has max_shifts_per_week_by_shift_name referencing non-existent shift: {shift_type}"
)
for week in self.weeks:
self.model.constraints.add(
sum(
self.model.works[worker.id, week, day, shift_type]
for day in days
if shift_type in self.get_shift_names_by_week_day(week, day)
) <= max_per_week
)
for week, day, shift in self.get_all_shiftclass_combinations():
# Only apply if this shift has force_assign_with set
if hasattr(shift, "force_assign_with") and shift.force_assign_with:
+9
View File
@@ -117,6 +117,7 @@ class Worker(BaseModel):
hard_day_dependencies: set[frozenset[str]] = set() # e.g. {frozenset({"Mon", "Tue"}), frozenset({"Fri", "Sun"})}
hard_day_exclusions: list[tuple[str, str]] = [] # e.g. [("Fri", "Sun")]
force_assign_with: dict[str, list[str]] = {} # e.g. {"oncall": ["weekend"]}
max_shifts_per_week_by_shift_name: dict[str, int] = {} # e.g. {"night": 2, "a": 3}
assign_as_block_preferences: dict[str, float] = {} # shift_name -> weight (positive = prefer block)
@@ -125,6 +126,8 @@ class Worker(BaseModel):
weekend_shift_target_number: int = 0
model_config = ConfigDict(
extra="allow",
)
@@ -422,3 +425,9 @@ class Worker(BaseModel):
if isinstance(with_shifts, str):
with_shifts = [with_shifts]
self.force_assign_with[shift] = with_shifts
def set_max_shifts_per_week(self, shift_name: str, max_per_week: int):
"""Helper to set the maximum number of shifts of a given type per week for this worker."""
if not hasattr(self, "max_shifts_per_week_by_shift_name"):
self.max_shifts_per_week_by_shift_name = {}
self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week
+75
View File
@@ -1381,3 +1381,78 @@ def test_workers_double_shifts_on_weekends():
Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
assert Rota.results.solver.status == "ok"
def test_max_shifts_per_week_by_shift_name_basic():
Rota = RotaBuilder(weeks_to_rota=2, start_date=datetime.date(2025, 6, 16))
w = Worker(name="A01", site="site1", grade=1)
w.max_shifts_per_week_by_shift_name = {"a": 1}
Rota.add_worker(w)
Rota.add_shifts(
SingleShift(
sites=("site1",),
name="a",
length=8,
days=days[:5],
workers_required=1,
),
)
Rota.build_and_solve(options={"ratio": 0.0})
# Worker should not be assigned more than 1 "a" shift per week
for week in Rota.weeks:
count = sum(
Rota.model.works[w.id, week, day, "a"].value > 0.5
for day in days[:5]
)
assert count <= 1
def test_max_shifts_per_week_by_shift_name_multiple_weeks():
Rota =generate_basic_rota(workers=4, weeks_to_rota=2)
for worker in Rota.get_workers():
worker.set_max_shifts_per_week("a", 1)
Rota.add_shifts(
SingleShift(
sites=("group1",),
name="a",
length=8,
days=days[:5],
workers_required=1,
),
)
Rota.build_and_solve(options={"ratio": 0.0})
assert Rota.results.solver.status == "error"
Rota.add_worker(
Worker(name="extraworker", site="group1", grade=1, max_shifts_per_week_by_shift_name={"a": 1})
)
Rota.build_and_solve(options={"ratio": 0.0})
assert Rota.results.solver.status == "ok"
Rota.export_rota_to_html("test_worker_double_shifts", folder="tests")
# Worker should not be assigned more than 2 "a" shifts per week
for worker in Rota.get_workers():
for week in Rota.weeks:
count = sum(
Rota.model.works[worker.id, week, day, "a"].value > 0.5
for day in days[:5]
)
assert count <= 2
def test_max_shifts_per_week_by_shift_name_invalid_shift():
Rota = RotaBuilder(weeks_to_rota=1, start_date=datetime.date(2025, 6, 16))
w = Worker(name="A03", site="site1", grade=1)
# "b" does not exist
w.max_shifts_per_week_by_shift_name = {"b": 1}
Rota.add_worker(w)
Rota.add_shifts(
SingleShift(
sites=("site1",),
name="a",
length=8,
days=days[:5],
workers_required=1,
),
)
with pytest.raises(InvalidShift):
Rota.build_and_solve(options={"ratio": 0.0})