[plots] Add option to diff and bin plot functions

flavio/plots/plotfunctions.py

@@ -68,7 +68,7 @@ def my_autopct(pct):
     return plt.pie(fracs, labels=labels, autopct=my_autopct, wedgeprops = {'linewidth':0.5}, colors=flavio.plots.colors.pastel)
 
 
-def diff_plot_th(obs_name, x_min, x_max, wc=None, steps=100, **kwargs):
+def diff_plot_th(obs_name, x_min, x_max, wc=None, steps=100, scale_factor=1, **kwargs):
     r"""Plot the central theory prediction of an observable dependending on
     a continuous parameter, e.g. $q^2$.
 
@@ -79,6 +79,8 @@ def diff_plot_th(obs_name, x_min, x_max, wc=None, steps=100, **kwargs):
       Wilson coefficients
     - `steps` (optional): number of steps in x. Defaults to 100. Less is
       faster but less precise.
+    - `scale_factor` (optional): factor by which all values will be multiplied.
+      Defaults to 1.
 
     Additional keyword arguments are passed to the matplotlib plot function,
     e.g. 'c' for colour.
@@ -95,12 +97,13 @@ def diff_plot_th(obs_name, x_min, x_max, wc=None, steps=100, **kwargs):
     ax = plt.gca()
     if 'c' not in kwargs and 'color' not in kwargs:
         kwargs['c'] = 'k'
-    ax.plot(x_arr, obs_arr, **kwargs)
+    ax.plot(x_arr, scale_factor * np.asarray(obs_arr), **kwargs)
 
 
 def diff_plot_th_err(obs_name, x_min, x_max, wc=None, steps=100,
                         steps_err=5, N=100, threads=1, label=None,
-                        plot_args=None, fill_args=None):
+                        plot_args=None, fill_args=None,
+                        scale_factor=1):
     r"""Plot the theory prediction of an observable dependending on
     a continuous parameter, e.g. $q^2$,
     with uncertainties as a function of this parameter.
@@ -123,6 +126,8 @@ def diff_plot_th_err(obs_name, x_min, x_max, wc=None, steps=100,
       to the matplotlib plot function, e.g. 'c' for colour.
     - `fill_args` (optional): dictionary with keyword arguments to be passed
       to the matplotlib fill_between function, e.g. 'facecolor'
+    - `scale_factor` (optional): factor by which all values will be multiplied.
+      Defaults to 1.
 
     A word of caution regarding the `steps_err` option. By default, the
     uncertainty is only computed at 10 steps and is interpolated in
@@ -154,13 +159,13 @@ def diff_plot_th_err(obs_name, x_min, x_max, wc=None, steps=100,
         plot_args['label'] = label
     if 'alpha' not in fill_args:
         fill_args['alpha'] = 0.5
-    ax.plot(x_arr, obs_arr, **plot_args)
+    ax.plot(x_arr, scale_factor * np.asarray(obs_arr), **plot_args)
     interp_err = scipy.interpolate.interp1d(x_err_arr, obs_err_arr,
                                             kind='quadratic')
     obs_err_arr_int = interp_err(x_arr)
     ax.fill_between(x_arr,
-                    obs_arr - obs_err_arr_int,
-                    obs_arr + obs_err_arr_int,
+                    scale_factor * np.asarray(obs_arr - obs_err_arr_int),
+                    scale_factor * np.asarray(obs_arr + obs_err_arr_int),
                     **fill_args)
 
 
@@ -218,6 +223,7 @@ def bin_plot_th(obs_name, bin_list, wc=None, divide_binwidth=False, N=50, thread
 
 def bin_plot_exp(obs_name, col_dict=None, divide_binwidth=False, include_measurements=None,
                 include_bins=None, exclude_bins=None,
+                scale_factor=1,
                 **kwargs):
     r"""Plot all existing binned experimental measurements of an observable
     dependending on a continuous parameter, e.g. $q^2$ (in the form of
@@ -244,6 +250,8 @@ def bin_plot_exp(obs_name, col_dict=None, divide_binwidth=False, include_measure
       boundaries) not to include in the plot. By default, all measured bins
       will be included. Should not be specified simultaneously with
       `include_bins`.
+    - `scale_factor` (optional): factor by which all values will be multiplied.
+      Defaults to 1.
 
     Additional keyword arguments are passed to the matplotlib errorbar function,
     e.g. 'c' for colour.
@@ -303,12 +311,16 @@ def bin_plot_exp(obs_name, col_dict=None, divide_binwidth=False, include_measure
                     # twice in the legend)
                     kwargs_m['label'] = m_obj.experiment
                     _experiment_labels.append(m_obj.experiment)
+            y = scale_factor * np.array(y)
+            dy_lower = scale_factor * np.array(dy_lower)
+            dy_upper = scale_factor * np.array(dy_upper)
             ax.errorbar(x, y, yerr=[dy_lower, dy_upper], xerr=dx, fmt='.', **kwargs_m)
     return y, bins
 
 
 def diff_plot_exp(obs_name, col_dict=None, include_measurements=None,
                 include_x=None, exclude_x=None,
+                scale_factor=1,
                 **kwargs):
     r"""Plot all existing experimental measurements of an observable
     dependending on a continuous parameter, e.g. $q^2$ (in the form of
@@ -329,6 +341,8 @@ def diff_plot_exp(obs_name, col_dict=None, include_measurements=None,
       not to include in the plot. By default, all measured values
       will be included. Should not be specified simultaneously with
       `include_x`.
+    - `scale_factor` (optional): factor by which all values will be multiplied.
+      Defaults to 1.
 
     Additional keyword arguments are passed to the matplotlib errorbar function,
     e.g. 'c' for colour.
@@ -382,6 +396,9 @@ def diff_plot_exp(obs_name, col_dict=None, include_measurements=None,
                     # twice in the legend)
                     kwargs_m['label'] = m_obj.experiment
                     _experiment_labels.append(m_obj.experiment)
+            y = scale_factor * np.array(y)
+            dy_lower = scale_factor * np.array(dy_lower)
+            dy_upper = scale_factor * np.array(dy_upper)
             ax.errorbar(x, y, yerr=[dy_lower, dy_upper], fmt='.', **kwargs_m)
     return y, xs
 

flavio/plots/test_plotfunctions.py

@@ -33,7 +33,7 @@ def test_error_budget_pie(self):
 
     def test_q2_th_diff(self):
         # without specifying WCs
-        diff_plot_th('dBR/dq2(B0->pienu)', 0, 25, steps=10)
+        diff_plot_th('dBR/dq2(B0->pienu)', 0, 25, steps=10, scale_factor=1000)
         # with WCs
         diff_plot_th('dBR/dq2(B+->pienu)', 0, 25,
                                    wc=flavio.WilsonCoefficients(), steps=10)
@@ -44,7 +44,8 @@ def test_q2_th_diff(self):
     def test_q2_th_diff_err(self):
         # without parallelization
         diff_plot_th_err('dBR/dq2(B0->pienu)', 1, 24, steps=5,
-                                                steps_err=3, N=10)
+                                                steps_err=3, N=10,
+                                                scale_factor=3)
         # with parallelization
         diff_plot_th_err('dBR/dq2(B0->pienu)', 1, 24, steps=5,
                                                  steps_err=3, N=10, threads=2)
@@ -65,7 +66,8 @@ def test_q2_plot_exp(self):
         bin_plot_exp('<dBR/dq2>(B0->K*mumu)')
         # with options
         bin_plot_exp('<dBR/dq2>(B0->K*mumu)', col_dict={'LHCb': 'r'},
-                                                        exclude_bins=[(1.1, 6)])
+                                                        exclude_bins=[(1.1, 6)],
+                                                        scale_factor=100)
         # check that observable not depending on q2 raises error
         with self.assertRaises(ValueError):
             bin_plot_exp('eps_K')
@@ -76,6 +78,7 @@ def test_q2_plot_exp(self):
         m.set_constraint(('dBR/dq2(B0->K*mumu)', 1), '1 +- 0.1 e-6')
         m.set_constraint(('dBR/dq2(B0->K*mumu)', 2), '2 +- 0.2 e-6')
         diff_plot_exp('dBR/dq2(B0->K*mumu)')
+        diff_plot_exp('dBR/dq2(B0->K*mumu)', scale_factor=10)
         # with options
         # check that observable not depending on q2 raises error
         with self.assertRaises(ValueError):