From ccfe4f80467ab2e3d7b07c979d2bc2d76436b3f0 Mon Sep 17 00:00:00 2001
From: fchinu <fachi2001@gmail.com>
Date: Mon, 21 Oct 2024 13:13:28 +0200
Subject: [PATCH] Add possibility of choosing legend location

---
 flarefly/fitter.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/flarefly/fitter.py b/flarefly/fitter.py
index c519055..96ca162 100644
--- a/flarefly/fitter.py
+++ b/flarefly/fitter.py
@@ -1463,6 +1463,9 @@ def plot_mass_fit(self, **kwargs):
             - extra_info_loc: list
                 location of extra info (one for chi2 and one for other info)
 
+            - legend_loc: str
+                location of the legend
+
             - num: int
                 number of bins to plot pdfs converted into histograms
 
@@ -1482,7 +1485,8 @@ def plot_mass_fit(self, **kwargs):
         mass_range = kwargs.get('extra_info_massrange', None)
         nhwhm = kwargs.get('extra_info_massnhwhm', None)
         nsigma = kwargs.get('extra_info_massnsigma', 3)
-        loc = kwargs.get('extra_info_loc', ['upper left', 'lower right'])
+        info_loc = kwargs.get('extra_info_loc', ['lower right', 'upper left'])
+        legend_loc = kwargs.get('legend_loc', 'best')
 
         mplhep.style.use(style)
 
@@ -1537,7 +1541,6 @@ def plot_mass_fit(self, **kwargs):
         plt.xlim(limits[0], limits[1])
         plt.xlabel(axis_title)
         plt.ylabel(rf'counts / {(limits[1]-limits[0])/bins*1000:0.1f} MeV/$c^2$')
-        plt.legend(loc='best')
         if logy:
             plt.yscale('log')
             plt.ylim(min(total_func) * norm_total_pdf / 5, max(total_func) * norm_total_pdf * 5)
@@ -1549,7 +1552,7 @@ def plot_mass_fit(self, **kwargs):
             chi2 = self.get_chi2()
             ndf = self.get_ndf()
             anchored_text_chi2 = AnchoredText(fr'$\chi^2 / \mathrm{{ndf}} =${chi2:.2f} / {ndf}',
-                                              loc = loc[0],
+                                              loc = info_loc[0],
                                               frameon=False)
             # signal and background info for all signals
             text = []
@@ -1597,11 +1600,13 @@ def plot_mass_fit(self, **kwargs):
                     extra_info += fr'  Signif.$({nsigma}\sigma)=${signif:.1f} $\pm$ {signif_err:.1f}'
                 text.append(extra_info)
             concatenated_text = '\n'.join(text)
-            anchored_text_signal = AnchoredText(concatenated_text, loc = loc[1], frameon=False)
+            anchored_text_signal = AnchoredText(concatenated_text, loc = info_loc[1], frameon=False)
 
             axs.add_artist(anchored_text_chi2)
             axs.add_artist(anchored_text_signal)
 
+        plt.legend(loc=legend_loc)
+
         Logger('plot_mass_fit now returns a tuple (fig, axs) !', 'WARNING')
 
         return fig, axs