z3_task.py

"""
    Solving Q1 in sheet 3 with Z3 the python API
    
    Author: Abdullah Khaled
    Supervised: Eng. Mai Shaaban
    
    How to run this code:
    1- install requirements
    pip install z3
    
    2- run the script
    python z3_task.py
"""

from z3 import Int, Solver, Int, Sum, Implies

######################### AUX #########################
"""
    These are auxiliary functions used to solve very specific problems 
    not related directly to the actual logic 
"""


def model_eval_col(model, X, col_index):
    res = 0
    for i in range(len(X)):
        res = res or model[X[i][col_index]].as_long()
    return res


def get_colums_as_list(matrix, col_index):
    res = []
    for i in range(len(matrix)):
        res.append(matrix[i][col_index])
    return res


####################### Analogy #######################
"""
    Let's Imagine our problem as a matrix where 
    X's rows indicate that the store is served by the warehouse in a particular location determined by the column index
"""

######################### I/P #########################
max_cost = 10  # the total cost must be less than this

k = 2

# SC: matrix store cost per location
SC = [
    [6, 1, 2, 1],
    [1, 9, 3, 9],
    [0, 1, 2, 9],
]

number_of_locations = len(SC[0])  # number of columns
number_of_stores = len(SC)  # number of rows

# CP: resembles the capacity constraint for each location in other words the number of stores that can be served from that location
CP = [2, 1, 1, 3]

######################### Definitions #########################
# X: matrix where columns are locations and rows are stores (the result i want)
# j: werehouse index
# i: store index
X = [[Int(f'x_{i+1}_{j+1}') for j in range(number_of_locations)]
     for i in range(number_of_stores)]

# A: matrix contains 1 if the wherehouse has supplied any shop
A = [Int(f'a_{i+1}') for i in range(number_of_locations)]
######################### Constraints #########################
# X_distinct: each row in x contains one and only one "1" value
X_distinct = [Sum(X[i]) == 1 for i in range(number_of_stores)]

# X_pos all values are positive X_lt2 less than 2
X_pos = [X[i][j] >= 0 for j in range(
    number_of_locations) for i in range(number_of_stores)]

X_lt2 = [X[i][j] < 2 for j in range(number_of_locations)
         for i in range(number_of_stores)]

# locations_lt_cp: sum of the output x must be less than or equal the capacity
current_l_cap = [Sum([X[i][j] for i in range(number_of_stores)])
                 for j in range(number_of_locations)]

location_lt_cp = [Sum(current_l_cap[i]) <= CP[i]
                  for i in range(number_of_locations)]

# lt_max_cost: Total cost is less than
total_cost_list = [SC[i][j] * X[i][j]
                   for j in range(number_of_locations) for i in range(number_of_stores)]

lt_max_cost = [Sum(total_cost_list) < max_cost]

# K constrain
# filling A
A_filled_hi = [Implies(Sum(get_colums_as_list(X, i)) > 0, A[i] == 1)
               for i in range(number_of_locations)]
A_filled_lo = [Implies(Sum(get_colums_as_list(X, i)) == 0, A[i] == 0)
               for i in range(number_of_locations)]
A_filled = A_filled_hi + A_filled_lo

A_sum_less_that_k = [Sum(A) <= k]
######################### Output #########################
print("constraints check:")
print('Sum(row) == 1 each shop must be supplied', X_distinct, 'All elements in X must be positive or zero', X_pos, 'All elements in x must be less than 2', X_lt2,
      'Total cost must be less than max cost', lt_max_cost, 'Total Capacity must be less than or equal the capacity constraint', location_lt_cp, "number of wherehouses must be less that k", A_filled, "and", A_sum_less_that_k, "", sep="\n-----\n")

solver = Solver()
solver.check(X_distinct + X_pos + X_lt2 + lt_max_cost +
             location_lt_cp + A_filled + A_sum_less_that_k)
X_res = solver.model()

# evaluated the selected locations
selected_locations = [model_eval_col(X_res, X, i)
                      for i in range(number_of_locations)]

print("-----")

print("resulted X:")
print(X_res)
print("-----")

print("selected locations:")
print(selected_locations)