MLP_mod.py

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

from BioBlender.table_values import values_fi


nstr = '_4GE.001'
surface_obj_name = ''


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

    ## 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)

    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)


vcols_from_nearest_fi(nstr, surface_obj_name)