1import numpy as np 

2from math import sqrt 

3from itertools import islice 

4 

5from ase.io.formats import string2index 

6from ase.utils import rotate 

7from ase.data import covalent_radii, atomic_numbers 

8from ase.data.colors import jmol_colors 

9 

10 

11class PlottingVariables: 

12 # removed writer - self 

13 def __init__(self, atoms, rotation='', show_unit_cell=2, 

14 radii=None, bbox=None, colors=None, scale=20, 

15 maxwidth=500, extra_offset=(0., 0.)): 

16 self.numbers = atoms.get_atomic_numbers() 

17 self.colors = colors 

18 if colors is None: 

19 ncolors = len(jmol_colors) 

20 self.colors = jmol_colors[self.numbers.clip(max=ncolors - 1)] 

21 

22 if radii is None: 

23 radii = covalent_radii[self.numbers] 

24 elif isinstance(radii, float): 

25 radii = covalent_radii[self.numbers] * radii 

26 else: 

27 radii = np.array(radii) 

28 

29 natoms = len(atoms) 

30 

31 if isinstance(rotation, str): 

32 rotation = rotate(rotation) 

33 

34 cell = atoms.get_cell() 

35 disp = atoms.get_celldisp().flatten() 

36 

37 if show_unit_cell > 0: 

38 L, T, D = cell_to_lines(self, cell) 

39 cell_vertices = np.empty((2, 2, 2, 3)) 

40 for c1 in range(2): 

41 for c2 in range(2): 

42 for c3 in range(2): 

43 cell_vertices[c1, c2, c3] = np.dot([c1, c2, c3], 

44 cell) + disp 

45 cell_vertices.shape = (8, 3) 

46 cell_vertices = np.dot(cell_vertices, rotation) 

47 else: 

48 L = np.empty((0, 3)) 

49 T = None 

50 D = None 

51 cell_vertices = None 

52 

53 nlines = len(L) 

54 

55 positions = np.empty((natoms + nlines, 3)) 

56 R = atoms.get_positions() 

57 positions[:natoms] = R 

58 positions[natoms:] = L 

59 

60 r2 = radii**2 

61 for n in range(nlines): 

62 d = D[T[n]] 

63 if ((((R - L[n] - d)**2).sum(1) < r2) & 

64 (((R - L[n] + d)**2).sum(1) < r2)).any(): 

65 T[n] = -1 

66 

67 positions = np.dot(positions, rotation) 

68 R = positions[:natoms] 

69 

70 if bbox is None: 

71 X1 = (R - radii[:, None]).min(0) 

72 X2 = (R + radii[:, None]).max(0) 

73 if show_unit_cell == 2: 

74 X1 = np.minimum(X1, cell_vertices.min(0)) 

75 X2 = np.maximum(X2, cell_vertices.max(0)) 

76 M = (X1 + X2) / 2 

77 S = 1.05 * (X2 - X1) 

78 w = scale * S[0] 

79 if w > maxwidth: 

80 w = maxwidth 

81 scale = w / S[0] 

82 h = scale * S[1] 

83 offset = np.array([scale * M[0] - w / 2, scale * M[1] - h / 2, 0]) 

84 else: 

85 w = (bbox[2] - bbox[0]) * scale 

86 h = (bbox[3] - bbox[1]) * scale 

87 offset = np.array([bbox[0], bbox[1], 0]) * scale 

88 

89 offset[0] = offset[0] - extra_offset[0] 

90 offset[1] = offset[1] - extra_offset[1] 

91 self.w = w + extra_offset[0] 

92 self.h = h + extra_offset[1] 

93 

94 positions *= scale 

95 positions -= offset 

96 

97 if nlines > 0: 

98 D = np.dot(D, rotation)[:, :2] * scale 

99 

100 if cell_vertices is not None: 

101 cell_vertices *= scale 

102 cell_vertices -= offset 

103 

104 cell = np.dot(cell, rotation) 

105 cell *= scale 

106 

107 self.cell = cell 

108 self.positions = positions 

109 self.D = D 

110 self.T = T 

111 self.cell_vertices = cell_vertices 

112 self.natoms = natoms 

113 self.d = 2 * scale * radii 

114 self.constraints = atoms.constraints 

115 

116 # extension for partial occupancies 

117 self.frac_occ = False 

118 self.tags = None 

119 self.occs = None 

120 

121 try: 

122 self.occs = atoms.info['occupancy'] 

123 self.tags = atoms.get_tags() 

124 self.frac_occ = True 

125 except KeyError: 

126 pass 

127 

128 

129def cell_to_lines(writer, cell): 

130 # XXX this needs to be updated for cell vectors that are zero. 

131 # Cannot read the code though! (What are T and D? nn?) 

132 nlines = 0 

133 nsegments = [] 

134 for c in range(3): 

135 d = sqrt((cell[c]**2).sum()) 

136 n = max(2, int(d / 0.3)) 

137 nsegments.append(n) 

138 nlines += 4 * n 

139 

140 positions = np.empty((nlines, 3)) 

141 T = np.empty(nlines, int) 

142 D = np.zeros((3, 3)) 

143 

144 n1 = 0 

145 for c in range(3): 

146 n = nsegments[c] 

147 dd = cell[c] / (4 * n - 2) 

148 D[c] = dd 

149 P = np.arange(1, 4 * n + 1, 4)[:, None] * dd 

150 T[n1:] = c 

151 for i, j in [(0, 0), (0, 1), (1, 0), (1, 1)]: 

152 n2 = n1 + n 

153 positions[n1:n2] = P + i * cell[c - 2] + j * cell[c - 1] 

154 n1 = n2 

155 

156 return positions, T, D 

157 

158 

159def make_patch_list(writer): 

160 from matplotlib.path import Path 

161 from matplotlib.patches import Circle, PathPatch, Wedge 

162 

163 indices = writer.positions[:, 2].argsort() 

164 patch_list = [] 

165 for a in indices: 

166 xy = writer.positions[a, :2] 

167 if a < writer.natoms: 

168 r = writer.d[a] / 2 

169 if writer.frac_occ: 

170 site_occ = writer.occs[str(writer.tags[a])] 

171 # first an empty circle if a site is not fully occupied 

172 if (np.sum([v for v in site_occ.values()])) < 1.0: 

173 # fill with white 

174 fill = '#ffffff' 

175 patch = Circle(xy, r, facecolor=fill, 

176 edgecolor='black') 

177 patch_list.append(patch) 

178 

179 start = 0 

180 # start with the dominant species 

181 for sym, occ in sorted(site_occ.items(), 

182 key=lambda x: x[1], 

183 reverse=True): 

184 if np.round(occ, decimals=4) == 1.0: 

185 patch = Circle(xy, r, facecolor=writer.colors[a], 

186 edgecolor='black') 

187 patch_list.append(patch) 

188 else: 

189 # jmol colors for the moment 

190 extent = 360. * occ 

191 patch = Wedge( 

192 xy, r, start, start + extent, 

193 facecolor=jmol_colors[atomic_numbers[sym]], 

194 edgecolor='black') 

195 patch_list.append(patch) 

196 start += extent 

197 

198 else: 

199 if ((xy[1] + r > 0) and (xy[1] - r < writer.h) and 

200 (xy[0] + r > 0) and (xy[0] - r < writer.w)): 

201 patch = Circle(xy, r, facecolor=writer.colors[a], 

202 edgecolor='black') 

203 patch_list.append(patch) 

204 else: 

205 a -= writer.natoms 

206 c = writer.T[a] 

207 if c != -1: 

208 hxy = writer.D[c] 

209 patch = PathPatch(Path((xy + hxy, xy - hxy))) 

210 patch_list.append(patch) 

211 return patch_list 

212 

213 

214class ImageChunk: 

215 """Base Class for a file chunk which contains enough information to 

216 reconstruct an atoms object.""" 

217 

218 def build(self, **kwargs): 

219 """Construct the atoms object from the stored information, 

220 and return it""" 

221 pass 

222 

223 

224class ImageIterator: 

225 """Iterate over chunks, to return the corresponding Atoms objects. 

226 Will only build the atoms objects which corresponds to the requested 

227 indices when called. 

228 Assumes ``ichunks`` is in iterator, which returns ``ImageChunk`` 

229 type objects. See extxyz.py:iread_xyz as an example. 

230 """ 

231 

232 def __init__(self, ichunks): 

233 self.ichunks = ichunks 

234 

235 def __call__(self, fd, index=None, **kwargs): 

236 if isinstance(index, str): 

237 index = string2index(index) 

238 

239 if index is None or index == ':': 

240 index = slice(None, None, None) 

241 

242 if not isinstance(index, (slice, str)): 

243 index = slice(index, (index + 1) or None) 

244 

245 for chunk in self._getslice(fd, index): 

246 yield chunk.build(**kwargs) 

247 

248 def _getslice(self, fd, indices): 

249 try: 

250 iterator = islice(self.ichunks(fd), 

251 indices.start, indices.stop, 

252 indices.step) 

253 except ValueError: 

254 # Negative indices. Go through the whole thing to get the length, 

255 # which allows us to evaluate the slice, and then read it again 

256 if not hasattr(fd, 'seekable') or not fd.seekable(): 

257 raise ValueError(('Negative indices only supported for ' 

258 'seekable streams')) 

259 

260 startpos = fd.tell() 

261 nchunks = 0 

262 for chunk in self.ichunks(fd): 

263 nchunks += 1 

264 fd.seek(startpos) 

265 indices_tuple = indices.indices(nchunks) 

266 iterator = islice(self.ichunks(fd), *indices_tuple) 

267 return iterator 

268 

269 

270def verify_cell_for_export(cell, check_orthorhombric=True): 

271 """Function to verify if the cell size is defined and if the cell is 

272 

273 Parameters: 

274 

275 cell: cell object 

276 cell to be checked. 

277 

278 check_orthorhombric: bool 

279 If True, check if the cell is orthorhombric, raise an ``ValueError`` if 

280 the cell is orthorhombric. If False, doesn't check if the cell is 

281 orthorhombric. 

282 

283 Raise a ``ValueError`` if the cell if not suitable for export to mustem xtl 

284 file or prismatic/computem xyz format: 

285 - if cell is not orthorhombic (only when check_orthorhombric=True) 

286 - if cell size is not defined 

287 """ 

288 

289 if check_orthorhombric and not cell.orthorhombic: 

290 raise ValueError('To export to this format, the cell needs to be ' 

291 'orthorhombic.') 

292 if cell.rank < 3: 

293 raise ValueError('To export to this format, the cell size needs ' 

294 'to be set: current cell is {}.'.format(cell)) 

295 

296 

297def verify_dictionary(atoms, dictionary, dictionary_name): 

298 """ 

299 Verify a dictionary have a key for each symbol present in the atoms object. 

300 

301 Parameters: 

302 

303 dictionary: dict 

304 Dictionary to be checked. 

305 

306 

307 dictionary_name: dict 

308 Name of the dictionary to be displayed in the error message. 

309 

310 cell: cell object 

311 cell to be checked. 

312 

313 

314 Raise a ``ValueError`` if the key doesn't match the atoms present in the 

315 cell. 

316 """ 

317 # Check if we have enough key 

318 for key in set(atoms.symbols): 

319 if key not in dictionary: 

320 raise ValueError('Missing the {} key in the `{}` dictionary.' 

321 ''.format(key, dictionary_name))