Coverage for /builds/kinetik161/ase/ase/io/x3d.py: 98.85%

1""" 

2Output support for X3D and X3DOM file types. 

3See http://www.web3d.org/x3d/specifications/ 

4X3DOM outputs to html that display 3-d manipulatable atoms in 

5modern web browsers and jupyter notebooks. 

6""" 

7 

8import xml.etree.ElementTree as ET 

9from xml.dom import minidom 

10 

11import numpy as np 

12 

13from ase.data import covalent_radii 

14from ase.data.colors import jmol_colors 

15from ase.utils import writer 

16 

17 

18@writer 

19def write_x3d(fd, atoms, format='X3D', style=None): 

20 """Writes to html using X3DOM. 

21 

22 Args: 

23 filename - str or file-like object, filename or output file object 

24 atoms - Atoms object to be rendered 

25 format - str, either 'X3DOM' for web-browser compatibility or 'X3D' 

26 to be readable by Blender. `None` to detect format based on file 

27 extension ('.html' -> 'X3DOM', '.x3d' -> 'X3D') 

28 style - dict, css style attributes for the X3D element 

29 """ 

30 X3D(atoms).write(fd, datatype=format, x3d_style=style) 

31 

32 

33@writer 

34def write_html(fd, atoms): 

35 """Writes to html using X3DOM. 

36 

37 Args: 

38 filename - str or file-like object, filename or output file object 

39 atoms - Atoms object to be rendered 

40 """ 

41 write_x3d(fd, atoms, format='X3DOM') 

42 

43 

44class X3D: 

45 """Class to write either X3D (readable by open-source rendering 

46 programs such as Blender) or X3DOM html, readable by modern web 

47 browsers. 

48 """ 

49 

50 def __init__(self, atoms): 

51 self._atoms = atoms 

52 

53 def write(self, fileobj, datatype, x3d_style=None): 

54 """Writes output to either an 'X3D' or an 'X3DOM' file, based on 

55 the extension. For X3D, filename should end in '.x3d'. For X3DOM, 

56 filename should end in '.html'. 

57 

58 Args: 

59 datatype - str, output format. 'X3D' or 'X3DOM' 

60 x3d_style - dict, css style attributes for the X3D element 

61 """ 

62 

63 # convert dictionary of style attributes to a css string 

64 if x3d_style is None: 

65 x3d_style = {} 

66 x3dstyle = " ".join(f'{k}="{v}";' for k, v in x3d_style.items()) 

67 

68 if datatype == 'X3DOM': 

69 template = X3DOM_template 

70 elif datatype == 'X3D': 

71 template = X3D_template 

72 else: 

73 raise ValueError(f'datatype not supported: {datatype}') 

74 

75 scene = x3d_atoms(self._atoms) 

76 document = template.format(scene=pretty_print(scene), style=x3dstyle) 

77 print(document, file=fileobj) 

78 

79 

80def x3d_atom(atom): 

81 """Represent an atom as an x3d, coloured sphere.""" 

82 

83 x, y, z = atom.position 

84 r, g, b = jmol_colors[atom.number] 

85 radius = covalent_radii[atom.number] 

86 

87 material = element('material', diffuseColor=f'{r} {g} {b}') 

88 

89 appearance = element('appearance', child=material) 

90 sphere = element('sphere', radius=f'{radius}') 

91 

92 shape = element('shape', children=(appearance, sphere)) 

93 return translate(shape, x, y, z) 

94 

95 

96def x3d_wireframe_box(box): 

97 """x3d wireframe representation of a box (3x3 array). 

98 

99 To draw a box, spanned by vectors a, b and c, it is necessary to 

100 draw 4 faces, each of which is a parallelogram. The faces are: 

101 (start from) , (vectors spanning the face) 

102 1. (0), (a, b) 

103 2. (c), (a, b) # opposite face to 1. 

104 3. (0), (a, c) 

105 4. (b), (a, c) # opposite face to 3.""" 

106 

107 # box may not be a cube, hence not just using the diagonal 

108 a, b, c = box 

109 faces = [ 

110 wireframe_face(a, b), 

111 wireframe_face(a, b, origin=c), 

112 wireframe_face(a, c), 

113 wireframe_face(a, c, origin=b), 

114 ] 

115 return group(faces) 

116 

117 

118def wireframe_face(vec1, vec2, origin=(0, 0, 0)): 

119 """x3d wireframe representation of a face spanned by vec1 and vec2.""" 

120 

121 x1, y1, z1 = vec1 

122 x2, y2, z2 = vec2 

123 

124 material = element('material', diffuseColor='0 0 0') 

125 appearance = element('appearance', child=material) 

126 

127 points = [ 

128 (0, 0, 0), 

129 (x1, y1, z1), 

130 (x1 + x2, y1 + y2, z1 + z2), 

131 (x2, y2, z2), 

132 (0, 0, 0), 

133 ] 

134 points = ' '.join(f'{x} {y} {z}' for x, y, z in points) 

135 

136 coordinates = element('coordinate', point=points) 

137 lineset = element('lineset', vertexCount='5', child=coordinates) 

138 shape = element('shape', children=(appearance, lineset)) 

139 

140 x, y, z = origin 

141 return translate(shape, x, y, z) 

142 

143 

144def x3d_atoms(atoms): 

145 """Convert an atoms object into an x3d representation.""" 

146 

147 atom_spheres = group([x3d_atom(atom) for atom in atoms]) 

148 wireframe = x3d_wireframe_box(atoms.cell) 

149 cell = group((wireframe, atom_spheres)) 

150 

151 # we want the cell to be in the middle of the viewport 

152 # so that we can (a) see the whole cell and (b) rotate around the center 

153 # therefore we translate so that the center of the cell is at the origin 

154 cell_center = atoms.cell.diagonal() / 2 

155 cell = translate(cell, *(-cell_center)) 

156 

157 # we want the cell, and all atoms, to be visible 

158 # - sometimes atoms appear outside the cell 

159 # - sometimes atoms only take up a small part of the cell 

160 # location of the viewpoint therefore takes both of these into account: 

161 # the scene is centered on the cell, so we find the furthest point away 

162 # from the cell center, and use this to determine the 

163 # distance of the viewpoint 

164 points = np.vstack((atoms.positions, atoms.cell[:])) 

165 max_xyz_extent = get_maximum_extent(points - cell_center) 

166 

167 # the largest separation between two points in any of x, y or z 

168 max_dim = max(max_xyz_extent) 

169 # put the camera twice as far away as the largest extent 

170 pos = f'0 0 {max_dim * 2}' 

171 # NB. viewpoint needs to contain an (empty) child to be valid x3d 

172 viewpoint = element('viewpoint', position=pos, child=element('group')) 

173 

174 return element('scene', children=(viewpoint, cell)) 

175 

176 

177def element(name, child=None, children=None, **attributes) -> ET.Element: 

178 """Convenience function to make an XML element. 

179 

180 If child is specified, it is appended to the element. 

181 If children is specified, they are appended to the element. 

182 You cannot specify both child and children.""" 

183 

184 # make sure we don't specify both child and children 

185 if child is not None: 

186 assert children is None, 'Cannot specify both child and children' 

187 children = [child] 

188 else: 

189 children = children or [] 

190 

191 element = ET.Element(name, **attributes) 

192 for child in children: 

193 element.append(child) 

194 return element 

195 

196 

197def translate(thing, x, y, z): 

198 """Translate a x3d element by x, y, z.""" 

199 return element('transform', translation=f'{x} {y} {z}', child=thing) 

200 

201 

202def group(things): 

203 """Group a (list of) x3d elements.""" 

204 return element('group', children=things) 

205 

206 

207def pretty_print(element: ET.Element, indent: int = 2): 

208 """Pretty print an XML element.""" 

209 

210 byte_string = ET.tostring(element, 'utf-8') 

211 parsed = minidom.parseString(byte_string) 

212 prettied = parsed.toprettyxml(indent=' ' * indent) 

213 # remove first line - contains an extra, un-needed xml declaration 

214 lines = prettied.splitlines()[1:] 

215 return '\n'.join(lines) 

216 

217 

218def get_maximum_extent(xyz): 

219 """Get the maximum extent of an array of 3d set of points.""" 

220 

221 return np.max(xyz, axis=0) - np.min(xyz, axis=0) 

222 

223 

224X3DOM_template = """\ 

225<html> 

226 <head> 

227 <title>ASE atomic visualization</title> 

228 <link rel="stylesheet" type="text/css" \ 

229 href="https://www.x3dom.org/release/x3dom.css"></link> 

230 <script type="text/javascript" \ 

231 src="https://www.x3dom.org/release/x3dom.js"></script> 

232 </head> 

233 <body> 

234 <X3D {style}> 

235 

236<!--Inserting Generated X3D Scene--> 

237{scene} 

238<!--End of Inserted Scene--> 

239 

240 </X3D> 

241 </body> 

242</html> 

243""" 

244 

245X3D_template = """\ 

246<?xml version="1.0" encoding="UTF-8"?> 

247<!DOCTYPE X3D PUBLIC "ISO//Web3D//DTD X3D 3.2//EN" \ 

248 "http://www.web3d.org/specifications/x3d-3.2.dtd"> 

249<X3D profile="Interchange" version="3.2" \ 

250 xmlns:xsd="http://www.w3.org/2001/XMLSchema-instance" \ 

251 xsd:noNamespaceSchemaLocation=\ 

252 "http://www.web3d.org/specifications/x3d-3.2.xsd" {style}> 

253 

254<!--Inserting Generated X3D Scene--> 

255{scene} 

256<!--End of Inserted Scene--> 

257 

258</X3D> 

259"""