@@ -166,6 +166,37 @@ def paths_to_3d_segments(paths, zs=0, zdir='z'):
166
166
segments .append (path_to_3d_segment (path , pathz , zdir ))
167
167
return segments
168
168
169
+ def path_to_3d_segment_with_codes (path , zs = 0 , zdir = 'z' ):
170
+ '''Convert a path to a 3D segment with path codes.'''
171
+
172
+ if not iterable (zs ):
173
+ zs = np .ones (len (path )) * zs
174
+
175
+ seg = []
176
+ codes = []
177
+ pathsegs = path .iter_segments (simplify = False , curves = False )
178
+ for (((x , y ), code ), z ) in zip (pathsegs , zs ):
179
+ seg .append ((x , y , z ))
180
+ codes .append (code )
181
+ seg3d = [juggle_axes (x , y , z , zdir ) for (x , y , z ) in seg ]
182
+ return seg3d , codes
183
+
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
+
189
+ if not iterable (zs ):
190
+ zs = np .ones (len (paths )) * zs
191
+
192
+ segments = []
193
+ codes_list = []
194
+ for path , pathz in zip (paths , zs ):
195
+ segs , codes = path_to_3d_segment_with_codes (path , pathz , zdir )
196
+ segments .append (segs )
197
+ codes_list .append (codes )
198
+ return segments , codes_list
199
+
169
200
class Line3DCollection (LineCollection ):
170
201
'''
171
202
A collection of 3D lines.
@@ -487,6 +518,7 @@ def __init__(self, verts, *args, **kwargs):
487
518
zsort = kwargs .pop ('zsort' , True )
488
519
PolyCollection .__init__ (self , verts , * args , ** kwargs )
489
520
self .set_zsort (zsort )
521
+ self ._codes3d = None
490
522
491
523
_zsort_functions = {
492
524
'average' : np .average ,
@@ -545,6 +577,14 @@ def set_verts(self, verts, closed=True):
545
577
# 2D verts will be updated at draw time
546
578
PolyCollection .set_verts (self , [], closed )
547
579
580
+ def set_verts_and_codes (self , verts , codes ):
581
+ '''Sets 3D vertices with path codes'''
582
+ # set vertices with closed=False to prevent PolyCollection from
583
+ # setting path codes
584
+ self .set_verts (verts , closed = False )
585
+ # and set our own codes instead.
586
+ self ._codes3d = codes
587
+
548
588
def set_3d_properties (self ):
549
589
# Force the collection to initialize the face and edgecolors
550
590
# just in case it is a scalarmappable with a colormap.
@@ -571,8 +611,8 @@ def do_3d_projection(self, renderer):
571
611
self ._facecolors3d = self ._facecolors
572
612
573
613
txs , tys , tzs = proj3d .proj_transform_vec (self ._vec , renderer .M )
574
- xyzlist = [(txs [si :ei ], tys [si :ei ], tzs [si :ei ]) \
575
- for si , ei in self ._segis ]
614
+ xyzlist = [(txs [si :ei ], tys [si :ei ], tzs [si :ei ])
615
+ for si , ei in self ._segis ]
576
616
577
617
# This extra fuss is to re-order face / edge colors
578
618
cface = self ._facecolors3d
@@ -586,18 +626,24 @@ def do_3d_projection(self, renderer):
586
626
587
627
# if required sort by depth (furthest drawn first)
588
628
if self ._zsort :
589
- z_segments_2d = [(self ._zsortfunc (zs ), list (zip (xs , ys )), fc , ec ) for
590
- (xs , ys , zs ), fc , ec in zip (xyzlist , cface , cedge )]
629
+ indices = range (len (xyzlist ))
630
+ z_segments_2d = [(self ._zsortfunc (zs ), list (zip (xs , ys )), fc , ec ,
631
+ idx ) for (xs , ys , zs ), fc , ec , idx in
632
+ zip (xyzlist , cface , cedge , indices )]
591
633
z_segments_2d .sort (key = lambda x : x [0 ], reverse = True )
592
634
else :
593
635
raise ValueError ("whoops" )
594
636
595
- segments_2d = [s for z , s , fc , ec in z_segments_2d ]
596
- PolyCollection .set_verts (self , segments_2d )
637
+ segments_2d = [s for z , s , fc , ec , idx in z_segments_2d ]
638
+ if self ._codes3d is not None :
639
+ codes = [self ._codes3d [idx ] for z , s , fc , ec , idx in z_segments_2d ]
640
+ PolyCollection .set_verts_and_codes (self , segments_2d , codes )
641
+ else :
642
+ PolyCollection .set_verts (self , segments_2d )
597
643
598
- self ._facecolors2d = [fc for z , s , fc , ec in z_segments_2d ]
644
+ self ._facecolors2d = [fc for z , s , fc , ec , idx in z_segments_2d ]
599
645
if len (self ._edgecolors3d ) == len (cface ):
600
- self ._edgecolors2d = [ec for z , s , fc , ec in z_segments_2d ]
646
+ self ._edgecolors2d = [ec for z , s , fc , ec , idx in z_segments_2d ]
601
647
else :
602
648
self ._edgecolors2d = self ._edgecolors3d
603
649
@@ -663,9 +709,10 @@ def draw(self, renderer):
663
709
664
710
def poly_collection_2d_to_3d (col , zs = 0 , zdir = 'z' ):
665
711
"""Convert a PolyCollection to a Poly3DCollection object."""
666
- segments_3d = paths_to_3d_segments (col .get_paths (), zs , zdir )
712
+ segments_3d , codes = paths_to_3d_segments_with_codes (col .get_paths (),
713
+ zs , zdir )
667
714
col .__class__ = Poly3DCollection
668
- col .set_verts (segments_3d )
715
+ col .set_verts_and_codes (segments_3d , codes )
669
716
col .set_3d_properties ()
670
717
671
718
def juggle_axes (xs , ys , zs , zdir ):
0 commit comments