2020from dwave .optimization .symbols import BinaryVariable
2121from dwave .system import LeapHybridCQMSampler , LeapHybridNLSampler
2222
23- from utils import DAYS , FULL_TIME_SHIFTS , SHIFTS , validate_nl_schedule
23+ from src . utils import DAYS , FULL_TIME_SHIFTS , SHIFTS , validate_nl_schedule
2424
2525
2626MSGS = {
4242}
4343
4444
45- def build_cqm (# params: ModelParams
45+ def build_cqm ( # params: ModelParams
4646 availability : dict [str , list [int ]],
4747 shifts : list [str ],
4848 min_shifts : int ,
@@ -62,7 +62,7 @@ def build_cqm(#params: ModelParams
6262 shift_forecast: A list of the number of expected employees needed per shift.
6363 allow_isolated_days_off: Whether on-off-on should be allowed in the schedule.
6464 max_consecutive_shifts: The maximum consectutive shifts to schedule a part-time employee for.
65- num_full_time: The number of full time employees.
65+ num_full_time: The number of full- time employees.
6666
6767 Returns:
6868 cqm: A Constrained Quadratic Model representing the problem.
@@ -118,7 +118,7 @@ def build_cqm(#params: ModelParams
118118 )
119119
120120 for employee in employees_ft :
121- # Schedule employees for at most max_shifts
121+ # Schedule full-time employees for all their shifts
122122 cqm .add_constraint (
123123 quicksum (x [employee , shift ] for shift in shifts ) <= FULL_TIME_SHIFTS ,
124124 label = f"overtime,{ employee } ," ,
@@ -129,7 +129,7 @@ def build_cqm(#params: ModelParams
129129 label = f"insufficient,{ employee } ," ,
130130 )
131131
132- # Every shift needs shift_min and shift_max employees working
132+ # Every shift needs shift_forecast employees working
133133 for i , shift in enumerate (shifts ):
134134 cqm .add_constraint (
135135 sum (x [employee , shift ] for employee in employees ) >= shift_forecast [i ],
@@ -234,16 +234,15 @@ def run_cqm(cqm: ConstrainedQuadraticModel):
234234 return feasible_sampleset , None
235235
236236
237- def build_nl (
237+ def build_nl ( # params: ModelParams
238238 availability : dict [str , list [int ]],
239239 shifts : list [str ],
240240 min_shifts : int ,
241241 max_shifts : int ,
242- shift_min : int ,
243- shift_max : int ,
244- requires_manager : bool ,
242+ shift_forecast : list ,
245243 allow_isolated_days_off : bool ,
246244 max_consecutive_shifts : int ,
245+ num_full_time : int ,
247246) -> tuple [Model , BinaryVariable ]:
248247 """Builds an employee scheduling nonlinear model.
249248
@@ -252,11 +251,10 @@ def build_nl(
252251 shifts (list[str]): Shift labels.
253252 min_shifts (int): Minimum shifts per employee.
254253 max_shifts (int): Maximum shifts per employee.
255- shift_min (int): Minimum employees per shift.
256- shift_max (int): Maximum employees per shift.
257- requires_manager (bool): Whether to require exactly one manager on every shift.
254+ shift_forecast (list[int]): A list of the number of expected employees needed per shift.
258255 allow_isolated_days_off (bool): Whether to allow isolated days off.
259256 max_consecutive_shifts (int): Maximum consecutive shifts per employee.
257+ num_full_time (int): The number of full-time employees.
260258
261259 Returns:
262260 tuple[Model, BinaryVariable]: the NL model and assignments decision variable
@@ -281,8 +279,8 @@ def build_nl(
281279 # Initialize model constants
282280 min_shifts_constant = model .constant (min_shifts )
283281 max_shifts_constant = model .constant (max_shifts )
284- shift_min_constant = model .constant (shift_min )
285- shift_max_constant = model .constant (shift_max )
282+ full_time_shifts_constant = model .constant (FULL_TIME_SHIFTS )
283+ shift_forecast_constant = model .constant (shift_forecast )
286284 max_consecutive_shifts_c = model .constant (max_consecutive_shifts )
287285 one_c = model .constant (1 )
288286
@@ -292,28 +290,32 @@ def build_nl(
292290
293291 # Objective: for infeasible solutions, focus on right number of shifts for employees
294292 target_shifts = model .constant ((min_shifts + max_shifts ) / 2 )
295- shift_difference_list = [
296- (assignments [e , :].sum () - target_shifts ) ** 2 for e in range (num_employees )
293+ shift_difference_list_pt = [
294+ (assignments [e , :].sum () - target_shifts ) ** 2 for e in range (num_full_time , num_employees )
297295 ]
298- obj += add (* shift_difference_list )
296+ shift_difference_list_ft = [
297+ (assignments [e , :].sum () - full_time_shifts_constant ) ** 2 for e in range (num_full_time )
298+ ]
299+ obj += add (* shift_difference_list_pt , * shift_difference_list_ft )
299300
300301 model .minimize (- obj )
301302
302303 # CONSTRAINTS:
303304 # Only schedule employees when they're available
304305 model .add_constraint ((availability_const >= assignments ).all ())
305306
306- for e in range (len (employees )):
307- # Schedule employees for at most max_shifts
308- model .add_constraint (assignments [e , :].sum () <= max_shifts_constant )
307+ # Schedule part-time employees for at most max_shifts
308+ model .add_constraint ((assignments [num_full_time :, :].sum (axis = 1 ) <= max_shifts_constant ).all ())
309309
310- # Schedule employees for at least min_shifts
311- model .add_constraint (assignments [e , :].sum () >= min_shifts_constant )
310+ # Schedule part-time employees for at least min_shifts
311+ model .add_constraint ((assignments [num_full_time :, :].sum (axis = 1 ) >= min_shifts_constant ).all ())
312+
313+ if num_full_time :
314+ # Schedule full-time employees for all their shifts
315+ model .add_constraint ((assignments [:num_full_time , :].sum (axis = 1 ) == full_time_shifts_constant ).all ())
312316
313- # Every shift needs shift_min and shift_max employees working
314- for s in range (num_shifts ):
315- model .add_constraint (assignments [:, s ].sum () <= shift_max_constant )
316- model .add_constraint (assignments [:, s ].sum () >= shift_min_constant )
317+ # Schedule only forecast number of employees per day
318+ model .add_constraint ((assignments .sum (axis = 0 ) == shift_forecast_constant ).all ())
317319
318320 managers_c = model .constant (
319321 [employees .index (e ) for e in employees if e [- 3 :] == "Mgr" ]
@@ -326,7 +328,7 @@ def build_nl(
326328 negthree_c = model .constant (- 3 )
327329 zero_c = model .constant (0 )
328330 # Adding many small constraints greatly improves feasibility
329- for e in range (len ( employees )):
331+ for e in range (num_full_time , num_employees ): # for part-time employees
330332 for s1 in range (len (shifts ) - 2 ):
331333 s2 , s3 = s1 + 1 , s1 + 2
332334 model .add_constraint (
@@ -337,12 +339,11 @@ def build_nl(
337339 <= zero_c
338340 )
339341
340- if requires_manager :
341- for shift in range (len (shifts )):
342- model .add_constraint (assignments [managers_c ][:, shift ].sum () == one_c )
342+ # At least 1 manager per shift
343+ model .add_constraint ((assignments [managers_c ].sum (axis = 0 ) >= one_c ).all ())
343344
344- # Don't exceed max_consecutive_shifts
345- for e in range (num_employees ):
345+ # Don't exceed max_consecutive_shifts for part-time employees
346+ for e in range (num_full_time , num_employees ):
346347 for s in range (num_shifts - max_consecutive_shifts + 1 ):
347348 s_window = s + max_consecutive_shifts + 1
348349 model .add_constraint (
@@ -368,12 +369,11 @@ def run_nl(
368369 shifts : list [str ],
369370 min_shifts : int ,
370371 max_shifts : int ,
371- shift_min : int ,
372- shift_max : int ,
373- requires_manager : bool ,
372+ shift_forecast : list [int ],
374373 allow_isolated_days_off : bool ,
375374 max_consecutive_shifts : int ,
376- time_limit : int | None = None ,
375+ num_full_time : int ,
376+ time_limit : Optional [int ] = None ,
377377 msgs : dict [str , tuple [str , str ]] = MSGS ,
378378) -> Optional [defaultdict [str , list [str ]]]:
379379 """Solves the NL scheduling model and detects any errors.
@@ -395,11 +395,10 @@ def run_nl(
395395 shifts ,
396396 min_shifts ,
397397 max_shifts ,
398- shift_min ,
399- shift_max ,
400- requires_manager ,
398+ shift_forecast ,
401399 allow_isolated_days_off ,
402400 max_consecutive_shifts ,
401+ num_full_time ,
403402 )
404403
405404 # Return errors if any error message list is populated
0 commit comments