changed to yield

wire_simulation.py

import bezier
import networkx as nx
from matplotlib import pyplot as plt

import numpy as np, numpy.random


# print(np.round(np.random.dirichlet(np.ones(10), size=1) * 170))

def distribute(value, nodes, max_number):
    while nodes >= 1:
        if nodes == 1:
            yield value
            return
        next = np.random.randint(0, min(value + 1, max_number))
        yield next
        value -= next
        nodes -= 1


def random_gen(fat_nodes, total, extra, max_number=3):
    assert extra >= fat_nodes, "number of extra values must be larger than fat_nodes"
    potential = np.random.choice(np.arange(total), fat_nodes, replace=False)
    rand_gen = distribute(extra - fat_nodes, fat_nodes, max_number - 1)  # maybe better this way.
    result = np.ones(total, dtype=np.int8)
    for idx, value in zip(potential, rand_gen):
        result[idx] = value + 2  # The fat_nodes get two at least.

    return result


def evenly_gen(fat_nodes, total, extra, max_number=3):
    assert extra >= fat_nodes, "number of extra values must be larger than fat_nodes"
    potential = np.random.choice(np.arange(total), fat_nodes, replace=False)
    if len(potential) > 0 and fat_nodes > 0:
        rand_gen = [len(i) for i in np.array_split(np.arange(extra), fat_nodes)]
    else:
        rand_gen = []

    result = np.ones(total, dtype=np.int8)
    for idx, value in zip(potential, rand_gen):
        result[idx] = value + 1

    return result


def create_nodes(prefix, ids):
    return [[f'{prefix}-{idx}-{k}' for k in range(value)] for idx, value in enumerate(ids)]


def neighbours(nodes, index):
    neighbour_index = [index]
    if index > 0:
        neighbour_index.append(index - 1)
    if index < len(nodes) - 1:
        neighbour_index.append(index + 1)

    for id in neighbour_index:
        for node in nodes[id]:
            yield node


def add_layer_nodes(graph, name, fat_nodes, total, extra, max_number=3):
    node_names = evenly_gen(fat_nodes, total, extra, max_number)
    nodes = create_nodes(name, node_names)
    for column in nodes:
        graph.add_nodes_from(column)
    return nodes


def flatten(nodes):
    return [n for layer in nodes for n in layer]


def filter(side, node):
    for grp in side:
        if node in grp:
            grp.remove(node)
    return side


def create_side_graph(pins, multilayer_pins, extra, max_number=3, length=4, alignment='up', move=10):
    G = nx.Graph()

    outer = add_layer_nodes(G, 'o', multilayer_pins, pins, extra, max_number)
    inner = add_layer_nodes(G, 'i', multilayer_pins, pins, extra, max_number)
    pos = create_layout(inner, outer, pins, 1.5, length=length, alignment=alignment, move=move)

    for i, nodes in enumerate(outer):
        for src in nodes:
            nbrs = list(neighbours(inner, i))
            if len(nbrs) > 0:
                tgt = nbrs[np.random.randint(0, len(nbrs))]
                G.add_edge(src, tgt)
                inner = filter(inner, tgt)

    return G, pos


def create_layout(inner, outer, nr_nodes, distance, length=4, alignment='right', move=10):
    shift = nr_nodes - (nr_nodes / 2) * distance
    shift /= distance

    direction = 1 if alignment == 'up' or alignment == 'right' else -1

    main_outer = []
    main_inner = []
    cross_outer = []
    cross_inner = []

    for i, grp in enumerate(inner):
        for idx, node in enumerate(grp):
            main_inner.append(i)
            cross_inner.append(idx + length)

    for i, grp in enumerate(outer):
        for idx, node in enumerate(grp):
            main_outer.append((i - shift) * distance)
            cross_outer.append(-idx - length)

    if alignment == 'right' or alignment == 'up':
        cross_outer = (-1 * np.array(cross_outer)).tolist()
        cross_inner = (-1 * np.array(cross_inner)).tolist()

    cross_inner = (np.array(cross_inner) + direction * move).tolist()
    cross_outer = (np.array(cross_outer) + direction * move).tolist()

    if alignment == 'up' or alignment == 'down':
        x_outer, y_outer = main_outer, cross_outer
        x_inner, y_inner = main_inner, cross_inner
    else:
        x_outer, y_outer = cross_outer, main_outer
        x_inner, y_inner = cross_inner, main_inner

    pos = dict(zip(flatten(inner), zip(x_inner, y_inner)))
    pos.update(dict(zip(flatten(outer), zip(x_outer, y_outer))))

    return pos


def plot(curve, num_pts=256, color=None, alpha=None, linewidth=1, ax=None):
    if curve._dimension != 2:
        raise NotImplementedError(
            "2D is the only supported dimension",
            "Current dimension",
            curve._dimension,
        )

    s_vals = np.linspace(0.0, 1.0, num_pts)
    points = curve.evaluate_multi(s_vals)

    if ax is None:
        figure = plt.figure()
        ax = figure.gca()

    ax.plot(points[0, :],
            points[1, :],
            linewidth=linewidth,
            solid_capstyle='round',
            color=color,
            alpha=alpha)
    return ax


def curves(G, pos, max_curve):
    edgelist = list(G.edges())
    edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
    for src, tgt in edge_pos:
        xs, ys = src
        xt, yt = tgt

        curve = np.random.uniform(0, max_curve, 2)
        nodes = np.asfortranarray([
            [xs, xt + curve[0], xt], [ys, yt + curve[1], yt]
        ])
        yield bezier.Curve.from_nodes(nodes)


def draw_graph(G, pos):
    fig = plt.figure()
    nx.draw(G, pos)
    plt.show()


def draw(curves, ax=None):
    for curve in curves:
        ax = plot(curve, ax=ax, color='black')
    return ax


# fig = plt.figure(dpi=500)
fig = plt.figure()
ax = fig.gca()
plt.axis('off')

G_up, pos_up = create_side_graph(25, 0, 0, 3, length=3, alignment='up', move=40)
G_down, pos_down = create_side_graph(25, 5, 15, 3, length=3, alignment='down', move=20)

G_right, pos_right = create_side_graph(25, 5, 15, 3, length=3, alignment='right', move=35)
G_left, pos_left = create_side_graph(25, 5, 15, 3, length=3, alignment='left', move=10)

ax = draw(curves(G_up, pos_up, 0.2), ax)
ax = draw(curves(G_down, pos_down, 0.2), ax)

ax = draw(curves(G_right, pos_right, 0.2), ax)
ax = draw(curves(G_left, pos_left, 0.2), ax)
# draw_graph(G, pos)
ax.margins(0.25, 0.25)  # Values in (-0.5, 0.0) zooms in to center

fig.savefig('amir.png')
plt.show()