Basically, I have done same changes in the XYZ importer as in the PDB IO.
[blender-addons-contrib.git] / io_mesh_xyz / import_xyz.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 #
20 #
21 #  Authors           : Clemens Barth (Blendphys@root-1.de), ...
22 #
23 #  Homepage(Wiki)    : http://development.root-1.de/Atomic_Blender.php
24 #  Tracker           : ... soon
25 #
26 #  Start of project              : 2011-12-01 by Clemens Barth
27 #  First publication in Blender  : 2011-12-18
28 #  Last modified                 : 2012-03-22
29 #
30 #  Acknowledgements: Thanks to ideasman, meta_androcto, truman, kilon,
31 #  dairin0d, PKHG, Valter, etc
32 #
33
34 import bpy
35 import io
36 import math
37 import os
38 from math import pi, cos, sin
39 from mathutils import Vector, Matrix
40
41 # These are variables, which contain the name of the XYZ file and
42 # the path of the XYZ file.
43 # They are used almost everywhere, which is the reason why they
44 # should stay global. First, they are empty and get 'filled' directly
45 # after having chosen the XYZ file (see 'class LoadXYZ' further below).
46
47 ATOM_XYZ_FILEPATH = ""
48
49 # Some string stuff for the console.
50 ATOM_XYZ_STRING = "Atomic Blender\n==================="
51
52
53 # -----------------------------------------------------------------------------
54 #                                                  Atom and element data
55
56
57 # This is a list that contains some data of all possible elements. The structure
58 # is as follows:
59 #
60 # 1, "Hydrogen", "H", [0.0,0.0,1.0], 0.32, 0.32, 0.32 , -1 , 1.54   means
61 #
62 # No., name, short name, color, radius (used), radius (covalent), radius (atomic),
63 #
64 # charge state 1, radius (ionic) 1, charge state 2, radius (ionic) 2, ... all
65 # charge states for any atom are listed, if existing.
66 # The list is fixed and cannot be changed ... (see below)
67
68 ATOM_XYZ_ELEMENTS_DEFAULT = (
69 ( 1,      "Hydrogen",        "H", (  1.0,   1.0,   1.0), 0.32, 0.32, 0.79 , -1 , 1.54 ),
70 ( 2,        "Helium",       "He", ( 0.85,   1.0,   1.0), 0.93, 0.93, 0.49 ),
71 ( 3,       "Lithium",       "Li", (  0.8,  0.50,   1.0), 1.23, 1.23, 2.05 ,  1 , 0.68 ),
72 ( 4,     "Beryllium",       "Be", ( 0.76,   1.0,   0.0), 0.90, 0.90, 1.40 ,  1 , 0.44 ,  2 , 0.35 ),
73 ( 5,         "Boron",        "B", (  1.0,  0.70,  0.70), 0.82, 0.82, 1.17 ,  1 , 0.35 ,  3 , 0.23 ),
74 ( 6,        "Carbon",        "C", ( 0.56,  0.56,  0.56), 0.77, 0.77, 0.91 , -4 , 2.60 ,  4 , 0.16 ),
75 ( 7,      "Nitrogen",        "N", ( 0.18,  0.31,  0.97), 0.75, 0.75, 0.75 , -3 , 1.71 ,  1 , 0.25 ,  3 , 0.16 ,  5 , 0.13 ),
76 ( 8,        "Oxygen",        "O", (  1.0,  0.05,  0.05), 0.73, 0.73, 0.65 , -2 , 1.32 , -1 , 1.76 ,  1 , 0.22 ,  6 , 0.09 ),
77 ( 9,      "Fluorine",        "F", ( 0.56,  0.87,  0.31), 0.72, 0.72, 0.57 , -1 , 1.33 ,  7 , 0.08 ),
78 (10,          "Neon",       "Ne", ( 0.70,  0.89,  0.96), 0.71, 0.71, 0.51 ,  1 , 1.12 ),
79 (11,        "Sodium",       "Na", ( 0.67,  0.36,  0.94), 1.54, 1.54, 2.23 ,  1 , 0.97 ),
80 (12,     "Magnesium",       "Mg", ( 0.54,   1.0,   0.0), 1.36, 1.36, 1.72 ,  1 , 0.82 ,  2 , 0.66 ),
81 (13,     "Aluminium",       "Al", ( 0.74,  0.65,  0.65), 1.18, 1.18, 1.82 ,  3 , 0.51 ),
82 (14,       "Silicon",       "Si", ( 0.94,  0.78,  0.62), 1.11, 1.11, 1.46 , -4 , 2.71 , -1 , 3.84 ,  1 , 0.65 ,  4 , 0.42 ),
83 (15,    "Phosphorus",        "P", (  1.0,  0.50,   0.0), 1.06, 1.06, 1.23 , -3 , 2.12 ,  3 , 0.44 ,  5 , 0.35 ),
84 (16,        "Sulfur",        "S", (  1.0,   1.0,  0.18), 1.02, 1.02, 1.09 , -2 , 1.84 ,  2 , 2.19 ,  4 , 0.37 ,  6 , 0.30 ),
85 (17,      "Chlorine",       "Cl", ( 0.12,  0.94,  0.12), 0.99, 0.99, 0.97 , -1 , 1.81 ,  5 , 0.34 ,  7 , 0.27 ),
86 (18,         "Argon",       "Ar", ( 0.50,  0.81,  0.89), 0.98, 0.98, 0.88 ,  1 , 1.54 ),
87 (19,     "Potassium",        "K", ( 0.56,  0.25,  0.83), 2.03, 2.03, 2.77 ,  1 , 0.81 ),
88 (20,       "Calcium",       "Ca", ( 0.23,   1.0,   0.0), 1.74, 1.74, 2.23 ,  1 , 1.18 ,  2 , 0.99 ),
89 (21,      "Scandium",       "Sc", ( 0.90,  0.90,  0.90), 1.44, 1.44, 2.09 ,  3 , 0.73 ),
90 (22,      "Titanium",       "Ti", ( 0.74,  0.76,  0.78), 1.32, 1.32, 2.00 ,  1 , 0.96 ,  2 , 0.94 ,  3 , 0.76 ,  4 , 0.68 ),
91 (23,      "Vanadium",        "V", ( 0.65,  0.65,  0.67), 1.22, 1.22, 1.92 ,  2 , 0.88 ,  3 , 0.74 ,  4 , 0.63 ,  5 , 0.59 ),
92 (24,      "Chromium",       "Cr", ( 0.54,   0.6,  0.78), 1.18, 1.18, 1.85 ,  1 , 0.81 ,  2 , 0.89 ,  3 , 0.63 ,  6 , 0.52 ),
93 (25,     "Manganese",       "Mn", ( 0.61,  0.47,  0.78), 1.17, 1.17, 1.79 ,  2 , 0.80 ,  3 , 0.66 ,  4 , 0.60 ,  7 , 0.46 ),
94 (26,          "Iron",       "Fe", ( 0.87,   0.4,   0.2), 1.17, 1.17, 1.72 ,  2 , 0.74 ,  3 , 0.64 ),
95 (27,        "Cobalt",       "Co", ( 0.94,  0.56,  0.62), 1.16, 1.16, 1.67 ,  2 , 0.72 ,  3 , 0.63 ),
96 (28,        "Nickel",       "Ni", ( 0.31,  0.81,  0.31), 1.15, 1.15, 1.62 ,  2 , 0.69 ),
97 (29,        "Copper",       "Cu", ( 0.78,  0.50,   0.2), 1.17, 1.17, 1.57 ,  1 , 0.96 ,  2 , 0.72 ),
98 (30,          "Zinc",       "Zn", ( 0.49,  0.50,  0.69), 1.25, 1.25, 1.53 ,  1 , 0.88 ,  2 , 0.74 ),
99 (31,       "Gallium",       "Ga", ( 0.76,  0.56,  0.56), 1.26, 1.26, 1.81 ,  1 , 0.81 ,  3 , 0.62 ),
100 (32,     "Germanium",       "Ge", (  0.4,  0.56,  0.56), 1.22, 1.22, 1.52 , -4 , 2.72 ,  2 , 0.73 ,  4 , 0.53 ),
101 (33,       "Arsenic",       "As", ( 0.74,  0.50,  0.89), 1.20, 1.20, 1.33 , -3 , 2.22 ,  3 , 0.58 ,  5 , 0.46 ),
102 (34,      "Selenium",       "Se", (  1.0,  0.63,   0.0), 1.16, 1.16, 1.22 , -2 , 1.91 , -1 , 2.32 ,  1 , 0.66 ,  4 , 0.50 ,  6 , 0.42 ),
103 (35,       "Bromine",       "Br", ( 0.65,  0.16,  0.16), 1.14, 1.14, 1.12 , -1 , 1.96 ,  5 , 0.47 ,  7 , 0.39 ),
104 (36,       "Krypton",       "Kr", ( 0.36,  0.72,  0.81), 1.31, 1.31, 1.24 ),
105 (37,      "Rubidium",       "Rb", ( 0.43,  0.18,  0.69), 2.16, 2.16, 2.98 ,  1 , 1.47 ),
106 (38,     "Strontium",       "Sr", (  0.0,   1.0,   0.0), 1.91, 1.91, 2.45 ,  2 , 1.12 ),
107 (39,       "Yttrium",        "Y", ( 0.58,   1.0,   1.0), 1.62, 1.62, 2.27 ,  3 , 0.89 ),
108 (40,     "Zirconium",       "Zr", ( 0.58,  0.87,  0.87), 1.45, 1.45, 2.16 ,  1 , 1.09 ,  4 , 0.79 ),
109 (41,       "Niobium",       "Nb", ( 0.45,  0.76,  0.78), 1.34, 1.34, 2.08 ,  1 , 1.00 ,  4 , 0.74 ,  5 , 0.69 ),
110 (42,    "Molybdenum",       "Mo", ( 0.32,  0.70,  0.70), 1.30, 1.30, 2.01 ,  1 , 0.93 ,  4 , 0.70 ,  6 , 0.62 ),
111 (43,    "Technetium",       "Tc", ( 0.23,  0.61,  0.61), 1.27, 1.27, 1.95 ,  7 , 0.97 ),
112 (44,     "Ruthenium",       "Ru", ( 0.14,  0.56,  0.56), 1.25, 1.25, 1.89 ,  4 , 0.67 ),
113 (45,       "Rhodium",       "Rh", ( 0.03,  0.49,  0.54), 1.25, 1.25, 1.83 ,  3 , 0.68 ),
114 (46,     "Palladium",       "Pd", (  0.0,  0.41,  0.52), 1.28, 1.28, 1.79 ,  2 , 0.80 ,  4 , 0.65 ),
115 (47,        "Silver",       "Ag", ( 0.75,  0.75,  0.75), 1.34, 1.34, 1.75 ,  1 , 1.26 ,  2 , 0.89 ),
116 (48,       "Cadmium",       "Cd", (  1.0,  0.85,  0.56), 1.48, 1.48, 1.71 ,  1 , 1.14 ,  2 , 0.97 ),
117 (49,        "Indium",       "In", ( 0.65,  0.45,  0.45), 1.44, 1.44, 2.00 ,  3 , 0.81 ),
118 (50,           "Tin",       "Sn", (  0.4,  0.50,  0.50), 1.41, 1.41, 1.72 , -4 , 2.94 , -1 , 3.70 ,  2 , 0.93 ,  4 , 0.71 ),
119 (51,      "Antimony",       "Sb", ( 0.61,  0.38,  0.70), 1.40, 1.40, 1.53 , -3 , 2.45 ,  3 , 0.76 ,  5 , 0.62 ),
120 (52,     "Tellurium",       "Te", ( 0.83,  0.47,   0.0), 1.36, 1.36, 1.42 , -2 , 2.11 , -1 , 2.50 ,  1 , 0.82 ,  4 , 0.70 ,  6 , 0.56 ),
121 (53,        "Iodine",        "I", ( 0.58,   0.0,  0.58), 1.33, 1.33, 1.32 , -1 , 2.20 ,  5 , 0.62 ,  7 , 0.50 ),
122 (54,         "Xenon",       "Xe", ( 0.25,  0.61,  0.69), 1.31, 1.31, 1.24 ),
123 (55,       "Caesium",       "Cs", ( 0.34,  0.09,  0.56), 2.35, 2.35, 3.35 ,  1 , 1.67 ),
124 (56,        "Barium",       "Ba", (  0.0,  0.78,   0.0), 1.98, 1.98, 2.78 ,  1 , 1.53 ,  2 , 1.34 ),
125 (57,     "Lanthanum",       "La", ( 0.43,  0.83,   1.0), 1.69, 1.69, 2.74 ,  1 , 1.39 ,  3 , 1.06 ),
126 (58,        "Cerium",       "Ce", (  1.0,   1.0,  0.78), 1.65, 1.65, 2.70 ,  1 , 1.27 ,  3 , 1.03 ,  4 , 0.92 ),
127 (59,  "Praseodymium",       "Pr", ( 0.85,   1.0,  0.78), 1.65, 1.65, 2.67 ,  3 , 1.01 ,  4 , 0.90 ),
128 (60,     "Neodymium",       "Nd", ( 0.78,   1.0,  0.78), 1.64, 1.64, 2.64 ,  3 , 0.99 ),
129 (61,    "Promethium",       "Pm", ( 0.63,   1.0,  0.78), 1.63, 1.63, 2.62 ,  3 , 0.97 ),
130 (62,      "Samarium",       "Sm", ( 0.56,   1.0,  0.78), 1.62, 1.62, 2.59 ,  3 , 0.96 ),
131 (63,      "Europium",       "Eu", ( 0.38,   1.0,  0.78), 1.85, 1.85, 2.56 ,  2 , 1.09 ,  3 , 0.95 ),
132 (64,    "Gadolinium",       "Gd", ( 0.27,   1.0,  0.78), 1.61, 1.61, 2.54 ,  3 , 0.93 ),
133 (65,       "Terbium",       "Tb", ( 0.18,   1.0,  0.78), 1.59, 1.59, 2.51 ,  3 , 0.92 ,  4 , 0.84 ),
134 (66,    "Dysprosium",       "Dy", ( 0.12,   1.0,  0.78), 1.59, 1.59, 2.49 ,  3 , 0.90 ),
135 (67,       "Holmium",       "Ho", (  0.0,   1.0,  0.61), 1.58, 1.58, 2.47 ,  3 , 0.89 ),
136 (68,        "Erbium",       "Er", (  0.0,  0.90,  0.45), 1.57, 1.57, 2.45 ,  3 , 0.88 ),
137 (69,       "Thulium",       "Tm", (  0.0,  0.83,  0.32), 1.56, 1.56, 2.42 ,  3 , 0.87 ),
138 (70,     "Ytterbium",       "Yb", (  0.0,  0.74,  0.21), 1.74, 1.74, 2.40 ,  2 , 0.93 ,  3 , 0.85 ),
139 (71,      "Lutetium",       "Lu", (  0.0,  0.67,  0.14), 1.56, 1.56, 2.25 ,  3 , 0.85 ),
140 (72,       "Hafnium",       "Hf", ( 0.30,  0.76,   1.0), 1.44, 1.44, 2.16 ,  4 , 0.78 ),
141 (73,      "Tantalum",       "Ta", ( 0.30,  0.65,   1.0), 1.34, 1.34, 2.09 ,  5 , 0.68 ),
142 (74,      "Tungsten",        "W", ( 0.12,  0.58,  0.83), 1.30, 1.30, 2.02 ,  4 , 0.70 ,  6 , 0.62 ),
143 (75,       "Rhenium",       "Re", ( 0.14,  0.49,  0.67), 1.28, 1.28, 1.97 ,  4 , 0.72 ,  7 , 0.56 ),
144 (76,        "Osmium",       "Os", ( 0.14,   0.4,  0.58), 1.26, 1.26, 1.92 ,  4 , 0.88 ,  6 , 0.69 ),
145 (77,       "Iridium",       "Ir", ( 0.09,  0.32,  0.52), 1.27, 1.27, 1.87 ,  4 , 0.68 ),
146 (78,     "Platinium",       "Pt", ( 0.81,  0.81,  0.87), 1.30, 1.30, 1.83 ,  2 , 0.80 ,  4 , 0.65 ),
147 (79,          "Gold",       "Au", (  1.0,  0.81,  0.13), 1.34, 1.34, 1.79 ,  1 , 1.37 ,  3 , 0.85 ),
148 (80,       "Mercury",       "Hg", ( 0.72,  0.72,  0.81), 1.49, 1.49, 1.76 ,  1 , 1.27 ,  2 , 1.10 ),
149 (81,      "Thallium",       "Tl", ( 0.65,  0.32,  0.30), 1.48, 1.48, 2.08 ,  1 , 1.47 ,  3 , 0.95 ),
150 (82,          "Lead",       "Pb", ( 0.34,  0.34,  0.38), 1.47, 1.47, 1.81 ,  2 , 1.20 ,  4 , 0.84 ),
151 (83,       "Bismuth",       "Bi", ( 0.61,  0.30,  0.70), 1.46, 1.46, 1.63 ,  1 , 0.98 ,  3 , 0.96 ,  5 , 0.74 ),
152 (84,      "Polonium",       "Po", ( 0.67,  0.36,   0.0), 1.46, 1.46, 1.53 ,  6 , 0.67 ),
153 (85,      "Astatine",       "At", ( 0.45,  0.30,  0.27), 1.45, 1.45, 1.43 , -3 , 2.22 ,  3 , 0.85 ,  5 , 0.46 ),
154 (86,         "Radon",       "Rn", ( 0.25,  0.50,  0.58), 1.00, 1.00, 1.34 ),
155 (87,      "Francium",       "Fr", ( 0.25,   0.0,   0.4), 1.00, 1.00, 1.00 ,  1 , 1.80 ),
156 (88,        "Radium",       "Ra", (  0.0,  0.49,   0.0), 1.00, 1.00, 1.00 ,  2 , 1.43 ),
157 (89,      "Actinium",       "Ac", ( 0.43,  0.67,  0.98), 1.00, 1.00, 1.00 ,  3 , 1.18 ),
158 (90,       "Thorium",       "Th", (  0.0,  0.72,   1.0), 1.65, 1.65, 1.00 ,  4 , 1.02 ),
159 (91,  "Protactinium",       "Pa", (  0.0,  0.63,   1.0), 1.00, 1.00, 1.00 ,  3 , 1.13 ,  4 , 0.98 ,  5 , 0.89 ),
160 (92,       "Uranium",        "U", (  0.0,  0.56,   1.0), 1.42, 1.42, 1.00 ,  4 , 0.97 ,  6 , 0.80 ),
161 (93,     "Neptunium",       "Np", (  0.0,  0.50,   1.0), 1.00, 1.00, 1.00 ,  3 , 1.10 ,  4 , 0.95 ,  7 , 0.71 ),
162 (94,     "Plutonium",       "Pu", (  0.0,  0.41,   1.0), 1.00, 1.00, 1.00 ,  3 , 1.08 ,  4 , 0.93 ),
163 (95,     "Americium",       "Am", ( 0.32,  0.36,  0.94), 1.00, 1.00, 1.00 ,  3 , 1.07 ,  4 , 0.92 ),
164 (96,        "Curium",       "Cm", ( 0.47,  0.36,  0.89), 1.00, 1.00, 1.00 ),
165 (97,     "Berkelium",       "Bk", ( 0.54,  0.30,  0.89), 1.00, 1.00, 1.00 ),
166 (98,   "Californium",       "Cf", ( 0.63,  0.21,  0.83), 1.00, 1.00, 1.00 ),
167 (99,   "Einsteinium",       "Es", ( 0.70,  0.12,  0.83), 1.00, 1.00, 1.00 ),
168 (100,       "Fermium",       "Fm", ( 0.70,  0.12,  0.72), 1.00, 1.00, 1.00 ),
169 (101,   "Mendelevium",       "Md", ( 0.70,  0.05,  0.65), 1.00, 1.00, 1.00 ),
170 (102,      "Nobelium",       "No", ( 0.74,  0.05,  0.52), 1.00, 1.00, 1.00 ),
171 (103,    "Lawrencium",       "Lr", ( 0.78,   0.0,   0.4), 1.00, 1.00, 1.00 ),
172 (104,       "Vacancy",      "Vac", (  0.5,   0.5,   0.5), 1.00, 1.00, 1.00),
173 (105,       "Default",  "Default", (  1.0,   1.0,   1.0), 1.00, 1.00, 1.00),
174 (106,         "Stick",    "Stick", (  0.5,   0.5,   0.5), 1.00, 1.00, 1.00),
175 )
176
177 # This list here contains all data of the elements and will be used during
178 # runtime. It is a list of classes.
179 # During executing Atomic Blender, the list will be initialized with the fixed
180 # data from above via the class structure below (CLASS_atom_xyz_Elements). We
181 # have then one fixed list (above), which will never be changed, and a list of
182 # classes with same data. The latter can be modified via loading a separate
183 # custom data file for instance.
184 ATOM_XYZ_ELEMENTS = []
185
186 # This is the list, which contains all atoms of all frames! Each item is a 
187 # list which contains the atoms of a single frame. It is a list of  
188 # 'CLASS_atom_xyz_atom'.
189 ALL_FRAMES = []
190 NUMBER_FRAMES = 0
191
192 # A list of ALL balls which are put into the scene
193 STRUCTURE = []
194
195 # This is the class, which stores the properties for one element.
196 class CLASS_atom_xyz_Elements(object):
197     __slots__ = ('number', 'name', 'short_name', 'color', 'radii', 'radii_ionic')
198     def __init__(self, number, name, short_name, color, radii, radii_ionic):
199         self.number = number
200         self.name = name
201         self.short_name = short_name
202         self.color = color
203         self.radii = radii
204         self.radii_ionic = radii_ionic
205
206 # This is the class, which stores the properties of one atom.
207 class CLASS_atom_xyz_atom(object):  
208     __slots__ = ('element', 'name', 'location', 'radius', 'color', 'material')
209     def __init__(self, element, name, location, radius, color, material):
210         self.element = element
211         self.name = name
212         self.location = location
213         self.radius = radius
214         self.color = color
215         self.material = material
216         
217
218 # -----------------------------------------------------------------------------
219 #                                                          Some small routines
220
221
222
223 # This function measures the distance between two objects (atoms),
224 # which are active.
225 def DEF_atom_xyz_distance():
226
227     if len(bpy.context.selected_bases) > 1:
228         object_1 = bpy.context.selected_objects[0]
229         object_2 = bpy.context.selected_objects[1]
230     else:
231         return "N.A."
232
233     dv = object_2.location - object_1.location
234     return str(dv.length)
235
236
237 # Routine to modify the radii via the type:
238 #
239 #        pre-defined, atomic or van der Waals
240 #
241 # Explanations here are also valid for the next 3 DEFs.
242 def DEF_atom_xyz_radius_type(rtype,how):
243
244     if how == "ALL_IN_LAYER":
245
246         # Note all layers that are active.
247         layers = []
248         for i in range(20):
249             if bpy.context.scene.layers[i] == True:
250                 layers.append(i)
251         # Put all objects, which are in the layers, into a list.
252         change_objects = []
253         for obj in bpy.context.scene.objects:
254             for layer in layers:
255                 if obj.layers[layer] == True:
256                     change_objects.append(obj)
257         # Consider all objects, which are in the list 'change_objects'.
258         for obj in change_objects:
259             if len(obj.children) != 0:
260                 if obj.children[0].type == "SURFACE" or obj.children[0].type  == "MESH":
261                     for element in ATOM_XYZ_ELEMENTS:
262                         if element.name in obj.name:
263                             obj.children[0].scale = (element.radii[int(rtype)],) * 3
264             else:
265                 if obj.type == "SURFACE" or obj.type == "MESH":
266                     for element in ATOM_XYZ_ELEMENTS:
267                         if element.name in obj.name:
268                             obj.scale = (element.radii[int(rtype)],) * 3
269
270     if how == "ALL_ACTIVE":
271         for obj in bpy.context.selected_objects:
272             if len(obj.children) != 0:
273                 if obj.children[0].type == "SURFACE" or obj.children[0].type  == "MESH":
274                     for element in ATOM_XYZ_ELEMENTS:
275                         if element.name in obj.name:
276                             obj.children[0].scale = (element.radii[int(rtype)],) * 3
277             else:
278                 if obj.type == "SURFACE" or obj.type == "MESH":
279                     for element in ATOM_XYZ_ELEMENTS:
280                         if element.name in obj.name:
281                             obj.scale = (element.radii[int(rtype)],) * 3
282
283
284 # Routine to modify the radii in picometer of a specific type of atom
285 def DEF_atom_xyz_radius_pm(atomname, radius_pm, how):
286
287     if how == "ALL_IN_LAYER":
288
289         layers = []
290         for i in range(20):
291             if bpy.context.scene.layers[i] == True:
292                 layers.append(i)
293         change_objects = []
294         for obj in bpy.context.scene.objects:
295             for layer in layers:
296                 if obj.layers[layer] == True:
297                     change_objects.append(obj)
298         for obj in change_objects:
299             if len(obj.children) != 0:
300                 if obj.children[0].type == "SURFACE" or obj.children[0].type  == "MESH":
301                     if atomname in obj.name:
302                         obj.children[0].scale = (radius_pm/100,) * 3
303             else:
304                 if obj.type == "SURFACE" or obj.type == "MESH":
305                     if atomname in obj.name:
306                         obj.scale = (radius_pm/100,) * 3
307
308     if how == "ALL_ACTIVE":
309         for obj in bpy.context.selected_objects:
310             if len(obj.children) != 0:
311                 if obj.children[0].type == "SURFACE" or obj.children[0].type  == "MESH":
312                     if atomname in obj.name:
313                         obj.children[0].scale = (radius_pm/100,) * 3
314             else:
315                 if obj.type == "SURFACE" or obj.type == "MESH":
316                     if atomname in obj.name:
317                         obj.scale = (radius_pm/100,) * 3
318
319
320 # Routine to scale the radii of all atoms
321 def DEF_atom_xyz_radius_all(scale, how):
322
323     if how == "ALL_IN_LAYER":
324
325         layers = []
326         for i in range(20):
327             if bpy.context.scene.layers[i] == True:
328                 layers.append(i)
329         change_objects = []
330         for obj in bpy.context.scene.objects:
331             for layer in layers:
332                 if obj.layers[layer] == True:
333                     change_objects.append(obj)
334         for obj in change_objects:
335             if len(obj.children) != 0:
336                 if obj.children[0].type == "SURFACE" or obj.children[0].type  == "MESH":
337                     obj.children[0].scale *= scale
338             else:
339                 if obj.type == "SURFACE" or obj.type == "MESH":
340                     obj.scale *= scale
341
342     if how == "ALL_ACTIVE":
343         for obj in bpy.context.selected_objects:
344             if len(obj.children) != 0:
345                 if obj.children[0].type == "SURFACE" or obj.children[0].type  == "MESH":
346                     obj.children[0].scale *= scale
347             else:
348                 if obj.type == "SURFACE" or obj.type == "MESH":
349                     obj.scale *= scale
350
351
352 def DEF_atom_xyz_read_elements():
353
354     ATOM_XYZ_ELEMENTS[:] = []
355
356     for item in ATOM_XYZ_ELEMENTS_DEFAULT:
357
358         # All three radii into a list
359         radii = [item[4],item[5],item[6]]
360         # The handling of the ionic radii will be done later. So far, it is an
361         # empty list.
362         radii_ionic = []
363
364         li = CLASS_atom_xyz_Elements(item[0],item[1],item[2],item[3],
365                                      radii,radii_ionic)
366         ATOM_XYZ_ELEMENTS.append(li)
367
368
369 def DEF_atom_xyz_read_xyz_file(filepath,radiustype):
370
371     global NUMBER_FRAMES
372
373     # Open the file ...
374     ATOM_XYZ_FILEPATH_p = io.open(ATOM_XYZ_FILEPATH, "r")
375
376     #Go through the whole file.
377     FLAG = False
378     for line in ATOM_XYZ_FILEPATH_p:
379
380         # ... the loop is broken here (EOF) ...
381         if line == "":
382             continue
383
384         split_list = line.rsplit()
385
386         if len(split_list) == 1:
387             number_atoms = int(split_list[0])
388             FLAG = True
389             
390         if FLAG == True:
391         
392             line = ATOM_XYZ_FILEPATH_p.readline()
393             line = line.rstrip()
394             comment = line
395             
396             all_atoms= []
397             for i in range(number_atoms):
398             
399                 line = ATOM_XYZ_FILEPATH_p.readline()
400                 line = line.rstrip()
401                 split_list = line.rsplit()
402                 short_name = str(split_list[0])
403                      
404                 # Go through all elements and find the element of the current atom.
405                 FLAG_FOUND = False
406                 for element in ATOM_XYZ_ELEMENTS:
407                     if str.upper(short_name) == str.upper(element.short_name):
408                         # Give the atom its proper names, color and radius:
409                         name = element.name
410                         # int(radiustype) => type of radius:
411                         # pre-defined (0), atomic (1) or van der Waals (2)
412                         radius = float(element.radii[int(radiustype)])
413                         color = element.color
414                         FLAG_FOUND = True
415                         break
416                 
417                 # Is it a vacancy or an 'unknown atom' ?
418                 if FLAG_FOUND == False:
419                     # Give this atom also a name. If it is an 'X' then it is a
420                     # vacancy. Otherwise ...
421                     if "X" in short_name:
422                         short_name = "VAC"
423                         name = "Vacancy"
424                         radius = float(ATOM_XYZ_ELEMENTS[-3].radii[int(radiustype)])
425                         color = ATOM_XYZ_ELEMENTS[-3].color
426                     # ... take what is written in the xyz file. These are somewhat
427                     # unknown atoms. This should never happen, the element list is
428                     # almost complete. However, we do this due to security reasons.
429                     else:
430                         name = str.upper(short_name)
431                         radius = float(ATOM_XYZ_ELEMENTS[-2].radii[int(radiustype)])
432                         color = ATOM_XYZ_ELEMENTS[-2].color
433               
434                 x = float(split_list[1])
435                 y = float(split_list[2])
436                 z = float(split_list[3])
437                 
438                 location = Vector((x,y,z))
439             
440                 all_atoms.append([short_name, name, location, radius, color])
441             
442             # We note here all elements. This needs to be done only once. 
443             if NUMBER_FRAMES == 0:
444                 elements = []
445                 for atom in all_atoms:
446                     FLAG_FOUND = False
447                     for element in elements:
448                         # If the atom name is already in the list, 
449                         # FLAG on 'True'.
450                         if element == atom[1]:
451                             FLAG_FOUND = True
452                             break
453                     # No name in the current list has been found? => New entry.
454                     if FLAG_FOUND == False:
455                         # Stored are: Atom label (e.g. 'Na'), the corresponding 
456                         # atom name (e.g. 'Sodium') and its color.
457                         elements.append(atom[1])
458             
459             structure = []
460             for element in elements:
461                 atoms_one_type = []
462                 for atom in all_atoms:
463                     if atom[1] == element:
464                         atoms_one_type.append(CLASS_atom_xyz_atom(
465                                                            atom[0],
466                                                            atom[1],
467                                                            atom[2],
468                                                            atom[3],
469                                                            atom[4],[]))
470                 structure.append(atoms_one_type)
471
472             ALL_FRAMES.append(structure)
473             NUMBER_FRAMES += 1
474             FLAG = False
475
476     ATOM_XYZ_FILEPATH_p.close()
477     
478     """
479     for frame in ALL_FRAMES:
480         for element in frame:
481             for atom in element:
482                 print(atom.element + "  " + str(atom.location))
483         print()    
484     """
485     return number_atoms
486
487
488 # This reads a custom data file.
489 def DEF_atom_xyz_custom_datafile(path_datafile):
490
491     if path_datafile == "":
492         return False
493
494     path_datafile = bpy.path.abspath(path_datafile)
495
496     if os.path.isfile(path_datafile) == False:
497         return False
498
499     # The whole list gets deleted! We build it new.
500     ATOM_XYZ_ELEMENTS[:] = []
501
502     # Read the data file, which contains all data
503     # (atom name, radii, colors, etc.)
504     data_file_p = io.open(path_datafile, "r")
505
506     for line in data_file_p:
507
508         if "Atom" in line:
509
510             line = data_file_p.readline()
511             # Number
512             line = data_file_p.readline()
513             number = line[19:-1]
514             # Name
515             line = data_file_p.readline()
516             name = line[19:-1]
517             # Short name
518             line = data_file_p.readline()
519             short_name = line[19:-1]
520             # Color
521             line = data_file_p.readline()
522             color_value = line[19:-1].split(',')
523             color = [float(color_value[0]),
524                      float(color_value[1]),
525                      float(color_value[2])]
526             # Used radius
527             line = data_file_p.readline()
528             radius_used = float(line[19:-1])
529             # Atomic radius
530             line = data_file_p.readline()
531             radius_atomic = float(line[19:-1])
532             # Van der Waals radius
533             line = data_file_p.readline()
534             radius_vdW = float(line[19:-1])
535             radii = [radius_used,radius_atomic,radius_vdW]
536             radii_ionic = []
537
538             element = CLASS_atom_xyz_Elements(number,name,short_name,color,
539                                               radii, radii_ionic)
540
541             ATOM_XYZ_ELEMENTS.append(element)
542
543     data_file_p.close()
544
545     return True
546
547 # -----------------------------------------------------------------------------
548 #                                                            The main routine
549
550 def DEF_atom_xyz_main(use_mesh,
551                       Ball_azimuth,
552                       Ball_zenith,
553                       Ball_radius_factor,
554                       radiustype,
555                       Ball_distance_factor,
556                       put_to_center, 
557                       use_camera,
558                       use_lamp,
559                       path_datafile):
560
561     # List of materials
562     atom_material_list = []
563
564     # ------------------------------------------------------------------------
565     # INITIALIZE THE ELEMENT LIST
566
567     DEF_atom_xyz_read_elements()
568
569     # ------------------------------------------------------------------------
570     # READING DATA OF ATOMS
571
572     Number_of_total_atoms = DEF_atom_xyz_read_xyz_file(ATOM_XYZ_FILEPATH, 
573                                                        radiustype)
574                                                
575     # We show the atoms of the first frame.
576     first_frame = ALL_FRAMES[0]
577
578     # ------------------------------------------------------------------------
579     # MATERIAL PROPERTIES FOR ATOMS
580
581     # Create first a new list of materials for each type of atom
582     # (e.g. hydrogen)
583     for atoms_of_one_type in first_frame:
584         # Take the first atom
585         atom = atoms_of_one_type[0]
586         material = bpy.data.materials.new(atom.name)
587         material.name = atom.name
588         material.diffuse_color = atom.color
589         atom_material_list.append(material)
590
591     # Now, we go through all atoms and give them a material. For all atoms ...
592     for atoms_of_one_type in first_frame:
593         for atom in atoms_of_one_type:
594             # ... and all materials ...
595             for material in atom_material_list:
596                 # ... select the correct material for the current atom via
597                 # comparison of names ...
598                 if atom.name in material.name:
599                     # ... and give the atom its material properties.
600                     # However, before we check, if it is a vacancy
601                     # The vacancy is represented by a transparent cube.
602                     if atom.name == "Vacancy":
603                         material.transparency_method = 'Z_TRANSPARENCY'
604                         material.alpha = 1.3
605                         material.raytrace_transparency.fresnel = 1.6
606                         material.raytrace_transparency.fresnel_factor = 1.6
607                         material.use_transparency = True
608                     # The atom gets its properties.
609                     atom.material = material
610
611     # ------------------------------------------------------------------------
612     # TRANSLATION OF THE STRUCTURE TO THE ORIGIN
613
614     # It may happen that the structure in a XYZ file already has an offset
615     # If chosen, the structure is first put into the center of the scene
616     # (the offset is substracted).
617
618     if put_to_center == True:
619
620         sum_vec = Vector((0.0,0.0,0.0))
621
622         # Sum of all atom coordinates
623         for atoms_of_one_type in first_frame:
624             sum_vec = sum([atom.location for atom in atoms_of_one_type], sum_vec)
625
626         # Then the average is taken
627         sum_vec = sum_vec / Number_of_total_atoms
628
629         # After, for each atom the center of gravity is substracted
630         for atoms_of_one_type in first_frame:
631             for atom in atoms_of_one_type:
632                 atom.location -= sum_vec
633     
634     # ------------------------------------------------------------------------
635     # SCALING
636
637     # Take all atoms and adjust their radii and scale the distances.
638     for atoms_of_one_type in first_frame:
639         for atom in atoms_of_one_type:
640             atom.location *= Ball_distance_factor
641     
642     # ------------------------------------------------------------------------
643     # DETERMINATION OF SOME GEOMETRIC PROPERTIES
644
645     # In the following, some geometric properties of the whole object are
646     # determined: center, size, etc.
647     sum_vec = Vector((0.0,0.0,0.0))
648
649     # First the center is determined. All coordinates are summed up ...
650     for atoms_of_one_type in first_frame:
651         sum_vec = sum([atom.location for atom in atoms_of_one_type], sum_vec)
652
653     # ... and the average is taken. This gives the center of the object.
654     object_center_vec = sum_vec / Number_of_total_atoms
655
656     # Now, we determine the size.The farest atom from the object center is
657     # taken as a measure. The size is used to place well the camera and light
658     # into the scene.
659
660     object_size_vec = []
661     for atoms_of_one_type in first_frame:
662         object_size_vec += [atom.location - object_center_vec for atom in atoms_of_one_type]
663
664     object_size = 0.0
665     object_size = max(object_size_vec).length
666
667     # ------------------------------------------------------------------------
668     # CAMERA AND LAMP
669     camera_factor = 15.0
670
671     # If chosen a camera is put into the scene.
672     if use_camera == True:
673
674         # Assume that the object is put into the global origin. Then, the
675         # camera is moved in x and z direction, not in y. The object has its
676         # size at distance math.sqrt(object_size) from the origin. So, move the
677         # camera by this distance times a factor of camera_factor in x and z.
678         # Then add x, y and z of the origin of the object.
679         object_camera_vec = Vector((math.sqrt(object_size) * camera_factor,
680                                     0.0,
681                                     math.sqrt(object_size) * camera_factor))
682         camera_xyz_vec = object_center_vec + object_camera_vec
683
684         # Create the camera
685         current_layers=bpy.context.scene.layers
686         bpy.ops.object.camera_add(view_align=False, enter_editmode=False,
687                                location=camera_xyz_vec,
688                                rotation=(0.0, 0.0, 0.0), layers=current_layers)
689         # Some properties of the camera are changed.
690         camera = bpy.context.scene.objects.active
691         camera.name = "A_camera"
692         camera.data.name = "A_camera"
693         camera.data.lens = 45
694         camera.data.clip_end = 500.0
695
696         # Here the camera is rotated such it looks towards the center of
697         # the object. The [0.0, 0.0, 1.0] vector along the z axis
698         z_axis_vec             = Vector((0.0, 0.0, 1.0))
699         # The angle between the last two vectors
700         angle                  = object_camera_vec.angle(z_axis_vec, 0)
701         # The cross-product of z_axis_vec and object_camera_vec
702         axis_vec               = z_axis_vec.cross(object_camera_vec)
703         # Rotate 'axis_vec' by 'angle' and convert this to euler parameters.
704         # 4 is the size of the matrix.
705         euler                  = Matrix.Rotation(angle, 4, axis_vec).to_euler()
706         camera.rotation_euler  = euler
707
708         # Rotate the camera around its axis by 90° such that we have a nice
709         # camera position and view onto the object.
710         bpy.ops.transform.rotate(value=(90.0*2*math.pi/360.0,),
711                                  axis=object_camera_vec,
712                                  constraint_axis=(False, False, False),
713                                  constraint_orientation='GLOBAL',
714                                  mirror=False, proportional='DISABLED',
715                                  proportional_edit_falloff='SMOOTH',
716                                  proportional_size=1, snap=False,
717                                  snap_target='CLOSEST', snap_point=(0, 0, 0),
718                                  snap_align=False, snap_normal=(0, 0, 0),
719                                  release_confirm=False)
720
721         # This does not work, I don't know why.
722         #
723         #for area in bpy.context.screen.areas:
724         #    if area.type == 'VIEW_3D':
725         #        area.spaces[0].region_3d.view_perspective = 'CAMERA'
726
727     # Here a lamp is put into the scene, if chosen.
728     if use_lamp == True:
729
730         # This is the distance from the object measured in terms of %
731         # of the camera distance. It is set onto 50% (1/2) distance.
732         lamp_dl = math.sqrt(object_size) * 15 * 0.5
733         # This is a factor to which extend the lamp shall go to the right
734         # (from the camera  point of view).
735         lamp_dy_right = lamp_dl * (3.0/4.0)
736
737         # Create x, y and z for the lamp.
738         object_lamp_vec = Vector((lamp_dl,lamp_dy_right,lamp_dl))
739         lamp_xyz_vec = object_center_vec + object_lamp_vec
740
741         # Create the lamp
742         current_layers=bpy.context.scene.layers
743         bpy.ops.object.lamp_add (type = 'POINT', view_align=False,
744                                  location=lamp_xyz_vec,
745                                  rotation=(0.0, 0.0, 0.0),
746                                  layers=current_layers)
747         # Some properties of the lamp are changed.
748         lamp = bpy.context.scene.objects.active
749         lamp.data.name = "A_lamp"
750         lamp.name = "A_lamp"
751         lamp.data.distance = 500.0
752         lamp.data.energy = 3.0
753         lamp.data.shadow_method = 'RAY_SHADOW'
754
755         bpy.context.scene.world.light_settings.use_ambient_occlusion = True
756         bpy.context.scene.world.light_settings.ao_factor = 0.2
757
758     # ------------------------------------------------------------------------
759     # SOME OUTPUT ON THE CONSOLE
760
761     print()
762     print()
763     print()
764     print(ATOM_XYZ_STRING)
765     print()
766     print("Total number of atoms       : " + str(Number_of_total_atoms))
767     print("Center of object (Angstrom) : ", object_center_vec)
768     print("Size of object (Angstrom)   : ", object_size)
769     print()
770
771     # ------------------------------------------------------------------------
772     # DRAWING THE ATOMS
773
774     bpy.ops.object.select_all(action='DESELECT')
775
776     # For each list of atoms of ONE type (e.g. Hydrogen)
777     for atoms_of_one_type in first_frame:
778
779         # Create first the vertices composed of the coordinates of all
780         # atoms of one type
781         atom_vertices = []
782         for atom in atoms_of_one_type:
783             # In fact, the object is created in the World's origin.
784             # This is why 'object_center_vec' is substracted. At the end
785             # the whole object is translated back to 'object_center_vec'.
786             atom_vertices.append( atom.location - object_center_vec )
787
788         # Build the mesh
789         atom_mesh = bpy.data.meshes.new("Mesh_"+atom.name)
790         atom_mesh.from_pydata(atom_vertices, [], [])
791         atom_mesh.update()
792         new_atom_mesh = bpy.data.objects.new(atom.name, atom_mesh)
793         bpy.context.scene.objects.link(new_atom_mesh)
794
795         # Now, build a representative sphere (atom)
796         current_layers=bpy.context.scene.layers
797
798         if atom.name == "Vacancy":
799             bpy.ops.mesh.primitive_cube_add(
800                             view_align=False, enter_editmode=False,
801                             location=(0.0, 0.0, 0.0),
802                             rotation=(0.0, 0.0, 0.0),
803                             layers=current_layers)
804         else:
805             # NURBS balls
806             if use_mesh == False:
807                 bpy.ops.surface.primitive_nurbs_surface_sphere_add(
808                             view_align=False, enter_editmode=False,
809                             location=(0,0,0), rotation=(0.0, 0.0, 0.0),
810                             layers=current_layers)
811             # UV balls
812             else:
813                 bpy.ops.mesh.primitive_uv_sphere_add(
814                             segments=Ball_azimuth, ring_count=Ball_zenith,
815                             size=1, view_align=False, enter_editmode=False,
816                             location=(0,0,0), rotation=(0, 0, 0),
817                             layers=current_layers)
818
819         ball = bpy.context.scene.objects.active
820         ball.scale  = (atom.radius*Ball_radius_factor,) * 3
821
822         if atom.name == "Vacancy":
823             ball.name = "Cube_"+atom.name
824         else:
825             ball.name = "Ball (NURBS)_"+atom.name
826         ball.active_material = atom.material
827         ball.parent = new_atom_mesh
828         new_atom_mesh.dupli_type = 'VERTS'
829         # The object is back translated to 'object_center_vec'.
830         new_atom_mesh.location = object_center_vec
831         STRUCTURE.append(new_atom_mesh)
832
833     print()
834
835     # ------------------------------------------------------------------------
836     # SELECT ALL LOADED OBJECTS
837     
838     bpy.ops.object.select_all(action='DESELECT')
839     obj = None
840     for obj in STRUCTURE:
841         obj.select = True
842     # activate the last selected object (perhaps another should be active?)
843     if obj:
844         bpy.context.scene.objects.active = obj
845     print("\n\nAll atoms (%d) have been drawn - finished.\n\n"
846            % (Number_of_total_atoms))
847
848     return Number_of_total_atoms
849     
850
851 def DEF_atom_xyz_build_frames(frame_delta, frame_skip):
852
853     scn = bpy.context.scene
854     current_layers = scn.layers
855
856     # Introduce the basis for all elements that appear in the structure.     
857     for element in STRUCTURE:
858      
859         bpy.ops.object.select_all(action='DESELECT')   
860         bpy.context.scene.objects.active = element
861         element.select = True
862         bpy.ops.object.shape_key_add(None)
863         
864     frame_skip += 1    
865
866     # Introduce the keys and reference the atom positions for each key.     
867     i = 0
868     for j, frame in enumerate(ALL_FRAMES):        
869            
870         if j % frame_skip == 0:
871            
872             for elements_frame, elements_structure in zip(frame,STRUCTURE):
873              
874                 key = elements_structure.shape_key_add()
875     
876                 for atom_frame, atom_structure in zip(elements_frame, key.data):
877     
878                     atom_structure.co = (atom_frame.location 
879                                        - elements_structure.location)
880     
881                 key.name = atom_frame.name + "_frame_" + str(i) 
882
883             i += 1
884
885     num_frames = i
886         
887     scn.frame_start = 0
888     scn.frame_end = frame_delta * num_frames
889
890     # Manage the values of the keys
891     for element in STRUCTURE:
892  
893         scn.frame_current = 0 
894
895         element.data.shape_keys.key_blocks[1].value = 1.0
896         element.data.shape_keys.key_blocks[2].value = 0.0
897         element.data.shape_keys.key_blocks[1].keyframe_insert("value")     
898         element.data.shape_keys.key_blocks[2].keyframe_insert("value")         
899
900         scn.frame_current += frame_delta
901     
902         for number in range(num_frames)[2:]:#-1]:
903     
904             element.data.shape_keys.key_blocks[number-1].value = 0.0
905             element.data.shape_keys.key_blocks[number].value = 1.0
906             element.data.shape_keys.key_blocks[number+1].value = 0.0
907             element.data.shape_keys.key_blocks[number-1].keyframe_insert("value")     
908             element.data.shape_keys.key_blocks[number].keyframe_insert("value")     
909             element.data.shape_keys.key_blocks[number+1].keyframe_insert("value")         
910                 
911             scn.frame_current += frame_delta
912             
913         number += 1    
914             
915         element.data.shape_keys.key_blocks[number].value = 1.0
916         element.data.shape_keys.key_blocks[number-1].value = 0.0
917         element.data.shape_keys.key_blocks[number].keyframe_insert("value")     
918         element.data.shape_keys.key_blocks[number-1].keyframe_insert("value")    
919         
920
921