ma.allequal testing

story645 · story645 · commit 3754dfcf39da · 2016-08-17T23:01:23.000-04:00
diff --git a/lib/matplotlib/category.py b/lib/matplotlib/category.py
@@ -20,9 +20,10 @@
 NP_NEW = (LooseVersion(np.version.version) >= LooseVersion('1.7'))
 
 
-def to_str_array(data, maxlen=100):
+def to_array(data, maxlen=100):
     if NP_NEW:
-        return np.array(data, dtype=np.unicode)
+        return np.array(data)
+    # dtype=np.unicode)
     if cbook.is_scalar_or_string(data):
         data = [data]
     try:
@@ -38,12 +39,11 @@ class StrCategoryConverter(munits.ConversionInterface):
 
     Conversion typically happens in the following order:
     1. default_units:
-        creates unit_data category-integer mapping and binds to axis
+        create unit_data category-integer mapping and binds to axis
     2. axis_info:
-        sets ticks/locator and label/formatter
+        set ticks/locator and labels/formatter
     3. convert:
-        maps input category data to integers using unit_data
-
+        map input category data to integers using unit_data
     """
     @staticmethod
     def convert(value, unit, axis):
@@ -55,7 +55,7 @@ def convert(value, unit, axis):
         if isinstance(value, six.string_types):
             return vmap.get(value, None)
 
-        vals = to_str_array(value)
+        vals = to_array(value)
         for lab, loc in vmap.items():
             vals[vals == lab] = loc
 
@@ -125,7 +125,7 @@ def __init__(self, data):
         self.unit_data = None
         self.units.default_units(data,
                                  self, sort=False)
-
+        self.loc2seq = dict(zip(self.unit_data.locs, self.unit_data.seq))
         self.vmin = min(self.unit_data.locs)
         self.vmax = max(self.unit_data.locs)
 
@@ -139,6 +139,12 @@ def __call__(self, value, clip=None):
         ret /= self.vmax
         return np.ma.array(ret, mask=~mask)
 
+    def inverse(self, value):
+        if not cbook.iterable(value):
+            value = np.asarray(value)
+        vscaled = np.asarray(value) * self.vmax
+        return [self.loc2seq[int(vs)] for vs in vscaled]
+
 
 def colors_from_categories(codings):
     """
@@ -209,7 +215,7 @@ def update(self, new_data, sort=True):
 
     def _set_seq_locs(self, data, value, sort):
         # magic to make it work under np1.6
-        strdata = to_str_array(data)
+        strdata = to_array(data)
 
         # np.unique makes dateframes work
         if sort:
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -313,6 +313,8 @@ def __init__(self, ax, cmap=None,
         if format is None:
             if isinstance(self.norm, colors.LogNorm):
                 self.formatter = ticker.LogFormatterMathtext()
+            elif isinstance(self.norm, category.CategoryNorm):
+                self.formatter = ticker.FixedFormatter(self.norm.unit_data.seq)
             else:
                 self.formatter = ticker.ScalarFormatter()
         elif cbook.is_string_like(format):
@@ -582,7 +584,7 @@ def _ticker(self):
                 elif isinstance(self.norm, colors.LogNorm):
                     locator = ticker.LogLocator()
                 elif isinstance(self.norm, category.CategoryNorm):
-                    locator = ticker.FixedLocator(self.norm.nvals + 0.5)
+                    locator = ticker.FixedLocator(self.norm.unit_data.locs)
                 else:
                     if mpl.rcParams['_internal.classic_mode']:
                         locator = ticker.MaxNLocator()
diff --git a/lib/matplotlib/tests/test_category.py b/lib/matplotlib/tests/test_category.py
@@ -138,7 +138,15 @@ class TestCategoryNorm(object):
     @pytest.mark.parametrize("data, nmap", testdata, ids=ids)
     def test_norm(self, data, nmap):
         norm = cat.CategoryNorm([205, 101, 302])
-        np.testing.assert_allclose(norm(data), nmap)
+        masked_nmap = np.ma.masked_equal(nmap, np.nan)
+        assert np.ma.allequal(norm(data), masked_nmap)
+
+    def test_invert(self):
+        data = [205, 302, 101]
+        strdata = ['205', '302', '101']
+        value = [0, .5, 1]
+        norm = cat.CategoryNorm(data)
+        assert norm.inverse(value) == strdata
 
 
 class TestColorsFromCategories(object):
