MLP_mod.py

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

from BioBlender.table_values import values_fi
from BioBlender.BioBlender import scale_vdw


nstr = '_4GE.001'
surface_obj_name = 'SURFACE'
radii_not_found = set()


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
    [x] step : get surface mesh.
    [ ] step : for every vertex on surface mesh find closest
               vertex in proxy_ob, and fi value

    '''
    normalize_and_make_numeric = True  # affects idx to fi

    ## storage
    atom_to_fi = defaultdict(list)
    proxy_obj = defaultdict(list)
    idx_to_fi = []
    mapper_obj = {}
    mapper_obj_invert = {}
    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):
        if block_name in texts:
            text = texts[block_name]
        else:
            text = texts.new(block_name)

        # if this isn't working it's because they are now floats.
        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)
            mapper_obj_invert[(start, end)] = key

            idx_to_fi.extend(atom_to_fi[key])

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

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

    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)
        max_dist = 2.7  # fudge, could test largest vdw in cloud.

        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.

            fi = surface_verts_fi.get(vidx)
            if fi:
                # return early, use a cached value
                return fi

            # -- if arrives here the value is not yet known.
            closest_elements = []

            # -- find_range returns in order of closest -> furthest
            for (co, index, dist) in kd.find_range(coordinate, mdist):
                # this is not efficient, but for testing will suffice.
                # --- translate index to element
                # k:        index ranges..(0, 20), (21, 50), etc
                # element:  the element associated that range
                for k, element in mapper_obj_invert.items():
                    if k[0] <= index <= k[1]:
                        radius = scale_vdw.get(element)
                        if not radius:
                            radii_not_found.add(element)
                        else:
                            radius = radius[0]
                            closest_elements.append([dist-radius, index])

            # -- establish fi for this surface vertex
            if not closest_elements:
                # search failure..
                fi = None
            else:
                if len(closest_elements) == 1:
                    idx = closest_elements[0][1]
                else:
                    # arrange nearest elements in order or dist-radius
                    # does not deal with interpolation of multiple close
                    list_member = sorted(closest_elements)[0]
                    idx = list_member[1]

                # at this point we have a closest idx, we know its fi.
                fi = idx_to_fi[idx]

            surface_verts_fi[vidx] = fi
            return fi

        # 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.vertices[vidx].co
                c = from_closest(vidx, coordinate, mdist=max_dist)
                rgb = (c, c, c) if isinstance(c, float) else (.5, .5, 1.)
                vcol_layer.data[i].color = rgb
                i += 1

vcols_from_nearest_fi(nstr, surface_obj_name)
print('radii not found:', radii_not_found)