Coverage for pesummary/core/plots/palette.py: 50.0%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3from pesummary.core.plots.seaborn import SEABORN 

4from pesummary.core.plots.seaborn._seaborn_palette import SEABORN_PALETTES 

5from matplotlib import colormaps 

6from matplotlib import colors as mcolors 

7import numpy as np 

8from itertools import cycle 

9 

10__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

11 

12 

13class ColorList(list): 

14 """Class inherited from list to add extra functionality to convert 

15 a list of colors to different formats 

16 

17 Methods 

18 ------- 

19 as_rgb 

20 return a ColorList object with colors in rgb format 

21 as_hex 

22 return a ColorList object with colors in hex format 

23 """ 

24 def as_rgb(self): 

25 return ColorList([mcolors.to_rgb(c) for c in self]) 

26 

27 def as_hex(self): 

28 return ColorList([mcolors.rgb2hex(rgb) for rgb in self]) 

29 

30 

31AVAILABLE_PALETTES = colormaps() + list(SEABORN_PALETTES.keys()) 

32# alternative for color palette from seaborn 

33if SEABORN: 

34 from seaborn import color_palette 

35else: 

36 def color_palette(palette, n_colors=1): 

37 # check to see if provided palette is in matplotlib 

38 if palette in colormaps(): 

39 cmap = colormaps[palette] 

40 colors = ColorList(cmap(np.linspace(0, 1, n_colors)).tolist()) 

41 return colors.as_rgb() 

42 # else check if provided in seaborn palettes 

43 elif palette in SEABORN_PALETTES.keys(): 

44 palette = SEABORN_PALETTES[palette] 

45 palette_cycle = cycle(palette) 

46 colors = ColorList([next(palette_cycle) for _ in range(n_colors)]) 

47 return colors.as_rgb() 

48 else: 

49 raise ValueError( 

50 "Unknown color palette: {}. Available palettes are: {}".format( 

51 palette, ", ".join(AVAILABLE_PALETTES) 

52 ) 

53 )