Coverage for pesummary/gw/plots/cmap.py: 88.2%

1# Copyright (C) 2014-2016 Leo Singer, Charlie Hoy 

2# 

3# This program is free software: you can redistribute it and/or modify 

4# it under the terms of the GNU General Public License as published by 

5# the Free Software Foundation, either version 3 of the License, or 

6# (at your option) any later version. 

7# 

8# This program is distributed in the hope that it will be useful, 

9# but WITHOUT ANY WARRANTY; without even the implied warranty of 

10# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

11# GNU General Public License for more details. 

12# 

13# You should have received a copy of the GNU General Public License 

14# along with this program. If not, see <http://www.gnu.org/licenses/>. 

15 

16from matplotlib import colormaps 

17from matplotlib import colors 

18import numpy as np 

19import pesummary 

20 

21 

22def register_cylon(): 

23 # Read in color map RGB data. 

24 path = pesummary.__file__[:-12] 

25 with open(path + "/gw/plots/cylon.csv") as f: 

26 data = np.loadtxt(f, delimiter=',') 

27 

28 # Create color map. 

29 cmap = colors.LinearSegmentedColormap.from_list("cylon", data) 

30 # Assign in module. 

31 locals().update({"cylon": cmap}) 

32 # Register with Matplotlib. 

33 try: 

34 colormaps.register(cmap=cmap) 

35 except ValueError: 

36 # colormap already registered 

37 pass 

38 # Create inverse color map 

39 cmap_r = colors.LinearSegmentedColormap.from_list("cylon_r", data[::-1]) 

40 locals().update({"cylon_r": cmap_r}) 

41 try: 

42 colormaps.register(cmap=cmap_r) 

43 except ValueError: 

44 # colormap already registered 

45 pass 

46 

47 

48def unregister_cylon(): 

49 colormaps.unregister("cylon") 

50 colormaps.unregister("cylon_r") 

51 

52 

53def colormap_with_fixed_hue(color, N=10): 

54 """Create a linear colormap with fixed hue 

55 

56 Parameters 

57 ---------- 

58 color: tuple 

59 color that determines the hue 

60 N: int, optional 

61 number of colors used in the palette 

62 """ 

63 from pesummary.core.plots.palette import color_palette 

64 from matplotlib.colors import LinearSegmentedColormap 

65 from matplotlib.colors import rgb_to_hsv, hsv_to_rgb, hex2color 

66 

67 color_hsv = rgb_to_hsv(hex2color(color)) 

68 base = color_palette("Blues", 10) 

69 base_hsv = np.array(list(map(rgb_to_hsv, base))) 

70 h, s, v = base_hsv.T 

71 

72 h_fixed = np.ones_like(h) * color_hsv[0] 

73 color_array = np.array(list(map( 

74 hsv_to_rgb, np.vstack([h_fixed, s * color_hsv[1], v]).T))) 

75 return LinearSegmentedColormap.from_list("mycmap", color_array)