66from matplotlib .colorbar import ColorbarBase
77
88
9- def _colorbar_extensions (spacing ):
9+ def _get_cmap_norms ():
10+ """
11+ Define a colormap and appropriate norms for each of the four
12+ possible settings of the extend keyword.
1013
14+ Helper function for _colorbar_extension_shape and
15+ colorbar_extension_length.
16+ """
1117 # Create a color map and specify the levels it represents.
1218 cmap = get_cmap ("RdBu" , lut = 5 )
1319 clevs = [- 5. , - 2.5 , - .5 , .5 , 1.5 , 3.5 ]
14-
1520 # Define norms for the color maps.
1621 norms = dict ()
1722 norms ['neither' ] = BoundaryNorm (clevs , len (clevs )- 1 )
1823 norms ['min' ] = BoundaryNorm ([- 10 ]+ clevs [1 :], len (clevs )- 1 )
1924 norms ['max' ] = BoundaryNorm (clevs [:- 1 ]+ [10 ], len (clevs )- 1 )
2025 norms ['both' ] = BoundaryNorm ([- 10 ]+ clevs [1 :- 1 ]+ [10 ], len (clevs )- 1 )
26+ return cmap , norms
27+
28+
29+ def _colorbar_extension_shape (spacing ):
30+ '''
31+ Produce 4 colorbars with rectangular extensions for either uniform
32+ or proportional spacing.
2133
34+ Helper function for test_colorbar_extension_shape.
35+ '''
36+ # Get a colormap and appropriate norms for each extension type.
37+ cmap , norms = _get_cmap_norms ()
38+ # Create a figure and adjust whitespace for subplots.
39+ fig = plt .figure ()
40+ fig .subplots_adjust (hspace = 4 )
41+ for i , extension_type in enumerate (('neither' , 'min' , 'max' , 'both' )):
42+ # Get the appropriate norm and use it to get colorbar boundaries.
43+ norm = norms [extension_type ]
44+ boundaries = values = norm .boundaries
45+ # Create a subplot.
46+ cax = fig .add_subplot (4 , 1 , i + 1 )
47+ # Turn off text and ticks.
48+ for item in cax .get_xticklabels () + cax .get_yticklabels () + \
49+ cax .get_xticklines () + cax .get_yticklines ():
50+ item .set_visible (False )
51+ # Generate the colorbar.
52+ cb = ColorbarBase (cax , cmap = cmap , norm = norm ,
53+ boundaries = boundaries , values = values ,
54+ extend = extension_type , extendrect = True ,
55+ orientation = 'horizontal' , spacing = spacing )
56+ # Return the figure to the caller.
57+ return fig
58+
59+
60+ def _colorbar_extension_length (spacing ):
61+ '''
62+ Produce 12 colorbars with variable length extensions for either
63+ uniform or proportional spacing.
64+
65+ Helper function for test_colorbar_extension_length.
66+ '''
67+ # Get a colormap and appropriate norms for each extension type.
68+ cmap , norms = _get_cmap_norms ()
2269 # Create a figure and adjust whitespace for subplots.
2370 fig = plt .figure ()
2471 fig .subplots_adjust (hspace = .6 )
25-
2672 for i , extension_type in enumerate (('neither' , 'min' , 'max' , 'both' )):
2773 # Get the appropriate norm and use it to get colorbar boundaries.
2874 norm = norms [extension_type ]
@@ -39,20 +85,33 @@ def _colorbar_extensions(spacing):
3985 boundaries = boundaries , values = values ,
4086 extend = extension_type , extendfrac = extendfrac ,
4187 orientation = 'horizontal' , spacing = spacing )
42-
4388 # Return the figure to the caller.
4489 return fig
4590
4691
92+ @image_comparison (
93+ baseline_images = ['colorbar_extensions_shape_uniform' ,
94+ 'colorbar_extensions_shape_proportional' ],
95+ extensions = ['png' ])
96+ def test_colorbar_extension_shape ():
97+ '''Test rectangular colorbar extensions.'''
98+ # Use default params so matplotlibrc doesn't cause the test to fail.
99+ rcParams .update (rcParamsDefault )
100+ # Create figures for uniform and proportionally spaced colorbars.
101+ fig1 = _colorbar_extension_shape ('uniform' )
102+ fig2 = _colorbar_extension_shape ('proportional' )
103+
104+
47105@image_comparison (
48106 baseline_images = ['colorbar_extensions_uniform' , 'colorbar_extensions_proportional' ],
49107 extensions = ['png' ])
50- def test_colorbar_extensions ():
51- # Use default params so .matplotlibrc doesn't cause the test to fail.
108+ def test_colorbar_extension_length ():
109+ '''Test variable length colorbar extensions.'''
110+ # Use default params so matplotlibrc doesn't cause the test to fail.
52111 rcParams .update (rcParamsDefault )
53112 # Create figures for uniform and proportionally spaced colorbars.
54- fig1 = _colorbar_extensions ('uniform' )
55- fig2 = _colorbar_extensions ('proportional' )
113+ fig1 = _colorbar_extension_length ('uniform' )
114+ fig2 = _colorbar_extension_length ('proportional' )
56115
57116
58117if __name__ == '__main__' :
0 commit comments