-
Notifications
You must be signed in to change notification settings - Fork 1
/
draw.py
170 lines (143 loc) · 6.56 KB
/
draw.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
#!/usr/bin/env python
"""
Loads a json molecule and draws atoms in Blender.
Blender scripts are weird. Either run this inside of Blender or in a shell with
blender foo.blend -P molecule_to_blender.py
The script expects an input file named "molecule.json" and should be in the
same directory as "atoms.json"
Written by Patrick Fuller, [email protected], 28 Nov 12
"""
import bpy
from math import acos
from mathutils import Vector
import json
import os
import sys
# Atomic radii from wikipedia, scaled to Blender radii (C = 0.4 units)
# http://en.wikipedia.org/wiki/Atomic_radii_of_the_elements_(data_page)
# Atomic colors from cpk
# http://jmol.sourceforge.net/jscolors/
PATH = os.path.dirname(os.path.realpath(__file__))
with open(os.path.join(PATH, "atoms.json")) as in_file:
atom_data = json.load(in_file)
def draw_molecule(molecule, center=(0, 0, 0), show_bonds=True, join=True):
"""Draws a JSON-formatted molecule in Blender.
This method uses a couple of tricks from [1] to improve rendering speed.
In particular, it minimizes the amount of unique meshes and materials,
and doesn't draw until all objects are initialized.
[1] https://blenderartists.org/forum/showthread.php
?273149-Generating-a-large-number-of-mesh-primitives
Args:
molecule: The molecule to be drawn, as a python object following the
JSON convention set in this project.
center: (Optional, default (0, 0, 0)) Cartesian center of molecule. Use
to draw multiple molecules in different locations.
show_bonds: (Optional, default True) Draws a ball-and-stick model if
True, and a space-filling model if False.
join: (Optional, default True) Joins the molecule into a single object.
Set to False if you want to individually manipulate atoms/bonds.
Returns:
If run in a blender context, will return a visual object of the
molecule.
"""
shapes = []
# If using space-filling model, scale up atom size and remove bonds
# Add atom primitive
bpy.ops.object.select_all(action='DESELECT')
bpy.ops.mesh.primitive_uv_sphere_add()
sphere = bpy.context.object
# Initialize bond material if it's going to be used.
if show_bonds:
bpy.data.materials.new(name='bond')
# print(bpy.data.materials['bond'])
bpy.data.materials['bond'].diffuse_color = atom_data['bond']['color'] + [1.0]
bpy.data.materials['bond'].specular_intensity = 0.2
bpy.ops.mesh.primitive_cylinder_add()
cylinder = bpy.context.object
cylinder.active_material = bpy.data.materials['bond']
# draw atoms
for atom in molecule['atoms']:
if atom['element'] not in atom_data:
atom['element'] = 'undefined'
if atom['element'] not in bpy.data.materials:
key = atom['element']
bpy.data.materials.new(name=key)
bpy.data.materials[key].diffuse_color = atom_data[key]['color'] + [1.0]
bpy.data.materials[key].specular_intensity = 0.2
atom_sphere = sphere.copy()
atom_sphere.data = sphere.data.copy()
atom_sphere.location = [l + c for l, c in
zip(atom['location'], center)]
scale = 1 if show_bonds else 2.5
atom_sphere.dimensions = [atom_data[atom['element']]['radius'] *
scale * 2] * 3
atom_sphere.active_material = bpy.data.materials[atom['element']]
bpy.context.collection.objects.link(atom_sphere)
shapes.append(atom_sphere)
# draw bonds
for bond in (molecule['bonds'] if show_bonds else []):
start = molecule['atoms'][bond['atoms'][0]]['location']
end = molecule['atoms'][bond['atoms'][1]]['location']
diff = [c2 - c1 for c2, c1 in zip(start, end)]
cent = [(c2 + c1) / 2 for c2, c1 in zip(start, end)]
mag = sum([(c2 - c1) ** 2 for c1, c2 in zip(start, end)]) ** 0.5
v_axis = Vector(diff).normalized()
v_obj = Vector((0, 0, 1))
v_rot = v_obj.cross(v_axis)
# This check prevents gimbal lock (ie. weird behavior when v_axis is
# close to (0, 0, 1))
if v_rot.length > 0.01:
v_rot = v_rot.normalized()
axis_angle = [acos(v_obj.dot(v_axis))] + list(v_rot)
else:
v_rot = Vector((1, 0, 0))
axis_angle = [0] * 4
if bond['order'] not in range(1, 4):
sys.stderr.write("Improper number of bonds! Defaulting to 1.\n")
bond['order'] = 1
if bond['order'] == 1:
trans = [[0] * 3]
elif bond['order'] == 2:
trans = [[1.4 * atom_data['bond']['radius'] * x for x in v_rot],
[-1.4 * atom_data['bond']['radius'] * x for x in v_rot]]
elif bond['order'] == 3:
trans = [[0] * 3,
[2.2 * atom_data['bond']['radius'] * x for x in v_rot],
[-2.2 * atom_data['bond']['radius'] * x for x in v_rot]]
for i in range(bond['order']):
bond_cylinder = cylinder.copy()
bond_cylinder.data = cylinder.data.copy()
bond_cylinder.dimensions = [atom_data['bond']['radius'] * scale *
2] * 2 + [mag]
bond_cylinder.location = [c + scale * v for c,
v in zip(cent, trans[i])]
bond_cylinder.rotation_mode = 'AXIS_ANGLE'
bond_cylinder.rotation_axis_angle = axis_angle
bpy.context.collection.objects.link(bond_cylinder)
shapes.append(bond_cylinder)
# Remove primitive meshes
bpy.ops.object.select_all(action='DESELECT')
sphere.select_set(state=True)
if show_bonds:
cylinder.select_set(state=True)
# If the starting cube is there, remove it
if 'Cube' in bpy.data.objects.keys():
bpy.data.objects.get('Cube').select_set(state=True)
bpy.ops.object.delete()
for shape in shapes:
shape.select_set(state=True)
bpy.context.view_layer.objects.active = shapes[0]
bpy.ops.object.shade_smooth()
if join:
bpy.ops.object.join()
bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='MEDIAN')
bpy.context.view_layer.update()
if __name__ == '__main__':
"""Uses Blender's limited argv interface to pass args from main script."""
# show_bonds, join = True, True
# PATH = os.path.dirname(os.path.realpath(__file__))
# json_path = os.path.join(PATH, 'mol.json')
# with open(json_path) as fid:
# molecule = json.load(fid)
# draw_molecule(molecule, show_bonds=show_bonds, join=False)
pass