1import warnings 

2from math import pi, sin, cos 

3import numpy as np 

4 

5 

6def bz_vertices(icell, dim=3): 

7 """See https://xkcd.com/1421 ...""" 

8 from scipy.spatial import Voronoi 

9 icell = icell.copy() 

10 if dim < 3: 

11 icell[2, 2] = 1e-3 

12 if dim < 2: 

13 icell[1, 1] = 1e-3 

14 

15 I = (np.indices((3, 3, 3)) - 1).reshape((3, 27)) 

16 G = np.dot(icell.T, I).T 

17 vor = Voronoi(G) 

18 bz1 = [] 

19 for vertices, points in zip(vor.ridge_vertices, vor.ridge_points): 

20 if -1 not in vertices and 13 in points: 

21 normal = G[points].sum(0) 

22 normal /= (normal**2).sum()**0.5 

23 bz1.append((vor.vertices[vertices], normal)) 

24 return bz1 

25 

26 

27def bz_plot(cell, vectors=False, paths=None, points=None, 

28 elev=None, scale=1, interactive=False, 

29 pointstyle=None, ax=None, show=False): 

30 import matplotlib.pyplot as plt 

31 

32 if ax is None: 

33 fig = plt.gcf() 

34 

35 dimensions = cell.rank 

36 assert dimensions > 0, 'No BZ for 0D!' 

37 

38 if dimensions == 3: 

39 from mpl_toolkits.mplot3d import Axes3D 

40 from mpl_toolkits.mplot3d import proj3d 

41 from matplotlib.patches import FancyArrowPatch 

42 Axes3D # silence pyflakes 

43 

44 class Arrow3D(FancyArrowPatch): 

45 def __init__(self, xs, ys, zs, *args, **kwargs): 

46 FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs) 

47 self._verts3d = xs, ys, zs 

48 

49 def draw(self, renderer): 

50 xs3d, ys3d, zs3d = self._verts3d 

51 xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, 

52 zs3d, ax.axes.M) 

53 self.set_positions((xs[0], ys[0]), (xs[1], ys[1])) 

54 FancyArrowPatch.draw(self, renderer) 

55 

56 # FIXME: Compatibility fix for matplotlib 3.5.0: Handling of 3D 

57 # artists have changed and all 3D artists now need 

58 # "do_3d_projection". Since this class is a hack that manually 

59 # projects onto the 3D axes we don't need to do anything in this 

60 # method. Ideally we shouldn't resort to a hack like this. 

61 def do_3d_projection(self, *_, **__): 

62 return 0 

63 

64 azim = pi / 5 

65 elev = elev or pi / 6 

66 x = sin(azim) 

67 y = cos(azim) 

68 view = [x * cos(elev), y * cos(elev), sin(elev)] 

69 

70 if ax is None: 

71 ax = fig.add_subplot(projection='3d') 

72 elif dimensions == 2: 

73 # 2d in xy 

74 assert all(abs(cell[2][0:2]) < 1e-6) and all(abs(cell.T[2] 

75 [0:2]) < 1e-6) 

76 ax = plt.gca() 

77 cell = cell.copy() 

78 else: 

79 # 1d in x 

80 assert (all(abs(cell[2][0:2]) < 1e-6) and 

81 all(abs(cell.T[2][0:2]) < 1e-6) and 

82 abs(cell[0][1]) < 1e-6 and abs(cell[1][0]) < 1e-6) 

83 ax = plt.gca() 

84 cell = cell.copy() 

85 

86 icell = cell.reciprocal() 

87 kpoints = points 

88 bz1 = bz_vertices(icell, dim=dimensions) 

89 

90 maxp = 0.0 

91 minp = 0.0 

92 if dimensions == 1: 

93 x = np.array([-0.5 * icell[0, 0], 

94 0.5 * icell[0, 0], 

95 -0.5 * icell[0, 0]]) 

96 y = np.array([0, 0, 0]) 

97 ax.plot(x, y, c='k', ls='-') 

98 maxp = icell[0, 0] 

99 else: 

100 for points, normal in bz1: 

101 x, y, z = np.concatenate([points, points[:1]]).T 

102 if dimensions == 3: 

103 if np.dot(normal, view) < 0 and not interactive: 

104 ls = ':' 

105 else: 

106 ls = '-' 

107 ax.plot(x, y, z, c='k', ls=ls) 

108 elif dimensions == 2: 

109 ax.plot(x, y, c='k', ls='-') 

110 maxp = max(maxp, points.max()) 

111 minp = min(minp, points.min()) 

112 

113 def draw_axis3d(ax, vector): 

114 ax.add_artist(Arrow3D( 

115 [0, vector[0]], 

116 [0, vector[1]], 

117 [0, vector[2]], 

118 mutation_scale=20, 

119 arrowstyle='-|>', 

120 color='k', 

121 )) 

122 

123 def draw_axis2d(ax, x, y): 

124 ax.arrow(0, 0, x, y, 

125 lw=1, color='k', 

126 length_includes_head=True, 

127 head_width=0.03, 

128 head_length=0.05) 

129 

130 if vectors: 

131 if dimensions == 3: 

132 for i in range(3): 

133 draw_axis3d(ax, vector=icell[i]) 

134 maxp = max(maxp, 0.6 * icell.max()) 

135 elif dimensions == 2: 

136 for i in range(2): 

137 draw_axis2d(ax, icell[i, 0], icell[i, 1]) 

138 maxp = max(maxp, icell.max()) 

139 else: 

140 draw_axis2d(ax, icell[0, 0], 0) 

141 maxp = max(maxp, icell.max()) 

142 

143 if paths is not None: 

144 for names, points in paths: 

145 x, y, z = np.array(points).T 

146 if dimensions == 3: 

147 ax.plot(x, y, z, c='r', ls='-', marker='.') 

148 elif dimensions in [1, 2]: 

149 ax.plot(x, y, c='r', ls='-') 

150 

151 for name, point in zip(names, points): 

152 x, y, z = point 

153 if name == 'G': 

154 name = '\\Gamma' 

155 elif len(name) > 1: 

156 import re 

157 m = re.match(r'^(\D+?)(\d*)$', name) 

158 if m is None: 

159 raise ValueError('Bad label: {}'.format(name)) 

160 name, num = m.group(1, 2) 

161 if num: 

162 name = '{}_{{{}}}'.format(name, num) 

163 if dimensions == 3: 

164 ax.text(x, y, z, '$\\mathrm{' + name + '}$', 

165 ha='center', va='bottom', color='g') 

166 elif dimensions == 2: 

167 ha_s = ['right', 'left', 'right'] 

168 va_s = ['bottom', 'bottom', 'top'] 

169 

170 ha = ha_s[int(np.sign(x))] 

171 va = va_s[int(np.sign(y))] 

172 if abs(z) < 1e-6: 

173 ax.text(x, y, '$\\mathrm{' + name + '}$', 

174 ha=ha, va=va, color='g', 

175 zorder=5) 

176 else: 

177 if abs(y) < 1e-6 and abs(z) < 1e-6: 

178 ax.text(x, y, '$\\mathrm{' + name + '}$', 

179 ha='center', va='bottom', color='g', 

180 zorder=5) 

181 

182 if kpoints is not None: 

183 kw = {'c': 'b'} 

184 if pointstyle is not None: 

185 kw.update(pointstyle) 

186 for p in kpoints: 

187 if dimensions == 3: 

188 ax.scatter(p[0], p[1], p[2], **kw) 

189 elif dimensions == 2: 

190 ax.scatter(p[0], p[1], zorder=4, **kw) 

191 else: 

192 ax.scatter(p[0], 0, zorder=4, **kw) 

193 

194 ax.set_axis_off() 

195 

196 if dimensions in [1, 2]: 

197 ax.autoscale_view(tight=True) 

198 s = maxp * 1.05 

199 ax.set_xlim(-s, s) 

200 ax.set_ylim(-s, s) 

201 ax.set_aspect('equal') 

202 

203 if dimensions == 3: 

204 # ax.set_aspect('equal') <-- won't work anymore in 3.1.0 

205 ax.view_init(azim=azim / pi * 180, elev=elev / pi * 180) 

206 # We want aspect 'equal', but apparently there was a bug in 

207 # matplotlib causing wrong behaviour. Matplotlib raises 

208 # NotImplementedError as of v3.1.0. This is a bit unfortunate 

209 # because the workarounds known to StackOverflow and elsewhere 

210 # all involve using set_aspect('equal') and then doing 

211 # something more. 

212 # 

213 # We try to get square axes here by setting a square figure, 

214 # but this is probably rather inexact. 

215 fig = ax.get_figure() 

216 xx = plt.figaspect(1.0) 

217 fig.set_figheight(xx[1]) 

218 fig.set_figwidth(xx[0]) 

219 

220 # oldlibs tests with matplotlib 2.0.0 say we have no set_proj_type: 

221 if hasattr(ax, 'set_proj_type'): 

222 ax.set_proj_type('ortho') 

223 

224 minp0 = 0.9 * minp # Here we cheat a bit to trim spacings 

225 maxp0 = 0.9 * maxp 

226 ax.set_xlim3d(minp0, maxp0) 

227 ax.set_ylim3d(minp0, maxp0) 

228 ax.set_zlim3d(minp0, maxp0) 

229 

230 if hasattr(ax, 'set_box_aspect'): 

231 ax.set_box_aspect([1, 1, 1]) 

232 else: 

233 msg = ('Matplotlib axes have no set_box_aspect() method. ' 

234 'Aspect ratio will likely be wrong. ' 

235 'Consider updating to Matplotlib >= 3.3.') 

236 warnings.warn(msg) 

237 

238 if show: 

239 plt.show() 

240 

241 return ax