NNを可視化するために頂点の配置を整数計画法で計算するスクリプト（田中くんのため）

nnopt.py

import logging
import pickle
import time
import numpy
from z3 import *


NODES, EDGES, LAYERS = None, None, None

def parse(edges, layers):
    global NODES, EDGES, LAYERS
    with open(edges,  'rb') as r: EDGES = pickle.load(r)
    with open(layers, 'rb') as r: LAYERS = pickle.load(r)

    nodes = [node for node, _ in LAYERS.items()]
    assert len(nodes) == len(set(nodes))
    NODES = set(nodes)


def Abs(x):
    return If(x >= 0, x, -x)

def constraint_system():
    solver = Solver()
    n_vertices = len(NODES)

    # ノードごとに Z3 の整数変数を用意し，そのリストを「代入 (Assignment)」と呼ぶ．
    # i番目の要素がノードiに対応する変数．
    Assignment = [Int('I{}'.format(i)) for i in NODES]
    for Position in Assignment: solver.add(Position >= 0)  # 制約の追加：各変数の値は非負整数

    layers = set([layer for _, layer in LAYERS.items()])
    nodes_of_layers = dict([(layer, set([])) for layer in layers])  # 層 -> 層に属するノード集合
    for node, layer in LAYERS.items(): nodes_of_layers[layer] |= set([node])
    for layer in layers:  # 制約の追加：層内のノードに対する変数の値は一意（縦方向の順序が重ならないこと）
        solver.add(Distinct([Assignment[node] for node in nodes_of_layers[layer]]))

    # 最適化の目的関数は「（辺の長さ・重み）の和」の最小化
    Distance = Sum([Abs(Assignment[_from] - Assignment[_to]) * int(abs(weight) * 1000 + 1) for _from, _to, weight in EDGES])

    return {'Assignment': Assignment, 'solver': solver, 'Distance': Distance}


TIMEOUT = 1 * 60 * 1000  # 1 min time out set in milliseconds


def solve(problem):
    info = dict()
    t = time.time()
    solver, Assignment, Distance = [problem[key] for key in ['solver', 'Assignment', 'Distance']]

    history = []
    max_distance = 1 << 31

    while True:
        print(f'Distance: {max_distance}')
        solver.add(Distance < max_distance)
        timeout = int((t + 60 - time.time()) * 1000)
        if timeout < 0: break
        solver.set('timeout', timeout)
        result = solver.check()
        if result != sat: break
        current_best = solver.model()
        max_distance = current_best.evaluate(Distance)
        history.append(max_distance.as_long())

    info['solution'] = [current_best[v].as_long() for v in Assignment]
    info['distance'] = max_distance.as_long()
    info['optimum'] = result == unsat
    info['history'] = history
    logging.info('{} Distance reduction sequence: {}'.format('+' if result == unsat else '-', history))
    return info

if __name__ == '__main__':
    parse('edges_filtered.pkl', 'layers.pkl')
    constraint_system()
    s = time.time()
    info = solve(constraint_system())
    print(f'Time: {time.time() - s}, 重みつき距離の合計: {info["distance"]}, 最適性: {info["optimum"]}')
    print('Assignment: {info["solution"]}')

でも，結局，目的関数が実数値なため田中くんの要望に答えられなかった．このコードでは，与えられた重みを1,000倍したものを整数化して解いている．計算速度が遅いため田中くんにとっては残念だろう．