diff --git a/rota/shifts.py b/rota/shifts.py
index 8670d94..ab58bac 100644
--- a/rota/shifts.py
+++ b/rota/shifts.py
@@ -2,7 +2,7 @@ from calendar import week
import datetime
from distutils.log import debug
import itertools
-from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type
+from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal
import time
from pydantic import BaseModel, Extra, constr
@@ -65,9 +65,21 @@ SHIFT_BOUNDS = {
"weekend_count": (0, 60),
}
+VALID_SHIFT_CONSTRAINTS = Literal[
+ "night",
+ "pre",
+ "post",
+ "balance_across_groups",
+ "limit_grade_number",
+ "minimum_grade_number",
+ "max_shifts_per_week",
+ "require_remote_site_presence",
+ "require_remote_site_presence_week",
+]
+
class ShiftConstraint(BaseModel):
- name: str
+ name: VALID_SHIFT_CONSTRAINTS
options: bool | int | Dict | tuple = False
@@ -79,7 +91,7 @@ class SingleShift(BaseModel):
balance_across_groups (requires block assignment)
- postclear(2) / preclear(2)
+ post / pre clear
require_remote_site_presence_week:
options: (site, required_number)
@@ -180,9 +192,16 @@ class RotaBuilder(object):
SHIFT_BOUNDS=SHIFT_BOUNDS,
):
- console.print(Panel(f"""[white]
+ console.print(
+ Panel(
+ f"""[white]
{locals()}
- """, style="green", title="Generating Rota"), justify="left")
+ """,
+ style="green",
+ title="Generating Rota",
+ ),
+ justify="left",
+ )
self.SHIFT_BOUNDS = SHIFT_BOUNDS
self.shifts: List[SingleShift] = [] # type List[SingleShift]
@@ -191,7 +210,6 @@ class RotaBuilder(object):
self.use_shift_balance_extra = use_shift_balance_extra
self.use_bank_holiday_extra = use_bank_holiday_extra
-
self.workers: list[Worker] = []
# self.night_blocks = ["weekday", "weekend", "none"]
@@ -231,14 +249,16 @@ class RotaBuilder(object):
"avoid_shifts_by_grades": [],
}
-
self.results = None
self.ignore_valid_shifts = False
self.set_rota_dates(start_date, weeks_to_rota)
- def set_rota_dates(self, start_date: datetime.datetime, weeks_to_rota: int):
+ self.run_start_time: datetime.datetime | None = None
+ self.run_end_time: datetime.datetime | None = None
+
+ def set_rota_dates(self, start_date: datetime.date, weeks_to_rota: int):
self.weeks_to_rota = weeks_to_rota
@@ -286,22 +306,24 @@ class RotaBuilder(object):
):
console.print("Setting up solver")
- #solver = "scip"
+ solver = "scip"
if solver == "scip":
self.opt = SolverFactory(solver, executable="scip")
else:
self.opt = SolverFactory(solver)
-
-
try:
console.print("Solving")
if use_neos:
solver_manager = SolverManagerFactory("neos") # Solve using neos server
# results = solver_manager.solve(Rota.model, opt=opt, logfile="test.log")
results = solver_manager.solve(
- self.model, keepfiles=True, tee=True, opt=self.opt, logfile="test.log"
+ self.model,
+ keepfiles=True,
+ tee=True,
+ opt=self.opt,
+ logfile="test.log",
)
else:
results = self.opt.solve(
@@ -326,7 +348,7 @@ class RotaBuilder(object):
console.print(f"Attempting each shift individually")
self.solve_shifts_individually(options)
- #if not results.solver.status:
+ # if not results.solver.status:
# sys.exit(0)
def solve_shifts_by_block(self, options, block_length: int, folder="block_shifts"):
@@ -341,19 +363,23 @@ class RotaBuilder(object):
for block in range(0, no_blocks):
start_time = time.time()
start_date = original_start_date + datetime.timedelta(weeks=block)
- self.set_rota_dates(start_date = start_date,weeks_to_rota=block_length)
- console.print(f"Testing block: start_date - {self.start_date} , length {self.weeks_to_rota} weeks")
+ self.set_rota_dates(start_date=start_date, weeks_to_rota=block_length)
+ console.print(
+ f"Testing block: start_date - {self.start_date} , length {self.weeks_to_rota} weeks"
+ )
self.clear_shifts()
self.add_shifts(*shifts)
self.build_and_solve(options)
- self.export_rota_to_html(filename=f"{start_date}_{self.weeks_to_rota}", folder=folder)
+ self.export_rota_to_html(
+ filename=f"{start_date}_{self.weeks_to_rota}", folder=folder
+ )
end_time = time.time()
time_taken = end_time - start_time
outcomes[f"{start_date}_{self.weeks}"] = (
self.results.solver.status,
self.results.solver.termination_condition,
- time_taken
+ time_taken,
)
console.print(outcomes)
@@ -385,6 +411,9 @@ class RotaBuilder(object):
solve=True,
debug_if_fail: bool = False,
):
+
+ self.run_start_time = datetime.datetime.now()
+
self.build_shifts()
self.build_workers()
self.build_model()
@@ -392,6 +421,8 @@ class RotaBuilder(object):
if solve:
self.solve_model(options=options, debug_if_fail=debug_if_fail)
+ self.run_end_time = datetime.datetime.now()
+
def build_model(self):
# Initialize model
self.model = ConcreteModel()
@@ -1015,7 +1046,9 @@ class RotaBuilder(object):
# - self.model.night_per_site_t2[week, block, site]
# == self.model.night_per_site[week, block, site] - 1
# )
- for site in track(self.sites, description="Generating site balance constraints..."):
+ for site in track(
+ self.sites, description="Generating site balance constraints..."
+ ):
for shift in self.get_shifts_with_constraint("balance_across_groups"):
if site in shift.sites:
block = shift.name
@@ -1127,7 +1160,9 @@ class RotaBuilder(object):
# Most of our constraints apply per worker
for worker in self.workers:
- console.rule(f"Generate worker constraints: [bold blue]{worker.name}[/bold blue]")
+ console.rule(
+ f"Generate worker constraints: [bold blue]{worker.name}[/bold blue]"
+ )
logging.debug(f"Generate worker constraints: {worker.name}")
if self.constraint_options["minimise_shift_diffs"]:
@@ -1153,7 +1188,9 @@ class RotaBuilder(object):
except ValueError:
# Occurs if there are no shifts within a defined block
# TODO: test if this breaks (and we should check rathar than except)
- logging.debug(f"Failed to constrain max_shifts_per_month for worker: {worker.name}")
+ logging.debug(
+ f"Failed to constrain max_shifts_per_month for worker: {worker.name}"
+ )
pass
for constraint_dict in self.constraint_options["avoid_shifts_by_grades"]:
@@ -1218,7 +1255,9 @@ class RotaBuilder(object):
# Balance shifts within required limits (set by balance_offset)
# If balance_offset is too restrictive (for the rota length)
# no solution will be possible
- for shift in track(self.get_shifts(), description="Generate shift balance constraints"):
+ for shift in track(
+ self.get_shifts(), description="Generate shift balance constraints"
+ ):
if (
worker.site in shift.sites
): # Each site specfies which sites self.workers can fullfill it
@@ -1364,6 +1403,9 @@ class RotaBuilder(object):
# As the objective is to minimise t1+t2 and t1 and t2 are positive reals
# t1+t2 approximates the absolute target (which otherwise requires a quadratic solver)
# TODO: quadratic implementation so perfect solutions will be chosen
+
+ # NOTE: as this balances across the sum of the worker's shifts
+ # they can cancel each other out if a high weighting is given to multiple shifts
if self.constraint_options["balance_shifts_over_workers"]:
self.model.constraints.add(
self.model.worker_shift_count_t1[worker.id]
@@ -1553,10 +1595,15 @@ class RotaBuilder(object):
self.constraint_options["max_night_frequency"]
):
# Ignore excluded weeks
- if set(week_blocks).intersection(set(self.constraint_options["max_night_frequency_week_exclusions"])):
+ if set(week_blocks).intersection(
+ set(
+ self.constraint_options[
+ "max_night_frequency_week_exclusions"
+ ]
+ )
+ ):
continue
-
if self.get_shifts_with_constraint("night"):
# Prevent nights more than once every n weeks
self.model.constraints.add(
@@ -1839,7 +1886,9 @@ class RotaBuilder(object):
# def get_pre_may(week_days, max_pre)
- for n in track(range(len(weeks_days)), description="Generate week/day constraints"):
+ for n in track(
+ range(len(weeks_days)), description="Generate week/day constraints"
+ ):
week, day = weeks_days[n]
@@ -2044,6 +2093,9 @@ class RotaBuilder(object):
)
)
+ # NOTE: due to the way that this is implemented, if a
+ # shift spans 7 days you may get >7 allocations in a row
+ # as it only checks for a different shift allocation
for constraint_shift in self.get_shifts_with_constraints(
"pre",
):
@@ -2469,7 +2521,7 @@ class RotaBuilder(object):
"""
self.shifts.append(shift)
- def add_shifts(self, *shifts: SingleShift ) -> None:
+ def add_shifts(self, *shifts: SingleShift) -> None:
"""Add multiple shifts
Returns:
@@ -2610,7 +2662,7 @@ class RotaBuilder(object):
Returns:
list: contains tuple of all possible week / day / shift combinations
- """
+ """
week_day_shifts = self.week_day_shift_product
if week is not None:
@@ -2973,6 +3025,13 @@ class RotaBuilder(object):
timetable = []
+ if self.run_start_time is not None and self.run_end_time is not None:
+ timetable.append(
+ f"Start time: {self.run_start_time:%Y-%m-%d %H:%M:%S%z}
"
+ f"End time: {self.run_end_time:%Y-%m-%d %H:%M:%S%z}
"
+ f"Elapsed: {self.run_end_time-self.run_start_time}
"
+ )
+
timetable.append(
f"Rota start date: {self.start_date.isoformat()} ({self.weeks[-1]} weeks)
"
)
diff --git a/test/test_general_balancing.py b/test/test_general_balancing.py
index eb9c015..18d9031 100644
--- a/test/test_general_balancing.py
+++ b/test/test_general_balancing.py
@@ -1,5 +1,6 @@
import datetime
import pytest
+from pytest import approx
from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
from rota.workers import Worker
@@ -61,9 +62,9 @@ class TestBalancing:
for worker_name in shift_summary:
worker_shifts = shift_summary[worker_name]
- worker = Rota.get_worker_by_name(worker_name)
- assert worker_shifts["a"] in (22, 23, 24, 25)
- assert worker_shifts["b"] in (22, 23, 24, 25)
+ #worker = Rota.get_worker_by_name(worker_name)
+ assert worker_shifts["a"] in (approx(22), approx(23), approx(24), approx(25))
+ assert worker_shifts["b"] in (approx(22), approx(23), approx(24), approx(25))
def test_weighted_shift_balancing(self):
Rota = generate_basic_rota(20)
@@ -74,10 +75,10 @@ class TestBalancing:
name="a",
length=12.5,
days=days,
- balance_weighting=4,
+ balance_weighting=5,
workers_required=1,
force_as_block=False,
- constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+ constraint=[{"name": "pre","options": "2"}, {"name": "post","options": "2"}],
),
SingleShift(
sites=("group1", "group2"),
@@ -105,7 +106,8 @@ class TestBalancing:
worker_shifts = shift_summary[worker_name]
worker = Rota.get_worker_by_name(worker_name)
- assert worker_shifts["a"] in (46, 47, 48)
+ assert worker_shifts["a"] in (approx(46), approx(47), approx(48))
+ #assert worker_shifts["b"] in (46, 47, 48)
# assert worker_shifts["b"] in (22, 23, 24, 25)
def test_weighted_shift_balancing2(self):
@@ -127,7 +129,7 @@ class TestBalancing:
name="b",
length=12.5,
days=days,
- balance_weighting=4,
+ balance_weighting=8,
workers_required=1,
force_as_block=False,
),
@@ -141,7 +143,7 @@ class TestBalancing:
),
)
- Rota.build_and_solve(options={"ratio": 0.001})
+ Rota.build_and_solve(options={"ratio": 0.000})
Rota.export_rota_to_html("basic_balancing_weighted_shifts2")
shift_summary = Rota.get_shift_summary_dict()
@@ -151,7 +153,7 @@ class TestBalancing:
worker = Rota.get_worker_by_name(worker_name)
# assert worker_shifts["a"] in (46,47,48)
# assert worker_shifts["c"] in (46,47,48)
- assert worker_shifts["b"] in (52, 53, 54, 55)
+ assert worker_shifts["b"] in (approx(53), approx(54), approx(55))
def test_weighted_shift_balancing3(self):
Rota = generate_basic_rota(23)
@@ -172,7 +174,7 @@ class TestBalancing:
name="b",
length=12.5,
days=days,
- balance_weighting=4,
+ balance_weighting=1,
workers_required=1,
force_as_block=False,
),
@@ -181,13 +183,13 @@ class TestBalancing:
name="c",
length=12.5,
days=days[0],
- balance_weighting=4,
+ balance_weighting=8,
workers_required=1,
force_as_block=False,
),
)
- Rota.build_and_solve(options={"ratio": 0.000})
+ Rota.build_and_solve(options={"ratio": 0.0001})
Rota.export_rota_to_html("basic_balancing_weighted_shifts3")
shift_summary = Rota.get_shift_summary_dict()
@@ -197,7 +199,7 @@ class TestBalancing:
worker = Rota.get_worker_by_name(worker_name)
# assert worker_shifts["a"] in (46,47,48)
assert worker_shifts["c"] in (7, 8)
- assert worker_shifts["b"] in (52, 53, 54, 55)
+ #assert worker_shifts["b"] in (53, 54)
def test_weighted_shift_balancing4(self):
Rota = generate_basic_rota(10)