Skip to content

Commit

Permalink
Added categorical data example (#932)
Browse files Browse the repository at this point in the history
* Added categorical data example

* Added description and changed code slightly
  • Loading branch information
Joao-Dionisio authored Dec 3, 2024
1 parent d3e8ed0 commit 083efc6
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

## Unreleased
### Added
- Added categorical data example
- Added printProblem to print problem to stdout
- Added stage checks to presolve, freereoptsolve, freetransform
- Added primal_dual_evolution recipe and a plot recipe
Expand Down
73 changes: 73 additions & 0 deletions examples/finished/categorical_data.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
"""
This example shows how one can optimize a model with categorical data by converting it into integers.
There are three employees (Alice, Bob, Charlie) and three shifts. Each shift is assigned an integer:
Morning - 0
Afternoon - 1
Night - 2
The employees have availabilities (e.g. Alice can only work in the Morning and Afternoon), and different
salary demands. These constraints, and an additional one stipulating that every shift must be covered,
allows us to model a MIP with the objective of minimizing the money spent on salary.
"""

from pyscipopt import Model

# Define categorical data
shift_to_int = {"Morning": 0, "Afternoon": 1, "Night": 2}
employees = ["Alice", "Bob", "Charlie"]

# Employee availability
availability = {
"Alice": ["Morning", "Afternoon"],
"Bob": ["Afternoon", "Night"],
"Charlie": ["Morning", "Night"]
}

# Transform availability into integer values
availability_int = {}
for emp, available_shifts in availability.items():
availability_int[emp] = [shift_to_int[shift] for shift in available_shifts]


# Employees have different salary demands
cost = {
"Alice": [2,4,1],
"Bob": [3,2,7],
"Charlie": [3,3,3]
}

# Create the model
model = Model("Shift Assignment")

# x[e, s] = 1 if employee e is assigned to shift s
x = {}
for e in employees:
for s in shift_to_int.values():
x[e, s] = model.addVar(vtype="B", name=f"x({e},{s})")

# Each shift must be assigned to exactly one employee
for s in shift_to_int.values():
model.addCons(sum(x[e, s] for e in employees) == 1)

# Employees can only work shifts they are available for
for e in employees:
for s in shift_to_int.values():
if s not in availability_int[e]:
model.addCons(x[e, s] == 0)

# Minimize shift assignment cost
model.setObjective(
sum(cost[e][s]*x[e, s] for e in employees for s in shift_to_int.values()), "minimize"
)

# Solve the problem
model.optimize()

# Display the results
print("\nOptimal Shift Assignment:")
for e in employees:
for s, s_id in shift_to_int.items():
if model.getVal(x[e, s_id]) > 0.5:
print("%s is assigned to %s" % (e, s))

0 comments on commit 083efc6

Please sign in to comment.