from calendar import week import datetime import itertools from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal, Optional import time from pydantic import BaseModel, ConfigDict, Extra, constr, field_validator, model_validator, StrictInt, Field from pathlib import Path import uuid import datetime from pyomo.environ import * from pyomo.opt import SolverFactory, TerminationCondition, SolverStatus import urllib.parse from rota_generator.workers import Worker from rota_generator.console import console from rota_generator.bank_holidays import bank_holiday_map import json import csv from collections import defaultdict from io import StringIO import sys from rich.pretty import pprint from rich.progress import track from rich.panel import Panel from rich import print from loguru import logger logger.add("rota.log", rotation="1 MB", level="DEBUG") ShiftName = str DayStr = str WeekInt = int days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"] full_weekend_count = 2 SHIFT_BOUNDS = { "bank_holiday": (0, 9), "shift_count": (0, 400), "night_shift_count": (0, 100), "weekend_count": (0, 60), } #VALID_SHIFT_CONSTRAINTS = Literal[ # "night", # "pre", # "post", # "balance_across_groups", # "limit_grade_number", # "minimum_grade_number", # "max_shifts_per_week", # "max_shifts_per_week_block", # "require_remote_site_presence", # "require_remote_site_presence_week", #] class BaseShiftConstraint(BaseModel): weeks: Optional[List[int]] = None start_date: Optional[datetime.date] = None end_date: Optional[datetime.date] = None model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") @field_validator("weeks") def weeks_none_if_dates_set(cls, weeks, values): data = values.data if (data.get("start_date") is not None or data.get("end_date") is not None) and weeks is not None: raise ValueError("If start_date or end_date is set, weeks must not be set") return weeks #class ShiftConstraint(BaseShiftConstraint): # name: str # options: Any = None class NightConstraint(BaseShiftConstraint): options: Any = None class PreShiftConstraint(BaseShiftConstraint): days: Any = None ignore_shifts: List[str] = [] exclude_days: List[str] = [] allow_self: bool = True class PostShiftConstraint(BaseShiftConstraint): days: Any = None ignore_shifts: List[str] = [] exclude_days: List[str] = [] allow_self: bool = True class BalanceAcrossGroupsConstraint(BaseShiftConstraint): """""" class MinSummedGradeByShiftsPerDayConstraint(BaseShiftConstraint): shifts: List[str] = Field(..., description="List of shift names to sum grades over") min_grade_sum: int = Field(..., description="Minimum total grade required") days: Optional[List[str]] = None class LimitGradeNumberConstraint(BaseShiftConstraint): grade: int max_number: int class MinimumGradeNumberConstraint(BaseShiftConstraint): grade: int min_number: int class RequireRemoteSitePresenceConstraint(BaseShiftConstraint): site: str required_number: int class MaxShiftsPerWeekConstraint(BaseShiftConstraint): max_shifts: int | None = None class MaxShiftsPerWeekBlockConstraint(BaseShiftConstraint): week_block : int max_shifts: int | None = None def _pydantic_to_dict(obj): if isinstance(obj, list): return [_pydantic_to_dict(x) for x in obj] if hasattr(obj, "model_dump"): return obj.model_dump() return obj class WorkerRequirement(BaseModel): """ start date is inclusive, end date is not """ # NOTE: if number is changed mid week block allocations may not work? number: int = 1 start_date: datetime.date | None = None end_date: datetime.date | None = None @model_validator(mode="after") def check_dates(self): if self.start_date is not None and self.end_date is not None: if self.start_date > self.end_date: raise ValueError("Start date must be before end date") return self class SingleShift(BaseModel): """Class to hold all details for a shift start_date: The day that the shift starts on end_date: The day that the shift ends on (the last day it should be rota'd for) Due to week block assigments these (probably) need to be week ba Valid constraints balance_across_groups (requires block assignment) post / pre clear require_remote_site_presence_week: options: (site, required_number) night limit_grade_number: options: {grade: max_number, grade2: max_number2} minimum_grade_number: options: (grade, min_number) """ sites: List[str] name: str length: float days: str | list[str] balance_offset: float | None = None balance_weighting: float = 1 workers_required: StrictInt | list[WorkerRequirement] = 1 rota_on_nwds: bool = False assign_as_block: bool = False force_as_block: bool = False force_as_block_unless_nwd: bool = False hard_constrain_shift: bool = True bank_holidays_only: bool = False #constraint: list[ShiftConstraint] = [] constraints: List[BaseModel] = Field(default_factory=list) start_date: datetime.date | None = None end_date: datetime.date | None = None pair_proxy: bool = False display_char: str | None = None # Add this line force_assign_with: list[str] = [] # List of shift names to force assign with this shift on the same day model_config = ConfigDict( extra="allow", # orm_mode = True ) def __init__(self, **data: Any): if "constraint" in data: raise ValueError( "The 'constraint' field on SingleShift has been deprecated and removed. " "Please use the 'constraints' field with Pydantic constraint models instead." ) super().__init__(**data) # self.site = sites # self.name = name # self.length = length # self.days = days # self.balance_by_site = balance_by_site # balance_offset defines the max difference in allocated shifts # versus target shifts (hard constraint) # if 0 exactly equal shifts must be assigend (unlikely to be possible) if self.display_char is not None and len(self.display_char) > 1: raise ValueError("display_char must be a single character or None") if self.balance_offset is None: self.balance_offset = len(self.days) if self.force_as_block_unless_nwd: self.assign_as_block = True # self.assign_as_block = assign_as_block #self.constraint_options = {} #self.constraints = [] #for c in self.constraint: # self.constraints.append(c.name) # if c.options: # self.constraint_options[c.name] = c.options # constraint = c.constraint # match c: # case (constraint, options): # self.constraint_options[constraint] = options # case constraint: # self.constraints.append(constraint) if isinstance(self.days, str): self.days = [self.days] def __str__(self): return ( f"name: {self.name}, " f"sites: {self.sites}, " f"length: {self.length}, " f"balance_weighting: {self.balance_weighting}, " f"balance_offset: {self.balance_offset}, " f"hard_constrain_shift: {self.hard_constrain_shift}, " f"workers_required: {self.workers_required}, " f"rota_on_nwds: {self.rota_on_nwds}, " f"assign_as_block: {self.assign_as_block}, " f"force_as_block: {self.force_as_block}, " f"constraints: {self.constraints}" ) def has_constraint(self, constraint: Type) -> bool: for c in self.constraints: if isinstance(c, constraint): return True return False def get_constraint(self, constraint: Type) -> BaseShiftConstraint | None: for c in self.constraints: if isinstance(c, constraint): return c return None def get_constraints(self, constraint: Type) -> list[BaseShiftConstraint]: return [c for c in self.constraints if isinstance(c, constraint)] def get_shift_summary(self): """ Returns a text summary of the shift """ return f"{self.name}: {self.workers_required} workers for sites ({', '.join(self.site)}) on days ({', '.join(self.days)})" def get_shift_number(self): return len(self.days) def get_worker_requirement_by_date(self, date: datetime.date): """Must be called after build_shifts Returns the number of workers required for a given date If none are required 0 is returned""" if isinstance(self.workers_required, int): return self.workers_required for wr in self.workers_required: # print(date, wr) if wr.start_date <= date < wr.end_date: return wr.number return 0 def get_display_char(self): """Returns the display character for the shift""" if self.display_char is not None: return self.display_char return self.name[0] class RotaBuilder(object): """Class to hold and manipulate shifts""" def __init__( self, start_date: datetime.date = datetime.date.today(), weeks_to_rota: int = 26, balance_offset_modifier: int = 1, ltft_balance_offset: int = 1, use_previous_shifts: bool = False, use_shift_balance_extra: bool = False, use_bank_holiday_extra: bool = False, allow_force_assignment_with_leave_conflict: bool = False, SHIFT_BOUNDS=SHIFT_BOUNDS, name: str = "", bank_holidays: dict[datetime.date, str] = bank_holiday_map, ): self.name = name console.print( Panel( f"""[white] {locals()} """, style="green", title=f"Generating Rota: {name}", ), justify="left", ) self.SHIFT_BOUNDS = SHIFT_BOUNDS self.shifts: List[SingleShift] = [] # type List[SingleShift] # self.allowed_multi_shift_sets: list[set[str]] = [] self.use_previous_shifts = use_previous_shifts 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"] self.sites: set = set() self.balance_offset_modifier = balance_offset_modifier self.ltft_balance_offset = ltft_balance_offset self.constraint_options = { "balance_nights": True, "constrain_time_off_after_nights": False, "balance_nights_across_sites": True, "balance_bank_holidays": True, "balance_blocks": True, "balance_shifts": True, # Does not use a quadratic function "balance_shifts_true_quadratic": False, # Does not use a quadratic function "balance_shifts_quadratic": False, # Will prevent spreading of spreading across different shifts "balance_shifts_over_workers": True, "minimise_shift_diffs": False, # less sophisticated version of balance_shifts_over_workers "maximum_allowed_shift_diff": None, "balance_weekends": True, "max_weekends": 100, # Don't assign multiple shifts every (n) weeks "max_weekend_frequency": 1, # Requires balance weekends # Don't assign multiple shifts every (n) weeks "max_night_frequency": 1, "max_night_frequency_week_exclusions": [], "max_shifts_per_week": 7, "max_shifts_per_month": 40, "max_days_per_week_block": [], #"max_days_per_week_block": [(3, 2)], # (max_days, week_block) # The following will cause an unsolvable problem if a shift # is forced as a block and spans the time peroid "prevent_monday_after_full_weekends": [], "prevent_monday_and_tuesday_after_full_weekends": [], "prevent_fridays_before_full_weekends": [], "prevent_thursdays_before_full_weekends": [], "hard_constrain_pair_separation": False, "avoid_shifts_by_grades": [], "avoid_shifts_by_worker_names": [], "allocate_locum_shifts": True, "distribute_locum_shifts": True, "balance_locum_shifts": False, "maximum_allowed_locum_shifts_per_worker": None, "minimum_allowed_locum_shifts_per_worker": None, } self.constraint_options["min_summed_grade_by_shifts_per_day"]: List[MinSummedGradeByShiftsPerDayConstraint] = [] self.terminate_on_warning = [ "Worker/duplicate id", "Worker/duplicate name", "Worker/no valid shifts", "Shift/invalid start date", "Shift/invalid end date", "Locum/no locum availability", #"Force assignment/leave conflict", # Currently will always be infeasible ] self.results = None self.warnings: list[tuple[str, str]] = [] self.ignore_valid_shifts = False self.set_rota_dates(start_date, weeks_to_rota) self.run_start_time: datetime.datetime | None = None self.run_end_time: datetime.datetime | None = None self.allow_force_assignment_with_leave_conflict = allow_force_assignment_with_leave_conflict self.bank_holidays = bank_holidays self.paired_shifts = [] self.exported_rota_file = None self.enable_unavailabilities = True self.MAX_SHIFTS_PER_DAY = 2 def set_rota_constraint(self, constraint: str, value: Any): if constraint not in self.constraint_options: raise ValueError(f"Constraint {constraint} not valid") self.constraint_options[constraint] = value def set_rota_dates(self, start_date: datetime.date, weeks_to_rota: int): self.weeks_to_rota = weeks_to_rota self.days = days self.weeks = [i for i in range(1, weeks_to_rota + 1)] self.weeks_days_product = list(itertools.product(self.weeks, days)) if start_date.weekday() != 0: raise ValueError( f"Start date {start_date.isoformat()} must be a Mon (not a {days[start_date.weekday()]})" ) self.start_date = start_date self.rota_days_length = len(self.weeks) * 7 self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length) console.print( f"Weeks to rota: {weeks_to_rota} ({self.start_date} - {self.rota_end_date})" ) # Generate a map for week, day combinations to a given date self.week_day_date_map = {} self.date_week_day_map = {} n: int = 0 for week, day in self.get_week_day_combinations(): d = self.start_date + datetime.timedelta(n) self.week_day_date_map[(week, day)] = d self.date_week_day_map[d] = (week, day) n = n + 1 self.unavailable_to_work = set() self.unavailable_to_work_reason = {} self.pref_not_to_work = {} self.pref_not_to_work_reason = {} self.work_requests: set[ Tuple[str | uuid.UUID | int, WeekInt, DayStr, ShiftName] ] = set() self.work_requests_map = {} self.locum_availability: set[ Tuple[str | uuid.UUID | int, WeekInt, DayStr, ShiftName] ] = set() self.locum_availability_map = {} def solve_model( self, solver: str = "appsi_highs", options: dict = {}, debug_if_fail: bool = False, ): console.print("Setting up solver") console.print(f"{options=}") if solver == "scip": if "seconds" in options: options["limits/time"] = options.pop("seconds") if "ratio" in options: options["limits/gap"] = options.pop("ratio") if "threads" in options: options.pop("threads") self.opt = SolverFactory(solver, executable="scip") elif solver == "appsi_highs": self.opt = SolverFactory(solver) if "seconds" in options: options["time_limit"] = options.pop("seconds") if "ratio" in options: options["mip_rel_gap"] = options.pop("ratio") else: self.opt = SolverFactory(solver) try: console.print("Solving") console.print(f"Options: {options}") log_file = f"logs/{self.name}_{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}.log" results = self.opt.solve( self.model, tee=True, options=options, # options={ # "threads": 10, # }, # logfile=log_file, keepfiles=True, load_solutions=False, ) except KeyboardInterrupt: return pass if ( not results.solver.termination_condition == TerminationCondition.infeasible # and not results.solver.status == SolverStatus.aborted ): self.model.solutions.load_from(results) self.results = results console.print(f"Complete - outcome: {results.solver.status}") console.print(f"Termination condition: {results.solver.termination_condition}") if results.solver.status != "ok" and debug_if_fail: console.print(f"Attempting each shift individually") self.solve_shifts_individually(options) # if not results.solver.status: # sys.exit(0) def solve_shifts_by_block(self, options, block_length: int, folder="block_shifts"): shifts = self.shifts outcomes = {} no_blocks = self.weeks_to_rota - block_length original_start_date = self.start_date 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.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 ) 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, ) console.print(outcomes) sys.exit(0) def solve_shifts_individually(self, options, folder="individual_shifts"): shifts = self.shifts outcomes = {} for shift in shifts: console.print(f"Testing shifts: {shift.name}") self.ignore_valid_shifts = True self.clear_shifts() self.add_shift(shift) self.build_and_solve(options) self.export_rota_to_html(filename=shift.name, folder=folder) outcomes[shift.name] = ( self.results.solver.status, self.results.solver.termination_condition, ) console.print(outcomes) def build_and_solve( self, options={"ratio": 0.1, "seconds": 1000, "threads": 10}, solve=True, debug_if_fail: bool = False, export=False, export_with_timestamp=False, solver="appsi_highs", ): self.run_start_time = datetime.datetime.now() self.build_shifts() self.build_workers() self.build_model() console.print("Building model") console.print(f"{options=}") if solve: self.solve_model( options=options, debug_if_fail=debug_if_fail, solver=solver ) self.run_end_time = datetime.datetime.now() if export: self.export_rota_to_html(self.name, timestamp_filename=export_with_timestamp) def build_model(self): # Initialize model self.model = ConcreteModel() # binary variables representing if a worker is scheduled somewhere # NOTE: this will assign to all possible shift combinations (which we probably don't want) self.model.works = Var( ( (worker.id, week, day, shiftname) for worker in self.workers for week, day, shiftname in self.get_all_shiftname_combinations() ), within=Binary, initialize=0, ) self.model.works_day = Var( ( (worker.id, week, day) for worker in self.workers for week, day in self.get_week_day_combinations() ), within=Binary, initialize=0, ) self.model.shift_week_worker_assigned = Var( ( (shift, week, worker.id) for worker in self.workers for week in self.weeks for shift in self.get_shift_names() ), within=Binary, initialize=0, ) self.model.blocks_worker_shift_assigned = Var( ( (worker.id, week, shift) for worker, week, shift in self.worker_week_shifts_to_assign_as_blocks() ), within=Binary, initialize=0, ) self.model.blocks_assigned = Var( ( (week, shift) for week, shift in self.week_shifts_to_assign_as_blocks() ), within=NonNegativeIntegers, initialize=0, ) if self.constraint_options["balance_bank_holidays"]: # Bank holidays self.model.bank_holiday_count = Var( ((worker.id) for worker in self.workers), within=NonNegativeIntegers, initialize=0, bounds=self.SHIFT_BOUNDS["bank_holiday"], ) self.model.bank_holiday_count_w = Var( ((worker.id) for worker in self.workers), within=NonNegativeIntegers, initialize=0, ) self.model.multi_shift_together = Var( ( (worker.id, week, day, idx) for worker in self.workers for week, day in self.get_week_day_combinations() for idx, allowed_set in enumerate( getattr(worker, "allowed_multi_shift_sets", []) ) if allowed_set.issubset(self.get_shift_names_by_week_day(week, day)) ), domain=Binary, initialize=0, ) # Used to limit number of workers on night shift per site # if we force a binary it will in effect hard constrain to < 2 # from a single site if self.constraint_options["balance_nights_across_sites"]: self.model.night_per_site = Var( ( (week, shift.name, site) for site in self.sites for week in self.weeks for shift in self.get_shifts_with_constraint(NightConstraint) ), # within=Binary, within=NonNegativeIntegers, initialize=0, ) self.model.night_per_site_t1 = Var( ( (week, shift.name, site) for site in self.sites for week in self.weeks for shift in self.get_shifts_with_constraint(NightConstraint) ), domain=NonNegativeIntegers, initialize=0, ) self.model.night_per_site_t2 = Var( ( (week, shift.name, site) for site in self.sites for week in self.weeks for shift in self.get_shifts_with_constraint(NightConstraint) ), domain=NonNegativeIntegers, initialize=0, ) # Replacement for balance_nights_across_sites self.model.shift_per_site = Var( ( (week, shift.name, site) for site in self.sites for week in self.weeks for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint) ), # within=Binary, within=NonNegativeIntegers, initialize=0, ) self.model.shift_per_site_t1 = Var( ( (week, shift.name, site) for site in self.sites for week in self.weeks for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint) ), domain=NonNegativeIntegers, initialize=0, ) self.model.shift_per_site_t2 = Var( ( (week, shift.name, site) for site in self.sites for week in self.weeks for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint) ), domain=NonNegativeIntegers, initialize=0, ) self.model.shift_count = Var( ( (worker.id, shift) for worker in self.workers for shift in self.get_shift_names() ), domain=NonNegativeReals, initialize=0, bounds=self.SHIFT_BOUNDS["shift_count"], ) self.model.shift_count_diff = Var( ( (worker.id, shift) for worker in self.workers for shift in self.get_shift_names() ), within=Reals, initialize=0, ) self.model.shift_count_diff_summed = Var( (worker.id for worker in self.workers), within=Reals, initialize=0, ) if self.constraint_options["balance_shifts_over_workers"]: self.model.worker_shift_count_t1 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) self.model.worker_shift_count_t2 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) if ( self.constraint_options["balance_shifts"] or self.constraint_options["balance_shifts_quadratic"] ): self.model.shift_count_t1 = Var( ( (worker.id, shift.name) for worker in self.workers for shift in self.get_shifts() ), domain=NonNegativeReals, initialize=0, ) self.model.shift_count_t2 = Var( ( (worker.id, shift.name) for worker in self.workers for shift in self.get_shifts() ), domain=NonNegativeReals, initialize=0, ) if self.constraint_options["balance_shifts_quadratic"]: self.model.shift_count_w = Var( ( (worker.id, shift.name) for worker in self.workers for shift in self.get_shifts() ), domain=NonNegativeReals, initialize=0, ) if self.constraint_options["balance_nights"]: # We also try to even out the night shifts seperately self.model.night_shift_count = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, bounds=self.SHIFT_BOUNDS["night_shift_count"], ) self.model.night_shift_count_t1 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) self.model.night_shift_count_t2 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) self.model.night_shift_count_w = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) # self.model.worker_weekend_assigned = Var( # ((worker.id, week) for worker in self.workers # for week in self.weeks), # domain=NonNegativeReals, # initialize=0, # ) # self.model.works_weekend_count = Var( # ((worker.id, week) for worker in self.workers # for week in self.weeks), # domain=NonNegativeIntegers, # initialize=0, # ) # self.model.works_saturday = Var( # ((worker.id, week) for worker in self.workers # for week in self.weeks), # domain=Binary, # initialize=0, # ) # self.model.works_sunday = Var( # ((worker.id, week) for worker in self.workers # for week in self.weeks), # domain=Binary, # initialize=0, # ) # self.model.works_whole_weekend = Var( # ((worker.id, week) for worker in self.workers # for week in self.weeks), # domain=Binary, # initialize=0, # ) self.model.works_weekend = Var( ((worker.id, week) for worker in self.workers for week in self.weeks), domain=Binary, initialize=0, ) self.model.worker_weekend_count = Var( ((worker.id) for worker in self.workers), domain=NonNegativeIntegers, initialize=0, bounds=self.SHIFT_BOUNDS["weekend_count"], ) if self.constraint_options["balance_weekends"]: self.model.weekend_shift_count_t1 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) self.model.weekend_shift_count_t2 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) self.model.weekend_shift_count_w = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, initialize=0, ) # self.model.weekend_count_t1 = Var( # ((worker.id) for worker in self.workers), # domain=NonNegativeReals, # initialize=0, # ) # self.model.weekend_count_t2 = Var( # ((worker.id) for worker in self.workers), # domain=NonNegativeReals, # initialize=0, # ) # self.model.unavailable = Var(((worker, week, day) for worker in self.workers for week in weeks for day in days), # within=Binary, initialize=0) for worker in self.workers: for week, day, shift_name in getattr(worker, "forced_assignments", []): # Only add if this shift is valid for this week/day if shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True): # Check for leave conflict if hasattr(self, "unavailable_to_work"): if (worker.id, week, day) in self.unavailable_to_work: if self.allow_force_assignment_with_leave_conflict: self.unavailable_to_work.remove((worker.id, week, day)) self.add_warning( "Force assignment/leave conflict", f"Worker '{worker.name}' has a forced assignment for shift '{shift_name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which conflicts with a leave request. The leave request has been overridden.", ) else: self.add_warning( "Force assignment/leave conflict", f"Worker '{worker.name}' has a forced assignment for shift '{shift_name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which conflicts with a leave request.", ) if not self.enable_unavailabilities: self.unavailable_to_work = set() self.unavailable_to_work_reason = {} def availability_init(model, wid, week, day): # This needs to be higher than the max number of possible shifts per day if (wid, week, day) in self.unavailable_to_work: # print((wid, week, day)) return 0 return self.MAX_SHIFTS_PER_DAY + 1 self.model.available = Param( ( (worker.id, week, day) for worker in self.workers for week, day in self.get_week_day_combinations() ), initialize=availability_init, ) locum_request_sets = set() for locum_request in self.locum_availability: if locum_request[3] == "*": pass else: locum_request_sets.add(locum_request[3]) # Validate locum shifts are valid invalid_shifts = locum_request_sets - set(self.get_shift_names()) if invalid_shifts and not self.ignore_valid_shifts: raise InvalidShift( f"Invalid locum shift request [shift(s): ${invalid_shifts}]" ) def locum_request_init(model, wid, week, day, shift): if (wid, week, day, shift) in self.locum_availability: self.locum_availability_map[(wid, week, day)] = shift return 1 return 0 self.model.locum_availability = Param( ( (worker.id, week, day, shift) for worker in self.workers if not worker.locum for week, day, shift in self.get_all_shiftname_combinations() ), initialize=locum_request_init, ) work_request_sets = set() for work_request in self.work_requests: if work_request[3] == "*": pass else: work_request_sets.add(work_request[3]) # Validate shifts are valid invalid_shifts = work_request_sets - set(self.get_shift_names()) if invalid_shifts and not self.ignore_valid_shifts: raise InvalidShift( f"Invalid worker shift request [shift(s): ${invalid_shifts}]" ) def work_request_init(model, wid, week, day, shift): if (wid, week, day, shift) in self.work_requests: self.work_requests_map[(wid, week, day)] = shift return 1 return 0 self.model.work_requests = Param( ( (worker.id, week, day, shift) for worker in self.workers for week, day, shift in self.get_all_shiftname_combinations() ), initialize=work_request_init, ) def pref_init(model, wid, week, day): if (wid, week, day) in self.pref_not_to_work: return self.pref_not_to_work[(wid, week, day)] return 0 # It would also be possible to assign sequential days as blocks # which would allow more weight to be given if the whole block # is not allocated self.model.pref_not_to_work = Param( ( (worker.id, week, day) for worker in self.workers for week, day in self.get_week_day_combinations() ), initialize=pref_init, ) # binary variables representing if a worker is necessary # self.model.needed = Var([worker.id for worker in self.workers], within=Binary, initialize=0) # binary variables representing if a worker worked on sunday but not on saturday (avoid if possible) # self.model.no_pref = Var([worker.id for worker in self.workers], within=Binary, initialize=0) if self.get_workers_who_require_locums(): # We use this as a time check that their is some valid locum ability if not self.get_all_locum_availability(): self.add_warning( "Locum/no locum availability", "No locum availability set" ) self.model.locum_required = Var( ( (week, day, shiftname) for week, day, shiftname in self.get_all_shiftname_combinations() ), within=Binary, initialize=0, ) self.model.locum_works = Var( ( (worker.id, week, day, shiftname) for worker in self.workers for week, day, shiftname in self.get_all_shiftname_combinations() ), within=Binary, initialize=0, ) self.model.locum_shifts_by_worker = Var( (worker.id for worker in self.workers), within=Integers, initialize=0, ) self.model.locum_shifts_t1 = Var( (worker.id for worker in self.workers), domain=NonNegativeReals, initialize=0, ) self.model.locum_shifts_t2 = Var( (worker.id for worker in self.workers), domain=NonNegativeReals, initialize=0, ) self.build_model_constraints() def build_model_constraints(self): self.model.constraints = ConstraintList() # Create a set of constraints for constraint in self.constraint_options.get("min_summed_grade_by_shifts_per_day", []): weeks = constraint.weeks if constraint.weeks is not None else self.weeks days_list = constraint.days if constraint.days is not None else self.days for week in weeks: for day in days_list: # Only consider if all shifts are present that day if all(shift in self.get_shift_names_by_week_day(week, day) for shift in constraint.shifts): expr_terms = [ worker.grade * self.model.works[worker.id, week, day, shift] for shift in constraint.shifts for worker in self.get_workers_for_shift(self.get_shift_by_name(shift)) ] if expr_terms: # Only add constraint if there are eligible workers self.model.constraints.add( sum(expr_terms) >= constraint.min_grade_sum ) # Ensure each shift has the requisit number of workers assigned for ( week, day, shift, workers_required, site_required, ) in self.get_required_workers_and_site_combinations(): # print(week, day, shift, workers_required, site_required) self.model.constraints.add( workers_required == sum( self.model.works[worker.id, week, day, shift] for worker in self.workers if worker.site in site_required ) ) if self.get_workers_who_require_locums(): self.model.constraints.add( self.model.locum_required[week, day, shift] == sum( self.model.works[worker.id, week, day, shift] for worker in self.workers if worker.locum ) ) self.model.constraints.add( self.model.locum_required[week, day, shift] == sum( self.model.locum_works[worker.id, week, day, shift] for worker in self.workers # if not worker.locum ) ) self.model.constraints.add( 0 == sum( self.model.locum_works[worker.id, week, day, shift] for worker in self.workers if worker.locum ) ) # And it is not assigned if the worker is from the wrong site if [worker for worker in self.workers if worker.site not in site_required]: self.model.constraints.add( 0 == sum( self.model.works[worker.id, week, day, shift] for worker in self.workers if worker.site not in site_required ) ) # # Ensure shifts are only assigned by workers from the allowed sites # self.model.constraints.add(0 == sum( # self.model.works[worker.id, week, day, shift] # for worker in self.workers if worker.site not in site_required)) # Constraint: total hours worked hours worked # for worker in self.workers: # self.model.constraints.add( # 1200 >= sum(shift_lengths[shift] * self.model.works[worker, week, day, shift] for week in weeks for day in days for shift in days_shifts[day]) # ) # for week in weeks: # for worker in self.workers: # self.model.constraints.add( # 48 >= sum(shift_lengths[shift] * self.model.works[worker, week, day, shift] for day in days for shift in days_shifts[day]) # ) # Define a maximum number hours that can be worked per week # NOTE: this works on a calander week not a 7 day period # TODO: consider fixing # def maxHoursPerWeekRule(model, wid, week): # return 72 >= sum( # shift.length * model.works[wid, week, day, shift.name] # for day, shift in self.get_day_shiftclass_combinations()) # self.model.max_hours_per_week_constraint = Constraint( # [worker.id for worker in self.workers], # self.weeks, # rule=maxHoursPerWeekRule) # Define an upper limit for hours worked over the rota period # def maxHoursRule(model, wid): # return 8000 >= sum( # shift.length * model.works[wid, week, day, shift.name] # for week, day, shift in self.get_all_shiftclass_combinations()) # self.model.max_hours_constraint = Constraint( # [worker.id for worker in self.workers], rule=maxHoursRule) # Limit grades def limitGradesRule(model, week, shift, grade, limit): workers = [w for w in self.workers if w.grade == grade] if not workers: return Constraint.Skip return ( sum( model.shift_week_worker_assigned[shift, week, w.id] for w in workers ) <= limit ) for shift in self.get_shifts_with_constraint(LimitGradeNumberConstraint): for constraint in shift.get_constraints(LimitGradeNumberConstraint): grade = constraint.grade limit = constraint.max_number # self.model.limit_grades_constraint = Constraint( setattr( self.model, f"limit_grades_constraint_{shift.name}_{grade}", Constraint( [week for week in self.weeks], # [shift for shift in self.get_shifts_with_constraint("limit_grade_number")], [shift.name], [grade], [limit], rule=limitGradesRule, # name=f"limit_grades_constraint_{shift.name}_{grade}" ), ) def minimumGradesRule(model, week, shift, grade, limit): workers = [w for w in self.workers if w.grade >= grade] if len(workers) < limit: # return Constraint.Skip raise ValueError("minimum_grade_number: not enough valid workers") return ( sum( model.shift_week_worker_assigned[shift, week, w.id] for w in workers ) >= limit ) for shift in self.get_shifts_with_constraint(MinimumGradeNumberConstraint): constraint = shift.get_constraints(MinimumGradeNumberConstraint)[0] grade, min_required = constraint.grade, constraint.min_number # self.model.limit_grades_constraint = Constraint( setattr( self.model, f"minimum_grade_number_{shift.name}", Constraint( [week for week in self.weeks], # [shift for shift in self.get_shifts_with_constraint("limit_grade_number")], [shift.name], [grade], [min_required], rule=minimumGradesRule, # name=f"limit_grades_constraint_{shift.name}_{grade}" ), ) def presenceAtRemoteSiteWeek( model, week, shift, required_site, required_number ): required_site_workers = [ w for w in self.workers if required_site == w.remote_site ] return ( sum( model.shift_week_worker_assigned[shift, week, w.id] for w in required_site_workers ) >= required_number ) for shift in self.get_shifts_with_constraint( RequireRemoteSitePresenceConstraint ): constraint = shift.get_constraints(RequireRemoteSitePresenceConstraint)[0] site, required_number = constraint.site, constraint.required_number setattr( self.model, f"require_remote_site_presence_week_{shift.name}", Constraint( [week for week in self.weeks], [shift.name], [site], [required_number], rule=presenceAtRemoteSiteWeek, ), ) #def presenceAtRemoteSite( # model, day, week, shift, required_site, required_number #): # required_site_workers = [ # w for w in self.workers if required_site == w.remote_site # ] # for w in required_site_workers: # if (w.id, week, day, shift) not in model.works: # return Constraint.Skip # return ( # sum(model.works[w.id, week, day, shift] for w in required_site_workers) # >= required_number # ) #for shift in self.get_shifts_with_constraint("require_remote_site_presence"): # site, required_number = shift.constraint_options[ # "require_remote_site_presence" # ] # # self.model.require_presence_at_site_overnight_rule = Constraint( # setattr( # self.model, # f"require_remote_site_presence_{shift.name}", # Constraint( # [day for day in self.days], # [week for week in self.weeks], # [shift.name], # [site], # [required_number], # rule=presenceAtRemoteSite, # ), # ) # # Count the number of workers from each site on each night shift # # As 1 or 0 is optimum we can simply subtract 1 from the number # if self.constraint_options["balance_nights_across_sites"]: # for site in self.sites: # for shift in self.get_shifts_with_constraint("night"): # block = shift.name # for week in self.weeks: # self.model.constraints.add( # self.model.night_per_site[week, block, site] # == sum( # self.model.shift_week_worker_assigned[ # block, week, worker.id # ] # for worker in self.workers # if worker.site == site # ) # # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site) # ) # self.model.constraints.add( # self.model.night_per_site_t1[week, block, site] # - 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 shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint): if site in shift.sites: block = shift.name for week in self.weeks: self.model.constraints.add( self.model.shift_per_site[week, block, site] == sum( self.model.shift_week_worker_assigned[ block, week, worker.id ] for worker in self.workers if worker.site == site ) # 1 >= sum(self.model.shift_week_worker_assigned[shift.name, week, worker.id] for worker in self.workers if worker.site == site) ) self.model.constraints.add( self.model.shift_per_site_t1[week, block, site] - self.model.shift_per_site_t2[week, block, site] == self.model.shift_per_site[week, block, site] - 1 ) weeks_days = self.get_week_day_combinations() #for week in track(self.weeks, description="Generating week constraints..."): # for shift_name in self.shifts_to_assign_or_force_as_blocks(): for week, shift_name in self.week_shifts_to_assign_as_blocks(): shift = self.get_shift_by_name(shift_name) 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 ) ) 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 ) ) 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 = [] 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) 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 ] for worker in self.workers ) <= workers_required ) else: self.model.constraints.add( sum( self.model.blocks_worker_shift_assigned[ worker.id, week, shift.name ] for worker in full_workers ) <= 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) worker: Worker for worker in self.workers: console.rule( f"Generate worker constraints: [bold blue]{worker.name}[/bold blue]" ) logger.debug(f"Generate worker constraints: {worker.name}") for week, day in self.get_week_day_combinations(): shifts_today = self.get_shift_names_by_week_day(week, day) total_assigned = sum( self.model.works[worker.id, week, day, shift] for shift in shifts_today ) # If assigned to any shift, works_day must be 1 self.model.constraints.add(total_assigned >= self.model.works_day[worker.id, week, day]) # If not assigned to any shift, works_day must be 0 self.model.constraints.add(total_assigned <= 1000 * self.model.works_day[worker.id, week, day]) # --- Max days worked per week constraint (global or per-worker) --- # Use per-worker value if set, otherwise use global constraint option if set max_days_worked_per_week = getattr(worker, "max_days_worked_per_week", None) if max_days_worked_per_week is None: max_days_worked_per_week = self.constraint_options.get("max_days_worked_per_week", None) if max_days_worked_per_week is not None: for week in self.weeks: # Create a binary variable for each day to indicate if the worker works any shift that day for day in self.days: var_name = f"works_any_{worker.id}_{week}_{day}" if not hasattr(self.model, "works_any_day_vars"): self.model.works_any_day_vars = {} if (worker.id, week, day) not in self.model.works_any_day_vars: self.model.works_any_day_vars[(worker.id, week, day)] = Var(within=Binary) setattr(self.model, var_name, self.model.works_any_day_vars[(worker.id, week, day)]) works_any = self.model.works_any_day_vars[(worker.id, week, day)] total_assigned = sum( self.model.works[worker.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day) ) # If assigned to any shift, works_any must be 1 self.model.constraints.add(total_assigned >= works_any) # If not assigned to any shift, works_any must be 0 self.model.constraints.add(total_assigned <= 1000 * works_any) # Now sum the binaries for the week self.model.constraints.add( sum(self.model.works_any_day_vars[(worker.id, week, day)] for day in self.days) <= max_days_worked_per_week ) for week, day, shift_name in getattr(worker, "forced_assignments", []): # Only add if this shift is valid for this week/day if shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True): # We have already checked for leave conflicts when building the model self.model.constraints.add( self.model.works[worker.id, week, day, shift_name] == 1 ) else: date = self.get_date_by_week_day(week, day) self.add_warning( "Invalid forced assignment", f"Worker '{worker.name}' has forced assignment for non-existent shift: {shift_name} on week {week}, day {day} (date: {date})", ) 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: for other_shift_name in shift.force_assign_with: # Only add constraint if the other shift is available on this day if other_shift_name in self.get_shift_names_by_week_day(week, day): for worker in self.get_workers_for_shift(shift): self.model.constraints.add( self.model.works[worker.id, week, day, shift.name] <= self.model.works[worker.id, week, day, other_shift_name] ) # Add hard shift dependencies for this worker if hasattr(worker, "hard_day_dependencies"): for dep in worker.hard_day_dependencies: day_group = dep.days weeks_to_apply = self.weeks if dep.weeks is not None: weeks_to_apply = dep.weeks if dep.start_date is not None or dep.end_date is not None: if dep.start_date is None: dep.start_date = self.start_date if dep.end_date is None: dep.end_date = self.rota_end_date weeks_to_apply = self.get_weeks_by_date_range( dep.start_date, dep.end_date ) for week in weeks_to_apply: assigned_any_shift = {} for day in day_group: shifts_today = self.get_shift_names_by_week_day(week, day) var_name = f"hard_day_dep_{worker.id}_{week}_{day}" if not hasattr(self.model, "hard_day_dep_vars"): self.model.hard_day_dep_vars = {} if (worker.id, week, day) not in self.model.hard_day_dep_vars: self.model.hard_day_dep_vars[(worker.id, week, day)] = Var(within=Binary) setattr(self.model, var_name, self.model.hard_day_dep_vars[(worker.id, week, day)]) assigned_any_shift[day] = self.model.hard_day_dep_vars[(worker.id, week, day)] total_assigned = sum(self.model.works[worker.id, week, day, shift] for shift in shifts_today) self.model.constraints.add(total_assigned >= assigned_any_shift[day]) self.model.constraints.add(total_assigned <= len(shifts_today) * assigned_any_shift[day]) # All-or-none: all days in the group must have the same assignment status vals = list(assigned_any_shift.values()) for i in range(1, len(vals)): self.model.constraints.add(vals[i] == vals[0]) # Add hard day exclusions for this worker if hasattr(worker, "hard_day_exclusions"): for exclusion in worker.hard_day_exclusions: days_to_exclude = exclusion.days weeks_to_apply = self.weeks if exclusion.weeks is not None: weeks_to_apply = exclusion.weeks if exclusion.start_date is not None or exclusion.end_date is not None: if exclusion.start_date is None: exclusion.start_date = self.start_date if exclusion.end_date is None: exclusion.end_date = self.rota_end_date weeks_to_apply = self.get_weeks_by_date_range( exclusion.start_date, exclusion.end_date ) for week in weeks_to_apply: # For each week, prevent any shift being assigned on both days self.model.constraints.add( sum( self.model.works_day[worker.id, week, day] for day in days_to_exclude #for shift in self.get_shift_names_by_week_day(week, day) ) <= 1 ) #self.model.constraints.add( # sum( # self.model.works[worker.id, week, day1, shift] # for shift in self.get_shift_names_by_week_day(week, day1) # ) # + sum( # self.model.works[worker.id, week, day2, shift] # for shift in self.get_shift_names_by_week_day(week, day2) # ) # <= 1 #) for week, day in self.get_week_day_combinations(): shifts_today = self.get_shift_names_by_week_day(week, day) worker_allowed_sets = getattr(worker, "allowed_multi_shift_sets", []) # for idx, allowed_set in enumerate(worker_allowed_sets): # if allowed_set.issubset(shifts_today): # var = self.model.multi_shift_together[worker.id, week, day, idx] # self.model.constraints.add( # sum(self.model.works[worker.id, week, day, shift] for shift in allowed_set) # == len(allowed_set) * var # ) for idx, allowed_set in enumerate(worker_allowed_sets): if allowed_set.issubset(shifts_today): var = self.model.multi_shift_together[worker.id, week, day, idx] # var is 1 iff all shifts in allowed_set are assigned for shift in allowed_set: self.model.constraints.add( var <= self.model.works[worker.id, week, day, shift] ) self.model.constraints.add( var >= sum( self.model.works[worker.id, week, day, shift] for shift in allowed_set ) - len(allowed_set) + 1 ) shifts_today = self.get_shift_names_by_week_day(week, day) for shift_name in shifts_today: if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with: for other_shift in worker.force_assign_with[shift_name]: if other_shift in shifts_today: self.model.constraints.add( self.model.works[worker.id, week, day, shift_name] <= self.model.works[worker.id, week, day, other_shift] ) if self.get_workers_who_require_locums(): for week, day, shift in self.get_all_shiftclass_combinations(): if not worker.locum: self.model.constraints.add( self.get_locum_availability(worker, week, day, shift) >= self.model.locum_works[worker.id, week, day, shift.name] ) self.model.constraints.add( self.model.locum_shifts_by_worker[worker.id] == sum( self.model.locum_works[worker.id, week, day, shift] for week, day, shift in self.get_all_shiftname_combinations() ) ) self.model.constraints.add( self.model.locum_shifts_t1[worker.id] - self.model.locum_shifts_t2[worker.id] == self.model.locum_shifts_by_worker[worker.id] ) if worker.locum_max_shifts: self.model.constraints.add( worker.locum_max_shifts >= sum( self.model.locum_works[worker.id, week, day, shiftname] for week, day, shiftname in self.get_all_shiftname_combinations() ) ) if worker.locum_max_shifts_per_week: for week_to_check in self.weeks: self.model.constraints.add( worker.locum_max_shifts_per_week >= sum( self.model.locum_works[worker.id, week, day, shiftname] for week, day, shiftname in self.get_all_shiftname_combinations( week=week_to_check ) ) ) if self.constraint_options["distribute_locum_shifts"]: self.model.constraints.add( sum( self.model.locum_required[week, day, shift] for week, day, shift in self.get_all_shiftname_combinations() ) / (len(self.get_locum_workers()) - 1) >= self.model.locum_shifts_by_worker[worker.id] ) if ( self.constraint_options["maximum_allowed_locum_shifts_per_worker"] is not None ): self.model.constraints.add( self.constraint_options[ "maximum_allowed_locum_shifts_per_worker" ] >= self.model.locum_shifts_by_worker[worker.id] ) if worker.locum_availability: if ( self.constraint_options[ "minimum_allowed_locum_shifts_per_worker" ] is not None ): self.model.constraints.add( self.constraint_options[ "minimum_allowed_locum_shifts_per_worker" ] <= self.model.locum_shifts_by_worker[worker.id] ) if ( self.constraint_options["minimise_shift_diffs"] or self.constraint_options["maximum_allowed_shift_diff"] is not None ): self.model.constraints.add( self.model.shift_count_diff_summed[worker.id] == sum( self.model.shift_count_diff[worker.id, shift.name] for shift in self.get_shifts() ) ) if self.constraint_options["maximum_allowed_shift_diff"] is not None: self.model.constraints.add( self.model.shift_count_diff_summed[worker.id] <= self.constraint_options["maximum_allowed_shift_diff"] ) self.model.constraints.add( self.model.shift_count_diff_summed[worker.id] >= -self.constraint_options["maximum_allowed_shift_diff"] ) for max_days, weeks in self.constraint_options["max_days_per_week_block"]: for week_blocks in self.get_week_block_iterator(weeks): self.model.constraints.add( sum( self.model.works_day[worker.id, week, day] for week, day in self.get_week_day_combinations() if week in week_blocks ) <= max_days ) for week_blocks in self.get_week_block_iterator(4): # Prevent more than n number shifts per 4 weeks try: self.model.constraints.add( self.constraint_options["max_shifts_per_month"] >= sum( self.model.works[worker.id, week, day, shiftname] for week, day, shiftname in self.get_all_shiftname_combinations() if week in week_blocks ) ) except ValueError: # Occurs if there are no shifts within a defined block # TODO: test if this breaks (and we should check rathar than except) self.add_warning( "max shifts per month constraint", f"Failed to constrain max_shifts_per_month for worker: {worker.name}", ) pass for constraint_dict in self.constraint_options[ "avoid_shifts_by_worker_names" ]: if invalid_shifts := set(constraint_dict["shifts"]).difference( set(self.get_shift_names()) ): message = f"avoid shifts by worker constraint contains a non existent shift: {invalid_shifts}" if self.ignore_valid_shifts: self.add_warning("Non existent shift", message) continue # Skip if we don't want to raise an error else: raise InvalidShift(message) if worker.name in constraint_dict["names"]: self.model.constraints.add( 0 == sum( self.model.works[worker.id, week, day, shift] for week, day, shift in self.get_all_shiftname_combinations() if week in constraint_dict["weeks"] and shift in constraint_dict["shifts"] ) ) for constraint_dict in self.constraint_options["avoid_shifts_by_grades"]: if invalid_shifts := set(constraint_dict["shifts"]).difference( set(self.get_shift_names()) ): message = f"avoid shifts by grade constraint contains a non existent shift: {invalid_shifts}" if self.ignore_valid_shifts: self.add_warning("Non existent shift", message) continue # Skip if we don't want to raise an error else: raise InvalidShift(message) if invalid_grades := set(constraint_dict["grades"]).difference( set(self.get_worker_grades()) ): message = f"avoid shifts by grade constraint contains a non existent worker grade: {invalid_grades}" self.add_warning("Non existent shift", message) raise ValueError(message) if worker.grade in constraint_dict["grades"]: self.model.constraints.add( 0 == sum( self.model.works[worker.id, week, day, shift] for week, day, shift in self.get_all_shiftname_combinations() if week in constraint_dict["weeks"] and shift in constraint_dict["shifts"] ) ) # Count number of weekends an worker works # if self.constraint_options["balance_weekends"]: self.model.constraints.add( self.model.worker_weekend_count[worker.id] == sum(self.model.works_weekend[worker.id, week] for week in self.weeks) ) if self.constraint_options["max_weekends"] > -1: self.model.constraints.add( self.constraint_options["max_weekends"] >= self.model.worker_weekend_count[worker.id] ) # 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" ): if ( worker.site in shift.sites ): # Each site specfies which sites self.workers can fullfill it # calaculate the total number of shifts that need assigning over the rota # total_shifts = (len(self.weeks) * len(shift.days) * # TODO: check if shift_counts is still required? # total_shifts = ( # self.shift_counts[shift.name] * shift.workers_required # ) total_shifts = self.shift_worker_counts[shift.name] full_time_equivalent_joined = sum( self.full_time_equivalent_sites[i][shift.name] for i in shift.sites ) target_shifts = ( total_shifts / full_time_equivalent_joined * worker.get_fte(shift=shift.name) ) if self.use_previous_shifts: if shift.name in worker.previous_shifts: worked, allocated = worker.previous_shifts[shift.name] target_shifts = ( target_shifts + float(allocated) - float(worked) ) if self.use_shift_balance_extra: if shift.name in worker.shift_balance_extra: extra = worker.shift_balance_extra[shift.name] # TODO look at how this affects allocation (how does it affect fte) match extra: case "double": target_shifts = target_shifts * 2 case "half": target_shifts = target_shifts / 2 case _: target_shifts = target_shifts + extra # print(worker.name, shift.name, target_shifts) worker.shift_target_number[shift.name] = target_shifts if shift.hard_constrain_shift: adjusted_balance_offset = shift.balance_offset if worker.get_fte(shift=shift.name) < 100: adjusted_balance_offset = ( adjusted_balance_offset * self.ltft_balance_offset ) max_shifts = ( target_shifts + adjusted_balance_offset * self.balance_offset_modifier ) min_shifts = ( target_shifts - adjusted_balance_offset * self.balance_offset_modifier ) # --- Store for export --- if not hasattr(worker, "hard_constrain_shift_limits"): worker.hard_constrain_shift_limits = {} worker.hard_constrain_shift_limits[shift.name] = { "min_shifts": min_shifts, "max_shifts": max_shifts, "target_shifts": target_shifts, } self.model.constraints.add( inequality( min_shifts, sum( self.model.works[worker.id, week, day, shift.name] for week, day in self.get_week_day_combinations_for_shift( shift ) ), max_shifts, ) ) else: # Make sure shifts aren't assigned to those from other sites # Not needed if there are no shifts if self.get_week_day_combinations_for_shift(shift): self.model.constraints.add( 0 == sum( self.model.works[worker.id, week, day, shift.name] for week, day in self.get_week_day_combinations_for_shift( shift ) # for week, day in self.get_week_day_combinations() ) ) # model.shift_count stores the number of shifts that a worker is assigned to worker by shift # this is used (by the objective) to balance the total number of shifts worked self.model.constraints.add( self.model.shift_count[worker.id, shift.name] == sum( self.model.works[worker.id, week, day, shift.name] for week, day in self.get_week_day_combinations_for_shift(shift) ) ) self.model.constraints.add( self.model.shift_count_diff[worker.id, shift.name] == self.model.shift_count[worker.id, shift.name] - worker.shift_target_number[shift.name] ) if self.constraint_options["balance_shifts"]: self.model.constraints.add( self.model.shift_count_t1[worker.id, shift.name] - self.model.shift_count_t2[worker.id, shift.name] == self.model.shift_count_diff[worker.id, shift.name] ) if self.constraint_options["balance_shifts_quadratic"]: # This may need to be updated xU = 25 xL = 1 self.model.constraints.add( inequality( xL, self.model.shift_count_t1[worker.id, shift.name] + self.model.shift_count_t2[worker.id, shift.name] + 1, xU, ) ) self.model.constraints.add( self.model.shift_count_w[worker.id, shift.name] >= xL * ( self.model.shift_count_t1[worker.id, shift.name] + self.model.shift_count_t2[worker.id, shift.name] + 1 ) * 2 - xL * xL ) self.model.constraints.add( self.model.shift_count_w[worker.id, shift.name] >= xU * ( self.model.shift_count_t1[worker.id, shift.name] + self.model.shift_count_t2[worker.id, shift.name] + 1 ) * 2 - xU * xU ) # Define worker_shift_count_t1 and worker_shift_count_t2 constraints for the object # Thus bypassing the need for a quadratic solver # t1-t2 is the target # 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] - self.model.worker_shift_count_t2[worker.id] == sum( (self.model.shift_count_diff[worker.id, shift.name]) * shift.balance_weighting for shift in self.get_shifts() ) ) if self.constraint_options["balance_bank_holidays"]: extra_bank_holiday = 0 if self.use_bank_holiday_extra: extra_bank_holiday = worker.bank_holiday_extra self.model.constraints.add( self.model.bank_holiday_count[worker.id] == sum( self.model.works[worker.id, week, day, shift] for week, day, shift in self.get_all_shiftname_combinations() if self.week_day_date_map[(week, day)] in self.bank_holidays ) + extra_bank_holiday ) xU = len(self.get_bank_holiday_week_days()) + 1 xL = 1 self.model.constraints.add( inequality( xL, self.model.bank_holiday_count[worker.id] + 1, xU, ) ) self.model.constraints.add( self.model.bank_holiday_count_w[worker.id] >= xL * (self.model.bank_holiday_count[worker.id] + 1) * 2 - xL * xL ) self.model.constraints.add( self.model.bank_holiday_count_w[worker.id] >= xU * (self.model.bank_holiday_count[worker.id] + 1) * 2 - xU * xU ) if self.constraint_options["balance_nights"]: self.model.constraints.add( self.model.night_shift_count[worker.id] == sum( self.model.works[worker.id, week, day, shift.name] for week, day, shift in self.get_week_day_shift_combinations_for_constraint( NightConstraint ) # for week, day in self.get_week_day_combinations_for_shift(shift) # for shift in self.get_shifts_with_constraint("night") ) ) night_shift_target_number = sum( worker.shift_target_number[shift.name] for shift in self.get_shifts_with_constraint(NightConstraint) ) self.model.constraints.add( self.model.night_shift_count_t1[worker.id] - self.model.night_shift_count_t2[worker.id] == self.model.night_shift_count[worker.id] - night_shift_target_number ) # This may need to be updated xU = 6 xL = 1 self.model.constraints.add( inequality( xL, self.model.night_shift_count_t1[worker.id] + self.model.night_shift_count_t2[worker.id] + 1, xU, ) ) self.model.constraints.add( self.model.night_shift_count_w[worker.id] >= xL * ( self.model.night_shift_count_t1[worker.id] + self.model.night_shift_count_t2[worker.id] + 1 ) * 2 - xL * xL ) self.model.constraints.add( self.model.night_shift_count_w[worker.id] >= xU * ( self.model.night_shift_count_t1[worker.id] + self.model.night_shift_count_t2[worker.id] + 1 ) * 2 - xU * xU ) # self.model.constraints.add( # self.model.night_shift_count_w[worker.id] >= 0) # We use a similar method to balance the number of weekends worked # This works as long as weekend shifts are assigned as blocks! weekend_shift_target_number = sum( worker.shift_target_number[shift.name] / len(shift.days) for shift in self.get_shifts() if "Sat" in shift.days or "Sun" in shift.days ) worker.weekend_shift_target_number = weekend_shift_target_number if self.constraint_options["balance_weekends"]: if weekend_shift_target_number > 0: self.model.constraints.add( self.model.weekend_shift_count_t1[worker.id] - self.model.weekend_shift_count_t2[worker.id] == self.model.worker_weekend_count[worker.id] - weekend_shift_target_number ) xU = self.constraint_options["max_weekends"] xL = 1 self.model.constraints.add( inequality( xL, self.model.weekend_shift_count_t1[worker.id] + self.model.weekend_shift_count_t2[worker.id] + 1, xU, ) ) self.model.constraints.add( self.model.weekend_shift_count_w[worker.id] >= xL * ( self.model.weekend_shift_count_t1[worker.id] + self.model.weekend_shift_count_t2[worker.id] + 1 ) * 2 - xL * xL ) self.model.constraints.add( self.model.weekend_shift_count_w[worker.id] >= xU * ( self.model.weekend_shift_count_t1[worker.id] + self.model.weekend_shift_count_t2[worker.id] + 1 ) * 2 - xU * xU ) # Ensure worker is not allocated shifts on non working days if worker.non_working_day_list: for week, day, shift in self.get_all_shiftclass_combinations(): for n, start, end in worker.non_working_day_list: if not shift.rota_on_nwds and day == n: # print(start, self.week_day_date_map[(week, day)], end) # If this shift is force assigned, allow it but add a warning if start <= self.week_day_date_map[(week, day)] < end: if (week, day, shift.name) in getattr(worker, "forced_assignments", []): self.add_warning( "Force assignment/NWD conflict", f"Worker '{worker.name}' has a forced assignment for shift '{shift.name}' on week {week}, day {day} (date: {self.get_date_by_week_day(week, day)}), which is a non-working day." ) else: self.model.constraints.add( 0 == self.model.works[ worker.id, week, day, shift.name ] ) if self.get_workers_who_require_locums(): if not worker.locum_on_nwds and day == n: if start <= self.week_day_date_map[(week, day)] < end: self.model.constraints.add( 0 == self.model.locum_works[ worker.id, week, day, shift.name ] ) if self.constraint_options["balance_blocks"]: if self.constraint_options["max_night_frequency"]: for week_blocks in self.get_week_block_iterator( self.constraint_options["max_night_frequency"] ): # Ignore excluded weeks if set(week_blocks).intersection( set( self.constraint_options[ "max_night_frequency_week_exclusions" ] ) ): continue if self.get_shifts_with_constraint(NightConstraint): # Prevent nights more than once every n weeks self.model.constraints.add( 1 >= sum( self.model.shift_week_worker_assigned[ shift.name, week, worker.id ] for week in week_blocks for shift in self.get_shifts_with_constraint( NightConstraint ) ) ) # if self.constraint_options["balance_weekends"]: if self.constraint_options["max_weekend_frequency"]: for week_blocks in self.get_week_block_iterator( self.constraint_options["max_weekend_frequency"] ): # Prevent weekend shifts more than once every n weeks self.model.constraints.add( 1 >= sum( self.model.works_weekend[worker.id, week] for week in week_blocks ) ) for constraint_shift in self.get_shifts_with_constraints( MaxShiftsPerWeekBlockConstraint ): constraint = constraint_shift.get_constraint(MaxShiftsPerWeekBlockConstraint) # If not supplied, default to the number of days per shift if constraint.max_shifts is not None: shift_number = constraint.max_shifts else: shift_number = constraint_shift.get_shift_number() for week_blocks in self.get_week_block_iterator( constraint.week_block ): # Prevent shifts more than once every n weeks self.model.constraints.add( shift_number >= sum( self.model.works[ worker.id, week, day, constraint_shift.name ] for week in week_blocks for day in constraint_shift.days ) ) for week in self.weeks: weekend_days = self.days[5:] # ["Sat", "Sun"] total_assigned = sum( self.model.works[worker.id, week, day, shiftname] for w, day, shiftname in self.get_all_shiftname_combinations(week=week) if day in weekend_days ) # If any shift is assigned on Sat or Sun, works_weekend is 1 self.model.constraints.add(total_assigned >= self.model.works_weekend[worker.id, week]) # If works_weekend is 1, at least one shift must be assigned self.model.constraints.add(total_assigned <= 1000 * self.model.works_weekend[worker.id, week]) for constraint_shift in self.get_shifts_with_constraints( MaxShiftsPerWeekConstraint ): # If not supplied, default to the number of days per shift constraint = constraint_shift.get_constraint(MaxShiftsPerWeekConstraint) if constraint.max_shifts is not None: shift_number = constraint.max_shifts else: shift_number = constraint_shift.get_shift_number() try: self.model.constraints.add( shift_number >= sum( self.model.works[ worker.id, w, day, constraint_shift.name ] # for shiftname in self.get_shift_names_by_week_day(week, day) # for day in self.days for w, day in self.get_week_day_combinations_for_shift( constraint_shift ) if w == week ) ) except ValueError: # This happens if a shift if not assigned on the week (should we test for this instead?) pass try: self.model.constraints.add( self.constraint_options["max_shifts_per_week"] >= sum( self.model.works[worker.id, w, day, shiftname] # for shiftname in self.get_shift_names_by_week_day(week, day) # for day in self.days for w, day, shiftname in self.get_all_shiftname_combinations( week=week ) ) ) except ValueError: pass # TODO: consider excluding shifts that span relevant period if ( worker.site in self.constraint_options["prevent_monday_after_full_weekends"] ): # Ignore last week if week + 1 in self.weeks: self.model.constraints.add( full_weekend_count >= sum( self.model.works[worker.id, week, "Sat", shift] for shift in self.get_shift_names_by_week_day( week, "Sat" ) ) + sum( self.model.works[worker.id, week, "Sun", shift] for shift in self.get_shift_names_by_week_day( week, "Sun" ) ) + sum( self.model.works[worker.id, week + 1, "Mon", shift] for shift in self.get_shift_names_by_week_day( week + 1, "Mon" ) ) ) if ( worker.site in self.constraint_options[ "prevent_monday_and_tuesday_after_full_weekends" ] ): # Ignore last week if week + 1 in self.weeks: self.model.constraints.add( full_weekend_count >= sum( self.model.works[worker.id, week, "Sat", shift] for shift in self.get_shift_names_by_week_day( week, "Sat" ) ) + sum( self.model.works[worker.id, week, "Sun", shift] for shift in self.get_shift_names_by_week_day( week, "Sun" ) ) + sum( self.model.works[worker.id, week + 1, "Mon", shift] for shift in self.get_shift_names_by_week_day( week + 1, "Mon" ) ) + sum( self.model.works[worker.id, week + 1, "Tue", shift] for shift in self.get_shift_names_by_week_day( week + 1, "Tue" ) ) ) if ( worker.site in self.constraint_options["prevent_fridays_before_full_weekends"] ): self.model.constraints.add( full_weekend_count >= sum( self.model.works[worker.id, week, "Sat", shift.name] for shift in self.get_shift_names_by_week_day(week, "Sat") ) + sum( self.model.works[worker.id, week, "Sun", shift.name] for shift in self.get_shift_names_by_week_day(week, "Sat") ) + sum( self.model.works[worker.id, week, "Fri", shift.name] for shift in self.get_shift_names_by_week_day(week, "Sat") ) ) if ( worker.site in self.constraint_options["prevent_thursdays_before_full_weekends"] ): self.model.constraints.add( full_weekend_count >= sum( self.model.works[worker.id, week, "Sat", shift.name] for shift in self.get_shift_names_by_week_day(week, "Sat") ) + sum( self.model.works[worker.id, week, "Sun", shift.name] for shift in self.get_shift_names_by_week_day(week, "Sun") ) + sum( self.model.works[worker.id, week, "Thu", shift.name] for shift in self.get_shift_names_by_week_day(week, "Thu") ) ) # # model.weekend_count stores the number of weekend shifts that a worker is assigned to wor # # this is used (by the objective) to balance the total number of weekend shifts worked # self.model.constraints.add( # # We could use a helper to get required shifts to check # self.model.works_whole_weekend[worker.id, week] >= sum( # self.model.works[worker.id, week, day, shift.name] # for day in self.days[5:] # for shift in self.get_shifts()) - 1) # self.model.constraints.add( # self.model.works_whole_weekend[worker.id, week] <= sum( # self.model.works[worker.id, week, day, shift.name] # for day in self.days[5:] # for shift in self.get_shifts()) / 2) # if self.constraint_options["balance_weekends"]: # Try disabling tihs to allow more than 2 shifts ot be assigned on a weekend #self.model.constraints.add( # self.model.works_weekend[worker.id, week] # >= sum( # self.model.works[worker.id, week, day, shiftname] # for w, day, shiftname in self.get_all_shiftname_combinations( # week=week # ) # if day in self.days[5:] # ) # / 2 #) #self.model.constraints.add( # self.model.works_weekend[worker.id, week] # <= sum( # self.model.works[worker.id, week, day, shiftname] # for w, day, shiftname in self.get_all_shiftname_combinations( # week=week # ) # if day in self.days[5:] # ) #) # self.model.constraints.add( # self.model.works_saturday[worker.id, week] == sum( # self.model.works[worker.id, week, "Sat", shift.name] # for shift in self.get_shifts())) # self.model.constraints.add( # self.model.works_sunday[worker.id, week] == sum( # self.model.works[worker.id, week, "Sun", shift.name] # for shift in self.get_shifts())) # self.model.constraints.add(1 == sum( # self.model.shift_week_worker_assigned[s.name, week, worker.id] # for s in self.get_shifts_with_constraint("night"))) # for shift in self.get_shifts_with_constraint("single_block_per_week"): # if shift.force_as_block: # # Force nights to be assigned in blocks # # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days) # # ) # # Force night shifts to be assigned in blocks # self.model.constraints.add( # self.model.shift_week_worker_assigned[ # shift.name, week, worker.id # ] # * len(shift.days) # == sum( # self.model.works[worker.id, w, d, s] # for w, d, s in self.get_all_shiftname_combinations(week=week) # #for day in self.days # ) # ) # else: # raise ValueError(f"Requires {shift.name} to have force_as_block") # for shift in self.get_shifts_with_constraint("night"): for shift in self.get_shifts(week=week): if shift.force_as_block: if shift.get_worker_requirement_by_date( self.get_week_start_date(week) ): # Force shifts to be assigned in blocks self.model.constraints.add( self.model.shift_week_worker_assigned[ shift.name, week, worker.id ] * len(shift.days) == sum( self.model.works[worker.id, week, day, shift.name] for day in shift.days ) ) # # def get_pre_may(week_days, max_pre) for n in track( range(len(weeks_days)), description="Generate week/day constraints" ): week, day = weeks_days[n] pre_map = [] for pre_n in range(1, self.max_pre + 1): try: pweek, pday = weeks_days[n - pre_n] p = 1 except: pweek, pday = weeks_days[n] p = 0 pre_map.append((p, pweek, pday)) post_map = [] for post_n in range(1, self.max_post + 1): try: pweek, pday = weeks_days[n + post_n] p = 1 except: pweek, pday = weeks_days[n] p = 0 post_map.append((p, pweek, pday)) # p1 = 1 # p2 = 1 # p3 = 1 # p4 = 1 # p5 = 1 # p6 = 1 # p7 = 1 # if n > 0: # pweek, pday = weeks_days[n - 1] # else: # p1 = 0 # pweek, pday = weeks_days[n] # try: # p2week, p2day = weeks_days[n - 2] # except IndexError: # p2week, p2day = weeks_days[n] # p2 = 0 # try: # p3week, p3day = weeks_days[n - 3] # except IndexError: # p3week, p3day = weeks_days[n] # p3 = 0 # try: # p4week, p4day = weeks_days[n - 4] # except IndexError: # p4week, p4day = weeks_days[n] # p4 = 0 # try: # p5week, p5day = weeks_days[n - 5] # except IndexError: # p5week, p5day = weeks_days[n] # p5 = 0 # try: # p6week, p6day = weeks_days[n - 6] # except IndexError: # p6week, p6day = weeks_days[n] # p6 = 0 # try: # p7week, p7day = weeks_days[n - 7] # except IndexError: # p7week, p7day = weeks_days[n] # p7 = 0 # pre_map = [ # (p1, pweek, pday), # (p2, p2week, p2day), # (p3, p3week, p3day), # (p4, p4week, p4day), # (p5, p5week, p5day), # (p6, p6week, p6day), # (p7, p7week, p7day), # ] # n1 = 1 # n2 = 1 # n3 = 1 # n4 = 1 # n5 = 1 # n6 = 1 # n7 = 1 # if n > 0: # nweek, nday = weeks_days[n - 1] # else: # n1 = 0 # nweek, nday = weeks_days[n] # try: # n2week, n2day = weeks_days[n - 2] # except IndexError: # n2week, n2day = weeks_days[n] # n2 = 0 # try: # n3week, n3day = weeks_days[n - 3] # except IndexError: # n3week, n3day = weeks_days[n] # n3 = 0 # try: # n4week, n4day = weeks_days[n - 4] # except IndexError: # n4week, n4day = weeks_days[n] # n4 = 0 # try: # n5week, n5day = weeks_days[n - 5] # except IndexError: # n5week, n5day = weeks_days[n] # n5 = 0 # try: # n6week, n6day = weeks_days[n - 6] # except IndexError: # n6week, n6day = weeks_days[n] # n6 = 0 # try: # n7week, n7day = weeks_days[n - 7] # except IndexError: # n7week, n7day = weeks_days[n] # n7 = 0 # post_map = [ # (n1, nweek, nday), # (n2, n2week, n2day), # (n3, n3week, n3day), # (n4, n4week, n4day), # (n5, n5week, n5day), # (n6, n6week, n6day), # (n7, n7week, n7day), # ] # n1 = 1 # n2 = 1 # try: # nweek, nday = weeks_days[n + 1] # except IndexError: # nweek, nday = weeks_days[n] # n1 = 0 # try: # n2week, n2day = weeks_days[n + 2] # except IndexError: # n2week, n2day = weeks_days[n] # n2 = 0 # try: # n3week, n3day = weeks_days[n + 3] # except IndexError: # n3week, n3day = weeks_days[n] # n3 = 0 # IF paired we check the following against both workers workers = [worker] if self.constraint_options["hard_constrain_pair_separation"]: for worker_pairs in self.worker_pairs: if worker_pairs[0] == worker: workers = worker_pairs if self.get_shift_names_by_week_day(week, day): # Unable to work (hard constraint not preference) self.model.constraints.add( self.model.available[worker.id, week, day] >= sum( self.model.works[worker.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day) ) ) if self.get_locum_workers(): self.model.constraints.add( self.model.available[worker.id, week, day] >= sum( self.model.locum_works[worker.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day) ) ) # single shift per day (unless multi-shift allowed) # This is signifantly slower so only enable if required shifts_today = self.get_shift_names_by_week_day(week, day) assigned_vars = [ self.model.works[worker.id, week, day, shift] for shift in shifts_today ] worker_allowed_sets = getattr( worker, "allowed_multi_shift_sets", [] ) force_assign_sets = [] for shift_name in shifts_today: if hasattr(worker, "force_assign_with") and shift_name in worker.force_assign_with: force_set = set([shift_name] + list(worker.force_assign_with[shift_name])) if force_set.issubset(shifts_today): force_assign_sets.append(force_set) shift_obj = self.get_shift_by_name(shift_name) if getattr(shift_obj, "force_assign_with", []): force_set = set([shift_name] + list(shift_obj.force_assign_with)) if force_set.issubset(shifts_today): force_assign_sets.append(force_set) # If no allowed sets, we can only assign one shift allowed_sets_today = [ allowed_set for allowed_set in worker_allowed_sets if allowed_set.issubset(shifts_today) ] + force_assign_sets if allowed_sets_today: # Forbid any multi-shift assignment that includes shifts from different allowed sets # Build all possible multi-shift assignments all_multi = [ set(c) for i in range(2, len(shifts_today) + 1) for c in itertools.combinations(shifts_today, i) ] # Build all allowed multi-shift patterns (all non-empty subsets of each allowed set) allowed_patterns = [{shift} for shift in shifts_today] for allowed_set in allowed_sets_today: allowed_patterns += [ set(s) for i in range(2, len(allowed_set) + 1) for s in itertools.combinations(allowed_set, i) ] # Forbid all other multi-shift combinations forbidden = [s for s in all_multi if s not in allowed_patterns] for forbidden_set in forbidden: self.model.constraints.add( sum( self.model.works[worker.id, week, day, shift] for shift in forbidden_set ) <= len(forbidden_set) - 1 ) # Still, at most all shifts per day self.model.constraints.add( sum( self.model.works[worker.id, week, day, shift] for shift in shifts_today ) <= len(shifts_today) ) else: # Only one shift per day if no allowed sets self.model.constraints.add(sum(assigned_vars) <= 1) # This applies to locums as well if self.get_workers_who_require_locums(): self.model.constraints.add( 1 >= sum( self.model.locum_works[worker.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day) ) ) self.model.constraints.add( 1 >= sum( self.model.works[worker.id, week, day, shift] + self.model.locum_works[worker.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day) ) ) if self.constraint_options["hard_constrain_pair_separation"]: for worker_pairs in self.worker_pairs: if worker_pairs[0] == worker: self.model.constraints.add( 1 >= sum( self.model.works[w.id, week, day, shift] for shift in self.get_shift_names_by_week_day( week, day ) for w in worker_pairs ) ) # 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( PreShiftConstraint ): constraint = constraint_shift.get_constraint(PreShiftConstraint) if constraint.allow_self: ignore_shifts = constraint.ignore_shifts + [constraint_shift.name] else: ignore_shifts = constraint.ignore_shifts if week not in constraint.weeks: continue for n in range(0, constraint.days): if day in constraint.exclude_days: continue if day in constraint_shift.days: # Only apply if worker can work this shift if worker.site not in constraint_shift.sites: continue try: works = self.model.works[ worker.id, week, day, constraint_shift.name ] except KeyError: continue self.model.constraints.add( 1 >= works + sum( pre_map[n][0] * self.model.works[ w.id, pre_map[n][1], pre_map[n][2], shiftname ] for shiftname in self.get_shift_names_by_week_day( pre_map[n][1], pre_map[n][2] ) if shiftname not in ignore_shifts for w in workers if w.site in constraint_shift.sites # Only workers who can work this shift ) ) for constraint_shift in self.get_shifts_with_constraints( PostShiftConstraint ): constraint = constraint_shift.get_constraint(PostShiftConstraint) if constraint.allow_self: ignore_shifts = constraint.ignore_shifts + [constraint_shift.name] else: ignore_shifts = constraint.ignore_shifts if week not in constraint.weeks: continue for n in range(0, constraint.days): if day in constraint.exclude_days: continue if day in constraint_shift.days: # Only apply if worker can work this shift if worker.site not in constraint_shift.sites: continue try: works = self.model.works[ worker.id, week, day, constraint_shift.name ] except KeyError: continue self.model.constraints.add( 1 >= works + sum( post_map[n][0] * self.model.works[ w.id, post_map[n][1], post_map[n][2], shiftname ] for shiftname in self.get_shift_names_by_week_day( post_map[n][1], post_map[n][2] ) if shiftname not in ignore_shifts for w in workers if w.site in constraint_shift.sites # Only workers who can work this shift ) ) # Night constraint means we won't assign a shift the day before # an unavailability for constraint_shift in self.get_shifts_with_constraint(NightConstraint): if ( worker.id, # pweek, pre_map[0][1], # pday, pre_map[0][2], constraint_shift.name, ) in self.model.works: # Ensure night prior to unavalibity is not assigned self.model.constraints.add( self.model.available[worker.id, week, day] >= self.model.works[ # worker.id, pweek, pday, constraint_shift.name worker.id, pre_map[0][1], pre_map[0][2], constraint_shift.name, ] ) ## Ensure each shift has the requisit number of workers assigned # for ( # week, # day, # shift, # workers_required, # site_required, # ) in self.get_required_workers_and_site_combinations(): # # print(week, day, shift, workers_required, site_required) # self.model.constraints.add( # workers_required # == sum( # self.model.works[worker.id, week, day, shift] # for worker in self.workers # if worker.locum # ) # ) # for worker in self.workers: # if worker.locum: # for week, day, shift in self.get_all_shiftname_combinations(): # self.model.constraints.add( # self.model.locum_works[worker.id, week, day, shift] # == self.model.works[worker.id, week, day, shift] # ) self.define_objectives() print("Building model completed") def define_objectives(self): # Define an objective function with model as input, to pass later def obj_rule(m): # c = len(workers) prefer_multi_shift_expr = sum( getattr(worker, "prefer_multi_shift_together", 0) * self.model.multi_shift_together[worker.id, week, day, idx] for worker in self.workers for week, day in self.get_week_day_combinations() for idx, allowed_set in enumerate( getattr(worker, "allowed_multi_shift_sets", []) ) if allowed_set.issubset(self.get_shift_names_by_week_day(week, day)) ) balance_modifier_constant = 10 balance_quadratic_shift_modifier_constant = 3 block_shift_balancing_constant = 1 locum_shift_balance_modifier_constant = 40 locum_shift_balancing = 0 if self.get_workers_who_require_locums(): if self.constraint_options["balance_locum_shifts"]: locum_shift_balancing = sum( locum_shift_balance_modifier_constant * ( self.model.locum_shifts_t1[(worker.id)] + self.model.locum_shifts_t2[(worker.id)] ) for worker in self.workers ) if self.constraint_options["balance_shifts_over_workers"]: worker_shift_balancing = sum( balance_modifier_constant * ( self.model.worker_shift_count_t1[(worker.id)] + self.model.worker_shift_count_t2[(worker.id)] ) for worker in self.workers ) else: worker_shift_balancing = 0 if self.constraint_options["balance_shifts"]: shift_balancing = sum( balance_modifier_constant * ( self.model.shift_count_t1[worker.id, shift.name] + self.model.shift_count_t2[worker.id, shift.name] ) for worker in self.workers for shift in self.get_shifts() ) else: shift_balancing = 0 if self.constraint_options["balance_shifts_quadratic"]: quadratic_shift_balancing = sum( balance_quadratic_shift_modifier_constant * ( self.model.shift_count_t1[worker.id, shift.name] + self.model.shift_count_t2[worker.id, shift.name] ) for worker in self.workers for shift in self.get_shifts() ) else: quadratic_shift_balancing = 0 true_quadratic_shift_balancing = 0 if self.constraint_options["balance_shifts_true_quadratic"]: shift_diff_modifier_constant = 10 true_quadratic_shift_balancing = sum( shift_diff_modifier_constant * self.model.shift_count_diff_summed[(worker.id)] for worker in self.workers ) if self.constraint_options["minimise_shift_diffs"]: shift_diff_modifier_constant = 10 shift_diff_balancing = sum( shift_diff_modifier_constant * self.model.shift_count_diff_summed[(worker.id)] for worker in self.workers ) else: shift_diff_balancing = 0 if self.constraint_options["balance_nights"]: night_balance_modifier_constant = 10 night_shift_balancing = sum( night_balance_modifier_constant * self.model.night_shift_count_w[(worker.id)] # (self.model.night_shift_count_t1[(worker.id)] + # self.model.night_shift_count_t2[(worker.id)]) for worker in self.workers ) else: night_shift_balancing = 0 if self.constraint_options["balance_bank_holidays"]: bank_holiday_balance_modifier_constant = 1000 bank_holiday_balancing = sum( bank_holiday_balance_modifier_constant * (self.model.bank_holiday_count_w[(worker.id)] - 1) # (self.model.night_shift_count_t1[(worker.id)] + # self.model.night_shift_count_t2[(worker.id)]) for worker in self.workers ) else: bank_holiday_balancing = 0 if self.constraint_options["balance_weekends"]: weekend_balance_modifier_constant = 5 weekend_shift_balancing = sum( weekend_balance_modifier_constant * self.model.weekend_shift_count_w[(worker.id)] for worker in self.workers ) else: weekend_shift_balancing = 0 preference_constant = 10 # Preferences (not to work) preferences = sum( self.model.pref_not_to_work[worker.id, week, day] * self.model.works[worker.id, week, day, shift] * preference_constant for worker in self.workers for week, day, shift in self.get_all_shiftname_combinations() ) # Work requsets work_request_constant = 10000 # Needs to be very big to override loss from bank holiday requests work_requests = sum( self.model.work_requests[worker.id, week, day, shift] * self.model.works[worker.id, week, day, shift] * work_request_constant for worker in self.workers for week, day, shift in self.get_all_shiftname_combinations() ) # # Spread nights if self.constraint_options["balance_nights_across_sites"]: nights_site_balancing = sum( ( self.model.night_per_site_t1[week, shift.name, site] + self.model.night_per_site_t2[week, shift.name, site] ) * 20 # self.model.night_per_site2[week, block, site] for week in self.weeks for shift in self.get_shifts_with_constraint(NightConstraint) for site in self.sites ) else: nights_site_balancing = 0 if self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint): block_site_balancing = sum( ( self.model.shift_per_site_t1[week, shift.name, site] + self.model.shift_per_site_t2[week, shift.name, site] ) * 2000 # self.model.night_per_site2[week, block, site] for week in self.weeks for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint) for site in self.sites ) else: block_site_balancing = 0 if self.constraint_options["balance_blocks"]: blocks_balancing = sum( block_shift_balancing_constant * self.model.blocks_assigned[week, shift] for week in self.weeks for shift in self.shifts_to_assign_as_blocks() ) 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] 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 worker.site in shift.sites ) # Quadratic :( # worker_pairs_balancing = 0 # worker_pairs_constant = 100000 # if self.worker_pairs: # for worker_a, worker_b in self.worker_pairs: # print(worker_a, worker_b) # for week, day in self.get_week_day_combinations(): # worker_pairs_balancing = worker_pairs_balancing + sum(self.model.works[worker_a.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * sum(self.model.works[worker_b.id, week, day, shift] for shift in self.get_shift_names_by_week_day(week, day)) * worker_pairs_constant return ( weekend_shift_balancing + quadratic_shift_balancing + shift_balancing + worker_shift_balancing + night_shift_balancing + shift_diff_balancing + preferences + nights_site_balancing + block_site_balancing + bank_holiday_balancing + blocks_balancing - work_requests + locum_shift_balancing + true_quadratic_shift_balancing - prefer_multi_shift_expr + prefer_block_expr ) # add objective function to the model. rule (pass function) or expr (pass expression directly) self.model.obj = Objective(rule=obj_rule, sense=minimize) # def allow_shifts_together(self, *shift_names: str): # """ # Allow the given set of shift names to be assigned together on a single day. # """ # self.allowed_multi_shift_sets.append(set(shift_names)) def add_warning(self, warning_type: str, message: str): print(f"[bold red]WARNING:[/bold red] {warning_type} - {message}") self.warnings.append((warning_type, message)) if warning_type in self.terminate_on_warning: raise WarningTermination(warning_type) def get_warnings(self, warning_type: None | str = None): if warning_type is None: return self.warnings else: return [warning for warning in self.warnings if warning[0] == warning_type] def add_worker(self, worker: Worker) -> None: """Add a worker to the rota Args: worker (Worker): """ self.workers.append(worker) def add_workers(self, workers: List) -> None: """Add multiple workers to the rota Args: workers (List(Worker)): """ for worker in workers: self.add_worker(worker) # self.workers.extend(workers) def build_workers(self) -> None: """Process loaded shifts and workers Must be called prior to attempting to solve """ if not self.workers: raise NoWorkers("Workers must be added prior to calling build_workers") self.workers_id_map: dict[str, Worker] = {} self.workers_name_map: dict[str, Worker] = {} for worker in track(self.workers, description="Building workers"): print(worker.name) worker.load_rota(self) wid = worker.id if wid in self.workers_id_map: message = f"Worker with id '{wid}' has been added twice" self.add_warning("Worker/duplicate id", message) self.workers_id_map[wid] = worker if worker.name in self.workers_name_map: message = f"Worker with name '{worker.name}' has been added twice" self.add_warning("Worker/duplicate name", message) self.workers_name_map[worker.name] = worker if worker.site not in self.sites: message = f"Worker with name '{worker.name}' ({worker.id}) has no valid shifts (site: {worker.site})" logger.warning(message) self.add_warning("Worker/no valid shifts", message) for date, shift_name in worker.forced_assignments_by_date: try: week, day = self.get_week_day_by_date(date) worker.force_assign_shift(week, day, shift_name) except WeekDayNotFound: self.add_warning( "Worker/forced assignment date not found", f"Worker {worker.name} ({worker.id}) has a forced assignment for {date} ({shift_name}) but this date is not in the rota", ) continue self.workers = sorted(self.workers) self.full_time_equivalent = sum( w.fte_adj for w in self.workers ) # TODO: rewrite as per shift self.full_time_equivalent_sites = {} self.workers_at_sites = {} for site in track(self.sites, description="Generating site workers"): self.workers_at_sites[site] = [w for w in self.workers if w.site == site] self.full_time_equivalent_sites[site] = {} for shift in self.get_shifts(): self.full_time_equivalent_sites[site][shift.name] = sum( w.get_fte(shift=shift.name) for w in self.workers if w.site == site ) self.worker_pairs = [] pairs = defaultdict(list) for w in self.workers: if w.pair is not None: pairs[w.pair].append(w) for p in pairs: self.worker_pairs.append(tuple(pairs[p])) def add_worker_name_constraint_by_week( self, names: Iterable[str], weeks: Iterable[int], shifts ): self.constraint_options["avoid_shifts_by_worker_names"].append( {"names": names, "weeks": weeks, "shifts": shifts} ) def add_grade_constraint_by_week( self, grades: Iterable[int], weeks: Iterable[int], shifts ): self.constraint_options["avoid_shifts_by_grades"].append( {"grades": grades, "weeks": weeks, "shifts": shifts} ) def add_min_summed_grade_by_shifts_per_day_constraint( self, shifts: list[str], min_grade_sum: int, weeks: list[int] = None, days: list[str] = None, ): """ Helper to add a minimum summed grade constraint for a set of shifts per day. """ constraint = MinSummedGradeByShiftsPerDayConstraint( shifts=shifts, min_grade_sum=min_grade_sum, weeks=weeks, days=days, ) self.constraint_options["min_summed_grade_by_shifts_per_day"].append(constraint) def add_shift(self, shift: SingleShift) -> None: """Add a shift to the collection :param SingleShift shift: Shift object """ self.shifts.append(shift) def add_shifts(self, *shifts: SingleShift) -> None: """Add multiple shifts Returns: None """ self.shifts.extend(shifts) def clear_shifts(self) -> None: self.shifts = [] def get_shift_names_by_week_day(self, week, day: DayStr, return_empty_if_week_day_not_found:bool=False) -> set: """Returns the shifts required for a specific day Returns: set: Set of ShiftName """ if (week, day) not in self.week_day_shifts_dict: if return_empty_if_week_day_not_found: return set() raise WeekDayNotFound(f"Week {week}, Day {day} not found in week_day_shifts_dict") return self.week_day_shifts_dict[(week, day)] def get_week_day_shift_names_by_week_day(self, week, day: DayStr) -> set: """Returns the shifts required for a specific day Returns: set: Set of ShiftName """ return set( (week, day, shift_name) for shift_name in self.week_day_shifts_dict[(week, day)] ) def get_shift_by_name(self, name: str) -> SingleShift: """ :param name: """ return self.shifts_by_name[name] def get_shift_length_by_name(self, name): """Get the length of a shift Args: name (str): Name of a shift Returns: int: Length of the shift """ return self.shifts_by_name[name].length def get_date_by_week_day(self, week: WeekInt, day: DayStr) -> datetime.date: try: return self.week_day_date_map[(week, day)] except KeyError: self.add_warning( "Rota/week day date map", f"Week {week}, Day {day} not found in week_day_date_map", ) return datetime.date(1900, 1, 1) def get_weeks_by_date_range(self, start_date: datetime.date | None, end_date: datetime.date | None) -> list[WeekInt]: """Get a list of weeks that fall within a date range. Args: start_date (datetime.date): The start date of the range. end_date (datetime.date): The end date of the range. Returns: list[WeekInt]: A list of week numbers that fall within the date range. """ # TODO: think about error if not monday start date? weeks = [] if start_date is None: start_date = self.start_date if end_date is None: end_date = self.rota_end_date for week in self.weeks: week_start = self.get_date_by_week_day(week, "Mon") week_end = self.get_date_by_week_day(week, "Sun") if week_start >= start_date and week_end <= end_date: weeks.append(week) return weeks def pair_shifts(self, *shifts: SingleShift): # NOTE: currently designed to pair two shifts # this could be extended... if len(shifts) < 2: raise ValueError(f"Must pair at least two shifts: {shifts}") self.paired_shifts.append([*shifts]) def get_date_range( self, start_date: datetime.date = None, end_date: datetime.date = None ): """Gets a range of dates If either start_date or end_date are not provided defaults to the rota dates""" if start_date is None: start_date = self.start_date if end_date is None: end_date = self.rota_end_date for n in range(int((end_date - start_date).days) + 1): yield start_date + datetime.timedelta(n) def allow_shifts_together_for_all_workers(self, *shift_names: str): """ Allow the given set of shift names to be assigned together on a single day for all workers. """ for worker in self.workers: worker.allow_shifts_together(*shift_names) def build_shifts(self): """ Process the added shifts """ self.shifts_by_name = {} self.shift_names = [] # type: List[ShiftName] # self.week_day_shifts_dict = {(week, day): set() for week in self.weeks for day in days} self.week_day_shifts_dict = {} self.sites = set() self.week_day_shift_product = [] self.week_day_shiftclass_product = [] for s in self.shifts: pass for s in self.shifts: if s.name in self.shift_names: raise InvalidShift(f"Duplicate shift: {s.name}") if s.start_date is not None: if s.start_date > self.rota_end_date or s.start_date < self.start_date: self.add_warning( "Shift/invalid start date", f"Shift '{s.name}' has a start date outside of the rota limits ({s.start_date} vs [{self.start_date}---{self.rota_end_date}])", ) else: s.start_date = self.start_date if s.end_date is not None: if s.end_date > self.rota_end_date or s.end_date < self.start_date: self.add_warning( "Shift/invalid end date", f"Shift '{s.name}' has an end date outside of the rota limits ({s.end_date} vs [{self.start_date}---{self.rota_end_date}])", ) else: s.end_date = self.rota_end_date if hasattr(s, "force_assign_with") and s.force_assign_with: missing = set(s.force_assign_with) - set(shift.name for shift in self.shifts) if missing: raise InvalidShift( f"Shift '{s.name}' has force_assign_with referencing non-existent shift(s): {missing}" ) all_shift_names = set(shift.name for shift in self.shifts) for worker in self.workers: if hasattr(worker, "force_assign_with"): for shift_name, with_shifts in getattr(worker, "force_assign_with", {}).items(): missing = set([shift_name] + list(with_shifts)) - all_shift_names if missing: raise InvalidShift( f"Worker '{worker.name}' has force_assign_with referencing non-existent shift(s): {missing}" ) # Check worker requirements if not isinstance(s.workers_required, int): # Validate the worker requirements for worker_requirement in s.workers_required: if not isinstance(worker_requirement, WorkerRequirement): raise InvalidShift( f"Invalid worker requirement: {worker_requirement}" ) if worker_requirement.start_date is None: worker_requirement.start_date = s.start_date elif worker_requirement.start_date > s.end_date: self.add_warning( "Shift/worker requirement start date", f"Worker requirement start date is after shift end date: {worker_requirement.start_date} > {s.end_date}", ) if worker_requirement.end_date is None: worker_requirement.end_date = s.end_date elif worker_requirement.end_date < s.start_date: self.add_warning( "Shift/worker requirement end date", f"Worker requirement end date is before shift start date: {worker_requirement.end_date} < {s.start_date}", ) self.shifts_by_name[s.name] = s self.shift_names.append(s.name) # for day in s.days: # self.week_day_shifts_dict[(week, day)].add(s.name) for site in list(s.sites): self.sites.add(site) for constraint in s.constraints: if constraint.weeks is None: constraint.weeks = self.get_weeks_by_date_range(constraint.start_date, constraint.end_date) self.shift_counts = defaultdict(int) self.shift_worker_counts = defaultdict(int) for week, day in self.weeks_days_product: self.week_day_shifts_dict[(week, day)] = set() current_date = self.get_date_by_week_day(week, day) for s in self.shifts: # Check if the shift has started or ended if current_date < s.start_date or current_date > s.end_date: continue # Shift not required if there are no worker requirements workers_required = s.get_worker_requirement_by_date(current_date) if not workers_required: self.add_warning( "Shift/worker requirement is zero", f"Shift '{s.name}' has no worker requirements on {current_date}", ) continue if s.bank_holidays_only: if current_date in self.bank_holidays: self.week_day_shift_product.append((week, day, s.name)) self.week_day_shiftclass_product.append((week, day, s)) self.week_day_shifts_dict[(week, day)].add(s.name) self.shift_counts[s.name] = self.shift_counts[s.name] + 1 self.shift_worker_counts[s.name] = ( self.shift_worker_counts[s.name] + workers_required ) else: if day in s.days: self.week_day_shift_product.append((week, day, s.name)) self.week_day_shiftclass_product.append((week, day, s)) self.week_day_shifts_dict[(week, day)].add(s.name) self.shift_counts[s.name] = self.shift_counts[s.name] + 1 self.shift_worker_counts[s.name] = ( self.shift_worker_counts[s.name] + workers_required ) self.max_pre = 1 self.max_post = 1 for shift in self.get_shifts_with_constraint(PreShiftConstraint): self.max_pre = max(shift.get_constraint(PreShiftConstraint).days, self.max_pre) for shift in self.get_shifts_with_constraint(PostShiftConstraint): self.max_post = max(shift.get_constraint(PostShiftConstraint).days, self.max_post) # --- Check allowed_multi_shift_sets validity --- all_shift_names = set(s.name for s in self.shifts) for worker in self.workers: for allowed_set in getattr(worker, "allowed_multi_shift_sets", []): missing = set(allowed_set) - all_shift_names if missing: raise InvalidShift( f"Worker '{worker.name}' has allowed_multi_shift_set(s) with non-existent shift(s): {missing}" ) ## For paired shifts we create proxy shifts # self.paired_shift_map = {} # for shifts in self.paired_shifts: # sites = shifts[0].sites # for shift in shifts: # self.paired_shift_map[shift] = shifts # # if shift.sites != sites: # # This will lead to odd/broken behaviour # self.add_warning( # "Paired shifts/site mismatch", # f"Shift '{shift.name}' has different sites to paired shifts", # ) # # Todo replace with week_day..... # self.day_shift_product = [] # self.day_shiftclass_product = [] # for s in self.shifts: # for day in days: # if day in s.days: # self.day_shift_product.append((day, s.name)) # self.day_shiftclass_product.append((day, s)) # # def get_day_shiftname_combinations(self): # """Returns a list of all required day / shift combinations # # Returns: # list: list of in the following format (str->day, str->shift) # """ # return self.day_shift_product # # def get_day_shiftclass_combinations(self) -> List[Tuple[DayStr, SingleShift]]: # """Returns a list of all required day / shift combinations # # Returns: # list: list of in the following format (str->day, shiftObject->shift) # """ # return self.day_shiftclass_product def get_all_shiftname_combinations( self, week: int = None, day: int = None ) -> List[Tuple[WeekInt, DayStr, ShiftName]]: """Returns a list of all possible week / day / shift combinations Returns: list: contains tuple of all possible week / day / shift combinations """ week_day_shifts = self.week_day_shift_product if week is not None: week_day_shifts = [t for t in week_day_shifts if t[0] == week] if day is not None: week_day_shifts = [t for t in week_day_shifts if t[1] == day] return week_day_shifts def get_all_shiftclass_combinations( self, week: int = None, day: int = None ) -> List[Tuple[WeekInt, DayStr, SingleShift]]: """Returns a list of all possible week / day / shift combinations Returns: list: contains tuple of all possible week / day / shift combinations """ week_day_shifts = self.week_day_shiftclass_product if week is not None: week_day_shifts = [t for t in week_day_shifts if t[0] == week] if day is not None: week_day_shifts = [t for t in week_day_shifts if t[1] == day] return week_day_shifts def get_week_day_combinations(self) -> list: """Returns a list of all week / day tuple combinations Returns: list: list of possible week day combinations [(1, "Mon"), (1, "Tue"), ... (n, "Fri")] """ return self.weeks_days_product def get_week_day_by_date(self, date: datetime.date) -> tuple[int, DayStr]: """Returns the week and day for a given date Args: date (datetime.date): Date to get week and day for Returns: tuple[int, DayStr]: Tuple containing week number and day string """ try: week, day = self.date_week_day_map[date] return week, day except KeyError: raise WeekDayNotFound(f"No week/day found for date {date}") def get_week_day_combinations_for_shift(self, shift: SingleShift) -> list: return [ (week, day) for week, day in self.get_week_day_combinations() if shift.name in self.get_shift_names_by_week_day(week, day) ] def get_week_day_shift_combinations_for_constraint(self, constraint) -> list: constraint_shifts = self.get_shifts_with_constraint(constraint) week_day_shifts = [] for tup in self.get_all_shiftclass_combinations(): if tup[2] in constraint_shifts: week_day_shifts.append(tup) return week_day_shifts def get_shifts(self, week: int | None = None) -> List[SingleShift]: """Returns a list of all the registered shifts Returns: List[SingleShift]: list of registered shifts (as SingleShift) """ if week is not None: date = self.get_week_start_date(week) shifts = [] for shift in self.shifts: if shift.start_date <= date < shift.end_date: if shift.get_worker_requirement_by_date(date): shifts.append(shift) return shifts else: return self.shifts def get_shifts_for_worker_site(self, worker_site): shifts = [shift for shift in self.shifts if worker_site in shift.sites] return shifts def get_shifts_for_worker(self, worker_id): worker = self.get_worker_by_id(worker_id) shifts = [shift for shift in self.shifts if worker.site in shift.sites] return shifts def get_shifts_with_constraint(self, constraint) -> List[SingleShift]: return [shift for shift in self.shifts if shift.has_constraint(constraint)] def get_shifts_with_constraints( self, *constraints, week: int | None = None ) -> List[SingleShift]: shift_names = set() for constraint in constraints: shift_names.update( [ shift.name for shift in self.get_shifts(week) if shift.has_constraint(constraint) ] ) return [self.get_shift_by_name(s) for s in shift_names] def get_shift_names(self) -> List[ShiftName]: ( """Returns a list of all the registered shift names Returns: List[ShiftName]: List of names of all available shifts """ """ """ ) return self.shift_names def get_week_block_iterator(self, block_length: int): """Gets a two dimensional list of week blocks in specified length e.g. block_length = 4 [ [ 1, 2, 3, 4], [ 2, 3, 4, 5], ... ] Args: block_length (int): length of blocks to create Returns: list: two dimensional list containing weeks in blocks """ blocks = [] for i in range(len(self.weeks) - block_length + 1): block = self.weeks[i : i + block_length] blocks.append(block) return blocks def get_required_workers_and_site_combinations(self): """Returns a list of all possible shifts, the workers required and site Returns: list: list (week, day, shift name, workers required, shift.sites) """ l = [] for week, day, shift in self.week_day_shiftclass_product: workers_required = shift.get_worker_requirement_by_date( self.get_date_by_week_day(week, day) ) # if day in shift.days: l.append((week, day, shift.name, workers_required, shift.sites)) return l def get_not_required_shifts(self): """Returns a set of all possible shifts combinations that are not required Includes those on days Returns: list: (week, day, shift) """ l = set() for week in self.weeks: for day in days: for shift in self.shifts: l.add((week, day, shift)) return l - set(self.get_all_shiftclass_combinations()) def shifts_to_assign_as_blocks(self) -> List[str]: return [shift.name for shift in self.get_shifts() if shift.assign_as_block] def shifts_to_force_as_blocks(self) -> List[str]: return [ shift.name for shift in self.get_shifts() if (shift.force_as_block or shift.force_as_block_unless_nwd) ] # def shifts_to_assign_or_force_as_blocks(self, week: int | None = None) -> List[str]: def shifts_to_assign_or_force_as_blocks(self) -> List[str]: s = self.shifts_to_assign_as_blocks() s.extend(self.shifts_to_force_as_blocks()) return s def worker_week_shifts_to_assign_as_blocks(self) -> List[Tuple[Worker, int, str]]: """Returns a list of tuples containing worker id, week and shift name for shifts to assign as blocks""" return [ (worker, week, shift.name) for worker in self.workers for week in self.weeks for shift in self.get_shifts() if ( shift.assign_as_block or shift.force_as_block or shift.force_as_block_unless_nwd or ( hasattr(worker, "assign_as_block_preferences") and shift.name in getattr(worker, "assign_as_block_preferences", {}) and getattr(worker, "assign_as_block_preferences")[shift.name] != 0 ) ) ] def week_shifts_to_assign_as_blocks(self) -> List[Tuple[int, str]]: """Returns a list of tuples containing week and shift name for shifts to assign as blocks""" return set([ (week, shift.name) for worker in self.workers for week in self.weeks for shift in self.get_shifts() if ( shift.assign_as_block or shift.force_as_block or shift.force_as_block_unless_nwd or ( hasattr(worker, "assign_as_block_preferences") and shift.name in getattr(worker, "assign_as_block_preferences", {}) and getattr(worker, "assign_as_block_preferences")[shift.name] != 0 ) ) ]) def get_all_locum_availability(self): return self.locum_availability_map def get_locum_availability(self, worker: Worker, week: int, day: DayStr, shift): if (worker.id, week, day) in self.locum_availability_map: if shift.name in self.locum_availability_map[(worker.id, week, day)]: return 1 return 0 def get_all_workers(self) -> List[Worker]: return self.workers def get_worker_by_id(self, worker_id: str) -> Worker: return self.workers_id_map[worker_id] def get_worker_by_name(self, name: str) -> Worker: return self.workers_name_map[name] def get_workers_by_group(self) -> dict[str, list[Worker]]: group_workers: dict[str, list[Worker]] = defaultdict(list) for worker in self.workers: group_workers[worker.site].append(worker) return group_workers def get_workers_by_remote_group(self) -> dict[str, list[Worker]]: group_workers: dict[str, list[Worker]] = defaultdict(list) for worker in self.workers: group_workers[worker.remote_site].append(worker) return group_workers def get_workers_by_site(self, site: str) -> list[Worker]: """ Returns a list of workers assigned to a specific site. Args: site (str): The site name. Returns: list[Worker]: List of workers at the given site. """ return [worker for worker in self.workers if worker.site == site] def get_workers_by_grade(self) -> Dict[int, list[Worker]]: group_workers: dict[int, list[Worker]] = defaultdict(list) for worker in self.workers: group_workers[worker.grade].append(worker) return group_workers def get_workers_for_shift(self, shift: SingleShift) -> List[Worker]: return [worker for worker in self.workers if worker.site in shift.sites] def get_workers_total_fte(self) -> float: """Does not take into account shift adjusted ftes""" return self.full_time_equivalent def get_worker_details(self): w: dict[str, Worker] = defaultdict(list) worker: Worker for worker in self.workers: w[worker.site].append(worker) l = [] for site in w: l.append(f"{site}: {sum(worker.get_fte() for worker in w[site]) / 100}") t = "\n".join(l) return f"Full time equivalent trainees by site:\n{t}" def get_workers(self): return self.workers def get_worker_grades(self) -> set[int]: return set([worker.grade for worker in self.workers]) def get_workers_who_require_locums(self): return [worker for worker in self.get_all_workers() if worker.locum] def get_locum_workers(self): return [ worker for worker in self.get_all_workers() if not worker.locum and worker.locum_availability ] def get_bank_holiday_week_days(self): return [ (week, day) for week, day in self.get_week_day_combinations() if self.week_day_date_map[(week, day)] in self.bank_holidays ] def get_week_start_date(self, week: int): """Week start date is 1 indexed""" week = week - 1 return self.start_date + datetime.timedelta(weeks=week) def force_assign_shift(self, worker_name: str, week: int, day: str, shift_name: str): """ Force a specific worker to be assigned to a shift on a specific week and day. """ worker = self.get_worker_by_name(worker_name) worker.forced_assignments.append((week, day, shift_name)) def force_assign_shift_by_date(self, worker_name: str, date: datetime.date, shift_name: str): """ Force a specific worker to be assigned to a shift on a specific date. """ worker = self.get_worker_by_name(worker_name) worker.forced_assignments_by_date.append((date, shift_name)) # RESULTS def export_rota_to_html( self, filename: str = "rota", folder=None, timestamp_filename=False ): if timestamp_filename: filename = f"{filename}_{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}" if folder is not None: output_file = Path("output", folder, f"{filename}.html") else: output_file = Path("output", f"{filename}.html") # Make sure path exits output_file.parent.mkdir(parents=True, exist_ok=True) with open(output_file, "w") as f: f.write(self.get_worker_timetable_html(True)) url = urllib.parse.quote(str(output_file.resolve())) print(f"Rota exported [bold blue]('[link url={url}]{url}')[/bold blue]") self.exported_rota_file = output_file def export_rota_to_csv(self, filename: str = "rota"): output_file = Path("output", f"{filename}.csv") works = self.model.works with open(output_file, "w", newline="") as f: wr = csv.writer(f, quoting=csv.QUOTE_ALL) l = ["Name"] l.extend([worker.name for worker in self.workers]) wr.writerow(l) l2 = ["Site"] l2.extend([worker.site for worker in self.workers]) wr.writerow(l2) l3 = ["Grade"] l3.extend([worker.grade for worker in self.workers]) wr.writerow(l3) l4 = ["FTE"] l4.extend([worker.fte for worker in self.workers]) wr.writerow(l4) for week, day in self.get_week_day_combinations(): d = [f"Week {week} Day {day}"] for worker in self.workers: i = "" for shift in self.get_shift_names_by_week_day(week, day): if works[worker.id, week, day, shift].value == 1: i = shift d.append(i) wr.writerow(d) def get_work_table(self): """Build a timetable of the week as a dictionary from the model's optimal solution.""" works = self.model.works week_table = { week: {day: {shift: [] for shift in self.get_shift_names()} for day in days} for week in self.weeks } for week in self.weeks: for worker in self.workers: for day, shift in self.get_day_shiftname_combinations(): if works[worker.id, week, day, shift].value == 1: week_table[week][day][shift].append(worker.get_details()) return week_table def get_shift_timetable_by_week(self): """Build a timetable of the week as a dictionary from the model's optimal solution.""" works = self.model.works shift_timetable = { shift.name: {week: {day: [] for day in days} for week in self.weeks} for shift in self.shifts } for shift in self.shifts: for week in self.weeks: for day in days: assigned_workers = [] for worker in self.workers: if works[worker.id, week, day, shift.name].value > 0.5: assigned_workers.append(worker) shift_timetable[shift.name][week][day] = assigned_workers return shift_timetable def get_worker_timetable(self): works = self.model.works timetable = { worker.name: {week: {day: [] for day in days} for week in self.weeks} for worker in self.workers } for worker in self.workers: for week in self.weeks: for day in days: assigned_shifts = [] for shift in self.get_shift_names_by_week_day(week, day): if works[worker.id, week, day, shift].value > 0.5: assigned_shifts.append(shift) timetable[worker.name][week][day] = assigned_shifts return timetable def get_worker_timetable_brief( self, show_prefs=False, marker_every=30, show_unavailable=False ): model = self.model week_string = f"{'-Week-':20}" + "".join( [7 * str(f"{w}")[-1:] for w in self.weeks] ) days_string = f"{'-Day-':20}" + "".join("MTWTFSS" * len(self.weeks)) timetable = [] for worker in self.workers: shifts = [] w = [f"{worker.name:20}"] for week, day in self.get_week_day_combinations(): shift_display_char = "-" shift_names = [] for shift in self.get_shift_names_by_week_day(week, day): if model.works[worker.id, week, day, shift].value > 0.5: shift_names.append( self.get_shift_by_name(shift).get_display_char() ) if shift_names: shift_display_char = "".join(shift_names) shifts.extend(shift_names) w.append(shift_display_char) shift_count = "" for s in set(shifts): shift_count = shift_count + f"{s}: {shifts.count(s)}, " shift_count = ( shift_count + f"#weekends_worked: {model.worker_weekend_count[worker.id].value}\\#" ) timetable.append(f"{''.join(w)} {shift_count}") if show_prefs: # prefs w = [f"{'Preferences':20}"] for week, day in self.get_week_day_combinations(): if model.pref_not_to_work[worker.id, week, day] > 0: w.append("Y") else: w.append("N") timetable.append("".join(w)) if show_unavailable: # prefs w = [f"{'Unavailable':20}"] for week, day in self.get_week_day_combinations(): if model.available[worker.id, week, day] > 0: w.append("A") else: w.append("U") timetable.append("".join(w)) i = marker_every while i < len(timetable): timetable.insert(i, week_string) timetable.insert(i + 1, days_string) i = i + marker_every + 2 timetable.append(week_string) timetable.insert(0, week_string) timetable.insert(0, days_string) return "\n".join(timetable) def get_worker_timetable_html( self, include_html_tag=False, table_name="rota-table" ): model = self.model 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)

" ) date_row = [""] n = 0 for week, day in self.get_week_day_combinations(): d = self.start_date + datetime.timedelta(n) th_class = "date" if d in self.bank_holidays: th_class = th_class + " bank-holiday" date_row.append( f"Week {week}: {day}" ) n = n + 1 timetable.append(f"{''.join(date_row)}") current_site = "" for worker in self.get_all_workers(): if worker.site != current_site: try: timetable.append( f"{worker.site} n={len(self.workers_at_sites[worker.site])} fte={self.full_time_equivalent_sites[worker.site]}" ) except KeyError as e: print(e) current_site = worker.site shifts = [] locum_shifts = [] nwds = json.dumps(worker.non_working_day_list, default=str) # if worker.nwd: # # TODO: limit to dates # nwds = ", ".join([i[0] for i in worker.nwd]) # else: # nwds = None worker_targets = json.dumps(worker.shift_target_number) tds = [] n = 0 shift_tds = [] for week, day in self.get_week_day_combinations(): d = self.start_date + datetime.timedelta(n) n = n + 1 shift_display_char = "-" shift_names = [] assigned_shift_names = [] locum = False shift_name = "" force_assigned = False for shift in self.get_shift_names_by_week_day(week, day): shift_obj = self.get_shift_by_name(shift) if model.works[worker.id, week, day, shift].value > 0.8: shift_names.append(shift_obj.get_display_char()) assigned_shift_names.append(shift) # Check if this shift was force assigned if (week, day, shift) in getattr(worker, "forced_assignments", []): force_assigned = True elif self.get_workers_who_require_locums(): if model.locum_works[worker.id, week, day, shift].value > 0.8: locum_shifts.append(shift) shift_names.append(shift_obj.get_display_char()) assigned_shift_names.append(shift) locum = True # Check if this shift was force assigned if (week, day, shift) in getattr(worker, "forced_assignments", []): force_assigned = True if shift_names: shift_display_char = "
".join( [ "" + name + "" for name in shift_names ] ) shifts.extend(assigned_shift_names) shift_name = assigned_shift_names[0] if assigned_shift_names else "" # Visual indication for multiple shifts multi_shift = len(shift_names) > 1 css_class = day if multi_shift: css_class += " multi-shift" title = f"{shift_name} ({d})" unavailable_reason = "" if model.available[worker.id, week, day] > 0: available = True else: available = False css_class = "unavailable" try: unavailable_reason = self.unavailable_to_work_reason[ (worker.id, week, day) ] except KeyError: print("Error getting reason") print(worker.id, worker.name) print(f"Week {week}, Day {day}") unavailable_reason = "????" title = f"{title} / {unavailable_reason}" # --- Highlight force assigned shifts --- if force_assigned: css_class += " force-assigned" title = f"{shift_name} ({d}) [FORCE ASSIGNED]" else: title = f"{shift_name} ({d})" bank_holiday = "" if d in self.bank_holidays: title = f"{title} [Bank Holiday]" css_class = " ".join((css_class, "bank-holiday")) bank_holiday = f" data-bank-holiday='{self.bank_holidays[d]}'" requests = "" if (worker.id, week, day) in self.work_requests_map: css_class = " ".join((css_class, "shift-requested")) title = " ".join((title, "[REQUESTED]")) requests = f" data-shift-request='{self.work_requests_map[worker.id, week, day]}'" if (worker.id, week, day) in self.locum_availability_map: css_class = " ".join((css_class, "locum-availability")) title = " ".join((title, "[LOCUM AVAILABILITY]")) requests = f" data-locum-request='{self.locum_availability_map[worker.id, week, day]}'" if locum: css_class = " ".join((css_class, "locum-shift")) remote_site = "" if shift_name: shift = self.get_shift_by_name(shift_name) if shift.has_constraint(RequireRemoteSitePresenceConstraint): remote_site = f" data-shift-remote-site='{shift.get_constraint(RequireRemoteSitePresenceConstraint).site}'" if shift.has_constraint(NightConstraint): css_class = " ".join((css_class, "night-shift")) shift_tds.append( f"{shift_display_char}" ) shift_count = "" shift_count_dict = {} for s in set(shifts): c = shifts.count(s) shift_count_dict[s] = c shift_count = shift_count + f"{s}: {c}, " locum_shift_count = "" locum_shift_count_dict = {} for s in set(locum_shifts): c = locum_shifts.count(s) locum_shift_count_dict[s] = c locum_shift_count = locum_shift_count + f"{s}: {c}, " shift_diff_dict = {} for shift in self.get_shifts(): diff = model.shift_count_diff[worker.id, shift.name].value shift_diff_dict[shift.name] = diff worker_td = """ {name} ({grade}) [{fte}]""".format( site=worker.site, nwds=nwds, worker_id=worker.id, name=worker.name, fte=worker.fte, fte_adj=worker.get_fte(), start_date=worker.calculated_start_date, end_date=worker.calculated_end_date, oops=json.dumps(json.dumps(worker.oop, default=str)), targets=worker_targets, worker_shift_counts=json.dumps(shift_count_dict), locum_shift_counts=json.dumps(locum_shift_count_dict), weekend_target=worker.weekend_shift_target_number, remote_site=worker.remote_site, grade=worker.grade, pair=worker.pair, shift_balance_extra=json.dumps(worker.shift_balance_extra), shift_diff=json.dumps(shift_diff_dict), shift_diff_summed=model.shift_count_diff_summed[worker.id].value, bank_holiday_extra=worker.bank_holiday_extra, ) shift_count = ( shift_count + f"#weekends_worked: {model.worker_weekend_count[worker.id].value}\\#" ) if self.constraint_options["balance_bank_holidays"]: bank_holiday_count = model.bank_holiday_count[worker.id].value bank_holiday_count_w = model.bank_holiday_count_w[worker.id].value - 1 else: bank_holiday_count = -1 bank_holiday_count_w = -1 # print(worker.name, bank_holiday_count, bank_holiday_count_w) timetable.append( f"{worker_td}{''.join(shift_tds)}" ) # timetable.append("{}".format("".join(w), shift_count)) # if show_prefs: # # prefs # w = ["{:20}".format("Preferences")] # for week, day in self.get_week_day_combinations(): # if model.pref_not_to_work[worker.id, week, day] > 0: # w.append("Y") # else: # w.append("N") # timetable.append("".join(w)) # if show_unavailable: # # prefs # w = ["{:20}".format("Unavailable")] # for week, day in self.get_week_day_combinations(): # if model.available[worker.id, week, day] > 0: # w.append("A") # else: # w.append("U") # timetable.append("".join(w)) result_stream = StringIO() if self.results is not None: self.results.write(ostream=result_stream) self.results.write_json(ostream=result_stream) result_string = result_stream.getvalue() else: result_string = "Rota not run" joined_timetable = "\n".join(timetable) hard_constrain_shift_info = [] for worker in self.workers: for shift_name, limits in getattr(worker, "hard_constrain_shift_limits", {}).items(): hard_constrain_shift_info.append({ "worker": worker.name, "shift": shift_name, **limits, }) # Add warnings export warnings_html = "

Warnings

\n"
        for warning_type, message in self.warnings:
            warnings_html += f"{warning_type}: {message}\n"
        warnings_html += "
" # Human-readable worker details export worker_details_human = [] for worker in self.get_all_workers(): details = [ f"Name: {worker.name}", f"ID: {worker.id}", f"Site: {worker.site}", f"Grade: {worker.grade}", f"FTE: {worker.fte}", f"Remote site: {getattr(worker, 'remote_site', '')}", f"Pair: {getattr(worker, 'pair', '')}", f"Start date: {getattr(worker, 'calculated_start_date', '')}", f"End date: {getattr(worker, 'calculated_end_date', '')}", f"Bank holiday extra: {getattr(worker, 'bank_holiday_extra', '')}", f"Shift balance extra: {getattr(worker, 'shift_balance_extra', '')}", f"Non-working days: {getattr(worker, 'non_working_day_list', '')}", f"OOP: {getattr(worker, 'oop', '')}", f"Unavailable: {''}", f"Shift targets: {json.dumps(getattr(worker, 'shift_target_number', {}))}", ] # Build unavailability list for this worker unavail_list = [] for entry in sorted(self.unavailable_to_work): wid, wk, d = entry if wid != worker.id: continue date = self.get_date_by_week_day(wk, d) reason = self.unavailable_to_work_reason.get((wid, wk, d), "") if reason: unavail_list.append(f"{date.isoformat()} ({d}) [week {wk}]: {reason}") else: # date may be a string if fallback used if isinstance(date, str): unavail_list.append(f"{date} ({d}) [week {wk}]") else: unavail_list.append(f"{date.isoformat()} ({d}) [week {wk}]") # replace placeholder Unavailable entry details = [line for line in details if not line.startswith("Unavailable:")] details.insert( -1, # before Shift targets f"Unavailable: {', '.join(unavail_list) if unavail_list else ''}", ) worker_details_human.append("\n".join(details)) worker_details_human_str = "\n\n".join(worker_details_human) # Collect force assigned shifts for export force_assigned_info = [] for worker in self.get_all_workers(): for assignment in getattr(worker, "forced_assignments", []): # Support both tuple and pydantic model if hasattr(assignment, "week") and hasattr(assignment, "day") and hasattr(assignment, "shift_name"): week = assignment.week day = assignment.day shift_name = assignment.shift_name else: week, day, shift_name = assignment date = self.get_date_by_week_day(week, day) force_assigned_info.append({ "worker": worker.name, "week": week, "day": day, "date": str(date), "shift": shift_name, }) html = f"""
{joined_timetable}
{warnings_html}

Rota settings

 
        {json.dumps({k: _pydantic_to_dict(v) for k, v in self.constraint_options.items()}, indent=4)}

Shifts settings
{"
".join(str(i) for i in self.shifts)}

Worker details

 
        {worker_details_human_str}

Output

 
        {result_string}

Export table

Shift timetables

Extra

Locum availability

{"TEST"} {self.locum_availability_map}
Force assigned shifts 
        {json.dumps(force_assigned_info, indent=4)}

Hard Constrain Shift Min/Max

 
     {json.dumps(hard_constrain_shift_info, indent=4)}
     "
""" if include_html_tag: html = f""" {html}""" return html def get_shift_summary_dict(self): """""" timetable = { worker.name: {shift: "" for shift in self.get_shift_names()} for worker in self.workers } for worker in self.workers: for shift in self.get_shift_names(): timetable[worker.name][shift] = round( self.model.shift_count[worker.id, shift].value ) return timetable def get_worker_shifts_by_date(self, worker: Worker) -> Dict[datetime.date, str]: shifts = {} n = 0 for week, day in self.get_week_day_combinations(): d = self.start_date + datetime.timedelta(n) n = n + 1 temp = "" for shift in self.get_shift_names_by_week_day(week, day): if self.model.works[worker.id, week, day, shift].value > 0.5: shifts[d] = shift return shifts def get_workers_total_shifts(self) -> Dict[str, int]: shifts = {} for worker in self.workers: shifts[worker.name] = len( [i for i in self.get_worker_shift_list(worker) if i != ""] ) return shifts def get_shift_worker_list(self, shift_name: str) -> List[set[Worker]]: """ Returns a list of workers assigned to a specific shift. Args: shift_name (str): The name of the shift. Returns: List[Worker]: List of workers assigned to the specified shift. """ shift_list = [] for week, day in self.get_week_day_combinations(): workers: set = set() for worker in self.workers: if self.model.works[worker.id, week, day, shift_name].value > 0.5: workers.add(worker) shift_list.append(workers) return shift_list def get_worker_shift_list(self, worker: Worker, include_locums=False, search_multiple_assignments=False) -> List: """ This function returns a list of shifts assigned to a worker for each week and day. It was originally designed to return a list of shift names (when only a single shift per day was supported), but it can also return a set of shift names if `search_multiple_assignments` is True. """ shifts = [] for week, day in self.get_week_day_combinations(): # d = self.start_date + datetime.timedelta(n) # n = n + 1 if search_multiple_assignments: temp = set() else: temp = "" for shift in self.get_shift_names_by_week_day(week, day): if self.model.works[worker.id, week, day, shift].value > 0.90: if search_multiple_assignments: temp.add(shift) else: temp = shift elif ( include_locums and self.model.locum_works[worker.id, week, day, shift].value > 0.90 ): if search_multiple_assignments: temp.add(shift) else: temp = shift shifts.append(temp) return shifts def get_worker_shift_count(self, worker: Worker, shift_name: str) -> int: """Returns the number of shifts of a specific type assigned to a worker.""" count = 0 for week, day in self.get_week_day_combinations(): if shift_name in self.get_shift_names_by_week_day(week, day): if self.model.works[worker.id, week, day, shift_name].value > 0.5: count += 1 return count def get_worker_days_worked(self, worker: Worker) -> int: """ Returns the number of unique days on which the worker is assigned at least one shift. """ days_worked = 0 for week, day in self.get_week_day_combinations(): for shift in self.get_shift_names_by_week_day(week, day): if self.model.works[worker.id, week, day, shift].value > 0.5: days_worked += 1 break # Only count each day once, even if multiple shifts return days_worked def get_worker_shift_list_string(self, worker: Worker) -> str: shifts = self.get_worker_shift_list(worker) # Convert shift to a string representation return "".join( [ self.get_shift_by_name(i).get_display_char() if i != "" else "-" for i in shifts ] ) def get_shift_summary(self): works = self.model.works timetable = { worker.get_details(): {shift: "" for shift in self.get_shift_names()} for worker in self.workers } # timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers } t = [] for worker in self.workers: l = [] total_shifts = 0 for shift in self.get_shift_names(): # for shift in ["truro_twilight"]: c = [ works[worker.id, week, day, shift].value for week in self.weeks for day in days ].count(1) if c > 0: l.append(f"{shift} ({c})") total_shifts = total_shifts + c timetable[worker.get_details()][shift] = c t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}") return "\n".join(t) def get_shift_summary_html(self): works = self.model.works timetable = { worker.get_details(): { shift.get_display_char(): "" for shift in self.get_shifts() } for worker in self.workers } # timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers } t = [] for worker in self.workers: l = [] total_shifts = 0 for shift in self.get_shift_names(): # for shift in ["truro_twilight"]: c = [ works[worker.id, week, day, shift] for week in self.weeks for day in days ].count(1) if c > 0: l.append(f"{shift} ({c})") total_shifts = total_shifts + c shift_obj = self.get_shift_by_name(shift) timetable[worker.get_details()][shift_obj.get_display_char()] = c t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}") return "\n".join(t) class NoActiveSites(Exception): """Raised when there are no active sites""" pass class NoWorkers(Exception): """Raised when there are no active sites""" pass class InvalidShift(Exception): """Raised when there are no active sites""" pass class WarningTermination(Exception): """Raised when a warning in the termination group is raised""" pass class WeekDayNotFound(Exception): """Raised when a week/day combination is not found""" pass