Merge pull request #431 from efiring/collection_log_offset

mdboom · mdboom · commit cb578e409c71 · 2011-08-22T11:46:06.000-07:00
collections: fix bug in scatter with log scale; handle masked offsets
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -101,7 +101,7 @@ def __init__(self,
         # Force _offsets to be Nx2
         self._offsets.shape = (0, 2)
         if offsets is not None:
-            offsets = np.asarray(offsets)
+            offsets = np.asanyarray(offsets)
             offsets.shape = (-1, 2)             # Make it Nx2
             if transOffset is not None:
                 self._offsets = offsets
@@ -160,7 +160,10 @@ def get_datalim(self, transData):
             offsets = transOffset.transform_non_affine(offsets)
             transOffset = transOffset.get_affine()
 
-        offsets = np.asarray(offsets, np.float_)
+        offsets = np.asanyarray(offsets, np.float_)
+        if np.ma.isMaskedArray(offsets):
+            offsets = offsets.filled(np.nan)
+            # get_path_collection_extents handles nan but not masked arrays
         offsets.shape = (-1, 2)                     # Make it Nx2
 
         result = mpath.get_path_collection_extents(
@@ -198,16 +201,22 @@ def _prepare_points(self):
                 ys = self.convert_yunits(offsets[:,1])
                 offsets = zip(xs, ys)
 
-        offsets = np.asarray(offsets, np.float_)
+        offsets = np.asanyarray(offsets, np.float_)
         offsets.shape = (-1, 2)             # Make it Nx2
 
         if not transform.is_affine:
             paths = [transform.transform_path_non_affine(path) for path in paths]
             transform = transform.get_affine()
         if not transOffset.is_affine :
             offsets = transOffset.transform_non_affine(offsets)
+            # This might have changed an ndarray into a masked array.
             transOffset = transOffset.get_affine()
 
+        if np.ma.isMaskedArray(offsets):
+            offsets = offsets.filled(np.nan)
+            # Changing from a masked array to nan-filled ndarray
+            # is probably most efficient at this point.
+
         return transform, transOffset, offsets, paths
 
     @allow_rasterization
@@ -289,7 +298,7 @@ def set_offsets(self, offsets):
 
         ACCEPTS: float or sequence of floats
         """
-        offsets = np.asarray(offsets, np.float_)
+        offsets = np.asanyarray(offsets, np.float_)
         offsets.shape = (-1, 2)             # Make it Nx2
         #This decision is based on how they are initialized above
         if self._uniform_offsets is None:
@@ -588,13 +597,13 @@ def __init__(self, paths, sizes=None, **kwargs):
 
     def set_paths(self, paths):
         self._paths = paths
-        
+
     def get_paths(self):
         return self._paths
-        
+
     def get_sizes(self):
         return self._sizes
-        
+
     @allow_rasterization
     def draw(self, renderer):
         if self._sizes is not None:
