feat(core): Support multiple people and weights

j3soon · Aug 10, 2024 · 02e6649 · 02e6649
1 parent 748f379
commit 02e6649
Show file tree

Hide file tree

Showing 9 changed files with 140 additions and 31 deletions.
diff --git a/core/nurse_scheduling/preference_types.py b/core/nurse_scheduling/preference_types.py
@@ -45,6 +45,7 @@ def assign_shifts_evenly(ctx: Context, preference, preference_idx):
         ctx.model.AddMultiplicationEquality(L2, diff, diff)
 
         # Add the objective
+        # TODO: Support custom weight
         weight = -1000000
         ctx.objective += weight * L2
         ctx.reports.append(Report(f"assign_shifts_evenly_L2_p_{p}", L2, lambda x: x == 0))
@@ -54,46 +55,49 @@ def shift_request(ctx: Context, preference, preference_idx):
     # For all people, try to fulfill the shift requests.
     # Note that a shift is represented as (d, r)
     # i.e., max(weight * shifts[(d, r, p)]), for all satisfying (d, r)
-    p = preference.person
-    preference_days = utils.parse_dates(preference.date, ctx.startdate, ctx.enddate)
-    for date in preference_days:
-        d = (date - ctx.startdate).days
-        r = ctx.map_rid_r[preference.shift]
-        # Add the objective
-        weight = 1
-        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))
+    people = utils.ensure_list(preference.person)
+    dates = utils.parse_dates(preference.date, ctx.startdate, ctx.enddate)
+    for p in people:
+        for date in dates:
+            d = (date - ctx.startdate).days
+            r = ctx.map_rid_r[preference.shift]
+            # Add the objective
+            weight = utils.one_or_value(preference.weight)
+            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
+    people = utils.ensure_list(preference.person)
+    # TODO: Check pattern is list
     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)]
+    for p in people:
+        # 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
-        )
+            # 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))
+            # Add the objective
+            weight = utils.neg_one_or_value(preference.weight)
+            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,

diff --git a/core/nurse_scheduling/utils.py b/core/nurse_scheduling/utils.py
@@ -2,6 +2,12 @@
 import re
 
 
+def one_or_value(val):
+    return 1 if val is None else val
+
+def neg_one_or_value(val):
+    return -1 if val is None else val
+
 def ensure_list(val):
     if val is None:
         return []

diff --git a/core/tests/testcases/or-tools-example_2_date-format.yaml b/core/tests/testcases/or-tools-example_2_date-format.yaml
@@ -1,5 +1,5 @@
 apiVersion: alpha
-description: OR-Tools Example 2 with shift requests. From <https://developers.google.com/optimization/scheduling/employee_scheduling>.
+description: OR-Tools Example 2 with shift requests with various date format. From <https://developers.google.com/optimization/scheduling/employee_scheduling>.
 startdate: 2023-09-02
 enddate: 2023-09-08
 people:

diff --git a/core/tests/testcases/shift-request_1_people-2_weight.csv b/core/tests/testcases/shift-request_1_people-2_weight.csv
@@ -0,0 +1,6 @@
+,18,19,20,21,22,23,24
+,Fri,Sat,Sun,Mon,Tue,Wed,Thu
+Nurse 0,,D,D,D,D,D,D
+Nurse 1,E,E,E,E,E,E,E
+Nurse 2,N,N,N,N,N,N,N
+Nurse 3,D,,,,,,
diff --git a/core/tests/testcases/shift-request_1_people-2_weight.yaml b/core/tests/testcases/shift-request_1_people-2_weight.yaml
@@ -0,0 +1,27 @@
+apiVersion: alpha
+description: Test shift request with multiple people and weights
+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: shift request
+    person: [0, 1]
+    date: [18, 19, 20, 21, 22, 23, 24]
+    shift: N
+    weight: -1
diff --git a/core/tests/testcases/unwanted-pattern_1_people-2.csv b/core/tests/testcases/unwanted-pattern_1_people-2.csv
@@ -0,0 +1,6 @@
+,18,19,20,21,22,23,24
+,Fri,Sat,Sun,Mon,Tue,Wed,Thu
+Nurse 0,,D,D,D,D,D,D
+Nurse 1,E,E,E,E,E,E,E
+Nurse 2,N,N,N,N,N,N,N
+Nurse 3,D,,,,,,
diff --git a/core/tests/testcases/unwanted-pattern_1_people-2.yaml b/core/tests/testcases/unwanted-pattern_1_people-2.yaml
@@ -0,0 +1,25 @@
+apiVersion: alpha
+description: Test unwanted pattern with multiple people
+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, 1]
+    pattern: [N]
diff --git a/core/tests/testcases/unwanted-pattern_1_people-2_weight.csv b/core/tests/testcases/unwanted-pattern_1_people-2_weight.csv
@@ -0,0 +1,6 @@
+,18,19,20,21,22,23,24
+,Fri,Sat,Sun,Mon,Tue,Wed,Thu
+Nurse 0,D,D,D,D,D,D,D
+Nurse 1,E,E,E,E,E,E,E
+Nurse 2,N,N,N,N,N,,
+Nurse 3,,,,,,N,N
diff --git a/core/tests/testcases/unwanted-pattern_1_people-2_weight.yaml b/core/tests/testcases/unwanted-pattern_1_people-2_weight.yaml
@@ -0,0 +1,29 @@
+apiVersion: alpha
+description: Test unwanted pattern with multiple people and weights
+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, 1]
+    pattern: [N]
+  - type: unwanted shift type successions
+    person: 0
+    pattern: [D]
+    weight: 1