Skip to content

Instantly share code, notes, and snippets.

@zeffii
Forked from anonymous/PDB_IMPORT.py
Last active August 29, 2015 14:14
Show Gist options
  • Save zeffii/21fdce98b73285621f83 to your computer and use it in GitHub Desktop.
Save zeffii/21fdce98b73285621f83 to your computer and use it in GitHub Desktop.
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment