feat(core): Add unwanted shift type successions preference

The `ortools_expression_to_bool_var` utility function is directly copied from the 2023/08/20 POC.

core/nurse_scheduling/context.py

@@ -12,6 +12,7 @@ def __init__(self) -> None:
         self.n_days = None
         self.n_requirements = None
         self.n_people = None
+        self.map_rid_r = None
         self.model: cp_model.CpModel = None
         self.model_vars = None
         self.shifts = None

core/nurse_scheduling/preference_types.py

@@ -2,6 +2,7 @@
 from .context import Context
 from .report import Report
 
+# Leave most parsing to the caller, keep the function here simple.
 
 def all_requirements_fulfilled(ctx: Context, preference, preference_idx):
     # Hard constraint
@@ -34,7 +35,7 @@ def assign_shifts_evenly(ctx: Context, preference, preference_idx):
         target_n_shifts = round(ctx.n_days * sum(requirement.required_people for requirement in ctx.requirements) / ctx.n_people)
         unique_var_prefix = f"pref_{preference_idx}_p_{p}_"
 
-        # Construct: L2 = actual_n_shifts - target_n_shifts) ** 2
+        # Construct: L2 = (actual_n_shifts - target_n_shifts) ** 2
         MAX = max(ctx.n_days - target_n_shifts, target_n_shifts)
         diff_var_name = f"{unique_var_prefix}diff"
         ctx.model_vars[diff_var_name] = diff = ctx.model.NewIntVar(0, MAX, diff_var_name)
@@ -65,9 +66,41 @@ def shift_request(ctx: Context, preference, preference_idx):
         ctx.objective += weight * ctx.shifts[(d, r, p)]
         ctx.reports.append(Report(f"shift_request_p_{p}_d_{d}_r_{r}", ctx.shifts[(d, r, p)], lambda x: x == 1))
 
+def unwanted_shift_type_successions(ctx: Context, preference, preference_idx):
+    # Soft constraint
+    # For all people, for all start date, try to avoid unwanted shift type successions.
+    # Note that a shift is represented as (d, r)
+    # i.e., max(weight * (actual_n_matched == target_n_matched)), for all p,
+    # where actual_n_matched = sum_{(d, r)}(shifts[(d, r, p)]), for all satisfying (d, r)
+    p = preference.person
+    pattern = preference.pattern
+    # TODO: Consider history
+    for d_begin in range(ctx.n_days - len(pattern) + 1):
+        actual_n_matched = 0
+        target_n_matched = len(pattern)
+        for i in range(len(pattern)):
+            d = d_begin + i
+            r = ctx.map_rid_r[pattern[i]]
+            actual_n_matched += ctx.shifts[(d, r, p)]
+
+        # Construct: is_match = (actual_n_matched == target_n_matched)
+        unique_var_prefix = f"pref_{preference_idx}_p_{p}_"
+        is_match_var_name = f"{unique_var_prefix}is_match"
+        ctx.model_vars[is_match_var_name] = is_match = utils.ortools_expression_to_bool_var(
+            ctx.model, is_match_var_name,
+            actual_n_matched == target_n_matched,
+            actual_n_matched != target_n_matched
+        )
+
+        # Add the objective
+        weight = -100
+        ctx.objective += weight * is_match
+        ctx.reports.append(Report(f"unwanted_shift_type_successions_p_{p}", is_match, lambda x: x != target_n_matched))
+
 PREFERENCE_TYPES_TO_FUNC = {
     "all requirements fulfilled": all_requirements_fulfilled,
     "all people work at most one shift per day": all_people_work_at_most_one_shift_per_day,
     "assign shifts evenly": assign_shifts_evenly,
     "shift request": shift_request,
+    "unwanted shift type successions": unwanted_shift_type_successions,
 }

core/nurse_scheduling/scheduler.py

@@ -29,6 +29,11 @@ def schedule(filepath: str, validate=True, deterministic=False):
     ctx.n_requirements = len(ctx.requirements)
     ctx.n_people = len(ctx.people)
     ctx.dates = [ctx.startdate + timedelta(days=d) for d in range(ctx.n_days)]
+    ctx.map_rid_r = {}
+    for r in range(ctx.n_requirements):
+        if ctx.requirements[r].id in ctx.map_rid_r:
+            raise ValueError(f"Duplicated requirement ID: {r.id}")
+        ctx.map_rid_r[ctx.requirements[r].id] = r
 
     logging.info("Initializing solver model...")
     ctx.model = cp_model.CpModel()
@@ -40,6 +45,9 @@ def schedule(filepath: str, validate=True, deterministic=False):
 
     logging.info("Creating shift variables...")
     # Ref: https://developers.google.com/optimization/scheduling/employee_scheduling
+    # In the following code, we always use the convention of (d, r, p)
+    # to represent the index of (day, requirement, person).
+    # The object will not be abbreviated as (d, r, p) to avoid confusion.
     for d in range(ctx.n_days):
         for r in range(ctx.n_requirements):
             # TODO(Optimize): Skip if no people is required in that day

core/nurse_scheduling/utils.py

@@ -2,3 +2,12 @@ def required_n_people(requirement):
     if not isinstance(requirement.required_people, int):
         raise NotImplementedError("required_people with type other than int is not supported yet")
     return requirement.required_people
+
+def ortools_expression_to_bool_var(
+        model, varname, true_expression, false_expression
+    ):
+    # Ref: https://stackoverflow.com/a/70571397
+    var = model.NewBoolVar(varname)
+    model.Add(true_expression).OnlyEnforceIf(var)
+    model.Add(false_expression).OnlyEnforceIf(var.Not())
+    return var

core/tests/test_all.py

@@ -1,15 +1,23 @@
 import glob
+import io
+import logging
 import os
 
 import nurse_scheduling
+import pandas
+import pytest
 
 
 def test_all():
     for filepath in glob.glob("tests/testcases/*.yaml"):
-        print(f"Testing '{filepath}' ...")
         base_filepath = os.path.splitext(filepath)[0]
+        logging.info(f"Testing '{base_filepath}' ...")
         df = nurse_scheduling.schedule(filepath, validate=False, deterministic=True)
-        print(df)
         with open(f"{base_filepath}.csv", 'r') as f:
             expected_csv = f.read()
-        assert df.to_csv(index=False, header=False) == expected_csv
+        actual_csv = df.to_csv(index=False, header=False)
+        if actual_csv != expected_csv:
+            logging.info(f"Actual CSV:\n{actual_csv}")
+            logging.info(f"Actual output:\n{df}")
+            logging.info(f"Expected output:\n{pandas.read_csv(io.StringIO(expected_csv))}")
+            pytest.fail(f"Output mismatch for '{base_filepath}'")

core/tests/testcases/unwanted-pattern_1_len-1.csv

@@ -0,0 +1,6 @@
+,18,19,20,21,22,23,24
+,Fri,Sat,Sun,Mon,Tue,Wed,Thu
+Nurse 0,E,E,E,E,D,D,D
+Nurse 1,N,N,N,N,N,N,N
+Nurse 2,D,D,D,D,,,
+Nurse 3,,,,,E,E,E

core/tests/testcases/unwanted-pattern_1_len-1.yaml

@@ -0,0 +1,25 @@
+apiVersion: alpha
+description: Test unwanted pattern with length 1 constraints
+startdate: 2023-08-18
+enddate: 2023-08-24
+people:
+  - description: Nurse 0
+  - description: Nurse 1
+  - description: Nurse 2
+  - description: Nurse 3
+requirements:
+  - id: D
+    description: Day shift requirement
+    required_people: 1
+  - id: E
+    description: Evening shift requirement
+    required_people: 1
+  - id: N
+    description: Night shift requirement
+    required_people: 1
+preferences:
+  - type: all requirements fulfilled
+  - type: all people work at most one shift per day
+  - type: unwanted shift type successions
+    person: 0
+    pattern: [N]

core/tests/testcases/unwanted-pattern_1_len-2.csv

@@ -0,0 +1,6 @@
+,18,19,20,21,22,23,24
+,Fri,Sat,Sun,Mon,Tue,Wed,Thu
+Nurse 0,N,D,N,E,N,E,N
+Nurse 1,D,N,E,N,D,N,E
+Nurse 2,,,D,D,,D,D
+Nurse 3,E,E,,,E,,

core/tests/testcases/unwanted-pattern_1_len-2.yaml

@@ -0,0 +1,25 @@
+apiVersion: alpha
+description: Test unwanted pattern with length 2 constraints
+startdate: 2023-08-18
+enddate: 2023-08-24
+people:
+  - description: Nurse 0
+  - description: Nurse 1
+  - description: Nurse 2
+  - description: Nurse 3
+requirements:
+  - id: D
+    description: Day shift requirement
+    required_people: 1
+  - id: E
+    description: Evening shift requirement
+    required_people: 1
+  - id: N
+    description: Night shift requirement
+    required_people: 1
+preferences:
+  - type: all requirements fulfilled
+  - type: all people work at most one shift per day
+  - type: unwanted shift type successions
+    person: 0
+    pattern: [N, N]

core/tests/testcases/unwanted-pattern_1_len-3.csv

@@ -0,0 +1,6 @@
+,18,19,20,21,22,23,24
+,Fri,Sat,Sun,Mon,Tue,Wed,Thu
+Nurse 0,N,D,N,N,D,N,N
+Nurse 1,D,N,E,E,N,E,E
+Nurse 2,,,D,D,,D,D
+Nurse 3,E,E,,,E,,

core/tests/testcases/unwanted-pattern_1_len-3.yaml

@@ -0,0 +1,25 @@
+apiVersion: alpha
+description: Test unwanted pattern with length 3 constraints
+startdate: 2023-08-18
+enddate: 2023-08-24
+people:
+  - description: Nurse 0
+  - description: Nurse 1
+  - description: Nurse 2
+  - description: Nurse 3
+requirements:
+  - id: D
+    description: Day shift requirement
+    required_people: 1
+  - id: E
+    description: Evening shift requirement
+    required_people: 1
+  - id: N
+    description: Night shift requirement
+    required_people: 1
+preferences:
+  - type: all requirements fulfilled
+  - type: all people work at most one shift per day
+  - type: unwanted shift type successions
+    person: 0
+    pattern: [N, N, N]

core/tests/testcases/unwanted-pattern_1_without.csv

@@ -0,0 +1,6 @@
+,18,19,20,21,22,23,24
+,Fri,Sat,Sun,Mon,Tue,Wed,Thu
+Nurse 0,N,N,N,N,E,E,E
+Nurse 1,E,E,E,E,D,D,D
+Nurse 2,D,D,D,D,,,
+Nurse 3,,,,,N,N,N

core/tests/testcases/unwanted-pattern_1_without.yaml

@@ -0,0 +1,22 @@
+apiVersion: alpha
+description: Test unwanted pattern without constraints
+startdate: 2023-08-18
+enddate: 2023-08-24
+people:
+  - description: Nurse 0
+  - description: Nurse 1
+  - description: Nurse 2
+  - description: Nurse 3
+requirements:
+  - id: D
+    description: Day shift requirement
+    required_people: 1
+  - id: E
+    description: Evening shift requirement
+    required_people: 1
+  - id: N
+    description: Night shift requirement
+    required_people: 1
+preferences:
+  - type: all requirements fulfilled
+  - type: all people work at most one shift per day