draw both graph and curvy ones.

wire_simulator.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)

    result = []
    for id in neighbour_index:
        result += nodes[id]

    return result


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):
    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)

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

    return G, pos


def create_layout(inner, outer, nr_nodes, distance):
    shift = nr_nodes - (nr_nodes / 2) * distance
    shift /= distance
    pos = dict()
    for i, grp in enumerate(inner):
        for idx, node in enumerate(grp):
            pos[node] = np.array([i, -idx - 1])
    for i, grp in enumerate(outer):
        for idx, node in enumerate(grp):
            pos[node] = np.array([(i - shift) * distance, idx + 2])
    return pos


def plot(curve, num_pts=256, color=None, alpha=None, 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=4,
            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):
    fig = plt.figure()
    ax = fig.gca()

    for curve in curves:
        ax = plot(curve, ax=ax, color='black')
    plt.axis('off')
    fig.savefig('amir.png')

    plt.show()


G, pos = create_side_graph(15, 0, 0, 3)
draw(curves(G, pos, 0.7))
draw_graph(G, pos)