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

87 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-08-02 00:12 +0000

1# fmt: off 

2 

3""" 

4Output support for X3D and X3DOM file types. 

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

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

7modern web browsers and jupyter notebooks. 

8""" 

9 

10import xml.etree.ElementTree as ET 

11from xml.dom import minidom 

12 

13import numpy as np 

14 

15from ase.data import covalent_radii 

16from ase.data.colors import jmol_colors 

17from ase.utils import writer 

18 

19 

20@writer 

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

22 """Writes to html using X3DOM. 

23 

24 Args: 

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

26 atoms - Atoms object to be rendered 

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

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

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

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

31 """ 

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

33 

34 

35@writer 

36def write_html(fd, atoms): 

37 """Writes to html using X3DOM. 

38 

39 Args: 

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

41 atoms - Atoms object to be rendered 

42 """ 

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

44 

45 

46class X3D: 

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

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

49 browsers. 

50 """ 

51 

52 def __init__(self, atoms): 

53 self._atoms = atoms 

54 

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

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

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

58 filename should end in '.html'. 

59 

60 Args: 

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

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

63 """ 

64 

65 # convert dictionary of style attributes to a css string 

66 if x3d_style is None: 

67 x3d_style = {} 

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

69 

70 if datatype == 'X3DOM': 

71 template = X3DOM_template 

72 elif datatype == 'X3D': 

73 template = X3D_template 

74 else: 

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

76 

77 scene = x3d_atoms(self._atoms) 

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

79 print(document, file=fileobj) 

80 

81 

82def x3d_atom(atom): 

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

84 

85 x, y, z = atom.position 

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

87 radius = covalent_radii[atom.number] 

88 

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

90 

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

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

93 

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

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

96 

97 

98def x3d_wireframe_box(box): 

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

100 

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

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

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

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

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

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

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

108 

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

110 a, b, c = box 

111 faces = [ 

112 wireframe_face(a, b), 

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

114 wireframe_face(a, c), 

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

116 ] 

117 return group(faces) 

118 

119 

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

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

122 

123 x1, y1, z1 = vec1 

124 x2, y2, z2 = vec2 

125 

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

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

128 

129 points = [ 

130 (0, 0, 0), 

131 (x1, y1, z1), 

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

133 (x2, y2, z2), 

134 (0, 0, 0), 

135 ] 

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

137 

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

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

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

141 

142 x, y, z = origin 

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

144 

145 

146def x3d_atoms(atoms): 

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

148 

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

150 wireframe = x3d_wireframe_box(atoms.cell) 

151 cell = group((wireframe, atom_spheres)) 

152 

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

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

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

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

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

158 

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

160 # - sometimes atoms appear outside the cell 

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

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

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

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

165 # distance of the viewpoint 

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

167 max_xyz_extent = get_maximum_extent(points - cell_center) 

168 

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

170 max_dim = max(max_xyz_extent) 

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

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

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

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

175 

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

177 

178 

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

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

181 

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

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

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

185 

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

187 if child is not None: 

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

189 children = [child] 

190 else: 

191 children = children or [] 

192 

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

194 for child in children: 

195 element.append(child) 

196 return element 

197 

198 

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

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

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

202 

203 

204def group(things): 

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

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

207 

208 

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

210 """Pretty print an XML element.""" 

211 

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

213 parsed = minidom.parseString(byte_string) 

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

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

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

217 return '\n'.join(lines) 

218 

219 

220def get_maximum_extent(xyz): 

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

222 

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

224 

225 

226X3DOM_template = """\ 

227<html> 

228 <head> 

229 <title>ASE atomic visualization</title> 

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

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

232 <script type="text/javascript" \ 

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

234 </head> 

235 <body> 

236 <X3D {style}> 

237 

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

239{scene} 

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

241 

242 </X3D> 

243 </body> 

244</html> 

245""" 

246 

247X3D_template = """\ 

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

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

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

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

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

253 xsd:noNamespaceSchemaLocation=\ 

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

255 

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

257{scene} 

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

259 

260</X3D> 

261"""