a few more fixes and improvements
This commit is contained in:
+50
-47
@@ -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)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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})
|
||||
Reference in New Issue
Block a user