Coverage for bilby/core/utils/plotting.py: 91%

1import functools 

2import os 

3import shutil 

4 

5from .log import logger 

6 

7 

8def latex_plot_format(func): 

9 """ 

10 Wrap the plotting function to set rcParams dependent on environment variables 

11 

12 The rcparams can be set directly from the env. variable `BILBY_STYLE` to 

13 point to a matplotlib style file. Or, if `BILBY_STYLE=default` (any case) a 

14 default setup is used, this is enabled by default. To not use any rcParams, 

15 set `BILBY_STYLE=none`. Occasionally, issues arrise with the latex 

16 `mathdefault` command. A fix is to define this command in the rcParams. An 

17 env. variable `BILBY_MATHDEFAULT` can be used to turn this fix on/off. 

18 Setting `BILBY_MATHDEFAULT=1` will enable the fix, all other choices 

19 (including undefined) will disable it. Additionally, the BILBY_STYLE and 

20 BILBY_MATHDEFAULT arguments can be passed into any 

21 latex_plot_format-wrapped plotting function and will be set directly. 

22 

23 """ 

24 @functools.wraps(func) 

25 def wrapper_decorator(*args, **kwargs): 

26 import matplotlib.pyplot as plt 

27 from matplotlib import rcParams 

28 

29 if "BILBY_STYLE" in kwargs: 

30 bilby_style = kwargs.pop("BILBY_STYLE") 

31 else: 

32 bilby_style = os.environ.get("BILBY_STYLE", "default") 

33 

34 if "BILBY_MATHDEFAULT" in kwargs: 

35 bilby_mathdefault = kwargs.pop("BILBY_MATHDEFAULT") 

36 else: 

37 bilby_mathdefault = int(os.environ.get("BILBY_MATHDEFAULT", "0")) 

38 

39 if bilby_mathdefault == 1: 

40 logger.debug("Setting mathdefault in the rcParams") 

41 rcParams['text.latex.preamble'] = r'\providecommand{\mathdefault}[1][]{}' 

42 

43 logger.debug("Using BILBY_STYLE={}".format(bilby_style)) 

44 if bilby_style.lower() == "none": 

45 return func(*args, **kwargs) 

46 elif os.path.isfile(bilby_style): 

47 plt.style.use(bilby_style) 

48 return func(*args, **kwargs) 

49 elif bilby_style in plt.style.available: 

50 plt.style.use(bilby_style) 

51 return func(*args, **kwargs) 

52 elif bilby_style.lower() == "default": 

53 _old_tex = rcParams["text.usetex"] 

54 _old_serif = rcParams["font.serif"] 

55 _old_family = rcParams["font.family"] 

56 if shutil.which("latex"): 

57 rcParams["text.usetex"] = True 

58 else: 

59 rcParams["text.usetex"] = False 

60 rcParams["font.serif"] = "Computer Modern Roman" 

61 rcParams["font.family"] = "serif" 

62 rcParams["text.usetex"] = _old_tex 

63 rcParams["font.serif"] = _old_serif 

64 rcParams["font.family"] = _old_family 

65 return func(*args, **kwargs) 

66 else: 

67 logger.debug( 

68 "Environment variable BILBY_STYLE={} not used" 

69 .format(bilby_style) 

70 ) 

71 return func(*args, **kwargs) 

72 return wrapper_decorator