Python to generate SMT-LIB for scheduling scouts in tents

scout-camp-gen.py

#!/usr/bin/python3
# Blog post here:
# https://aeshirey.github.io/code/2022/07/08/smt-for-scheduling-scouts-in-tents.html

# Run this as:
#    ./scout-camp-gen.py > scout-camp.smt2
# Then check satisfiability with:
#    z3 scout-camp.smt2 > solved.txt
# This will output something like:
# ...
#  (define-fun tent0 ((x!0 Int)) Int
#    (ite (= x!0 2) 5
#    (ite (= x!0 3) 6
#      (- 1))))
# 
# This indicates that for tent0,
#     position 2 (= x!0 2) is occupied by scout 5
#     position 3 (= x!0 3) is occupied by scout 6
#     other positions are -1 (unoccupied)
# (Positions are irrelevant in this implementation)

# If any of the constraints can't be met, this version simply returns unsatisfiable.
# In other words, all constraints are treated as HARD requirements and not simply lofty
# goals. Saying we don't want Abe & Dave to be together REQUIRES that they not be put 
# together. Maybe a future version will allow these to be goals but not requirements.


NUM_TENTS = 3
TENT_CAPACITY = 4
MAX_AGE_DIFF = 3

print(f'; Config:')
print(f'; NUM_TENTS = {NUM_TENTS}')
print(f'; TENT_CAPACITY = {TENT_CAPACITY}')
print(f'; MAX_AGE_DIFF = {MAX_AGE_DIFF}')
print()


# The list of scouts with their age and binary gender.
scouts = [
        ('Abe', 14, 'm'),     # 0
        ('Brian', 13, 'm'),   # 1
        ('Charlie', 14, 'm'), # 2
        ('Dave', 13, 'm'),    # 3
        ('Eddie', 14, 'm'),   # 4
        ('Lily', 15, 'f'),    # 5
        ('Megan', 14, 'f'),   # 6
        #('Navin', 13, 'm') # 7 -- Navin is the Jerk -- no one wants to be with him
        ]


# Who we *want* to be together
want_together = [
        ('Abe', 'Charlie'),
        ('Brian', 'Dave'),
        #('Eddie', 'Lily'), # this causes unsat b/c it requires they be together, but the gender requirement requires they not be
        ]

# Who we *don't want* to be together
dont_want_together = [
        ('Abe', 'Dave'),
        #('Navin', 'Abe'),
        #('Navin', 'Brian'),
        #('Navin', 'Charlie'),
        #('Navin', 'Dave'),
        #('Navin', 'Eddie'),
        ]


print(f'(define-fun name ((i Int)) String ',)
for (i,scout) in enumerate(scouts):
    print(f'(ite (= i {i}) "{scout[0]}" ',)
print(' "(empty)"')
print(')'*len(scouts))
print(')')
print()

print(f'(declare-const num_scouts Int)')
print(f'(assert (= num_scouts {len(scouts)}))')
print('(define-fun toInt ((b Bool)) Int (ite b 1 0)) ; convert bool to int')
print()


# Generate the tent functions
print(f'; Tent# is provided with the position and returns which scout is there or -1 for empty')
print(f'; Where position is 0-{TENT_CAPACITY}')
for i in range(NUM_TENTS):
    print(f'(declare-fun tent{i} (Int) Int)')
print()

print('(declare-fun tent_contains (Int Int) Bool)')
for scout in range(len(scouts)):
    for tent in range(NUM_TENTS):
        conds = ' '.join(f'(= (tent{tent} {i}) {scout})' for i in range(TENT_CAPACITY))
        print(f'(assert (= (tent_contains {tent} {scout}) (or {conds}))); {scouts[scout][0]}')


print('; Every position in every tent must be empty or a valid scout')
print('; Also, for simpler SMT-LIB code, the values are kept sorted ASC (so empty slots come earlier)')
for tent in range(NUM_TENTS):
    for position in range(TENT_CAPACITY):
        print(f'''(assert (and (<= -1 (tent{tent} {position})) (< (tent{tent} {position}) num_scouts)))''')
        if position > 0:
            print(f'(assert (<= (tent{tent} {position-1}) (tent{tent} {position})))')
print()


print('; Define gender functions')
print('(declare-fun is_male (Int) Bool)')
print('(declare-fun is_female (Int) Bool)')
for (i, scout) in enumerate(scouts):
    if scout[2] == 'm':
        print(f'(assert (is_male {i}))')
        print(f'(assert (not (is_female {i})))')
    elif scout[2] == 'f':
        print(f'(assert (is_female {i}))')
        print(f'(assert (not (is_male {i})))')
    else:
        raise ValueError("Expected binary gender")
print()


print('; Define age functions')
print('(declare-fun age (Int) Int)')
for (i, scout) in enumerate(scouts):
    print(f'(assert (= (age {i}) {scout[1]})) ; {scout[0]}')
print()


print('; Each scout must be represented exactly once')
for (i, scout) in enumerate(scouts):
    print(f'; {scout[0]} is scout #{i}')
    slots = [ f'(toInt (= {i} (tent{tent} {position})))' for tent in range(NUM_TENTS) for position in range(TENT_CAPACITY) ]
    slots = ' '.join(slots)
    print(f'(assert (= 1 (+ {slots})))')
print()




print('; Every tent must be single gender')
for tent in range(NUM_TENTS):
    print('(assert (or')
    # Either they're all empty/male
    print('    (and')
    for position in range(1, TENT_CAPACITY):
        print(f'        (not (is_female (tent{tent} {position})))')
    print('    )')

    # Or they're all empty/female
    print('    (and')
    for position in range(1, TENT_CAPACITY):
        print(f'        (not (is_male (tent{tent} {position})))')
    print('    )')
    print('  ))')
print()


print(f'; Max age difference is {MAX_AGE_DIFF}')
for tent in range(NUM_TENTS):
    for position1 in range(TENT_CAPACITY-1):
        for position2 in range(position1+1, TENT_CAPACITY):
            print(f'(assert (or (= -1 (tent{tent} {position1})) (= -1 (tent{tent} {position2})) (< (abs (- (age (tent{tent} {position1})) (age (tent{tent} {position2})))) 3)))')
print()

print(f'; Separate individuals')
for (name1, name2) in dont_want_together:
    # Get the number for these scouts
    sep1 = next(i for (i,s) in enumerate(scouts) if s[0] == name1)
    sep2 = next(i for (i,s) in enumerate(scouts) if s[0] == name2)

    # For simpler SMT-LIB code, make sure sep1 comes before sep2
    if sep1 > sep2:
        sep1, sep2 = sep2, sep1
        name1, name2 = name2, name1

    print(f'; {name1} & {name2}')
    for tent in range(NUM_TENTS):
        print(f'(assert (not (and (tent_contains {tent} {sep1}) (tent_contains {tent} {sep2}))))')
print()

print(f'; Combine individuals')
for (name1, name2) in want_together:
    # Get the number for these scouts
    sep1 = next(i for (i,s) in enumerate(scouts) if s[0] == name1)
    sep2 = next(i for (i,s) in enumerate(scouts) if s[0] == name2)

    # For simpler SMT-LIB code, make sure sep1 comes before sep2
    if sep1 > sep2:
        sep1, sep2 = sep2, sep1
        name1, name2 = name2, name1

    print(f'; {name1} & {name2}')
    for tent in range(NUM_TENTS):
        print(f'(assert (= (tent_contains {tent} {sep1}) (tent_contains {tent} {sep2})))')
    print()
print()


print('(check-sat)')
print('(get-model)')