added binning and regression to scatter plots

track_analyzer/plotting.py

@@ -644,13 +644,14 @@ def plot_MSD(data_dir, track, track_groups=None, df=None, df_out=None, fit_model
 
 
 def plot_param_vs_param(data_dir, x_param, y_param, df=None, hue=None, hue_order=None, set_axis_lim=None,
-                        plot_config=None,
+                        plot_config=None, x_bin_num=None, ci=None, fit_reg=False, scatter=True,
                         plot_dir=None, prefix='', suffix=''):
     """Plot a parameter of df (y_param) against another parameter (x_param). Optional: compare datasets with hue as datasets identifier."""
 
     plot_config = make_plot_config() if plot_config is None else plot_config
     color_list = plot_config['color_list']
     export_data_pts = plot_config['export_data_pts']
+    figsize = plot_config['figsize']
 
     # get df
     if df is None:
@@ -669,16 +670,37 @@ def plot_param_vs_param(data_dir, x_param, y_param, df=None, hue=None, hue_order
         print("Warning: parameter {} does not exist".format(y_param))
         return -1
 
+    # make labels
     info = tpr.get_info(data_dir)
     x_lab = tpr.make_param_label(x_param, l_unit=info['length_unit'], t_unit=info['time_unit'])
     y_lab = tpr.make_param_label(y_param, l_unit=info['length_unit'], t_unit=info['time_unit'])
 
+    # make sure data is float and finite
+    for p in [x_param,y_param]:
+        df[p] = df[p].astype(np.float)
+        df = df[np.isfinite(df[p])]
+
+    # make sure that sns.lmplot does not use the continuous colormap
     if hue is not None:
-        df[hue] = df[hue].astype('category')  # make sure that sns.scatterplot does not use the continuous colormap
+        df[hue] = df[hue].astype('category')  
 
-    fig, ax = plt.subplots(1, 1, figsize=plot_config['figsize'])
-    # sns.scatterplot(x=x_param,y=y_param,ax=ax,facecolors='none',edgecolor=color_list[0],data=df)
-    sns.scatterplot(x=x_param, y=y_param, ax=ax, hue=hue, hue_order=hue_order, data=df)
+    # make x_bins
+    if x_bin_num is None:
+        x_bins = None 
+    else: 
+        # make evenly spaced bins along x_param and exclude min and max
+        x_bins = np.linspace(df[x_param].min(),df[x_param].max(),x_bin_num+2)  # +2 so len(x_bins)=x_bin_num after removing min and max
+        x_bins = x_bins[1:-1]
+
+    # plot
+    g = sns.lmplot(x=x_param, y=y_param, hue=hue, hue_order=hue_order, data=df, ci=ci, fit_reg=fit_reg, x_bins=x_bins, scatter=scatter,
+        facet_kws={'despine':plot_config['despine']})
+    if hue is not None:
+        sns.move_legend(g, "right", frameon=False, title=None, bbox_to_anchor=(1.05, 0.5))
+
+    fig = g.figure
+    ax = g.ax
+    fig.set_size_inches(figsize[0], figsize[1])
     ax.set_xlabel(x_lab)
     ax.set_ylabel(y_lab)
     if set_axis_lim is not None:
@@ -695,9 +717,6 @@ def plot_param_vs_param(data_dir, x_param, y_param, df=None, hue=None, hue_order
     if plot_config['despine']:
         sns.despine(fig)
 
-    if hue is not None:
-        ax.legend(frameon=False)
-
     filename = osp.join(plot_dir, prefix + '{}_vs_{}{}{}'.format(y_param, x_param, suffix, plot_config['format']))
     fig.savefig(filename, dpi=plot_config['dpi'], bbox_inches='tight')
     plt.close(fig)
@@ -741,9 +760,9 @@ def plot_param_hist(data_dir, param, df=None, hue=None, hue_order=None, hist=Tru
     df = df[np.isfinite(df[param])]
 
     kind = "hist" if hist else "kde"
-    g = sns.displot(data=df, x=param, hue=hue, kind=kind, kde=kde,facet_kws={'legend_out':False,'despine':plot_config['despine']})
+    g = sns.displot(data=df, x=param, hue=hue, kind=kind, kde=kde,facet_kws={'despine':plot_config['despine']})
     if hue is not None:
-        sns.move_legend(g, "upper right", frameon=False, bbox_to_anchor=[0.95, 0.95],title=None)
+        sns.move_legend(g, "right", frameon=False, title=None, bbox_to_anchor=(1.05, 0.5))
     fig = g.figure
     ax = g.ax
     fig.set_size_inches(figsize[0], figsize[1])

track_analyzer/prepare.py

@@ -1334,10 +1334,12 @@ def make_traj_config(data_dir=None, export_config=True):
                   }
 
     scatter_config = {'run': True,  # run param_vs_param
-                      'couple_list': [['x', 'v'], ['y', 'v']],
-                      # list of [x,y] couples of variables to be plotted in scatter
-                      'mean_couple_list': [['x', 'v'], ['y', 'v']]
-                      # list of [x,y] couples of variables averaged along the whole track to be plotted in scatter
+                      'couple_list': [['x', 'v'], ['y', 'v']],  # list of [x,y] couples of variables to be plotted in scatter
+                      'mean_couple_list': [['x', 'v'], ['y', 'v']],  # list of [x,y] couples of variables averaged along the whole track to be plotted in scatter
+                      'x_bin_num': None,  #  number of evenly spaced bins along x axis
+                      'ci': None,  # show confidence interval [0,100]
+                      'fit_reg': False,  # show regression fit
+                      'scatter': True,  # show scatter 
                       }
 
     hist_config = {'run': True,  # run plot_param_hist

track_analyzer/scripts/analyze_tracks.py

@@ -184,7 +184,7 @@ def traj_analysis(data_dir, data=None, image=None, refresh=False, parallelize=Fa
             if quiet < 2:
                 print("Voronoi analysis...")
             vor_data = tca.compute_all_Voronoi(data_dir, df, outdir=sub_dir,compute_local_area=voronoi_config['compute_local_area'],
-                                            area_threshold=voronoi_config['area_threshold'],df_mean=df_prop)
+                                            area_threshold=voronoi_config['area_threshold'], df_mean=df_prop)
             # update csv files with area data
             df_prop.to_csv(mean_fn)
             df.to_csv(csv_fn) 
@@ -231,7 +231,8 @@ def traj_analysis(data_dir, data=None, image=None, refresh=False, parallelize=Fa
                 for param_vs_param in scatter_config['couple_list']:
                     x_param, y_param = param_vs_param
                     tpl.plot_param_vs_param(data_dir, x_param, y_param, df, plot_dir=sub_dir, plot_config=plot_config,
-                                            hue=hue, hue_order=hue_order)
+                                            hue=hue, hue_order=hue_order, x_bin_num=scatter_config["x_bin_num"], ci=scatter_config["ci"], 
+                                            fit_reg=scatter_config["fit_reg"], scatter=scatter_config["scatter"])
 
             if 'mean_couple_list' in scatter_config.keys():
                 if len(scatter_config['mean_couple_list']) > 0:
@@ -240,7 +241,8 @@ def traj_analysis(data_dir, data=None, image=None, refresh=False, parallelize=Fa
                 for param_vs_param in scatter_config['mean_couple_list']:
                     x_param, y_param = param_vs_param
                     tpl.plot_param_vs_param(data_dir, x_param, y_param, df_prop, plot_dir=sub_dir, plot_config=plot_config,
-                                            prefix='track_', hue=hue, hue_order=hue_order)
+                                            prefix='track_', hue=hue, hue_order=hue_order, x_bin_num=scatter_config["x_bin_num"], ci=scatter_config["ci"], 
+                                            fit_reg=scatter_config["fit_reg"], scatter=scatter_config["scatter"])
 
     return df_list