interim code

testing kdtree.py

# kdArrange v.003
import bmesh
import bpy
import mathutils
from collections import defaultdict

from BioBlender.table_values import values_fi


nstr = '_4GE.001'
surface_obj_name = 'SURFACE'


def vcols_from_nearest_fi(nstr, surface_obj_name):

    '''
    [x] step : first store atoms as {element_name: [co,..], }
    [x] step : generate singular ordered mesh by adding vertices
               in clumps of element types. (H,H,H,H,H,H,O,O,O,O,O,O..)
    [x] step : track start and end index for each element into mapper_obj
    [ ] step : get surface mesh.
    [ ] step : for every vertex on surface mesh find closest
               vertex in proxy_ob, and MLP value

    '''
    ## storage
    atom_to_fi = defaultdict(list)
    proxy_obj = defaultdict(list)
    idx_to_fi = []
    mapper_obj = {}
    verts = []
    surface_verts_fi = {}

    ## aliasing
    objs = bpy.data.objects
    texts = bpy.data.texts
    scene = bpy.context.scene
    meshes = bpy.data.meshes
    children = objs[nstr].children

    ## helper functions
    def bmesh_from_pyverts(verts):
        bm = bmesh.new()
        add_vert = bm.verts.new
        bm_verts = [add_vert(co) for co in verts]
        bm.verts.index_update()
        return bm

    def write_fi_to_textblock(block_name, idx_to_fi, normalize=True):
        if block_name in texts:
            text = texts[block_name]
        else:
            text = texts.new(block_name)

        if normalize:
            n = lambda fi: round(((float(fi) + 3) * 0.25), 4)
            idx_to_fi = [str(n(fi)) for fi in idx_to_fi]

        fi_string = '\n'.join(idx_to_fi)
        text.from_string(fi_string)

    def generate_or_update(verts):
        # -- get or create mesh
        if "proxy_mesh" in meshes:
            mesh = meshes["proxy_mesh"]
        else:
            mesh = meshes.new("proxy_mesh")

        # -- inject mesh with verts
        bm = bmesh_from_pyverts(verts)
        bm.to_mesh(mesh)
        bm.free()

        # -- create or update object with new mesh data
        if "proxy_obj" in objs:
            obj = objs['proxy_obj']
            obj.data = mesh
        else:
            obj = objs.new("proxy_obj", mesh)
            scene.objects.link(obj)

    def build_ktree(v):
        # documentation/blender_python_api_2_70_release/mathutils.kdtree.html
        size = len(v)
        kd = mathutils.kdtree.KDTree(size)

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

    def get_vcol_layer(obj):
        vcols = obj.data.vertex_colors
        if not ('fi_cols' in vcols):
            vcol_layer = obj.data.vertex_colors.new('fi_cols')
        else:
            vcol_layer = vcols.get('fi_cols')
        return vcol_layer

    if not verts:
        # fill `proxy_obj` and `atom_fo_fi`
        for o in children:
            _name = o.BBInfo[12:16].strip()
            _amino = o.BBInfo[17:20].strip()

            fi = values_fi[_amino][_name]
            co = o.location[:]

            proxy_obj[_name].append(co)   # element name
            atom_to_fi[_name].append(fi)

        # fills `idx_to_fi` and `mapper_obj`
        idx = 0
        for key in sorted(proxy_obj.keys()):
            start = len(verts)
            verts.extend(proxy_obj[key])
            end = len(verts)-1
            mapper_obj[key] = (start, end)
            idx_to_fi.extend(atom_to_fi[key])

        # print(mapper_obj)
        # {'C': (0, 97), 'N': (98, 124), 'O': (125, 153) ...}

    if verts:
        write_fi_to_textblock('fi_values.txt', idx_to_fi)
        generate_or_update(verts)
    else:
        print('no verts! - ending early')
        return

    if surface_obj_name and (surface_obj_name in objs):
        obj = objs.get(surface_obj_name)
        vcol_layer = get_vcol_layer(obj)

        kd = build_ktree(verts)

        def from_closest(vidx, coordinate, mdist):
            # Each surface vertex is present in at least 3 polygons.
            # To avoid repeating proximity search with kdtree this
            # algorithm memoizes the index and resulting fi value.
            # - this assumes that a lookup in a hashtable is more
            #   efficient than kdtree + tests..for a second and third
            #   time the vertex appears in a different polygon.

            for (co, index, dist) in kd.find_range(coordinate, mdist):
                pass

        # mesh data of surface
        # surface_mesh = obj.data
        # i = 0
        # for poly in surface_mesh.polygons:
        #     for idx, vidx in zip(poly.loop_indices, poly.vertices[:]):
        #         coordinate = surface_mesh.verts[vidx].co
        #         c = from_closest(vidx, coordinate, mdist=max_dist)
        #         vcol_layer.data[i].color = (c, c, c)
        #         i += 1



vcols_from_nearest_fi(nstr, surface_obj_name)