tree4.py

import math
from math import radians, degrees, tan

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 = []

    final_verts = intermediate_vertices + additional

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

    # for every edge check it's angle from the ground.
    # if angle less than 40 degrees, put a point on the ground to
    # make the edge.
    V = Vector
    get_angle = lambda v1, v2, v3: (V(v1) - V(v2)).angle(V(v3) - V(v2))

    for edge in edges:
        v1 = final_verts[edge[0]]  # high
        v2 = final_verts[edge[1]]  # low
        v3 = (v1[0], v1[1], v2[2])
        angle_radians = get_angle(v1, v2, v3)

        if (degrees(angle_radians)) <= 50:
            h = v1[2] - v3[2]
            v4 = Vector((v1[0], v1[1], v3[2]))
            adj = (v4-Vector(v2)).length  # current distance x, y
            new_distance = h/tan(radians(70))
            # [ ---------------------adj----------------]
            # [-----------------|---new_distance------- ]
            C = adj - new_distance
            ratio = C / adj
            new_vector = Vector(v2).lerp(v4, ratio)
            final_verts.append(new_vector[:])

    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