art3d.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886
  1. # art3d.py, original mplot3d version by John Porter
  2. # Parts rewritten by Reinier Heeres <reinier@heeres.eu>
  3. # Minor additions by Ben Axelrod <baxelrod@coroware.com>
  4. """
  5. Module containing 3D artist code and functions to convert 2D
  6. artists into 3D versions which can be added to an Axes3D.
  7. """
  8. import math
  9. import numpy as np
  10. from matplotlib import (
  11. artist, colors as mcolors, lines, text as mtext, path as mpath)
  12. from matplotlib.collections import (
  13. LineCollection, PolyCollection, PatchCollection, PathCollection)
  14. from matplotlib.colors import Normalize
  15. from matplotlib.patches import Patch
  16. from . import proj3d
  17. def _norm_angle(a):
  18. """Return the given angle normalized to -180 < *a* <= 180 degrees."""
  19. a = (a + 360) % 360
  20. if a > 180:
  21. a = a - 360
  22. return a
  23. def _norm_text_angle(a):
  24. """Return the given angle normalized to -90 < *a* <= 90 degrees."""
  25. a = (a + 180) % 180
  26. if a > 90:
  27. a = a - 180
  28. return a
  29. def get_dir_vector(zdir):
  30. """
  31. Return a direction vector.
  32. Parameters
  33. ----------
  34. zdir : {'x', 'y', 'z', None, 3-tuple}
  35. The direction. Possible values are:
  36. - 'x': equivalent to (1, 0, 0)
  37. - 'y': equivalent to (0, 1, 0)
  38. - 'z': equivalent to (0, 0, 1)
  39. - *None*: equivalent to (0, 0, 0)
  40. - an iterable (x, y, z) is returned unchanged.
  41. Returns
  42. -------
  43. x, y, z : array-like
  44. The direction vector. This is either a numpy.array or *zdir* itself if
  45. *zdir* is already a length-3 iterable.
  46. """
  47. if zdir == 'x':
  48. return np.array((1, 0, 0))
  49. elif zdir == 'y':
  50. return np.array((0, 1, 0))
  51. elif zdir == 'z':
  52. return np.array((0, 0, 1))
  53. elif zdir is None:
  54. return np.array((0, 0, 0))
  55. elif np.iterable(zdir) and len(zdir) == 3:
  56. return zdir
  57. else:
  58. raise ValueError("'x', 'y', 'z', None or vector of length 3 expected")
  59. class Text3D(mtext.Text):
  60. """
  61. Text object with 3D position and direction.
  62. Parameters
  63. ----------
  64. x, y, z
  65. The position of the text.
  66. text : str
  67. The text string to display.
  68. zdir : {'x', 'y', 'z', None, 3-tuple}
  69. The direction of the text. See `.get_dir_vector` for a description of
  70. the values.
  71. Other Parameters
  72. ----------------
  73. **kwargs
  74. All other parameters are passed on to `~matplotlib.text.Text`.
  75. """
  76. def __init__(self, x=0, y=0, z=0, text='', zdir='z', **kwargs):
  77. mtext.Text.__init__(self, x, y, text, **kwargs)
  78. self.set_3d_properties(z, zdir)
  79. def set_3d_properties(self, z=0, zdir='z'):
  80. x, y = self.get_position()
  81. self._position3d = np.array((x, y, z))
  82. self._dir_vec = get_dir_vector(zdir)
  83. self.stale = True
  84. @artist.allow_rasterization
  85. def draw(self, renderer):
  86. proj = proj3d.proj_trans_points(
  87. [self._position3d, self._position3d + self._dir_vec], renderer.M)
  88. dx = proj[0][1] - proj[0][0]
  89. dy = proj[1][1] - proj[1][0]
  90. angle = math.degrees(math.atan2(dy, dx))
  91. self.set_position((proj[0][0], proj[1][0]))
  92. self.set_rotation(_norm_text_angle(angle))
  93. mtext.Text.draw(self, renderer)
  94. self.stale = False
  95. def get_tightbbox(self, renderer):
  96. # Overwriting the 2d Text behavior which is not valid for 3d.
  97. # For now, just return None to exclude from layout calculation.
  98. return None
  99. def text_2d_to_3d(obj, z=0, zdir='z'):
  100. """Convert a Text to a Text3D object."""
  101. obj.__class__ = Text3D
  102. obj.set_3d_properties(z, zdir)
  103. class Line3D(lines.Line2D):
  104. """
  105. 3D line object.
  106. """
  107. def __init__(self, xs, ys, zs, *args, **kwargs):
  108. """
  109. Keyword arguments are passed onto :func:`~matplotlib.lines.Line2D`.
  110. """
  111. lines.Line2D.__init__(self, [], [], *args, **kwargs)
  112. self._verts3d = xs, ys, zs
  113. def set_3d_properties(self, zs=0, zdir='z'):
  114. xs = self.get_xdata()
  115. ys = self.get_ydata()
  116. zs = np.broadcast_to(zs, xs.shape)
  117. self._verts3d = juggle_axes(xs, ys, zs, zdir)
  118. self.stale = True
  119. def set_data_3d(self, *args):
  120. """
  121. Set the x, y and z data
  122. Parameters
  123. ----------
  124. x : array-like
  125. The x-data to be plotted.
  126. y : array-like
  127. The y-data to be plotted.
  128. z : array-like
  129. The z-data to be plotted.
  130. Notes
  131. -----
  132. Accepts x, y, z arguments or a single array-like (x, y, z)
  133. """
  134. if len(args) == 1:
  135. self._verts3d = args[0]
  136. else:
  137. self._verts3d = args
  138. self.stale = True
  139. def get_data_3d(self):
  140. """
  141. Get the current data
  142. Returns
  143. -------
  144. verts3d : length-3 tuple or array-like
  145. The current data as a tuple or array-like.
  146. """
  147. return self._verts3d
  148. @artist.allow_rasterization
  149. def draw(self, renderer):
  150. xs3d, ys3d, zs3d = self._verts3d
  151. xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
  152. self.set_data(xs, ys)
  153. lines.Line2D.draw(self, renderer)
  154. self.stale = False
  155. def line_2d_to_3d(line, zs=0, zdir='z'):
  156. """Convert a 2D line to 3D."""
  157. line.__class__ = Line3D
  158. line.set_3d_properties(zs, zdir)
  159. def _path_to_3d_segment(path, zs=0, zdir='z'):
  160. """Convert a path to a 3D segment."""
  161. zs = np.broadcast_to(zs, len(path))
  162. pathsegs = path.iter_segments(simplify=False, curves=False)
  163. seg = [(x, y, z) for (((x, y), code), z) in zip(pathsegs, zs)]
  164. seg3d = [juggle_axes(x, y, z, zdir) for (x, y, z) in seg]
  165. return seg3d
  166. def _paths_to_3d_segments(paths, zs=0, zdir='z'):
  167. """Convert paths from a collection object to 3D segments."""
  168. zs = np.broadcast_to(zs, len(paths))
  169. segs = [_path_to_3d_segment(path, pathz, zdir)
  170. for path, pathz in zip(paths, zs)]
  171. return segs
  172. def _path_to_3d_segment_with_codes(path, zs=0, zdir='z'):
  173. """Convert a path to a 3D segment with path codes."""
  174. zs = np.broadcast_to(zs, len(path))
  175. pathsegs = path.iter_segments(simplify=False, curves=False)
  176. seg_codes = [((x, y, z), code) for ((x, y), code), z in zip(pathsegs, zs)]
  177. if seg_codes:
  178. seg, codes = zip(*seg_codes)
  179. seg3d = [juggle_axes(x, y, z, zdir) for (x, y, z) in seg]
  180. else:
  181. seg3d = []
  182. codes = []
  183. return seg3d, list(codes)
  184. def _paths_to_3d_segments_with_codes(paths, zs=0, zdir='z'):
  185. """
  186. Convert paths from a collection object to 3D segments with path codes.
  187. """
  188. zs = np.broadcast_to(zs, len(paths))
  189. segments_codes = [_path_to_3d_segment_with_codes(path, pathz, zdir)
  190. for path, pathz in zip(paths, zs)]
  191. if segments_codes:
  192. segments, codes = zip(*segments_codes)
  193. else:
  194. segments, codes = [], []
  195. return list(segments), list(codes)
  196. class Line3DCollection(LineCollection):
  197. """
  198. A collection of 3D lines.
  199. """
  200. def set_sort_zpos(self, val):
  201. """Set the position to use for z-sorting."""
  202. self._sort_zpos = val
  203. self.stale = True
  204. def set_segments(self, segments):
  205. """
  206. Set 3D segments.
  207. """
  208. self._segments3d = segments
  209. LineCollection.set_segments(self, [])
  210. def do_3d_projection(self, renderer):
  211. """
  212. Project the points according to renderer matrix.
  213. """
  214. # see _update_scalarmappable docstring for why this must be here
  215. _update_scalarmappable(self)
  216. xyslist = [
  217. proj3d.proj_trans_points(points, renderer.M) for points in
  218. self._segments3d]
  219. segments_2d = [np.column_stack([xs, ys]) for xs, ys, zs in xyslist]
  220. LineCollection.set_segments(self, segments_2d)
  221. # FIXME
  222. minz = 1e9
  223. for xs, ys, zs in xyslist:
  224. minz = min(minz, min(zs))
  225. return minz
  226. @artist.allow_rasterization
  227. def draw(self, renderer, project=False):
  228. if project:
  229. self.do_3d_projection(renderer)
  230. LineCollection.draw(self, renderer)
  231. def line_collection_2d_to_3d(col, zs=0, zdir='z'):
  232. """Convert a LineCollection to a Line3DCollection object."""
  233. segments3d = _paths_to_3d_segments(col.get_paths(), zs, zdir)
  234. col.__class__ = Line3DCollection
  235. col.set_segments(segments3d)
  236. class Patch3D(Patch):
  237. """
  238. 3D patch object.
  239. """
  240. def __init__(self, *args, zs=(), zdir='z', **kwargs):
  241. Patch.__init__(self, *args, **kwargs)
  242. self.set_3d_properties(zs, zdir)
  243. def set_3d_properties(self, verts, zs=0, zdir='z'):
  244. zs = np.broadcast_to(zs, len(verts))
  245. self._segment3d = [juggle_axes(x, y, z, zdir)
  246. for ((x, y), z) in zip(verts, zs)]
  247. self._facecolor3d = Patch.get_facecolor(self)
  248. def get_path(self):
  249. return self._path2d
  250. def get_facecolor(self):
  251. return self._facecolor2d
  252. def do_3d_projection(self, renderer):
  253. s = self._segment3d
  254. xs, ys, zs = zip(*s)
  255. vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
  256. self._path2d = mpath.Path(np.column_stack([vxs, vys]))
  257. # FIXME: coloring
  258. self._facecolor2d = self._facecolor3d
  259. return min(vzs)
  260. class PathPatch3D(Patch3D):
  261. """
  262. 3D PathPatch object.
  263. """
  264. def __init__(self, path, *, zs=(), zdir='z', **kwargs):
  265. Patch.__init__(self, **kwargs)
  266. self.set_3d_properties(path, zs, zdir)
  267. def set_3d_properties(self, path, zs=0, zdir='z'):
  268. Patch3D.set_3d_properties(self, path.vertices, zs=zs, zdir=zdir)
  269. self._code3d = path.codes
  270. def do_3d_projection(self, renderer):
  271. s = self._segment3d
  272. xs, ys, zs = zip(*s)
  273. vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
  274. self._path2d = mpath.Path(np.column_stack([vxs, vys]), self._code3d)
  275. # FIXME: coloring
  276. self._facecolor2d = self._facecolor3d
  277. return min(vzs)
  278. def _get_patch_verts(patch):
  279. """Return a list of vertices for the path of a patch."""
  280. trans = patch.get_patch_transform()
  281. path = patch.get_path()
  282. polygons = path.to_polygons(trans)
  283. if len(polygons):
  284. return polygons[0]
  285. else:
  286. return []
  287. def patch_2d_to_3d(patch, z=0, zdir='z'):
  288. """Convert a Patch to a Patch3D object."""
  289. verts = _get_patch_verts(patch)
  290. patch.__class__ = Patch3D
  291. patch.set_3d_properties(verts, z, zdir)
  292. def pathpatch_2d_to_3d(pathpatch, z=0, zdir='z'):
  293. """Convert a PathPatch to a PathPatch3D object."""
  294. path = pathpatch.get_path()
  295. trans = pathpatch.get_patch_transform()
  296. mpath = trans.transform_path(path)
  297. pathpatch.__class__ = PathPatch3D
  298. pathpatch.set_3d_properties(mpath, z, zdir)
  299. class Patch3DCollection(PatchCollection):
  300. """
  301. A collection of 3D patches.
  302. """
  303. def __init__(self, *args, zs=0, zdir='z', depthshade=True, **kwargs):
  304. """
  305. Create a collection of flat 3D patches with its normal vector
  306. pointed in *zdir* direction, and located at *zs* on the *zdir*
  307. axis. 'zs' can be a scalar or an array-like of the same length as
  308. the number of patches in the collection.
  309. Constructor arguments are the same as for
  310. :class:`~matplotlib.collections.PatchCollection`. In addition,
  311. keywords *zs=0* and *zdir='z'* are available.
  312. Also, the keyword argument "depthshade" is available to
  313. indicate whether or not to shade the patches in order to
  314. give the appearance of depth (default is *True*).
  315. This is typically desired in scatter plots.
  316. """
  317. self._depthshade = depthshade
  318. super().__init__(*args, **kwargs)
  319. self.set_3d_properties(zs, zdir)
  320. def set_sort_zpos(self, val):
  321. """Set the position to use for z-sorting."""
  322. self._sort_zpos = val
  323. self.stale = True
  324. def set_3d_properties(self, zs, zdir):
  325. # Force the collection to initialize the face and edgecolors
  326. # just in case it is a scalarmappable with a colormap.
  327. self.update_scalarmappable()
  328. offsets = self.get_offsets()
  329. if len(offsets) > 0:
  330. xs, ys = offsets.T
  331. else:
  332. xs = []
  333. ys = []
  334. self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir)
  335. self._facecolor3d = self.get_facecolor()
  336. self._edgecolor3d = self.get_edgecolor()
  337. self.stale = True
  338. def do_3d_projection(self, renderer):
  339. # see _update_scalarmappable docstring for why this must be here
  340. _update_scalarmappable(self)
  341. xs, ys, zs = self._offsets3d
  342. vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
  343. fcs = (_zalpha(self._facecolor3d, vzs) if self._depthshade else
  344. self._facecolor3d)
  345. fcs = mcolors.to_rgba_array(fcs, self._alpha)
  346. self.set_facecolors(fcs)
  347. ecs = (_zalpha(self._edgecolor3d, vzs) if self._depthshade else
  348. self._edgecolor3d)
  349. ecs = mcolors.to_rgba_array(ecs, self._alpha)
  350. self.set_edgecolors(ecs)
  351. PatchCollection.set_offsets(self, np.column_stack([vxs, vys]))
  352. if vzs.size > 0:
  353. return min(vzs)
  354. else:
  355. return np.nan
  356. class Path3DCollection(PathCollection):
  357. """
  358. A collection of 3D paths.
  359. """
  360. def __init__(self, *args, zs=0, zdir='z', depthshade=True, **kwargs):
  361. """
  362. Create a collection of flat 3D paths with its normal vector
  363. pointed in *zdir* direction, and located at *zs* on the *zdir*
  364. axis. 'zs' can be a scalar or an array-like of the same length as
  365. the number of paths in the collection.
  366. Constructor arguments are the same as for
  367. :class:`~matplotlib.collections.PathCollection`. In addition,
  368. keywords *zs=0* and *zdir='z'* are available.
  369. Also, the keyword argument "depthshade" is available to
  370. indicate whether or not to shade the patches in order to
  371. give the appearance of depth (default is *True*).
  372. This is typically desired in scatter plots.
  373. """
  374. self._depthshade = depthshade
  375. super().__init__(*args, **kwargs)
  376. self.set_3d_properties(zs, zdir)
  377. def set_sort_zpos(self, val):
  378. """Set the position to use for z-sorting."""
  379. self._sort_zpos = val
  380. self.stale = True
  381. def set_3d_properties(self, zs, zdir):
  382. # Force the collection to initialize the face and edgecolors
  383. # just in case it is a scalarmappable with a colormap.
  384. self.update_scalarmappable()
  385. offsets = self.get_offsets()
  386. if len(offsets) > 0:
  387. xs, ys = offsets.T
  388. else:
  389. xs = []
  390. ys = []
  391. self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir)
  392. self._facecolor3d = self.get_facecolor()
  393. self._edgecolor3d = self.get_edgecolor()
  394. self._sizes3d = self.get_sizes()
  395. self._linewidth3d = self.get_linewidth()
  396. self.stale = True
  397. def do_3d_projection(self, renderer):
  398. # see _update_scalarmappable docstring for why this must be here
  399. _update_scalarmappable(self)
  400. xs, ys, zs = self._offsets3d
  401. vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
  402. fcs = (_zalpha(self._facecolor3d, vzs) if self._depthshade else
  403. self._facecolor3d)
  404. ecs = (_zalpha(self._edgecolor3d, vzs) if self._depthshade else
  405. self._edgecolor3d)
  406. sizes = self._sizes3d
  407. lws = self._linewidth3d
  408. # Sort the points based on z coordinates
  409. # Performance optimization: Create a sorted index array and reorder
  410. # points and point properties according to the index array
  411. z_markers_idx = np.argsort(vzs)[::-1]
  412. # Re-order items
  413. vzs = vzs[z_markers_idx]
  414. vxs = vxs[z_markers_idx]
  415. vys = vys[z_markers_idx]
  416. if len(fcs) > 1:
  417. fcs = fcs[z_markers_idx]
  418. if len(ecs) > 1:
  419. ecs = ecs[z_markers_idx]
  420. if len(sizes) > 1:
  421. sizes = sizes[z_markers_idx]
  422. if len(lws) > 1:
  423. lws = lws[z_markers_idx]
  424. vps = np.column_stack((vxs, vys))
  425. fcs = mcolors.to_rgba_array(fcs, self._alpha)
  426. ecs = mcolors.to_rgba_array(ecs, self._alpha)
  427. self.set_edgecolors(ecs)
  428. self.set_facecolors(fcs)
  429. self.set_sizes(sizes)
  430. self.set_linewidth(lws)
  431. PathCollection.set_offsets(self, vps)
  432. return np.min(vzs) if vzs.size else np.nan
  433. def _update_scalarmappable(sm):
  434. """
  435. Update a 3D ScalarMappable.
  436. With ScalarMappable objects if the data, colormap, or norm are
  437. changed, we need to update the computed colors. This is handled
  438. by the base class method update_scalarmappable. This method works
  439. by, detecting if work needs to be done, and if so stashing it on
  440. the ``self._facecolors`` attribute.
  441. With 3D collections we internally sort the components so that
  442. things that should be "in front" are rendered later to simulate
  443. having a z-buffer (in addition to doing the projections). This is
  444. handled in the ``do_3d_projection`` methods which are called from the
  445. draw method of the 3D Axes. These methods:
  446. - do the projection from 3D -> 2D
  447. - internally sort based on depth
  448. - stash the results of the above in the 2D analogs of state
  449. - return the z-depth of the whole artist
  450. the last step is so that we can, at the Axes level, sort the children by
  451. depth.
  452. The base `draw` method of the 2D artists unconditionally calls
  453. update_scalarmappable and rely on the method's internal caching logic to
  454. lazily evaluate.
  455. These things together mean you can have the sequence of events:
  456. - we create the artist, do the color mapping and stash the results
  457. in a 3D specific state.
  458. - change something about the ScalarMappable that marks it as in
  459. need of an update (`ScalarMappable.changed` and friends).
  460. - We call do_3d_projection and shuffle the stashed colors into the
  461. 2D version of face colors
  462. - the draw method calls the update_scalarmappable method which
  463. overwrites our shuffled colors
  464. - we get a render that is wrong
  465. - if we re-render (either with a second save or implicitly via
  466. tight_layout / constrained_layout / bbox_inches='tight' (ex via
  467. inline's defaults)) we again shuffle the 3D colors
  468. - because the CM is not marked as changed update_scalarmappable is
  469. a no-op and we get a correct looking render.
  470. This function is an internal helper to:
  471. - sort out if we need to do the color mapping at all (has data!)
  472. - sort out if update_scalarmappable is going to be a no-op
  473. - copy the data over from the 2D -> 3D version
  474. This must be called first thing in do_3d_projection to make sure that
  475. the correct colors get shuffled.
  476. Parameters
  477. ----------
  478. sm : ScalarMappable
  479. The ScalarMappable to update and stash the 3D data from
  480. """
  481. if sm._A is None:
  482. return
  483. copy_state = sm._update_dict['array']
  484. ret = sm.update_scalarmappable()
  485. if copy_state:
  486. if sm._is_filled:
  487. sm._facecolor3d = sm._facecolors
  488. elif sm._is_stroked:
  489. sm._edgecolor3d = sm._edgecolors
  490. def patch_collection_2d_to_3d(col, zs=0, zdir='z', depthshade=True):
  491. """
  492. Convert a :class:`~matplotlib.collections.PatchCollection` into a
  493. :class:`Patch3DCollection` object
  494. (or a :class:`~matplotlib.collections.PathCollection` into a
  495. :class:`Path3DCollection` object).
  496. Parameters
  497. ----------
  498. za
  499. The location or locations to place the patches in the collection along
  500. the *zdir* axis. Default: 0.
  501. zdir
  502. The axis in which to place the patches. Default: "z".
  503. depthshade
  504. Whether to shade the patches to give a sense of depth. Default: *True*.
  505. """
  506. if isinstance(col, PathCollection):
  507. col.__class__ = Path3DCollection
  508. elif isinstance(col, PatchCollection):
  509. col.__class__ = Patch3DCollection
  510. col._depthshade = depthshade
  511. col.set_3d_properties(zs, zdir)
  512. class Poly3DCollection(PolyCollection):
  513. """
  514. A collection of 3D polygons.
  515. .. note::
  516. **Filling of 3D polygons**
  517. There is no simple definition of the enclosed surface of a 3D polygon
  518. unless the polygon is planar.
  519. In practice, Matplotlib fills the 2D projection of the polygon. This
  520. gives a correct filling appearance only for planar polygons. For all
  521. other polygons, you'll find orientations in which the edges of the
  522. polygon intersect in the projection. This will lead to an incorrect
  523. visualization of the 3D area.
  524. If you need filled areas, it is recommended to create them via
  525. `~mpl_toolkits.mplot3d.axes3d.Axes3D.plot_trisurf`, which creates a
  526. triangulation and thus generates consistent surfaces.
  527. """
  528. def __init__(self, verts, *args, zsort='average', **kwargs):
  529. """
  530. Parameters
  531. ----------
  532. verts : list of array-like Nx3
  533. Each element describes a polygon as a sequence of ``N_i`` points
  534. ``(x, y, z)``.
  535. zsort : {'average', 'min', 'max'}, default: 'average'
  536. The calculation method for the z-order.
  537. See `~.Poly3DCollection.set_zsort` for details.
  538. *args, **kwargs
  539. All other parameters are forwarded to `.PolyCollection`.
  540. Notes
  541. -----
  542. Note that this class does a bit of magic with the _facecolors
  543. and _edgecolors properties.
  544. """
  545. super().__init__(verts, *args, **kwargs)
  546. self.set_zsort(zsort)
  547. self._codes3d = None
  548. _zsort_functions = {
  549. 'average': np.average,
  550. 'min': np.min,
  551. 'max': np.max,
  552. }
  553. def set_zsort(self, zsort):
  554. """
  555. Set the calculation method for the z-order.
  556. Parameters
  557. ----------
  558. zsort : {'average', 'min', 'max'}
  559. The function applied on the z-coordinates of the vertices in the
  560. viewer's coordinate system, to determine the z-order.
  561. """
  562. self._zsortfunc = self._zsort_functions[zsort]
  563. self._sort_zpos = None
  564. self.stale = True
  565. def get_vector(self, segments3d):
  566. """Optimize points for projection."""
  567. if len(segments3d):
  568. xs, ys, zs = np.row_stack(segments3d).T
  569. else: # row_stack can't stack zero arrays.
  570. xs, ys, zs = [], [], []
  571. ones = np.ones(len(xs))
  572. self._vec = np.array([xs, ys, zs, ones])
  573. indices = [0, *np.cumsum([len(segment) for segment in segments3d])]
  574. self._segslices = [*map(slice, indices[:-1], indices[1:])]
  575. def set_verts(self, verts, closed=True):
  576. """Set 3D vertices."""
  577. self.get_vector(verts)
  578. # 2D verts will be updated at draw time
  579. PolyCollection.set_verts(self, [], False)
  580. self._closed = closed
  581. def set_verts_and_codes(self, verts, codes):
  582. """Set 3D vertices with path codes."""
  583. # set vertices with closed=False to prevent PolyCollection from
  584. # setting path codes
  585. self.set_verts(verts, closed=False)
  586. # and set our own codes instead.
  587. self._codes3d = codes
  588. def set_3d_properties(self):
  589. # Force the collection to initialize the face and edgecolors
  590. # just in case it is a scalarmappable with a colormap.
  591. self.update_scalarmappable()
  592. self._sort_zpos = None
  593. self.set_zsort('average')
  594. self._facecolor3d = PolyCollection.get_facecolor(self)
  595. self._edgecolor3d = PolyCollection.get_edgecolor(self)
  596. self._alpha3d = PolyCollection.get_alpha(self)
  597. self.stale = True
  598. def set_sort_zpos(self, val):
  599. """Set the position to use for z-sorting."""
  600. self._sort_zpos = val
  601. self.stale = True
  602. def do_3d_projection(self, renderer):
  603. """
  604. Perform the 3D projection for this object.
  605. """
  606. # see _update_scalarmappable docstring for why this must be here
  607. _update_scalarmappable(self)
  608. txs, tys, tzs = proj3d._proj_transform_vec(self._vec, renderer.M)
  609. xyzlist = [(txs[sl], tys[sl], tzs[sl]) for sl in self._segslices]
  610. # This extra fuss is to re-order face / edge colors
  611. cface = self._facecolor3d
  612. cedge = self._edgecolor3d
  613. if len(cface) != len(xyzlist):
  614. cface = cface.repeat(len(xyzlist), axis=0)
  615. if len(cedge) != len(xyzlist):
  616. if len(cedge) == 0:
  617. cedge = cface
  618. else:
  619. cedge = cedge.repeat(len(xyzlist), axis=0)
  620. # sort by depth (furthest drawn first)
  621. z_segments_2d = sorted(
  622. ((self._zsortfunc(zs), np.column_stack([xs, ys]), fc, ec, idx)
  623. for idx, ((xs, ys, zs), fc, ec)
  624. in enumerate(zip(xyzlist, cface, cedge))),
  625. key=lambda x: x[0], reverse=True)
  626. zzs, segments_2d, self._facecolors2d, self._edgecolors2d, idxs = \
  627. zip(*z_segments_2d)
  628. if self._codes3d is not None:
  629. codes = [self._codes3d[idx] for idx in idxs]
  630. PolyCollection.set_verts_and_codes(self, segments_2d, codes)
  631. else:
  632. PolyCollection.set_verts(self, segments_2d, self._closed)
  633. if len(self._edgecolor3d) != len(cface):
  634. self._edgecolors2d = self._edgecolor3d
  635. # Return zorder value
  636. if self._sort_zpos is not None:
  637. zvec = np.array([[0], [0], [self._sort_zpos], [1]])
  638. ztrans = proj3d._proj_transform_vec(zvec, renderer.M)
  639. return ztrans[2][0]
  640. elif tzs.size > 0:
  641. # FIXME: Some results still don't look quite right.
  642. # In particular, examine contourf3d_demo2.py
  643. # with az = -54 and elev = -45.
  644. return np.min(tzs)
  645. else:
  646. return np.nan
  647. def set_facecolor(self, colors):
  648. PolyCollection.set_facecolor(self, colors)
  649. self._facecolor3d = PolyCollection.get_facecolor(self)
  650. def set_edgecolor(self, colors):
  651. PolyCollection.set_edgecolor(self, colors)
  652. self._edgecolor3d = PolyCollection.get_edgecolor(self)
  653. def set_alpha(self, alpha):
  654. # docstring inherited
  655. artist.Artist.set_alpha(self, alpha)
  656. try:
  657. self._facecolor3d = mcolors.to_rgba_array(
  658. self._facecolor3d, self._alpha)
  659. except (AttributeError, TypeError, IndexError):
  660. pass
  661. try:
  662. self._edgecolors = mcolors.to_rgba_array(
  663. self._edgecolor3d, self._alpha)
  664. except (AttributeError, TypeError, IndexError):
  665. pass
  666. self.stale = True
  667. def get_facecolor(self):
  668. return self._facecolors2d
  669. def get_edgecolor(self):
  670. return self._edgecolors2d
  671. def poly_collection_2d_to_3d(col, zs=0, zdir='z'):
  672. """Convert a PolyCollection to a Poly3DCollection object."""
  673. segments_3d, codes = _paths_to_3d_segments_with_codes(
  674. col.get_paths(), zs, zdir)
  675. col.__class__ = Poly3DCollection
  676. col.set_verts_and_codes(segments_3d, codes)
  677. col.set_3d_properties()
  678. def juggle_axes(xs, ys, zs, zdir):
  679. """
  680. Reorder coordinates so that 2D xs, ys can be plotted in the plane
  681. orthogonal to zdir. zdir is normally x, y or z. However, if zdir
  682. starts with a '-' it is interpreted as a compensation for rotate_axes.
  683. """
  684. if zdir == 'x':
  685. return zs, xs, ys
  686. elif zdir == 'y':
  687. return xs, zs, ys
  688. elif zdir[0] == '-':
  689. return rotate_axes(xs, ys, zs, zdir)
  690. else:
  691. return xs, ys, zs
  692. def rotate_axes(xs, ys, zs, zdir):
  693. """
  694. Reorder coordinates so that the axes are rotated with zdir along
  695. the original z axis. Prepending the axis with a '-' does the
  696. inverse transform, so zdir can be x, -x, y, -y, z or -z
  697. """
  698. if zdir == 'x':
  699. return ys, zs, xs
  700. elif zdir == '-x':
  701. return zs, xs, ys
  702. elif zdir == 'y':
  703. return zs, xs, ys
  704. elif zdir == '-y':
  705. return ys, zs, xs
  706. else:
  707. return xs, ys, zs
  708. def _get_colors(c, num):
  709. """Stretch the color argument to provide the required number *num*."""
  710. return np.broadcast_to(
  711. mcolors.to_rgba_array(c) if len(c) else [0, 0, 0, 0],
  712. (num, 4))
  713. def _zalpha(colors, zs):
  714. """Modify the alphas of the color list according to depth."""
  715. # FIXME: This only works well if the points for *zs* are well-spaced
  716. # in all three dimensions. Otherwise, at certain orientations,
  717. # the min and max zs are very close together.
  718. # Should really normalize against the viewing depth.
  719. if len(colors) == 0 or len(zs) == 0:
  720. return np.zeros((0, 4))
  721. norm = Normalize(min(zs), max(zs))
  722. sats = 1 - norm(zs) * 0.7
  723. rgba = np.broadcast_to(mcolors.to_rgba_array(colors), (len(zs), 4))
  724. return np.column_stack([rgba[:, :3], rgba[:, 3] * sats])