tree3.py

from collections import defaultdict
import mathutils
from mathutils import Vector, kdtree


def grow(bv, tv, br, tr, ratio):

    edges = []
    radii = []

    def get_median(v_list):
        x, y, z = 0, 0, 0
        n = len(v_list)
        for v in v_list:
            x += v[0]
            y += v[1]
            z += v[2]
        return (x/n, y/n, z/n)

    base_size = len(bv)
    tips_size = len(tv)
    kd5 = defaultdict(list)

    # add base vectors to the tree
    kd = kdtree.KDTree(base_size)
    for i, vtx in enumerate(bv):
        kd.insert(Vector(vtx), i)
    kd.balance()

    # find the closet base vector to every tip vector
    for i, vtx in enumerate(tv):
        co, index, dist = kd.find(vtx)
        kd5[index].append(i)

    # find which base vectors are not keys of kd5.
    new_base_size = len(kd5)
    set_to_remove = set(kd5.keys()) ^ set(range(len(bv)))

    # pop vertices and reassign  kd5 keys with offset so they refer to
    # existing vertices.
    bv2 = [v for idx, v in enumerate(bv) if not (idx in set_to_remove)]
    kd52 = defaultdict(list)
    for idx, key in enumerate(sorted(kd5.keys())):
        kd52[idx] = kd5[key]

    intermediate_vertices = bv2 + tv

    # info, stick more
    additional = []
    # palms = 0
    # for key, values in kd5.items():
    #     if len(values) > 1:
    #         palms += 1
    #         vecs = [intermediate_vertices[i] for i in values + [key]]
    #         new_vec = get_median(vecs)
    #         v2 = Vector(new_vec)
    #         v1 = Vector(bv[key])
    #         v3 = v1.lerp(v2, ratio)
    #         additional.append(v3[:])

    final_verts = intermediate_vertices + additional

    # # add leafs
    for key, values in kd52.items():
        for idx in values:
            edges.append((idx + new_base_size, key))

    # # add stems
    # for key in range(len(bv)):
    #     edges.append((key, key + tipssize + palms))

    # test edges
    # for key, values in kd5.items():
    #     for idx in values:
    #         edges.append((idx, key))

    # # clunky for now
    # for v in bv:
    #     radii.append(abs(br))
    # for v in tv:
    #     radii.append(abs(tr))
    # for v in additional:
    #     # maybe some extra magic replated to ratio
    #     radius_notch = (tr + br) / 2
    #     radii.append(abs(radius_notch))

    return final_verts, edges, radii


def sv_main(base_verts=[[]], base_r=1.0, tip_verts=[[]], tip_r=0.5, ratio=1.0):

    verts_out = []
    edges_out = []
    radii_out = []

    in_sockets = [
        ['v', 'base_verts', base_verts],
        ['s', 'base_radius', base_r],
        ['v', 'tip_verts', tip_verts],
        ['s', 'tip_radius', tip_r],
        ['s', 'ratio', ratio]
    ]

    out_sockets = [
        ['v', 'verts', [verts_out]],
        ['s', 'edges', [edges_out]],
        ['s', 'radii', [radii_out]]
    ]

    if base_verts and base_verts[0]:
        bv = base_verts[0]
        if tip_verts and tip_verts[0]:
            tv = tip_verts[0]
            v, e, r = grow(bv, tv, base_r, tip_r, ratio)
            verts_out.extend(v)
            edges_out.extend(e)
            radii_out.extend(r)

    return in_sockets, out_sockets