Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

# -*- coding: utf-8 -*- 

# Copyright 2011-2018 Kwant authors. 

# 

# This file is part of Kwant. It is subject to the license terms in the file 

# LICENSE.rst found in the top-level directory of this distribution and at 

# http://kwant-project.org/license. A list of Kwant authors can be found in 

# the file AUTHORS.rst at the top-level directory of this distribution and at 

# http://kwant-project.org/authors. 

 

# This module is imported by plotter.py. It contains all the expensive imports 

# that we want to remove from plotter.py 

 

# All matplotlib imports must be isolated in a try, because even without 

# matplotlib iterators remain useful. Further, mpl_toolkits used for 3D 

# plotting are also imported separately, to ensure that 2D plotting works even 

# if 3D does not. 

 

import warnings 

from math import sqrt, pi 

import numpy as np 

 

try: 

import matplotlib 

import matplotlib.colors 

import matplotlib.cm 

from matplotlib.figure import Figure 

from matplotlib import collections 

from . import _colormaps 

mpl_available = True 

try: 

from mpl_toolkits import mplot3d 

has3d = True 

except ImportError: 

warnings.warn("3D plotting not available.", RuntimeWarning) 

has3d = False 

except ImportError: 

warnings.warn("matplotlib is not available, only iterator-providing " 

"functions will work.", RuntimeWarning) 

mpl_available = False 

 

 

# Collections that allow for symbols and linewiths to be given in data space 

# (not for general use, only implement what's needed for plotter) 

def isarray(var): 

if hasattr(var, '__getitem__') and not isinstance(var, str): 

return True 

else: 

return False 

 

 

def nparray_if_array(var): 

return np.asarray(var) if isarray(var) else var 

 

 

55 ↛ exitline 55 didn't exit the module, because the condition on line 55 was never falseif mpl_available: 

class LineCollection(collections.LineCollection): 

def __init__(self, segments, reflen=None, **kwargs): 

super().__init__(segments, **kwargs) 

self.reflen = reflen 

 

def set_linewidths(self, linewidths): 

self.linewidths_orig = nparray_if_array(linewidths) 

 

def draw(self, renderer): 

65 ↛ 71line 65 didn't jump to line 71, because the condition on line 65 was never false if self.reflen is not None: 

# Note: only works for aspect ratio 1! 

# 72.0 - there is 72 points in an inch 

factor = (self.axes.transData.frozen().to_values()[0] * 72.0 * 

self.reflen / self.figure.dpi) 

else: 

factor = 1 

 

super().set_linewidths(self.linewidths_orig * 

factor) 

return super().draw(renderer) 

 

 

class PathCollection(collections.PathCollection): 

def __init__(self, paths, sizes=None, reflen=None, **kwargs): 

super().__init__(paths, sizes=sizes, **kwargs) 

 

self.reflen = reflen 

self.linewidths_orig = nparray_if_array(self.get_linewidths()) 

 

self.transforms = np.array( 

[matplotlib.transforms.Affine2D().scale(x).get_matrix() 

for x in sizes]) 

 

def get_transforms(self): 

return self.transforms 

 

def get_transform(self): 

Affine2D = matplotlib.transforms.Affine2D 

94 ↛ 100line 94 didn't jump to line 100, because the condition on line 94 was never false if self.reflen is not None: 

# For the paths, use the data transformation but strip the 

# offset (will be added later with offsets) 

args = self.axes.transData.frozen().to_values()[:4] + (0, 0) 

return Affine2D().from_values(*args).scale(self.reflen) 

else: 

return Affine2D().scale(self.figure.dpi / 72.0) 

 

def draw(self, renderer): 

103 ↛ 109line 103 didn't jump to line 109, because the condition on line 103 was never false if self.reflen: 

# Note: only works for aspect ratio 1! 

factor = (self.axes.transData.frozen().to_values()[0] / 

self.figure.dpi * 72.0 * self.reflen) 

self.set_linewidths(self.linewidths_orig * factor) 

 

return collections.Collection.draw(self, renderer) 

 

 

112 ↛ exitline 112 didn't exit the module, because the condition on line 112 was never false if has3d: 

# Sorting is optional. 

sort3d = True 

 

# Compute the projection of a 3D length into 2D data coordinates 

# for this we use 2 3D half-circles that are projected into 2D. 

# (This gives the same length as projecting the full unit sphere.) 

 

phi = np.linspace(0, pi, 21) 

xyz = np.c_[np.cos(phi), np.sin(phi), 0 * phi].T.reshape(-1, 1, 21) 

# TODO: use np.block once we depend on numpy >= 1.13. 

unit_sphere = np.vstack([ 

np.hstack([xyz[0], xyz[2]]), 

np.hstack([xyz[1], xyz[0]]), 

np.hstack([xyz[2], xyz[1]]), 

]) 

 

def projected_length(ax, length): 

rc = np.array([ax.get_xlim3d(), ax.get_ylim3d(), ax.get_zlim3d()]) 

rc = np.apply_along_axis(np.sum, 1, rc) / 2. 

 

rs = unit_sphere * length + rc.reshape(-1, 1) 

 

transform = mplot3d.proj3d.proj_transform 

rp = np.asarray(transform(*(list(rs) + [ax.get_proj()]))[:2]) 

rc[:2] = transform(*(list(rc) + [ax.get_proj()]))[:2] 

 

coords = rp - np.repeat(rc[:2].reshape(-1, 1), len(rs[0]), axis=1) 

return sqrt(np.sum(coords**2, axis=0).max()) 

 

 

# Auxiliary array for calculating corners of a cube. 

corners = np.zeros((3, 8, 6), np.float_) 

corners[0, [0, 1, 2, 3], 0] = corners[0, [4, 5, 6, 7], 1] = \ 

corners[0, [0, 1, 4, 5], 2] = corners[0, [2, 3, 6, 7], 3] = \ 

corners[0, [0, 2, 4, 6], 4] = corners[0, [1, 3, 5, 7], 5] = 1.0 

 

 

class Line3DCollection(mplot3d.art3d.Line3DCollection): 

def __init__(self, segments, reflen=None, zorder=0, **kwargs): 

super().__init__(segments, **kwargs) 

self.reflen = reflen 

self.zorder3d = zorder 

 

def set_linewidths(self, linewidths): 

self.linewidths_orig = nparray_if_array(linewidths) 

 

def do_3d_projection(self, renderer): 

super().do_3d_projection(renderer) 

# The whole 3D ordering is flawed in mplot3d when several 

# collections are added. We just use normal zorder. Note the 

# "-" due to the different logic in the 3d plotting, we still 

# want larger zorder values to be plotted on top of smaller 

# ones. 

return -self.zorder3d 

 

def draw(self, renderer): 

169 ↛ 180line 169 didn't jump to line 180, because the condition on line 169 was never false if self.reflen: 

proj_len = projected_length(self.axes, self.reflen) 

args = self.axes.transData.frozen().to_values() 

# Note: unlike in the 2D case, where we can enforce equal 

# aspect ratio, this (currently) does not work with 

# 3D plots in matplotlib. As an approximation, we 

# thus scale with the average of the x- and y-axis 

# transformation. 

factor = proj_len * (args[0] + 

args[3]) * 0.5 * 72.0 / self.figure.dpi 

else: 

factor = 1 

 

super().set_linewidths( 

self.linewidths_orig * factor) 

super().draw(renderer) 

 

 

class Path3DCollection(mplot3d.art3d.Patch3DCollection): 

def __init__(self, paths, sizes, reflen=None, zorder=0, 

offsets=None, **kwargs): 

paths = [matplotlib.patches.PathPatch(path) for path in paths] 

 

192 ↛ 195line 192 didn't jump to line 195, because the condition on line 192 was never false if offsets is not None: 

kwargs['offsets'] = offsets[:, :2] 

 

super().__init__(paths, **kwargs) 

 

197 ↛ 200line 197 didn't jump to line 200, because the condition on line 197 was never false if offsets is not None: 

self.set_3d_properties(zs=offsets[:, 2], zdir="z") 

 

self.reflen = reflen 

self.zorder3d = zorder 

 

self.paths_orig = np.array(paths, dtype='object') 

self.linewidths_orig = nparray_if_array(self.get_linewidths()) 

self.linewidths_orig2 = self.linewidths_orig 

self.array_orig = nparray_if_array(self.get_array()) 

self.facecolors_orig = nparray_if_array(self.get_facecolors()) 

self.edgecolors_orig = nparray_if_array(self.get_edgecolors()) 

 

Affine2D = matplotlib.transforms.Affine2D 

self.orig_transforms = np.array( 

[Affine2D().scale(x).get_matrix() for x in sizes]) 

self.transforms = self.orig_transforms 

 

def set_array(self, array): 

self.array_orig = nparray_if_array(array) 

super().set_array(array) 

 

def set_color(self, colors): 

self.facecolors_orig = nparray_if_array(colors) 

self.edgecolors_orig = self.facecolors_orig 

super().set_color(colors) 

 

def set_edgecolors(self, colors): 

colors = matplotlib.colors.colorConverter.to_rgba_array(colors) 

self.edgecolors_orig = nparray_if_array(colors) 

super().set_edgecolors(colors) 

 

def get_transforms(self): 

# this is exact only for an isometric projection, for the 

# perspective projection used in mplot3d it's an approximation 

return self.transforms 

 

def get_transform(self): 

Affine2D = matplotlib.transforms.Affine2D 

236 ↛ 244line 236 didn't jump to line 244, because the condition on line 236 was never false if self.reflen: 

proj_len = projected_length(self.axes, self.reflen) 

 

# For the paths, use the data transformation but strip the 

# offset (will be added later with the offsets). 

args = self.axes.transData.frozen().to_values()[:4] + (0, 0) 

return Affine2D().from_values(*args).scale(proj_len) 

else: 

return Affine2D().scale(self.figure.dpi / 72.0) 

 

def do_3d_projection(self, renderer): 

xs, ys, zs = self._offsets3d 

 

# numpy complains about zero-length index arrays 

250 ↛ 251line 250 didn't jump to line 251, because the condition on line 250 was never true if len(xs) == 0: 

return -self.zorder3d 

 

proj = mplot3d.proj3d.proj_transform_clip 

vs = np.array(proj(xs, ys, zs, renderer.M)[:3]) 

 

256 ↛ 296line 256 didn't jump to line 296, because the condition on line 256 was never false if sort3d: 

indx = vs[2].argsort()[::-1] 

 

self.set_offsets(vs[:2, indx].T) 

 

261 ↛ 262line 261 didn't jump to line 262, because the condition on line 261 was never true if len(self.paths_orig) > 1: 

paths = np.resize(self.paths_orig, (vs.shape[1],)) 

self.set_paths(paths[indx]) 

 

265 ↛ 266line 265 didn't jump to line 266, because the condition on line 265 was never true if len(self.orig_transforms) > 1: 

self.transforms = self.transforms[indx] 

 

lw_orig = self.linewidths_orig 

269 ↛ 270line 269 didn't jump to line 270, because the condition on line 269 was never true if (isinstance(lw_orig, np.ndarray) and len(lw_orig) > 1): 

self.linewidths_orig2 = np.resize(lw_orig, 

(vs.shape[1],))[indx] 

 

# Note: here array, facecolors and edgecolors are 

# guaranteed to be 2d numpy arrays or None. (And 

# array is the same length as the coordinates) 

 

if self.array_orig is not None: 

super(Path3DCollection, 

self).set_array(self.array_orig[indx]) 

 

if (self.facecolors_orig is not None and 

self.facecolors_orig.shape[0] > 1): 

shape = list(self.facecolors_orig.shape) 

shape[0] = vs.shape[1] 

super().set_facecolors( 

np.resize(self.facecolors_orig, shape)[indx]) 

 

288 ↛ 290line 288 didn't jump to line 290, because the condition on line 288 was never true if (self.edgecolors_orig is not None and 

self.edgecolors_orig.shape[0] > 1): 

shape = list(self.edgecolors_orig.shape) 

shape[0] = vs.shape[1] 

super().set_edgecolors( 

np.resize(self.edgecolors_orig, 

shape)[indx]) 

else: 

self.set_offsets(vs[:2].T) 

 

# the whole 3D ordering is flawed in mplot3d when several 

# collections are added. We just use normal zorder, but correct 

# by the projected z-coord of the "center of gravity", 

# normalized by the projected z-coord of the world coordinates. 

# In doing so, several Path3DCollections are plotted probably 

# in the right order (it's not exact) if they have the same 

# zorder. Still, smaller and larger integer zorders are plotted 

# below or on top. 

 

bbox = np.asarray(self.axes.get_w_lims()) 

 

proj = mplot3d.proj3d.proj_transform_clip 

cz = proj(*(list(np.dot(corners, bbox)) + [renderer.M]))[2] 

 

return -self.zorder3d + vs[2].mean() / cz.ptp() 

 

def draw(self, renderer): 

315 ↛ 323line 315 didn't jump to line 323, because the condition on line 315 was never false if self.reflen: 

proj_len = projected_length(self.axes, self.reflen) 

args = self.axes.transData.frozen().to_values() 

factor = proj_len * (args[0] + 

args[3]) * 0.5 * 72.0 / self.figure.dpi 

 

self.set_linewidths(self.linewidths_orig2 * factor) 

 

super().draw(renderer)