#40 add title to the legend

rmsandu · Nov 30, 2019 · 23e828a · 23e828a
1 parent 7e7681c
commit 23e828a
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 18 deletions.
diff --git a/E_radiomics_stats.py b/E_radiomics_stats.py
@@ -70,8 +70,8 @@
           'x_label': 'Energy [kJ]',
           'y_label': 'Effective Ablation Volume [ml] for 3 Microwave Devices',
           'title': "Ablation Volume [ml] for 3 MWA devices. " + flag_subcapsular_title,
-          'x_lim': 120,
-          'y_lim': 120,
+          'x_lim': 100,
+          'y_lim': 100,
           'lin_reg': 1}
 scatter_plot(df, **kwargs)
 
@@ -80,23 +80,26 @@
 df_radiomics_covidien = df[df['Device_name'] == 'Covidien (Covidien MWA)']
 
 kwargs = {'x_data': 'Energy [kj]', 'y_data': 'Ablation Volume [ml]',
-          'y_label': 'Effective Ablation Volume [ml] for Amica',
+          'y_label': 'Effective Ablation Volume [ml]',
+          'legend_title': 'Amica',
           'title': "Ablation Volume [ml] for Amica " + flag_subcapsular_title,
-          'x_lim': 120,
-          'y_lim': 120,
+          'x_lim': 100,
+          'y_lim': 100,
           'lin_reg': 1}
 scatter_plot(df_radiomics_amica, **kwargs)
 
 kwargs = {'x_data': 'Energy [kj]', 'y_data': 'Ablation Volume [ml]',
-          'y_label': 'Effective Ablation Volume [ml] for Angyiodinamics (Solero)',
+          'y_label': 'Effective Ablation Volume [ml] ',
+          'legend_title': 'Solero',
           'title': "Ablation Volume [ml] for Angyiodinamics (Solero) " + flag_subcapsular_title,
-          'x_lim': 120,
-          'y_lim': 120,
+          'x_lim': 100,
+          'y_lim': 100,
           'lin_reg': 1}
 scatter_plot(df_radiomics_angyodinamics, **kwargs)
 
 kwargs = {'x_data': 'Energy [kj]', 'y_data': 'Ablation Volume [ml]',
           'y_label': 'Effective Ablation Volume [ml] for Covidien',
+          'legend_title': 'Covidien',
           'title': "Ablation Volume [ml] for Covidien " + flag_subcapsular_title,
           'x_lim': 120,
           'y_lim': 120,

diff --git a/utils/scatter_plot.py b/utils/scatter_plot.py
@@ -29,7 +29,7 @@ def scatter_plot(df1,  **kwargs):
     :return:
     """
     fig, ax = plt.subplots()
-    fontsize = 22
+    fontsize = 24
     if kwargs.get('x_data') is None:
         print('No X input data to plot')
         return
@@ -48,7 +48,7 @@ def scatter_plot(df1,  **kwargs):
         cbar = plt.colorbar()
         cbar.ax.set_title(kwargs['colormap'], fontsize=8)
     else:
-        df_to_plot.plot.scatter(x=kwargs["x_data"], y=kwargs["y_data"], s=100, alpha=0.7)
+        df_to_plot.plot.scatter(x=kwargs["x_data"], y=kwargs["y_data"], s=100, alpha=0.7, color='purple', marker='*')
     if kwargs.get('size') is not None:
         size = np.asarray(50 * (df_to_plot[kwargs['size']] + 1)).reshape(len(df_to_plot), 1)
         size_mask = ~np.isnan(size)
@@ -57,7 +57,7 @@ def scatter_plot(df1,  **kwargs):
         fig, ax = plt.subplots()
         x = df_to_plot[kwargs['x_data']]
         y = df_to_plot[kwargs['y_data']]
-        sc = ax.scatter(x, y, s=size, alpha=0.6,  marker='o', color='steelblue')
+        sc = ax.scatter(x, y, s=size, alpha=0.6,  marker='*', color='purple')
         legend_1 = ax.legend(*sc.legend_elements("sizes", num=6, func=lambda x: x / 50 - 1, color='steelblue'),
                              title=kwargs['size'], labelspacing=1.5, borderpad=1.5, handletextpad=3.5,
                              fontsize=fontsize, loc='upper right')
@@ -79,19 +79,24 @@ def scatter_plot(df1,  **kwargs):
         label_r2 = r'$R^2:{0:.2f}$'.format(r_squared)
         label_r = r'$r: {0:.2f}$'.format(correlation_coef)
         plt.plot(X, regr.predict(X), color='orange', linewidth=1.5, label='Linear Regression')
+        plt.plot([], [], ' ', label='n = ' + str(len(kwargs['x_data'])))
         plt.plot([], [], ' ', label=label_r)
         plt.plot([], [], ' ', label=label_r2)
-        plt.legend(fontsize=fontsize, loc='upper left')
+        if kwargs.get('legend_title'):
+            plt.legend(fontsize=fontsize, loc='best', title=kwargs['legend_title'], title_fontsize=fontsize)
+        else:
+            plt.legend(fontsize=fontsize, loc='upper left')
+
     if kwargs.get('x_lim') is not None and kwargs.get('y_lim') is not None:
         plt.xlim([0, kwargs['x_lim']])
         plt.ylim([0, kwargs['y_lim']])
     title = kwargs['title']
-    plt.tick_params(labelsize=fontsize, color='black')
-    ax.tick_params(axis='y', labelsize=fontsize, color='k')
-    ax.tick_params(axis='x', labelsize=fontsize, color='k')
-    ax.xaxis.label.set_color('black')
-    ax.yaxis.label.set_color('black')
-    matplotlib.rc('axes', labelcolor='black')
+    # plt.tick_params(labelsize=fontsize, color='black')
+    # ax.tick_params(axis='y', labelsize=fontsize, color='k')
+    # ax.tick_params(axis='x', labelsize=fontsize, color='k')
+    # ax.xaxis.label.set_color('black')
+    # ax.yaxis.label.set_color('black')
+    # matplotlib.rc('axes', labelcolor='black')
     figpathHist = os.path.join("figures", title)
 
     if kwargs.get('x_label') is not None and kwargs.get('y_label') is None: