@@ -598,16 +598,33 @@ def set_zsort(self, zsort):
598598 self .stale = True
599599
600600 def get_vector (self , segments3d ):
601- """Optimize points for projection."""
602- if len (segments3d ):
603- xs , ys , zs = np .row_stack (segments3d ).T
604- else : # row_stack can't stack zero arrays.
605- xs , ys , zs = [], [], []
606- ones = np .ones (len (xs ))
607- self ._vec = np .array ([xs , ys , zs , ones ])
601+ """
602+ Optimize points for projection.
608603
609- indices = [0 , * np .cumsum ([len (segment ) for segment in segments3d ])]
610- self ._segslices = [* map (slice , indices [:- 1 ], indices [1 :])]
604+ Parameters
605+ ----------
606+ segments3d : NumPy array or list of NumPy arrays
607+ List of vertices of the boundary of every segment. If all paths are
608+ of equal length and this argument is a NumPy arrray, then it should
609+ be of shape (num_faces, num_vertices, 3).
610+ """
611+ if isinstance (segments3d , np .ndarray ):
612+ if segments3d .ndim != 3 or segments3d .shape [- 1 ] != 3 :
613+ raise ValueError ("segments3d must be a MxNx3 array, but got " +
614+ "shape {}" .format (segments3d .shape ))
615+ self ._segments = segments3d
616+ else :
617+ num_faces = len (segments3d )
618+ num_verts = np .fromiter (map (len , segments3d ), dtype = np .intp )
619+ max_verts = num_verts .max (initial = 0 )
620+ padded = np .empty ((num_faces , max_verts , 3 ))
621+ for i , face in enumerate (segments3d ):
622+ padded [i , :len (face )] = face
623+ mask = np .arange (max_verts ) >= num_verts [:, None ]
624+ mask = mask [..., None ] # add a component axis
625+ # ma.array does not broadcast the mask for us
626+ mask = np .broadcast_to (mask , padded .shape )
627+ self ._segments = np .ma .array (padded , mask = mask )
611628
612629 def set_verts (self , verts , closed = True ):
613630 """Set 3D vertices."""
@@ -649,37 +666,39 @@ def do_3d_projection(self, renderer):
649666 self .update_scalarmappable ()
650667 self ._facecolors3d = self ._facecolors
651668
652- txs , tys , tzs = proj3d ._proj_transform_vec (self ._vec , renderer .M )
653- xyzlist = [( txs [ sl ], tys [ sl ], tzs [ sl ]) for sl in self . _segslices ]
669+ psegments = proj3d ._proj_transform_vectors (self ._segments , renderer .M )
670+ num_faces = len ( psegments )
654671
655672 # This extra fuss is to re-order face / edge colors
656673 cface = self ._facecolors3d
657674 cedge = self ._edgecolors3d
658- if len (cface ) != len ( xyzlist ) :
659- cface = cface .repeat (len ( xyzlist ) , axis = 0 )
660- if len (cedge ) != len ( xyzlist ) :
675+ if len (cface ) != num_faces :
676+ cface = cface .repeat (num_faces , axis = 0 )
677+ if len (cedge ) != num_faces :
661678 if len (cedge ) == 0 :
662679 cedge = cface
663680 else :
664- cedge = cedge .repeat (len ( xyzlist ) , axis = 0 )
681+ cedge = cedge .repeat (num_faces , axis = 0 )
665682
666- # sort by depth (furthest drawn first )
667- z_segments_2d = sorted (
668- (( self . _zsortfunc ( zs ), np . column_stack ([ xs , ys ]), fc , ec , idx )
669- for idx , (( xs , ys , zs ), fc , ec )
670- in enumerate ( zip ( xyzlist , cface , cedge ))),
671- key = lambda x : x [ 0 ], reverse = True )
683+ face_z = self . _zsortfunc ( psegments [..., 2 ], axis = - 1 )
684+ if isinstance ( face_z , np . ma . MaskedArray ):
685+ # NOTE: Unpacking .data is safe here, because every face has to
686+ # contain a valid vertex.
687+ face_z = face_z . data
688+ face_order = np . argsort ( face_z , axis = - 1 )[:: - 1 ]
672689
673- segments_2d = [s for z , s , fc , ec , idx in z_segments_2d ]
674690 if self ._codes3d is not None :
675- codes = [self ._codes3d [idx ] for z , s , fc , ec , idx in z_segments_2d ]
691+ segments_2d = [s .data [~ s .mask .all (axis = - 1 )]
692+ for s in psegments [face_order , :, :2 ]]
693+ codes = [self ._codes3d [idx ] for idx in face_order ]
676694 PolyCollection .set_verts_and_codes (self , segments_2d , codes )
677695 else :
696+ segments_2d = psegments [face_order , :, :2 ]
678697 PolyCollection .set_verts (self , segments_2d , self ._closed )
679698
680- self ._facecolors2d = [ fc for z , s , fc , ec , idx in z_segments_2d ]
699+ self ._facecolors2d = cface [ face_order ]
681700 if len (self ._edgecolors3d ) == len (cface ):
682- self ._edgecolors2d = [ ec for z , s , fc , ec , idx in z_segments_2d ]
701+ self ._edgecolors2d = cedge [ face_order ]
683702 else :
684703 self ._edgecolors2d = self ._edgecolors3d
685704
@@ -688,11 +707,11 @@ def do_3d_projection(self, renderer):
688707 zvec = np .array ([[0 ], [0 ], [self ._sort_zpos ], [1 ]])
689708 ztrans = proj3d ._proj_transform_vec (zvec , renderer .M )
690709 return ztrans [2 ][0 ]
691- elif tzs .size > 0 :
710+ elif psegments .size > 0 :
692711 # FIXME: Some results still don't look quite right.
693712 # In particular, examine contourf3d_demo2.py
694713 # with az = -54 and elev = -45.
695- return np .min (tzs )
714+ return np .min (psegments [..., 2 ] )
696715 else :
697716 return np .nan
698717
0 commit comments