tree2.py

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


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

    basesize = len(bv)
    size = len(tv)
    kd = kdtree.KDTree(size)

    for i, vtx in enumerate(bv):
        kd.insert(Vector(vtx), i)
    kd.balance()

    kd5 = defaultdict(list)

    for i, vtx in enumerate(tv):
        co, index, dist = kd.find(vtx)
        kd5[index].append(i+basesize)

    print(kd5)
    intermediate_vertices = bv + tv

    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)

    # info, stick more
    additional = []
    for key, values in kd5.items():
        if len(values) > 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)
            print(new_vec)
            additional.append(v3[:])

    final_verts = intermediate_vertices + additional

    edges = []

    # add leafs
    for key, values in kd5.items():
        for idx in values:
            edges.append((idx, key + size + basesize))

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

    # clunky for now
    radii = []
    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