PDB_IMPORT.py

import bpy
import bmesh

from collections import defaultdict

''' Aliases and variables and constants '''

meshes = bpy.data.meshes
objects = bpy.data.objects
scene = bpy.context.scene

filepath = r'C:\Users\dealga\Documents\GitHub\BioBlender\Test_molecules\02_4IHV_dsDNA.pdb'

# index 0 is visual, 1 is for game engine (i think..)
scale_cov = {
    'C': [1.1, 0.6], 'CA': [0.99, 0.6], 'N': [1.07, 0.6], 'O': [1.04, 0.6],
    'S': [1.46, 0.6], 'P': [1.51, 0.6], 'FE': [0.64, 0.6], 'MG': [0.65, 0.6],
    'ZN': [0.74, 0.6], 'CU': [0.72, 0.6], 'NA': [0.95, 0.6], 'K': [1.33, 0.6],
    'CL': [1.81, 0.6], 'MN': [0.46, 0.6], 'H': [0.53, 0.3], 'F': [1.36, 0.6]}

colors = {
    'C': [.1, .1, .1], 'CA': [.4, 1, .14], 'N': [.24, .41, .7], 'O': [.46, .1, .1],
    'S': [1, .75, .17], 'P': [1, .37, .05], 'FE': [1, .5, 0], 'MG': [.64, 1, .05],
    'ZN': [.32, .42, 1], 'CU': [1, .67, 0], 'NA': [.8, .48, 1], 'K': [.72, .29, 1],
    'CL': [.1, 1, .6], 'MN': [.67, .6, 1], 'H': [.9, .9, .9], 'F': [.27, .8, .21]}

scalex = 0.2
scaley = 0.2
scalez = 0.2

''' Handle PDB import '''


def import_pdb(filepath):

    atom_dict = defaultdict(list)

    def inspect_line_by_line(pdb_file):
        for line in pdb_file:
            # protein data line: pdl
            pdl = line.strip()
            if not pdl.startswith('ATOM'):
                continue

            values = pdl.split()
            x, y, z = values[6:9]
            xyz = (float(x) * scalex, float(y) * scaley, float(z) * scalez)
            atom_type = values[-1]
            atom_dict[atom_type].append(xyz)

    def pdb_import_intermediate():
        with open(filepath) as pdb_file:
            inspect_line_by_line(pdb_file)

    pdb_import_intermediate()
    return atom_dict

atom_dict = import_pdb(filepath)


''' Handle mesh construction '''


def pydata_from_bmesh(bm):
    v = [v.co[:] for v in bm.verts]
    e = [[i.index for i in e.verts] for e in bm.edges[:]]
    p = [[i.index for i in p.verts] for p in bm.faces[:]]
    return v, e, p


def create_icospehere(subdiv, d):
    bm = bmesh.new()
    bmesh.ops.create_icosphere(bm, subdivisions=subdiv, diameter=d)
    v, e, p = pydata_from_bmesh(bm)
    bm.free()
    return v, e, p


for atom_type, verts in atom_dict.items():

    ''' parent '''
    # generate a new mesh, ever atom of certain type goes into the same mesh
    mesh = meshes.new('atoms_' + atom_type)
    mesh.from_pydata(verts, [], [])
    mesh.update()

    # create an object to assign the mesh data to
    obj = objects.new('atoms_' + atom_type, mesh)
    scene.objects.link(obj)

    ''' child '''
    # create a donor mesh (child) to represent each of these atoms
    # this mesh is duplicated onto each vertex of the atoms_* parent mesh
    diameter = scale_cov[atom_type][0] * scalex  # unified
    verts, edges, faces = create_icospehere(2, diameter)
    mesh = meshes.new('repr_atom_' + atom_type)
    mesh.from_pydata(verts, [], faces)
    mesh.update()

    # child must be an object too
    obj_child_repr = objects.new('repr_atoms_' + atom_type, mesh)
    scene.objects.link(obj_child_repr)

    # make material and assign to child as active material.
    mat = bpy.data.materials.new('bb_material_' + atom_type)
    mat.use_nodes = True
    mat.use_fake_user = True  # usually handy
    obj_child_repr.active_material = mat

    # set the diffuse color for the node-based material
    atom_color = colors[atom_type] + [1.0]
    node = mat.node_tree.nodes['Diffuse BSDF']
    node.inputs[0].default_value = atom_color

    # attach the child to the parent object's vertices
    obj.dupli_type = 'VERTS'
    obj_child_repr.parent = obj

    def set_autosmooth(obj):
        mesh = obj.data
        smooth_states = [True] * len(mesh.polygons)
        mesh.polygons.foreach_set('use_smooth', smooth_states)
        mesh.update()

    set_autosmooth(obj_child_repr)