Merge branch 'add-rcparams-setting' into 'master'

Adds an rcparams configuration for plotting

See merge request !832

bilby/core/utils.py

@@ -19,6 +19,7 @@ import numpy as np
 from scipy.interpolate import interp2d
 from scipy.special import logsumexp
 import pandas as pd
+import matplotlib.pyplot as plt
 
 logger = logging.getLogger('bilby')
 
@@ -1160,27 +1161,67 @@ def safe_file_dump(data, filename, module):
 
 def latex_plot_format(func):
     """
-    Wrap a plotting function to set rcParams so that text renders nicely with
-    latex and Computer Modern Roman font.
+    Wrap the plotting function to set rcParams dependent on environment variables
+
+    The rcparams can be set directly from the env. variable `BILBY_STYLE` to
+    point to a matplotlib style file. Or, if `BILBY_STYLE=default` (any case) a
+    default setup is used, this is enabled by default. To not use any rcParams,
+    set `BILBY_STYLE=none`. Occasionally, issues arrise with the latex
+    `mathdefault` command. A fix is to define this command in the rcParams. An
+    env. variable `BILBY_MATHDEFAULT` can be used to turn this fix on/off.
+    Setting `BILBY_MATHDEFAULT=1` will enable the fix, all other choices
+    (including undefined) will disable it. Additionally, the BILBY_STYLE and
+    BILBY_MATHDEFAULT arguments can be passed into any
+    latex_plot_format-wrapped plotting function and will be set directly.
+
     """
     @functools.wraps(func)
     def wrapper_decorator(*args, **kwargs):
         from matplotlib import rcParams
-        _old_tex = rcParams["text.usetex"]
-        _old_serif = rcParams["font.serif"]
-        _old_family = rcParams["font.family"]
-        if find_executable("latex"):
-            rcParams["text.usetex"] = True
+
+        if "BILBY_STYLE" in kwargs:
+            bilby_style = kwargs.pop("BILBY_STYLE")
+        else:
+            bilby_style = os.environ.get("BILBY_STYLE", "default")
+
+        if "BILBY_MATHDEFAULT" in kwargs:
+            bilby_mathdefault = kwargs.pop("BILBY_MATHDEFAULT")
+        else:
+            bilby_mathdefault = int(os.environ.get("BILBY_MATHDEFAULT", "0"))
+
+        if bilby_mathdefault == 1:
+            logger.debug("Setting mathdefault in the rcParams")
             rcParams['text.latex.preamble'] = r'\newcommand{\mathdefault}[1][]{}'
+
+        logger.debug("Using BILBY_STYLE={}".format(bilby_style))
+        if bilby_style.lower() == "none":
+            return func(*args, **kwargs)
+        elif os.path.isfile(bilby_style):
+            plt.style.use(bilby_style)
+            return func(*args, **kwargs)
+        elif bilby_style in plt.style.available:
+            plt.style.use(bilby_style)
+            return func(*args, **kwargs)
+        elif bilby_style.lower() == "default":
+            _old_tex = rcParams["text.usetex"]
+            _old_serif = rcParams["font.serif"]
+            _old_family = rcParams["font.family"]
+            if find_executable("latex"):
+                rcParams["text.usetex"] = True
+            else:
+                rcParams["text.usetex"] = False
+            rcParams["font.serif"] = "Computer Modern Roman"
+            rcParams["font.family"] = "serif"
+            rcParams["text.usetex"] = _old_tex
+            rcParams["font.serif"] = _old_serif
+            rcParams["font.family"] = _old_family
+            return func(*args, **kwargs)
         else:
-            rcParams["text.usetex"] = False
-        rcParams["font.serif"] = "Computer Modern Roman"
-        rcParams["font.family"] = "serif"
-        value = func(*args, **kwargs)
-        rcParams["text.usetex"] = _old_tex
-        rcParams["font.serif"] = _old_serif
-        rcParams["font.family"] = _old_family
-        return value
+            logger.debug(
+                "Environment variable BILBY_STYLE={} not used"
+                .format(bilby_style)
+            )
+            return func(*args, **kwargs)
     return wrapper_decorator
 
 

docs/faq.txt

+=========================
+Frequency Asked Questions
+=========================
+
+Plotting questions
+------------------
+
+I'm running into latex errors when :code:`bilby` tries to create plots, what should I do?
+
+Matplotlib can be a little finicky. We wrap plotting commands in a function
+which can set up the rcParams and we use environment variables to allow
+configuration of this. See the docstring of this function for the allowed
+configuration options
+
+
+.. autofunction:: bilby.core.utils.latex_plot_format
+
+
+

docs/index.txt

@@ -27,5 +27,6 @@ Welcome to bilby's documentation!
    writing-documentation
    hyperparameters
    containers
+   faq
 
 

test/test.mplstyle

+figure.figsize: 3.4, 2.5

test/utils_test.py

 from __future__ import absolute_import, division
-
 import unittest
+import os
+
 import numpy as np
 from astropy import constants
 import lal
+import matplotlib.pyplot as plt
 
 import bilby
 from bilby.core import utils
@@ -253,5 +255,63 @@ class TestReflect(unittest.TestCase):
         self.assertTrue(np.testing.assert_allclose(utils.reflect(xprime), x) is None)
 
 
+class TestLatexPlotFormat(unittest.TestCase):
+    def setUp(self):
+        self.x = np.linspace(0, 1)
+        self.y = np.sin(self.x)
+        self.filename = "test_plot.png"
+
+    def tearDown(self):
+        if os.path.isfile(self.filename):
+            os.remove(self.filename)
+
+    def test_default(self):
+        @bilby.core.utils.latex_plot_format
+        def plot():
+            fig, ax = plt.subplots()
+            ax.plot(self.x, self.y)
+            fig.savefig(self.filename)
+        plot()
+        self.assertTrue(os.path.isfile(self.filename))
+
+    def test_mathedefault_one(self):
+        @bilby.core.utils.latex_plot_format
+        def plot():
+            fig, ax = plt.subplots()
+            ax.plot(self.x, self.y)
+            fig.savefig(self.filename)
+        plot(BILBY_MATHDEFAULT=1)
+        self.assertTrue(os.path.isfile(self.filename))
+
+    def test_mathedefault_zero(self):
+        @bilby.core.utils.latex_plot_format
+        def plot():
+            fig, ax = plt.subplots()
+            ax.plot(self.x, self.y)
+            fig.savefig(self.filename)
+        plot(BILBY_MATHDEFAULT=0)
+        self.assertTrue(os.path.isfile(self.filename))
+
+    def test_matplotlib_style(self):
+        @bilby.core.utils.latex_plot_format
+        def plot():
+            fig, ax = plt.subplots()
+            ax.plot(self.x, self.y)
+            fig.savefig(self.filename)
+
+        plot(BILBY_STYLE="fivethirtyeight")
+        self.assertTrue(os.path.isfile(self.filename))
+
+    def test_user_style(self):
+        @bilby.core.utils.latex_plot_format
+        def plot():
+            fig, ax = plt.subplots()
+            ax.plot(self.x, self.y)
+            fig.savefig(self.filename)
+
+        plot(BILBY_STYLE="test/test.mplstyle")
+        self.assertTrue(os.path.isfile(self.filename))
+
+
 if __name__ == "__main__":
     unittest.main()