glpk -> scip

magisch_corvee_script.py

@@ -1,17 +1,20 @@
 #! /usr/bin/env nix-shell
-#!nix-shell -i python3 -p python3 python3Packages.pyyaml glpk -I nixpkgs=flake:nixpkgs
+#!nix-shell -i python3 -p python3 python3Packages.pyyaml -I nixpkgs=flake:nixpkgs
+# pip install pyyaml pyscipopt pyright
 import sys
 import yaml
 import re
-import subprocess
+import argparse
 from collections import OrderedDict
-import glpm
+from pyscipopt import Model, quicksum
 from typing import Any, Tuple, TypeVar
 from dataclasses import dataclass, field
 conf = yaml.safe_load(open('config.yaml', 'r'))
 config = conf['config']
 config['ignore'].append('')
 
+QUADRATIC = False
+
 @dataclass
 class TaskConfig:
     personen: list[str]
@@ -24,16 +27,15 @@ tasks: dict[str, TaskConfig] = OrderedDict({k: TaskConfig(**t) for k, t in conf[
 X = TypeVar("X")
 def index(x: dict[X, Any]) -> dict[X, int]:
     return {v: k for k, v in enumerate(x)}
-
 daily_workloads = \
      [sum(task.workload * task.req[d] for task in tasks.values()) for d in range(config['days'])]
 ALL_DAYS: set[int] = set(range(config['days']))
 class Person(object):
     def __init__(self, conf={"dagen":ALL_DAYS}):
-        self.can = set()
+        self.can: set[str] = set()
-        self.loves = set()
+        self.loves: set[str] = set()
-        self.hates = set()
+        self.hates: set[str] = set()
-        self.does = set() # hardcoded
+        self.does: set[Tuple[int, str]] = set() # hardcoded
         self.has_prefs = False
         self.conf = conf
         self.conf['dagen'] = set(conf['dagen'])
@@ -46,7 +48,7 @@ class Person(object):
                 res -= config['weights']['hates']
         return res
     def workload(self, tasks):
-        return sum(tasks[t]['workload'] for (_,t) in self.does)
+        return sum(tasks[t].workload for (_,t) in self.does)
     def cost(self, num_people):
         return round(sum((daily_workloads[d] for d in self.conf['dagen'])) / num_people)
         # probabilistic round: int(math.floor(x + random.random()))
@@ -54,7 +56,7 @@ def read_people(conf_ppl) -> dict[str, Person]:
     people = OrderedDict()
     for x in conf_ppl:
         val = {"dagen": ALL_DAYS}
-        if type(x) == dict:
+        if isinstance(x, dict):
             x,val = x.popitem()
         people[x.lower()] = Person(val)
     return people
@@ -66,7 +68,7 @@ def make_task_lut(tasks: dict[str, TaskConfig]):
             task_lut[lookup] = t
     task_re = re.compile(config['task_re'])
     def lookup_tasks(tasks):
-        return (task_lut[x] for x in task_re.split(tasks) if not x in config['ignore'])
+        return (task_lut[x] for x in task_re.split(tasks) if x not in config['ignore'])
     return lookup_tasks
 def read_prefs(pref_file, tasks, people):
     lookup_tasks = make_task_lut(tasks)
@@ -95,89 +97,110 @@ def set_capabilities(tasks: dict[str, TaskConfig], people: dict[str, Person]):
         if conf.hardcode is not None:
             for day, pers in enumerate(conf.hardcode):
                 people[pers.lower()].does.add((day, task))
-# format as matrices
-def matrices(people: dict[str, Person], tasks: dict[str, TaskConfig]) -> Tuple[list[list[int]], list[list[int]], list[list[int]], dict[str, int], dict[Tuple[int, int], int], list[list[int]], dict[int, int], dict[int, int]]:
-    mat = lambda a,b: [[0 for j in b] for i in a]
-    loves = mat(people, tasks)
-    hates = mat(people, tasks)
-    capab = mat(people, tasks)
-    tsk_idx = index(tasks)
-    hardcode = {}
-    max_loads = {}
-    costs = {}
-    for i,p in enumerate(people.values()):
-        for t in p.loves: loves[i][tsk_idx[t]] = 1
-        for t in p.hates: hates[i][tsk_idx[t]] = 1
-        for t in p.can: capab[i][tsk_idx[t]] = 1
-        for (d,t) in p.does: hardcode[(i,tsk_idx[t],d)] = 1
-        if 'max_load' in p.conf: # max_load override for Pol
-            for d,l in enumerate(p.conf['max_load']):
-                max_loads[(i,d)] = l
-        # filter days that the person does not exist
-        for d in ALL_DAYS - p.conf['dagen']:
-            max_loads[(i,d)] = 0
-        costs[i] = p.cost(len(people))
-    req = mat(range(config['days']), tasks)
-    for di in range(config['days']):
-        for ti,t in enumerate(tasks.values()):
-            req[di][ti] = t.req[di]
-    workload = {tsk_idx[t]: tasks[t].workload for t in tasks}
-    return (loves, hates, capab, hardcode, max_loads, req, workload, costs)
-def read_assignment(file, people, tasks):
-    def between_the_lines(f, a=">>>>\n", b="<<<<\n"):
-        for l in f:
-            if l == a: break
-        for l in f:
-            if l == b: break
-            yield map(int, l.strip().split())
 
-    for p in people.values():
-        p.does = set()
-    person_vl = list(people.values())
-    task_nm = list(tasks.keys())
-    for [p,d,j,W,l] in between_the_lines(file):
-        person_vl[p-1].does.add((d-1, task_nm[j-1]))
-def write_data(people, tasks, file=sys.stdout):
-    [loves, hates, capab, hardcode, max_loads, req, workload, costs] = matrices(people, tasks)
-    print(glpm.matrix("L", loves), file=file)
-    print(glpm.matrix("H", hates), file=file)
-    print(glpm.matrix("C", capab, 1), file=file)
-    print(glpm.matrix("R", req, 0), file=file)
-    print(glpm.dict("Q", hardcode), file=file)
-    print(glpm.dict("Wl", workload), file=file)
-    print(glpm.dict("max_load", max_loads), file=file)
-    print(glpm.dict("Costs", costs), file=file)
-    print(glpm.param("D_count", config['days']), file=file)
-    print(glpm.param("P_count", len(people)), file=file)
-    print(glpm.param("J_count", len(tasks)), file=file)
-    print(glpm.param("ML", 6), file=file) # CHANGE THIS
-    print(glpm.param("WL", config['weights']['likes']), file=file)
-    print(glpm.param("WH", config['weights']['hates']), file=file)
 def write_tasks(people, tasks, file=sys.stdout):
     for name, p in people.items():
         days = [[] for i in range(config['days'])]
         for (d,t) in p.does:
             days[d].append((t, t in p.loves, t in p.hates))
-        q = lambda w: ",".join([t + (" <3" if l else "") + (" :(" if h else "") for (t,l,h) in w])
+        def q(w):
+            return ",".join([t + (" <3" if love else "") + (" :(" if hate else "") for (t,love,hate) in w])
         days_fmt = " {} ||" * len(days)
         days_filled = days_fmt.format(*map(q, days))
         print("| {} ||{} {} || {}".format(name, days_filled, p.vrolijkheid(), p.workload(tasks)), file=file)
         print("|-")
 
+def scipsol(people: dict[str, Person], tasks: dict[str, TaskConfig]):
+    max_loads: dict[Tuple[str, int], int] = {}
+    for i,p in people.items():
+        if 'max_load' in p.conf: # max_load override for Pol
+            for d,load in enumerate(p.conf['max_load']):
+                max_loads[(i,d)] = load
+        # filter days that the person does not exist
+        for d in ALL_DAYS - p.conf['dagen']:
+            max_loads[(i,d)] = 0
+
+    m = Model()
+    does = {}
+    happiness = []
+    stdevs = []
+    errors = []
+    for pname, person in people.items():
+        workloads = []
+        p_error = m.addVar(vtype="I", name=f"{pname}_error", lb=0, ub=None)
+        for d in ALL_DAYS:
+            pdt = []
+            for task in tasks:
+                var = m.addVar(vtype="B", name=f"{pname}_does_{task}@{d}")
+                pdt.append(var)
+                does[(pname, d, task)] = var
+                # a person only does what (s)he is capable of
+                if task not in person.can:
+                    m.addCons(var == 0)
+                workloads.append(var * tasks[task].workload)
+            for task in person.loves:
+                happiness.append(does[(pname, d, task)] * config['weights']['likes'])
+            for task in person.hates:
+                happiness.append(does[(pname, d, task)] * (config['weights']['hates'] * -1))
+
+            # max_load_person: a person only has one task per day at most
+            m.addCons(quicksum(pdt) <= max_loads.get((pname, d), 1))
+
+        m.addCons(p_error >= person.cost(len(people)) - quicksum(workloads))
+        m.addCons(p_error >= quicksum(workloads) - person.cost(len(people)))
+        errors.append(p_error)
+
+        stdevs.append((person.cost(len(people)) - quicksum(workloads)) ** 2)
+
+        # min_load_person: person has at least 1 task
+        m.addCons(quicksum([does[(pname, d, task)] for d in ALL_DAYS for task in tasks]) >= 1)
+
+        # duplicate_jobs: a person does not perform the same job on all days
+        for task in tasks:
+            m.addCons(quicksum(does[(pname, d, task)] for d in ALL_DAYS) <= len(ALL_DAYS) - 1)
+
+        # max_load_person_total
+        m.addCons(quicksum([does[(pname, d, task)] * tasks[task].workload for d in ALL_DAYS for task in tasks]) <= 6)
+
+        # hardcode constraint
+        for d, task in person.does:
+            m.addCons(does[(pname, d, task)] == 1)
+
+
+    # all_allocated: each task in allocated
+    for j in tasks:
+        for d in ALL_DAYS:
+            m.addCons(quicksum(does[(p, d, j)] for p in people) == tasks[j].req[d])
+
+    objective = m.addVar(name="objvar", vtype="C", lb=None, ub=None)
+    m.setObjective(objective, "maximize")
+    if QUADRATIC:
+        m.addCons(objective <= quicksum(happiness) - quicksum(stdevs) / 2)
+    else:
+        if args.max_total_error is not None:
+            m.addCons(quicksum(errors) <= args.max_total_error)
+        m.addCons(objective <= quicksum(happiness) - quicksum(errors))
+        m.setObjective(objective, "maximize")
+    m.solveConcurrent()
+    for pname, person in people.items():
+        for d in ALL_DAYS:
+            for task in tasks:
+                if m.getVal(does[(pname, d, task)]):
+                    person.does.add((d, task))
+    write_tasks(people, tasks)
+
+    print("Totale vrolijkheid:", sum(p.vrolijkheid() for p in people.values()))
+    print("workload deviation:", m.getVal(quicksum(stdevs)))
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-q", "--quadratic", action="store_true")
+parser.add_argument("--max_total_error", type=int, default=None)
+
+args = parser.parse_args()
+QUADRATIC = args.quadratic
+
 people = read_people(conf['people'])
 with open('prefs_table', 'r') as pref_file:
     read_prefs(pref_file, tasks, people)
 set_capabilities(tasks, people)
-if len(sys.argv)>1 and sys.argv[1] == 'in':
+scipsol(people, tasks)
-    write_data(people, tasks)
-elif len(sys.argv)>1 and sys.argv[1] == 'out':
-    with open('output', 'r') as out_file:
-        read_assignment(out_file, people, tasks)
-    write_tasks(people, tasks)
-else:
-    with open('data', 'w') as out:
-        write_data(people, tasks, file=out)
-    subprocess.call(["glpsol", "--model", "model.glpm", "-d", "data", "--tmlim", "15", "--log", "output"], stdout=sys.stderr, stderr=sys.stdout)
-    with open('output', 'r') as file:
-        read_assignment(file, people, tasks)
-        write_tasks(people, tasks)