examples/notebook/school_scheduling_sat.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ortools.sat.python import cp_model\n",
    "\n",
    "\n",
    "class SchoolSchedulingProblem(object):\n",
    "\n",
    "  def __init__(self, subjects, teachers, curriculum, specialties, working_days,\n",
    "               periods, levels, sections, teacher_work_hours):\n",
    "    self.subjects = subjects\n",
    "    self.teachers = teachers\n",
    "    self.curriculum = curriculum\n",
    "    self.specialties = specialties\n",
    "    self.working_days = working_days\n",
    "    self.periods = periods\n",
    "    self.levels = levels\n",
    "    self.sections = sections\n",
    "    self.teacher_work_hours = teacher_work_hours\n",
    "\n",
    "\n",
    "class SchoolSchedulingSatSolver(object):\n",
    "\n",
    "  def __init__(self, problem):\n",
    "    # Problem\n",
    "    self.problem = problem\n",
    "\n",
    "    # Utilities\n",
    "    self.timeslots = [\n",
    "        '{0:10} {1:6}'.format(x, y)\n",
    "        for x in problem.working_days\n",
    "        for y in problem.periods\n",
    "    ]\n",
    "    self.num_days = len(problem.working_days)\n",
    "    self.num_periods = len(problem.periods)\n",
    "    self.num_slots = len(self.timeslots)\n",
    "    self.num_teachers = len(problem.teachers)\n",
    "    self.num_subjects = len(problem.subjects)\n",
    "    self.num_levels = len(problem.levels)\n",
    "    self.num_sections = len(problem.sections)\n",
    "    self.courses = [\n",
    "        x * self.num_levels + y\n",
    "        for x in problem.levels\n",
    "        for y in problem.sections\n",
    "    ]\n",
    "    self.num_courses = self.num_levels * self.num_sections\n",
    "\n",
    "    all_courses = range(self.num_courses)\n",
    "    all_teachers = range(self.num_teachers)\n",
    "    all_slots = range(self.num_slots)\n",
    "    all_sections = range(self.num_sections)\n",
    "    all_subjects = range(self.num_subjects)\n",
    "    all_levels = range(self.num_levels)\n",
    "\n",
    "    self.model = cp_model.CpModel()\n",
    "\n",
    "    self.assignment = {}\n",
    "    for c in all_courses:\n",
    "      for s in all_subjects:\n",
    "        for t in all_teachers:\n",
    "          for slot in all_slots:\n",
    "            if t in self.problem.specialties[s]:\n",
    "              name = 'C:{%i} S:{%i} T:{%i} Slot:{%i}' % (c, s, t, slot)\n",
    "              self.assignment[c, s, t, slot] = self.model.NewBoolVar(name)\n",
    "            else:\n",
    "              name = 'NO DISP C:{%i} S:{%i} T:{%i} Slot:{%i}' % (c, s, t, slot)\n",
    "              self.assignment[c, s, t, slot] = self.model.NewIntVar(0, 0, name)\n",
    "\n",
    "    # Constraints\n",
    "\n",
    "    # Each course must have the quantity of classes specified in the curriculum\n",
    "    for level in all_levels:\n",
    "      for section in all_sections:\n",
    "        course = level * self.num_sections + section\n",
    "        for subject in all_subjects:\n",
    "          required_slots = self.problem.curriculum[self.problem.levels[\n",
    "              level], self.problem.subjects[subject]]\n",
    "          self.model.Add(\n",
    "              sum(self.assignment[course, subject, teacher, slot]\n",
    "                  for slot in all_slots\n",
    "                  for teacher in all_teachers) == required_slots)\n",
    "\n",
    "    # Teacher can do at most one class at a time\n",
    "    for teacher in all_teachers:\n",
    "      for slot in all_slots:\n",
    "        self.model.Add(\n",
    "            sum([\n",
    "                self.assignment[c, s, teacher, slot]\n",
    "                for c in all_courses\n",
    "                for s in all_subjects\n",
    "            ]) <= 1)\n",
    "\n",
    "    # Maximum work hours for each teacher\n",
    "    for teacher in all_teachers:\n",
    "      self.model.Add(\n",
    "          sum([\n",
    "              self.assignment[c, s, teacher, slot] for c in all_courses\n",
    "              for s in all_subjects for slot in all_slots\n",
    "          ]) <= self.problem.teacher_work_hours[teacher])\n",
    "\n",
    "    # Teacher makes all the classes of a subject's course\n",
    "    teacher_courses = {}\n",
    "    for level in all_levels:\n",
    "      for section in all_sections:\n",
    "        course = level * self.num_sections + section\n",
    "        for subject in all_subjects:\n",
    "          for t in all_teachers:\n",
    "            name = 'C:{%i} S:{%i} T:{%i}' % (course, subject, teacher)\n",
    "            teacher_courses[course, subject, t] = self.model.NewBoolVar(name)\n",
    "            temp_array = [\n",
    "                self.assignment[course, subject, t, slot] for slot in all_slots\n",
    "            ]\n",
    "            self.model.AddMaxEquality(teacher_courses[course, subject, t],\n",
    "                                      temp_array)\n",
    "          self.model.Add(\n",
    "              sum(teacher_courses[course, subject, t]\n",
    "                  for t in all_teachers) == 1)\n",
    "\n",
    "  def solve(self):\n",
    "    print('Solving')\n",
    "    solver = cp_model.CpSolver()\n",
    "    solution_printer = SchoolSchedulingSatSolutionPrinter()\n",
    "    status = solver.Solve(self.model)\n",
    "    print()\n",
    "    print('status', status)\n",
    "    print('Branches', solver.NumBranches())\n",
    "    print('Conflicts', solver.NumConflicts())\n",
    "    print('WallTime', solver.WallTime())\n",
    "\n",
    "\n",
    "class SchoolSchedulingSatSolutionPrinter(cp_model.CpSolverSolutionCallback):\n",
    "\n",
    "  def __init__(self):\n",
    "    cp_model.CpSolverSolutionCallback.__init__(self)\n",
    "    self.__solution_count = 0\n",
    "\n",
    "  def OnSolutionCallback(self):\n",
    "    print('Found Solution!')\n",
    "\n",
    "\n",
    "# DATA\n",
    "subjects = ['English', 'Math', 'History']\n",
    "levels = ['1-', '2-', '3-']\n",
    "sections = ['A']\n",
    "teachers = ['Mario', 'Elvis', 'Donald', 'Ian']\n",
    "teachers_work_hours = [18, 12, 12, 18]\n",
    "working_days = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']\n",
    "periods = ['08:00-09:30', '09:45-11:15', '11:30-13:00']\n",
    "curriculum = {\n",
    "    ('1-', 'English'): 3,\n",
    "    ('1-', 'Math'): 3,\n",
    "    ('1-', 'History'): 2,\n",
    "    ('2-', 'English'): 4,\n",
    "    ('2-', 'Math'): 2,\n",
    "    ('2-', 'History'): 2,\n",
    "    ('3-', 'English'): 2,\n",
    "    ('3-', 'Math'): 4,\n",
    "    ('3-', 'History'): 2\n",
    "}\n",
    "\n",
    "# Subject -> List of teachers who can teach it\n",
    "specialties_idx_inverse = [\n",
    "    [1, 3],  # English   -> Elvis & Ian\n",
    "    [0, 3],  # Math      -> Mario & Ian\n",
    "    [2, 3]  # History   -> Donald & Ian\n",
    "]\n",
    "\n",
    "problem = SchoolSchedulingProblem(\n",
    "    subjects, teachers, curriculum, specialties_idx_inverse, working_days,\n",
    "    periods, levels, sections, teachers_work_hours)\n",
    "solver = SchoolSchedulingSatSolver(problem)\n",
    "solver.solve()\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}