def test_datalevels_output(self): # Test basic stuffs c = DataLevels(nlevels=2) assert_array_equal(c.levels, [0, 1]) c.set_data([1, 2, 3, 4]) assert_array_equal(c.levels, [1, 4]) c = DataLevels(levels=[1, 2, 3]) assert_array_equal(c.levels, [1, 2, 3]) c = DataLevels(nlevels=10, data=[0, 9]) assert_array_equal(c.levels, np.linspace(0, 9, num=10)) self.assertTrue(c.extend == 'neither') c = DataLevels(nlevels=10, data=[0, 9], vmin=2, vmax=3) assert_array_equal(c.levels, np.linspace(2, 3, num=10)) self.assertTrue(c.extend == 'both') c.set_extend('neither') self.assertTrue(c.extend == 'neither') with warnings.catch_warnings(record=True) as w: # Cause all warnings to always be triggered. warnings.simplefilter("always") # Trigger a warning. out = c.to_rgb() # Verify some things assert len(w) == 2 assert issubclass(w[0].category, RuntimeWarning) assert issubclass(w[1].category, RuntimeWarning) c = DataLevels(nlevels=10, data=[2.5], vmin=2, vmax=3) assert_array_equal(c.levels, np.linspace(2, 3, num=10)) self.assertTrue(c.extend == 'neither') c.update(dict(extend='both')) self.assertTrue(c.extend == 'both') self.assertRaises(AttributeError, c.update, dict(dummy='t')) c = DataLevels(nlevels=10, data=[0, 9], vmax=3) assert_array_equal(c.levels, np.linspace(0, 3, num=10)) self.assertTrue(c.extend == 'max') c = DataLevels(nlevels=10, data=[0, 9], vmin=1) assert_array_equal(c.levels, np.linspace(1, 9, num=10)) self.assertTrue(c.extend == 'min') c = DataLevels(nlevels=10, data=[0, 9], vmin=-1) assert_array_equal(c.levels, np.linspace(-1, 9, num=10)) self.assertTrue(c.extend == 'neither') c.set_plot_params() self.assertTrue(c.extend == 'neither') assert_array_equal(c.vmin, 0) assert_array_equal(c.vmax, 9) c.set_plot_params(vmin=1) assert_array_equal(c.vmin, 1) c.set_data([-12, 8]) assert_array_equal(c.vmin, 1) self.assertTrue(c.extend == 'min') c.set_data([2, 8]) self.assertTrue(c.extend == 'neither') c.set_extend('both') self.assertTrue(c.extend == 'both') c.set_data([3, 3]) self.assertTrue(c.extend == 'both') c.set_extend() self.assertTrue(c.extend == 'neither') # Test the conversion cm = mpl.colors.ListedColormap(['white', 'blue', 'red', 'black']) x = [-1, 0.9, 1.2, 2, 999, 0.8] c = DataLevels(levels=[0, 1, 2], data=x, cmap=cm) r = c.to_rgb() self.assertTrue(len(x) == len(r)) self.assertTrue(c.extend == 'both') assert_array_equal(r, cm([0, 1, 2, 3, 3, 1])) x = [0.9, 1.2] c = DataLevels(levels=[0, 1, 2], data=x, cmap=cm, extend='both') r = c.to_rgb() self.assertTrue(len(x) == len(r)) self.assertTrue(c.extend == 'both') assert_array_equal(r, cm([1, 2])) cm = mpl.colors.ListedColormap(['white', 'blue', 'red']) c = DataLevels(levels=[0, 1, 2], data=x, cmap=cm, extend='min') r = c.to_rgb() self.assertTrue(len(x) == len(r)) assert_array_equal(r, cm([1, 2])) cm = mpl.colors.ListedColormap(['blue', 'red', 'black']) c = DataLevels(levels=[0, 1, 2], data=x, cmap=cm, extend='max') r = c.to_rgb() self.assertTrue(len(x) == len(r)) assert_array_equal(r, cm([0, 1]))