ross/proc-rota
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: implement sub-block handling for shifts and update constraints in RotaBuilder

Browse Source
This commit is contained in:
Ross
2026-07-10 08:08:51 +01:00
parent d1fbbe2dc9
commit 5ff78faa07
3 changed files with 191 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gen_proc.py
+4 -3
View File
@@ -74,6 +74,7 @@ def main(
Rota.constraint_options["max_weekend_frequency"] = 3
Rota.constraint_options["max_days_per_week_block"] = [(8,3)]
Rota.constraint_options["maximum_allowed_shift_diff"] = 1
# wr = [
# WorkerRequirement(
@@ -174,11 +175,11 @@ def main(
length=12.5,
days=days[4:],
# balance_offset=3,
rota_on_nwds=True,
force_as_block=True,
#rota_on_nwds=True,
#force_as_block=True,
# assign_as_block=True,
constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
# force_as_block_unless_nwd=True
force_as_block_unless_nwd=[["Fri"], ["Sat", "Sun"]],
),
SingleShift(
sites=("torbay", "torbay twilights and weekends"),
rota_generator/shifts.py
+128 -92
View File
@@ -417,7 +417,7 @@ class SingleShift(BaseModel):
rota_on_nwds: bool = False
assign_as_block: bool = False
force_as_block: bool = False
force_as_block_unless_nwd: bool = False
force_as_block_unless_nwd: bool | List[List[str]] = False
hard_constrain_shift: bool = True
bank_holidays_only: bool = False
#constraint: list[ShiftConstraint] = []
@@ -1050,20 +1050,28 @@ class RotaBuilder(object):
initialize=0,
)
blocks_worker_keys = []
for worker, week, shift_name in self.worker_week_shifts_to_assign_as_blocks():
shift = self.get_shift_by_name(shift_name)
sub_blocks = self.get_shift_sub_blocks(shift)
for sb_idx in range(len(sub_blocks)):
blocks_worker_keys.append((worker.id, week, shift_name, sb_idx))
self.model.blocks_worker_shift_assigned = Var(
(
(worker.id, week, shift)
for worker, week, shift in self.worker_week_shifts_to_assign_as_blocks()
),
blocks_worker_keys,
within=Binary,
initialize=0,
)
blocks_assigned_keys = []
for week, shift_name in self.week_shifts_to_assign_as_blocks():
shift = self.get_shift_by_name(shift_name)
sub_blocks = self.get_shift_sub_blocks(shift)
for sb_idx in range(len(sub_blocks)):
blocks_assigned_keys.append((week, shift_name, sb_idx))
self.model.blocks_assigned = Var(
(
(week, shift)
for week, shift in self.week_shifts_to_assign_as_blocks()
),
blocks_assigned_keys,
within=NonNegativeIntegers,
initialize=0,
)
@@ -1900,99 +1908,103 @@ class RotaBuilder(object):
for week, shift_name in self.week_shifts_to_assign_as_blocks():
shift = self.get_shift_by_name(shift_name)
sub_blocks = self.get_shift_sub_blocks(shift)
for worker in self.get_workers_for_shift(shift):
if (worker.id, week, shift.name) in self.model.blocks_worker_shift_assigned:
try:
self.model.constraints.add(
8
* self.model.blocks_worker_shift_assigned[
worker.id, week, shift_name
]
>= sum(
self.model.works[worker.id, week, day, shift_name]
for day in shift.days
for sb_idx, sb in enumerate(sub_blocks):
for worker in self.get_workers_for_shift(shift):
if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned:
try:
self.model.constraints.add(
8
* self.model.blocks_worker_shift_assigned[
worker.id, week, shift_name, sb_idx
]
>= sum(
self.model.works[worker.id, week, day, shift_name]
for day in sb
)
)
)
except KeyError:
pass
self.model.constraints.add(
self.model.blocks_assigned[week, shift_name]
== sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift_name
]
for worker in self.workers
if (worker.id, week, shift.name) in self.model.blocks_worker_shift_assigned
except KeyError:
pass
self.model.constraints.add(
self.model.blocks_assigned[week, shift_name, sb_idx]
== sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift_name, sb_idx
]
for worker in self.workers
if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
)
)
)
workers_required = shift.get_worker_requirement_by_date(
self.get_week_start_date(week)
)
workers_required = shift.get_worker_requirement_by_date(
self.get_week_start_date(week)
)
if shift.force_as_block_unless_nwd:
workers = self.get_workers_for_shift(shift)
# Get workers who have a nwd on the shift
nwd_workers = []
full_workers = []
if shift.force_as_block_unless_nwd:
workers = self.get_workers_for_shift(shift)
# Get workers who have a nwd on the sub-block
nwd_workers = []
full_workers = []
for w in workers:
if w.non_working_day_list:
l = []
for (
nwd,
start_nwd_date,
end_nwd_date,
) in w.non_working_day_list:
if nwd in shift.days:
if start_nwd_date > self.get_week_start_date(week):
continue
if end_nwd_date < self.get_week_start_date(week):
continue
l.append(w)
for w in workers:
if w.non_working_day_list:
l = []
for (
nwd,
start_nwd_date,
end_nwd_date,
) in w.non_working_day_list:
if nwd in sb:
if start_nwd_date > self.get_week_start_date(week):
continue
if end_nwd_date < self.get_week_start_date(week):
continue
l.append(w)
if l:
nwd_workers.append(w)
if l:
nwd_workers.append(w)
else:
full_workers.append(w)
else:
full_workers.append(w)
if not nwd_workers: # Just do the usual
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name, sb_idx
]
for worker in self.workers
if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
)
<= workers_required
)
else:
full_workers.append(w)
if not nwd_workers: # Just do the usual
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name
]
for worker in self.workers
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name, sb_idx
]
for worker in full_workers
if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
)
<= workers_required + 1
)
<= workers_required
)
else:
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name
]
for worker in full_workers
elif shift.force_as_block:
if workers_required:
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name, sb_idx
]
for worker in self.workers
if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
)
<= workers_required
)
<= workers_required + 1
)
elif shift.force_as_block:
if workers_required:
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name
]
for worker in self.workers
)
<= workers_required
)
# Most of our constraints apply per worker
# Worker constraint loop (worker loop)
@@ -4185,24 +4197,27 @@ class RotaBuilder(object):
if self.constraint_options_model.balance_blocks:
blocks_balancing = sum(
block_shift_balancing_constant
* self.model.blocks_assigned[week, shift]
* self.model.blocks_assigned[week, shift, sb_idx]
for week in self.weeks
for shift in self.shifts_to_assign_as_blocks()
for sb_idx in range(len(self.get_shift_sub_blocks(self.get_shift_by_name(shift))))
if (week, shift, sb_idx) in self.model.blocks_assigned
)
else:
blocks_balancing = 0
prefer_block_expr = sum(
worker.assign_as_block_preferences.get(shift.name, 0)
* self.model.blocks_worker_shift_assigned[worker.id, week, shift.name]
* self.model.blocks_worker_shift_assigned[worker.id, week, shift.name, sb_idx]
for worker in self.workers
for week in self.weeks
for shift in self.shifts
if hasattr(worker, "assign_as_block_preferences")
and shift.name in worker.assign_as_block_preferences
and (worker.id, week, shift.name)
in self.model.blocks_worker_shift_assigned
and self.is_worker_eligible_for_shift(worker, shift)
for sb_idx in range(len(self.get_shift_sub_blocks(shift)))
if (worker.id, week, shift.name, sb_idx)
in self.model.blocks_worker_shift_assigned
)
# Quadratic :(
@@ -5046,6 +5061,27 @@ class RotaBuilder(object):
)
])
def get_shift_sub_blocks(self, shift: SingleShift) -> List[List[str]]:
if getattr(shift, "force_as_block_unless_nwd", None):
val = shift.force_as_block_unless_nwd
if isinstance(val, (list, tuple)) and all(isinstance(x, (list, tuple, set)) for x in val):
return [list(x) for x in val]
return [list(shift.days)]
# Check if it is a block shift globally or per-worker
is_block = (
shift.assign_as_block
or shift.force_as_block
or any(
hasattr(worker, "assign_as_block_preferences")
and getattr(worker, "assign_as_block_preferences", {}).get(shift.name, 0) != 0
for worker in self.workers
)
)
if is_block:
return [list(shift.days)]
return []
def get_all_locum_availability(self):
return self.locum_availability_map
test/test_nwds.py
+59 -2
View File
@@ -163,7 +163,7 @@ def test_nwd_force_as_block_force_split():
for worker in Rota.workers:
shifts = Rota.get_worker_shift_list(worker)
shifts_string = "".join([i if i != "" else "-" for i in shifts])
for week in weeks_from_list(shifts_string[7 * 5 :]):
for week in weeks_from_list(shifts_string[7 * 6 :]):
assert week in ("-----ww", "dddddww")
if worker.name == "worker1":
assert shifts_string[: 7 * 5].count("ww-") == 4
@@ -213,4 +213,61 @@ def test_nwd_testing():
if worker.name == "worker1":
assert week == "dddddww"
else:
assert week == "-----ww"
assert week == "-----ww"
def test_force_as_block_unless_nwd_with_subblocks():
# Test that a Fri/Sat/Sun shift can be split into Fri and Sat/Sun blocks
# if a worker does not work on Fridays.
Rota = setup_rota(weeks_to_rota=2)
start_date = Rota.start_date
# worker1 does not work on Fridays (Fri is NWD)
worker1 = Worker(
name="worker1", site="group1", grade=1,
nwds=[{"day": "Fri", "start_date": start_date, "end_date": start_date + datetime.timedelta(weeks=2)}],
)
# worker2 works normally
worker2 = Worker(name="worker2", site="group1", grade=1)
Rota.add_workers((worker1, worker2))
# Fri/Sat/Sun shift, requiring 1 worker, forced as block unless NWD
# We specify sub-blocks: [["Fri"], ["Sat", "Sun"]]
Rota.add_shifts(
SingleShift(
sites=("group1",), name="weekend", length=12.5,
days=["Fri", "Sat", "Sun"],
workers_required=1,
force_as_block_unless_nwd=[["Fri"], ["Sat", "Sun"]],
),
)
# We remove no valid shifts warning
Rota.terminate_on_warning.remove("Worker/no valid shifts")
Rota.build_and_solve()
assert Rota.results.solver.status == "ok"
# Check assignments:
# Since worker1 has NWD on Friday:
# - worker1 cannot work Friday.
# - But worker1 can work Sat and Sun as a block!
# So on Saturday and Sunday, worker1 should be assigned.
# On Friday, worker2 should be assigned.
for week in Rota.weeks:
w1_fri = Rota.model.works[worker1.id, week, "Fri", "weekend"].value
w1_sat = Rota.model.works[worker1.id, week, "Sat", "weekend"].value
w1_sun = Rota.model.works[worker1.id, week, "Sun", "weekend"].value
w2_fri = Rota.model.works[worker2.id, week, "Fri", "weekend"].value
w2_sat = Rota.model.works[worker2.id, week, "Sat", "weekend"].value
w2_sun = Rota.model.works[worker2.id, week, "Sun", "weekend"].value
assert w1_fri == 0
# Either worker1 works Sat/Sun and worker2 works Fri (split block)
# OR worker2 works Fri/Sat/Sun (full block) and worker1 works nothing
is_split = (w1_sat == 1 and w1_sun == 1 and w2_fri == 1 and w2_sat == 0 and w2_sun == 0)
is_full_w2 = (w1_sat == 0 and w1_sun == 0 and w2_fri == 1 and w2_sat == 1 and w2_sun == 1)
assert is_split or is_full_w2