plotting.py 60.1 KB
Newer Older
amichaut's avatar
amichaut committed
1
2
3
4
5
##########################################################################
# Track Analyzer - Quantification and visualization of tracking data     #
# Authors: Arthur Michaut                                                #
# Copyright 2016-2019 Harvard Medical School and Brigham and             #
#                          Women's Hospital                              #
6
# Copyright 2019-202é Institut Pasteur and CNRS–UMR3738                  #
amichaut's avatar
amichaut committed
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# See the COPYRIGHT file for details                                     #
#                                                                        #
# This file is part of Track Analyzer package.                           #
#                                                                        #
# Track Analyzer is free software: you can redistribute it and/or modify #
# it under the terms of the GNU General Public License as published by   #
# the Free Software Foundation, either version 3 of the License, or      #
# (at your option) any later version.                                    #
#                                                                        #
# Track Analyzer is distributed in the hope that it will be useful,      #
# but WITHOUT ANY WARRANTY; without even the implied warranty of         #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the           #
# GNU General Public License for more details .                          #
#                                                                        #
# You should have received a copy of the GNU General Public License      #
# along with Track Analyzer (COPYING).                                   #
# If not, see <https://www.gnu.org/licenses/>.                           #
##########################################################################

26
27
28
29
import sys
import os.path as osp
import datetime
import itertools
30
import copy
31

amichaut's avatar
amichaut committed
32
33
34
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
amichaut's avatar
amichaut committed
35
36
37
38
import pandas as pd
from skimage import io
import seaborn as sns
from scipy import stats
amichaut's avatar
amichaut committed
39
from scipy.spatial import voronoi_plot_2d
amichaut's avatar
amichaut committed
40
41
42
import napari
import tifffile as tifff

amichaut's avatar
amichaut committed
43
44
from track_analyzer import calculate as tca
from track_analyzer import prepare as tpr
amichaut's avatar
amichaut committed
45

amichaut's avatar
amichaut committed
46

47
def make_plot_config(data_dir=None, export_config=False):
48
49
    """ Generate config parameters for plotting """

50
    color_list = [c['color'] for c in list(plt.rcParams['axes.prop_cycle'])] + sns.color_palette("Set1",n_colors=9,desat=.5)
51
    # Plotting config
amichaut's avatar
amichaut committed
52
53
    plot_config = {'figsize': (5, 5),
                   'dpi': 300,
54
                   'figsize_factor':1,  # factor to modulate figsize on traj and map plots
55
                   'color_list': color_list,
amichaut's avatar
amichaut committed
56
57
58
59
60
61
62
                   'format': '.png',  # plot image format
                   'despine': True,  # use seaborn despine function
                   'logx': False,  # use log for x axis
                   'logy': False,  # use log for y axis
                   'invert_yaxis': True,  # to flip plot towards bottom as convential image orientation (origin at the top)
                   'export_data_pts': True,  # for numerical plots, export data points to csv
                   'save_as_stack': True,  # for movies, save as a tiff file, if false save as a series of files in a folder
amichaut's avatar
amichaut committed
63
                   }
64

65
66
67
68
69
70
71
72
73
74
    if export_config:
        if data_dir is None:
            raise Exception("ERROR: no data_dir given")
        else:
            config_dir = osp.join(data_dir, 'config')
            tpr.safe_mkdir(config_dir)

        fn = osp.join(config_dir, 'plot_config.csv')
        tpr.write_dict(plot_config, fn)

75
    return plot_config
amichaut's avatar
amichaut committed
76
77


amichaut's avatar
amichaut committed
78
def stack_max_proj(image_fn, z_dim, t_dim=None):
amichaut's avatar
amichaut committed
79
    """Perform a maximum projection of a 3D or 4D image. The dimension of z and time are given by z_dim and t_dim. """
amichaut's avatar
amichaut committed
80
    im = io.imread(image_fn)
amichaut's avatar
amichaut committed
81

amichaut's avatar
amichaut committed
82
    if t_dim is None:
amichaut's avatar
amichaut committed
83
        new_im = np.zeros((im.shape[1], im.shape[2]), 'uint8')
amichaut's avatar
amichaut committed
84
85
        new_im = np.max(im, axis=0)
    else:
amichaut's avatar
amichaut committed
86
87
        if t_dim == 1 and z_dim == 0:  # if the z dimension is along dimension 0 transpose
            im_ = im.transpose(1, 0, 2, 3)
amichaut's avatar
amichaut committed
88
        new_im = np.zeros((im.shape[0], im.shape[2], im.shape[3]), 'uint8')
amichaut's avatar
amichaut committed
89
90
91
        for i in range(im.shape[0]):
            new_im[i] = np.max(im[i], axis=0)
    fn, file_ext = osp.splitext(image_fn)
amichaut's avatar
amichaut committed
92
    out_fn = fn + '_maxproj.tif'
amichaut's avatar
amichaut committed
93
94
    tifff.imsave(out_fn, new_im)

amichaut's avatar
amichaut committed
95

amichaut's avatar
amichaut committed
96
def plot_cmap(plot_dir, label, cmap, vmin, vmax, plot_config=None, suffix=''):
97
    """ Plot colormap given by cmap with boundaries vmin and vmax."""
98
99

    plot_config = make_plot_config() if plot_config is None else plot_config
amichaut's avatar
amichaut committed
100

amichaut's avatar
amichaut committed
101
102
    fig = plt.figure(figsize=(8, 3))
    ax = fig.add_axes([0.05, 0.80, 0.9, 0.15])
amichaut's avatar
amichaut committed
103
    norm = plt.Normalize(vmin=vmin, vmax=vmax)
104
    cb = mpl.colorbar.ColorbarBase(ax, cmap=plt.get_cmap(cmap), norm=norm, orientation='horizontal')
amichaut's avatar
amichaut committed
105
    ax.tick_params(labelsize=16)
amichaut's avatar
amichaut committed
106
    cb.set_label(label=label, size=24)
amichaut's avatar
amichaut committed
107
    filename = osp.join(plot_dir, 'colormap'+suffix+plot_config['format'])
108
    fig.savefig(filename, dpi=plot_config['dpi'], bbox_inches='tight')
amichaut's avatar
amichaut committed
109
110
    plt.close(fig)

amichaut's avatar
amichaut committed
111

112
def plot_traj(df, frame, data_dir, groups=None, image={'image_fn': None, 't_dim': None, 'z_dim': None}, plot_dir=None,
amichaut's avatar
amichaut committed
113
              show_plot=False, dim=3, plot_config=None,traj_parameters=None,quiet=0):
114
115
116
117
118
    """ 
    Plot all trajectories of a given frame on an image if traj_parameters['no_bkg'] is False and an image is given.
    Plots can be color coded z value, by groups, or with random colors (traj_parameters['color_code']='z' or 'group' or 'random' or 'none')
    The trajectory path can be removed to keep only the dots if traj_parameters['traj'] is False.
    It can be plotted in 3D with plot3D, elevation and angle set the 3D view
amichaut's avatar
amichaut committed
119
    """
120

amichaut's avatar
amichaut committed
121
    if quiet < 1:
122
123
        sys.stdout.write("\033[K")  # go back to previous line
        print('plotting frame {}'.format(int(frame)), flush=True, end='\r')
amichaut's avatar
amichaut committed
124

125
    # get config parameters
126
    plot_config = make_plot_config() if plot_config is None else plot_config
amichaut's avatar
amichaut committed
127

128
    # unpack config
129
    color_list = plot_config['color_list']
130
    figsize_factor = plot_config['figsize_factor']
131
    show_tail = traj_parameters['show_tail']
amichaut's avatar
amichaut committed
132
    color_code = traj_parameters['color_code']
133
    cmap = traj_parameters['cmap']
amichaut's avatar
amichaut committed
134
135
136
137
    hide_labels = traj_parameters['hide_labels']
    no_bkg = traj_parameters['no_bkg']
    size_factor = traj_parameters['size_factor']
    plot3D = traj_parameters['plot3D']
138
    cmap_lim = traj_parameters['cmap_lim']
amichaut's avatar
amichaut committed
139
140
141
142
    show_axis = traj_parameters['show_axis']
    elevation = traj_parameters['elevation']
    angle = traj_parameters['angle']
    lab_size = traj_parameters['lab_size']
143
    invert_yaxis = plot_config['invert_yaxis']
amichaut's avatar
amichaut committed
144
    save_as_stack = plot_config['save_as_stack']
amichaut's avatar
amichaut committed
145

amichaut's avatar
amichaut committed
146

147
    # get image size
amichaut's avatar
amichaut committed
148
    info = tpr.get_info(data_dir)
149
150
151
152
    if 'image_width' not in info.keys() or 'image_height' not in info.keys():
        image_size = None
    else: 
        image_size = [info['image_width'], info['image_height']]
amichaut's avatar
amichaut committed
153
154
155
156
157

    # traj size
    ms_ref = plt.rcParams['lines.markersize']
    ms = ms_ref * size_factor
    lw = ms / 8
amichaut's avatar
amichaut committed
158
159

    if plot_dir is None:
amichaut's avatar
amichaut committed
160
        plot_dir = osp.join(data_dir, 'traj')
amichaut's avatar
amichaut committed
161
        tpr.safe_mkdir(plot_dir)
amichaut's avatar
amichaut committed
162

163
164
165
166
167
168
169
170
    # 3D PLOTTING. not supported anymore
    # if plot3D:
    #     fig = plt.figure()
    #     ax = fig.add_subplot(111, projection='3d')
    #     xmin, xmax, ymin, ymax = ax.axis('off')
    # else:
    # Get background image
    bkgd = tpr.get_background(image=image, frame=frame, df=df, no_bkg=no_bkg,
171
172
                        image_size=image_size,axis_on=show_axis, dpi=plot_config['dpi'],
                        figsize_factor=figsize_factor)
173
174
175
176
177
178
179
    fig = bkgd['fig']
    ax = bkgd['ax']
    xmin = bkgd['xmin']
    ymin = bkgd['ymin']
    xmax = bkgd['xmax']
    ymax = bkgd['ymax']
    no_bkg = bkgd['no_bkg']
180

181
    # get frame dataframe
182
    groups = df.groupby('frame') if groups is None else groups
amichaut's avatar
amichaut committed
183
    group = groups.get_group(frame).reset_index(drop=True)
amichaut's avatar
amichaut committed
184
185

    #### plotting positions as dots
amichaut's avatar
amichaut committed
186
187
    x = group['x'].values
    y = group['y'].values
188
189
    z = group['z_scaled'].values if dim == 3 else np.zeros(group.shape[0])
    t = group['t'].values
amichaut's avatar
amichaut committed
190

amichaut's avatar
amichaut committed
191
    ## color code
192
    if color_code == "z" or color_code == "t":
amichaut's avatar
amichaut committed
193
        if cmap_lim is None:
194
195
196
            if color_code == "z":
                if dim == 3:
                    cmap_lim = [df['z_scaled'].min(), df['z_scaled'].max()]
amichaut's avatar
amichaut committed
197
198
                else:
                    cmap_lim = [0, 0]
199
200
201
202
203
204
            elif color_code == "t":
                cmap_lim = [df['t'].min(), df['t'].max()]
        if color_code == "z":
            colors = tpr.get_cmap_color(z, cmap, vmin=cmap_lim[0], vmax=cmap_lim[1])
        elif color_code == "t":
            colors = tpr.get_cmap_color(t, cmap, vmin=cmap_lim[0], vmax=cmap_lim[1])
amichaut's avatar
amichaut committed
205

206
    elif color_code == "group":
207
208
209
        # check there are subsets in df
        if 'subset' in group.columns:
            colors = [color_list[i % len(color_list)] for i in group['subset_order'].values] # if too many colors repeat cycle
amichaut's avatar
amichaut committed
210
211
        else:
            colors = color_list[0]  # if color_coded by group but there's none, use only one color
amichaut's avatar
amichaut committed
212

amichaut's avatar
amichaut committed
213
    elif color_code == "random":
amichaut's avatar
amichaut committed
214
        colors = [color_list[i % len(color_list)] for i in group['track'].values]
amichaut's avatar
amichaut committed
215

amichaut's avatar
amichaut committed
216
217
    elif color_code == "none":
        colors = color_list[0]
amichaut's avatar
amichaut committed
218

amichaut's avatar
amichaut committed
219
220
221
    else:
        colors = color_list[0]

amichaut's avatar
amichaut committed
222
223
224
225
226
227
    # plotting labels
    if not hide_labels:
        group['track_lab'] = group['track'].map(lambda num: '{}'.format(int(num)))
        color_ = 'k' if no_bkg else 'w'
        ax.text(x, y, group['track_lab'].values, fontsize=lab_size, color=color_)

228
229
    # plot points
    ax.scatter(x, y, s=ms, color=colors)
amichaut's avatar
amichaut committed
230
231

    #### plotting trajectories as lines
232
    if show_tail:
amichaut's avatar
amichaut committed
233
234
        track_list = group['track'].values
        track_groups = df.groupby(['track'])
amichaut's avatar
amichaut committed
235
        for track in track_list:
amichaut's avatar
amichaut committed
236
237
238
239
            traj = tpr.get_traj(track_groups, track, max_frame=frame)
            traj_length = traj.shape[0]
            X = traj['x'].values
            Y = traj['y'].values
240
            Z = traj['z_scaled'].values if dim == 3 else np.zeros(traj.shape[0])
amichaut's avatar
amichaut committed
241
            t = traj['t'].values
amichaut's avatar
amichaut committed
242
243

            ## color code
244
245
246
            if color_code == "z":
                colors = tpr.get_cmap_color(Z, cmap, vmin=cmap_lim[0], vmax=cmap_lim[1])
            elif color_code == "t":
amichaut's avatar
amichaut committed
247
                colors = tpr.get_cmap_color(t, cmap, vmin=cmap_lim[0], vmax=cmap_lim[1])
248
            elif color_code == "group":
249
250
                if 'subset' in traj.columns:
                    colors = color_list[traj['subset_order'].values[0] % len(color_list)]
amichaut's avatar
amichaut committed
251
252
                else:
                    colors = color_list[0]  # if color_coded by group but there's none, use only one color
amichaut's avatar
amichaut committed
253
            elif color_code == "random":
amichaut's avatar
amichaut committed
254
                colors = color_list[track % len(color_list)]
amichaut's avatar
amichaut committed
255
256
257
258
259
            elif color_code == "none":
                colors = color_list[0]
            else:
                colors = color_list[0]

amichaut's avatar
amichaut committed
260
            if traj_length > 1:
261
262
263
                if color_code == "z" or color_code == "t":
                    for j in range(1, traj_length):
                        ax.plot([X[j - 1], X[j]], [Y[j - 1], Y[j]], lw=lw, ls='-', color=colors[j])
amichaut's avatar
amichaut committed
264
                else:
265
                    ax.plot(X, Y, lw=lw, ls='-', color=colors)
amichaut's avatar
amichaut committed
266
267

    if invert_yaxis:
amichaut's avatar
amichaut committed
268
        ax.axis([xmin, xmax, ymax, ymin])
amichaut's avatar
amichaut committed
269
270

    if show_plot:
271
        fig.show()
amichaut's avatar
amichaut committed
272
273
274
275
        return

    if show_axis:
        fig.tight_layout()
amichaut's avatar
amichaut committed
276

amichaut's avatar
amichaut committed
277
278
279
280
    if save_as_stack: 
        return fig
    else: 
        filename = osp.join(plot_dir, '{:04d}{}'.format(int(frame), plot_config['format']))
281
        fig.savefig(filename)
amichaut's avatar
amichaut committed
282
        plt.close(fig)
amichaut's avatar
amichaut committed
283

amichaut's avatar
amichaut committed
284
285

def plot_scalar_field(data_dir, df, data, field, frame, image={'image_fn': None, 't_dim': None, 'z_dim': None},
286
                      map_param={'no_bkg': False, 'vlim': None, 'show_axis': False, 'cmap': 'plasma'},
amichaut's avatar
amichaut committed
287
                      plot_dir=None, plot_config=None, quiet=0, dont_save=False):
amichaut's avatar
amichaut committed
288
289
    """Plot scalar field as colormap. The data needs to be generated before. """

amichaut's avatar
amichaut committed
290
    if quiet < 1:
291
292
        sys.stdout.write("\033[K")  # go back to previous line
        print('plotting {} {}'.format(field, int(frame)), flush=True, end='\r')
amichaut's avatar
amichaut committed
293
294

    if plot_dir is None:
amichaut's avatar
amichaut committed
295
        plot_dir = osp.join(data_dir, field)
amichaut's avatar
amichaut committed
296
    tpr.safe_mkdir(plot_dir)
amichaut's avatar
amichaut committed
297

amichaut's avatar
amichaut committed
298
    # misc param
299
300

    plot_config = make_plot_config() if plot_config is None else plot_config
amichaut's avatar
amichaut committed
301
    save_as_stack = plot_config['save_as_stack']
302
    figsize_factor = plot_config['figsize_factor']
amichaut's avatar
amichaut committed
303
304
305
306
307
    no_bkg = map_param['no_bkg']
    show_axis = map_param['show_axis']
    cmap = map_param['cmap']
    vlim = map_param['vlim']
    info = tpr.get_info(data_dir)
308
309
310
311
    if 'image_width' not in info.keys() or 'image_height' not in info.keys():
        image_size = None
    else: 
        image_size = [info['image_width'], info['image_height']]
312
    invert_yaxis = plot_config['invert_yaxis']
313
    cmap = copy.copy(plt.get_cmap(cmap))  # make a shallow copy of cmap as modifying Matplotlib colormaps is deprecated
amichaut's avatar
amichaut committed
314

amichaut's avatar
amichaut committed
315
316
317
318
    # extract data
    X = data[frame]['X']
    Y = data[frame]['Y']
    val = data[frame][field]
amichaut's avatar
amichaut committed
319
320
321
322
323
324
325

    # #remove edges for div and curl
    # if field=='div' or field=='curl':
    #     X=X[1:-1,1:-1]
    #     Y=Y[1:-1,1:-1]

    if image['image_fn'] is None:
amichaut's avatar
amichaut committed
326
        no_bkg = True
amichaut's avatar
amichaut committed
327

328
    bkgd = tpr.get_background(image=image, frame=frame, df=df, no_bkg=no_bkg,
329
330
                            image_size=image_size, axis_on=show_axis,dpi=plot_config['dpi'],
                            figsize_factor=figsize_factor)
331
332
333
334
335
336
337
    fig = bkgd['fig']
    ax = bkgd['ax']
    xmin = bkgd['xmin']
    ymin = bkgd['ymin']
    xmax = bkgd['xmax']
    ymax = bkgd['ymax']
    no_bkg = bkgd['no_bkg']
amichaut's avatar
amichaut committed
338
339

    val_masked = np.ma.array(val, mask=np.isnan(val))
amichaut's avatar
amichaut committed
340
341
    [vmin, vmax] = [val_masked.min(), val_masked.max()] if vlim is None else vlim
    cmap.set_bad('w', alpha=0)  # set NAN transparent
342
343
344
345

    # shading=nearest so color value is centered on grid points
    # for more info on pcolormesh behavior, see https://matplotlib.org/stable/gallery/images_contours_and_fields/pcolormesh_grids.html#sphx-glr-gallery-images-contours-and-fields-pcolormesh-grids-py
    C = ax.pcolormesh(X, Y, val_masked, cmap=cmap, alpha=0.5, vmin=vmin, vmax=vmax, shading='nearest')  
amichaut's avatar
amichaut committed
346
347
348
349
350
351
352

    if show_axis:
        ax.grid(False)
        ax.patch.set_visible(False)
        fig.set_tight_layout(True)

    if invert_yaxis:
amichaut's avatar
amichaut committed
353
        ax.axis([xmin, xmax, ymax, ymin])
amichaut's avatar
amichaut committed
354

amichaut's avatar
amichaut committed
355
356
357
    if save_as_stack or dont_save: 
        return fig
    else: 
amichaut's avatar
amichaut committed
358
        filename = osp.join(plot_dir, field + '_{:04d}.png'.format(int(frame)))
359
        fig.savefig(filename)
amichaut's avatar
amichaut committed
360
        plt.close(fig)
amichaut's avatar
amichaut committed
361

amichaut's avatar
amichaut committed
362

amichaut's avatar
amichaut committed
363
364
def plot_vector_field(data_dir, df, data, field, frame, plot_on_field=None, dim=3,
                      image={'image_fn': None, 't_dim': None, 'z_dim': None},
365
                      map_param={'no_bkg': False, 'vlim': None, 'show_axis': False, 'cmap': 'plasma',
amichaut's avatar
amichaut committed
366
                                 'size_factor': 1},
amichaut's avatar
amichaut committed
367
                      plot_dir=None, plot_config=None, quiet=0):
amichaut's avatar
amichaut committed
368
369
    """ Plot vector field"""

amichaut's avatar
amichaut committed
370
    if quiet < 1:
371
372
        sys.stdout.write("\033[K")  # go back to previous line
        print('plotting {} {}'.format(field, int(frame)), flush=True, end='\r')
amichaut's avatar
amichaut committed
373
374

    if plot_dir is None:
amichaut's avatar
amichaut committed
375
        plot_dir = osp.join(data_dir, field)
amichaut's avatar
amichaut committed
376
    tpr.safe_mkdir(plot_dir)
amichaut's avatar
amichaut committed
377

amichaut's avatar
amichaut committed
378
    # misc param
379
    plot_config = make_plot_config() if plot_config is None else plot_config
amichaut's avatar
amichaut committed
380
    save_as_stack = plot_config['save_as_stack']
381
    figsize_factor = plot_config['figsize_factor']
amichaut's avatar
amichaut committed
382
383
384
385
386
387
    no_bkg = map_param['no_bkg']
    show_axis = map_param['show_axis']
    # cmap=map_param['cmap']
    # vlim=map_param['vlim']
    size_factor = map_param['size_factor']
    info = tpr.get_info(data_dir)
388
389
390
391
    if 'image_width' not in info.keys() or 'image_height' not in info.keys():
        image_size = None
    else: 
        image_size = [info['image_width'], info['image_height']]
392
    invert_yaxis = plot_config['invert_yaxis']
amichaut's avatar
amichaut committed
393
394

    # import image
amichaut's avatar
amichaut committed
395
    if image['image_fn'] is None:
amichaut's avatar
amichaut committed
396
        no_bkg = True
amichaut's avatar
amichaut committed
397

amichaut's avatar
amichaut committed
398
    no_plot_on_field = False
amichaut's avatar
amichaut committed
399
    if plot_on_field is not None:
400
401
        if 'plot_on' not in plot_on_field.keys(): 
            plot_on_field['plot_on'] = None
amichaut's avatar
amichaut committed
402
        if plot_on_field['plot_on'] is not None:
amichaut's avatar
amichaut committed
403
404
405
406
            map_param_ = map_param
            map_param_['cmap'] = plot_on_field['cmap']
            map_param_['vlim'] = plot_on_field['vlim']
            dim = 2  # to ensure that arrows are plotted in black and the z data is not use
amichaut's avatar
amichaut committed
407
            fig = plot_scalar_field(data_dir, df, data, plot_on_field['plot_on'], frame, image=image,
amichaut's avatar
amichaut committed
408
                                        map_param=map_param_, plot_dir=plot_dir, plot_config=None,
409
                                        quiet=2, dont_save=True)
amichaut's avatar
amichaut committed
410
            ax = fig.gca()
amichaut's avatar
amichaut committed
411
            invert_yaxis = False  # to ensure it's not inverted a second time
amichaut's avatar
amichaut committed
412
        else:
amichaut's avatar
amichaut committed
413
414
415
            no_plot_on_field = True
    else:
        no_plot_on_field = True
amichaut's avatar
amichaut committed
416
417

    if no_plot_on_field:
418
        bkgd = tpr.get_background(image=image, frame=frame, df=df, no_bkg=no_bkg,
419
420
                                image_size=image_size, axis_on=show_axis,
                                dpi=plot_config['dpi'],figsize_factor=figsize_factor)
421
        fig = bkgd['fig']
422
423
424
425
426
427
        ax = bkgd['ax']
        xmin = bkgd['xmin']
        ymin = bkgd['ymin']
        xmax = bkgd['xmax']
        ymax = bkgd['ymax']
        no_bkg = bkgd['no_bkg']
428

amichaut's avatar
amichaut committed
429
430
431
432
    # extract data
    dimensions = ['x', 'y', 'z'] if dim == 3 else ['x', 'y']
    vdata = [field + d for d in dimensions]  # eg ['vx','vy'] or ['ax','ay','az']
    val = [data[frame]['X'], data[frame]['Y']] + [data[frame][vd] for vd in vdata]  # eg ['X','Y','vx','vy']
amichaut's avatar
amichaut committed
433

amichaut's avatar
amichaut committed
434
435
436
    # norm=plt.Normalize(vlim[0],vlim[1]) if vlim is not None else None
    # Q=ax.quiver(*val,units='inches',cmap=cmap,norm=norm,width=0.005)
    Q = ax.quiver(*val, units='inches', width=0.005 * size_factor, angles='xy')
amichaut's avatar
amichaut committed
437
438
439
440
441
442
443

    if show_axis:
        ax.grid(False)
        ax.patch.set_visible(False)
        fig.set_tight_layout(True)

    if invert_yaxis:
amichaut's avatar
amichaut committed
444
445
        ylim = ax.get_ylim()
        ax.set_ylim(ylim[1], ylim[0])
amichaut's avatar
amichaut committed
446

amichaut's avatar
amichaut committed
447
448
449
450
    if save_as_stack: 
        return fig
    else: 
        filename = osp.join(plot_dir, 'vector_' + field + '_{:04d}.png'.format(int(frame)))
451
        fig.savefig(filename)
amichaut's avatar
amichaut committed
452
        plt.close(fig)
amichaut's avatar
amichaut committed
453

amichaut's avatar
amichaut committed
454

amichaut's avatar
amichaut committed
455
def plot_Voronoi(data_dir, df, frame, data, show_local_area=True,
amichaut's avatar
amichaut committed
456
                 image={'image_fn': None, 't_dim': None, 'z_dim': None},
457
                 map_param={'no_bkg': False, 'vlim': None, 'show_axis': False,'cmap':'plasma','size_factor': 1,'line_width':1.},
458
                 plot_dir=None, plot_config=None,quiet=0):
amichaut's avatar
amichaut committed
459
460
461
462
463
464
465
466
467
468
469
470
471
472
    """
    Plot Voronoi tesselation and local area in 2D only.
    :param data_dir:
    :param df:
    :param frame:
    :param data:
    :param show_local_area:
    :param image:
    :param map_param:
    :param plot_dir:
    :param plot_config:
    :return:
    """

amichaut's avatar
amichaut committed
473
    if quiet < 1:
474
475
        sys.stdout.write("\033[K")  # go back to previous line
        print('plotting {} {}'.format('voronoi', int(frame)), flush=True, end='\r')
amichaut's avatar
amichaut committed
476
477
478
479
480
481
482

    if plot_dir is None:
        plot_dir = osp.join(data_dir, 'voronoi')
    tpr.safe_mkdir(plot_dir)

    # misc param
    plot_config = make_plot_config() if plot_config is None else plot_config
amichaut's avatar
amichaut committed
483
    save_as_stack = plot_config['save_as_stack']
484
    figsize_factor = plot_config['figsize_factor']
amichaut's avatar
amichaut committed
485
486
487
488
    no_bkg = map_param['no_bkg']
    show_axis = map_param['show_axis']
    cmap = map_param['cmap']
    vlim = map_param['vlim']
489
    line_width = map_param['line_width']
amichaut's avatar
amichaut committed
490
    info = tpr.get_info(data_dir)
491
492
493
494
    if 'image_width' not in info.keys() or 'image_height' not in info.keys():
        image_size = None
    else: 
        image_size = [info['image_width'], info['image_height']]
amichaut's avatar
amichaut committed
495
496
497
498
499
500
    invert_yaxis = plot_config['invert_yaxis']

    # import image
    if image['image_fn'] is None:
        no_bkg = True

501
    bkgd = tpr.get_background(image=image, frame=frame, df=df, no_bkg=no_bkg,
502
503
                            image_size=image_size, axis_on=show_axis, dpi=plot_config['dpi'],
                            figsize_factor=figsize_factor)
504
505
506
507
508
509
510
    fig = bkgd['fig']
    ax = bkgd['ax']
    xmin = bkgd['xmin']
    ymin = bkgd['ymin']
    xmax = bkgd['xmax']
    ymax = bkgd['ymax']
    no_bkg = bkgd['no_bkg']
amichaut's avatar
amichaut committed
511
512
513

    # plot tesselation
    vor = data[frame]['vor']
amichaut's avatar
amichaut committed
514
    if vor is not None:
515
        voronoi_plot_2d(vor, show_points=False, show_vertices=False, ax=ax, line_width=line_width)
amichaut's avatar
amichaut committed
516
517
518
519
520
521
522
523
524
525

        # plot local area on top
        if show_local_area:
            areas = data[frame]['areas']
            if areas is not None:
                for pt_id, reg_num in enumerate(vor.point_region):
                    indices = vor.regions[reg_num]
                    area = areas[pt_id]
                    if not np.isnan(area):
                        color = tpr.get_cmap_color(area, cmap, vmin=vlim[0], vmax=vlim[1])
526
                        ax.fill(*zip(*vor.vertices[indices]), color=color, alpha=0.5, linewidth=0)
amichaut's avatar
amichaut committed
527

528
529
530
531
    # ensure axis limits are constant
    ax.set_xlim(xmin,xmax)
    ax.set_ylim(ymin,ymax)

amichaut's avatar
amichaut committed
532
533
534
535
536
537
538
539
    if show_axis:
        ax.grid(False)
        ax.patch.set_visible(False)
        fig.set_tight_layout(True)

    if invert_yaxis:
        ylim = ax.get_ylim()
        ax.set_ylim(ylim[1], ylim[0])
amichaut's avatar
amichaut committed
540
541
542
543
544
    
    if save_as_stack: 
        return fig
    else: 
        filename = osp.join(plot_dir, 'voronoi_{:04d}.png'.format(int(frame)))
545
        fig.savefig(filename)
amichaut's avatar
amichaut committed
546
        plt.close(fig)
amichaut's avatar
amichaut committed
547
548


amichaut's avatar
amichaut committed
549
550
def plot_hist_persistence_length(data_dir, track_groups, tracks, minimal_traj_length=40, normalize=True, dim=3,
                                 plot_config=None):
amichaut's avatar
amichaut committed
551
    plt.close('all')
552
553
554

    plot_config = make_plot_config() if plot_config is None else plot_config

amichaut's avatar
amichaut committed
555
    pers_length_dict = {}
amichaut's avatar
amichaut committed
556
    for track in tracks:
amichaut's avatar
amichaut committed
557
558
559
560
        traj = tpr.get_traj(track_groups, track)
        traj_length, c = traj.shape
        if traj_length > minimal_traj_length:
            pers_length_dict[track] = tca.get_obj_persistence_length(track_groups, track, traj, dim=dim)
amichaut's avatar
amichaut committed
561

amichaut's avatar
amichaut committed
562
    pers_lengths = pd.Series(pers_length_dict)
amichaut's avatar
amichaut committed
563
    fig, ax = plt.subplots()
amichaut's avatar
amichaut committed
564
    if normalize:
amichaut's avatar
amichaut committed
565
        pers_lengths.plot.hist(weights=np.ones_like(pers_lengths * 100) / len(pers_lengths), ax=ax)
566
        ax.set_ylabel('trajectories proportion ')
amichaut's avatar
amichaut committed
567
    else:
568
569
570
        pers_lengths.plot.hist(ax=ax)
        ax.set_ylabel('trajectories count')
    ax.set_xlabel(r'persistence length ($\mu m$) ')
amichaut's avatar
amichaut committed
571
    filename = osp.join(data_dir, 'persistence_lenght.svg')
572
    fig.savefig(filename, dpi=plot_config['dpi'], bbox_inches='tight')
amichaut's avatar
amichaut committed
573

amichaut's avatar
amichaut committed
574
575

def plot_MSD(data_dir, track, track_groups=None, df=None, df_out=None, fit_model="biased_diff", dim=2, save_plot=True,
amichaut's avatar
amichaut committed
576
             print_traj_info=True, frame_subset=None, fitrange=None, plot_dir=None, plot_config=None, logx=True,
amichaut's avatar
amichaut committed
577
             logy=True):
amichaut's avatar
amichaut committed
578
579
    """Compute MSD of a trajectory and fit it with a random walk model. If df_out is given, save the output of a fit. If save_plot is False, don't plot the MSD (useful for large number of tracks to analyze)."""

580
581
    plot_config = make_plot_config() if plot_config is None else plot_config
    color_list = plot_config['color_list']
amichaut's avatar
amichaut committed
582
583
584

    data = tpr.get_data(data_dir)
    timescale = data['timescale']
amichaut's avatar
amichaut committed
585
    if df is None:
amichaut's avatar
amichaut committed
586
        df = data['df']
amichaut's avatar
amichaut committed
587
588

    if track_groups is None:
amichaut's avatar
amichaut committed
589
590
        track_groups = df.groupby(['track'])
    traj = tpr.get_traj(track_groups, track)
amichaut's avatar
amichaut committed
591
592

    if frame_subset is not None:
amichaut's avatar
amichaut committed
593
        ind = ((traj['frame'] >= frame_subset[0]) & (traj['frame'] <= frame_subset[1]))
amichaut's avatar
amichaut committed
594
595
        traj = traj[ind]

amichaut's avatar
amichaut committed
596
597
598
599
    if dim == 2:
        dimensions = ['x_scaled', 'y_scaled']
    elif dim == 3:
        dimensions = ['x_scaled', 'y_scaled', 'z_scaled']
amichaut's avatar
amichaut committed
600

amichaut's avatar
amichaut committed
601
602
    # compute and fit MSD
    msd = tca.compute_msd(traj, timescale, dimensions)
amichaut's avatar
amichaut committed
603
    if fit_model is not None:
604
        results = tca.fit_msd(msd, mean_vel=traj['v'].mean(), dim=dim, model=fit_model, fitrange=fitrange)
amichaut's avatar
amichaut committed
605
        if df_out is not None and results['success']:
amichaut's avatar
amichaut committed
606
            ind = df_out['track'] == track
amichaut's avatar
amichaut committed
607
            for param in results['param'].keys():
amichaut's avatar
amichaut committed
608
609
                df_out.loc[ind, param] = results['param'][param]
            df_out.loc[ind, 'redchi'] = results['redchi']
amichaut's avatar
amichaut committed
610
611

    if save_plot:
612
613
614
615
616
617
        if plot_dir is None:
            plot_dir = osp.join(data_dir, 'MSD')
        else: 
            plot_dir = osp.join(plot_dir, 'MSD')  #  save in a separate folder
        tpr.safe_mkdir(plot_dir)

amichaut's avatar
amichaut committed
618
        info = tpr.get_info(data_dir)
amichaut's avatar
amichaut committed
619
        D_unit = tpr.make_param_label('D', l_unit=info['length_unit'], t_unit=info['time_unit'], only_unit=True)
amichaut's avatar
amichaut committed
620

621
        fig, ax = plt.subplots(1, 1, figsize=plot_config['figsize'])
amichaut's avatar
amichaut committed
622
        msd.plot.scatter(x="tau", y="msd", logx=logx, logy=logy, ax=ax)
amichaut's avatar
amichaut committed
623
624
        if fit_model is not None:
            if results['success']:
amichaut's avatar
amichaut committed
625
626
627
628
629
630
631
632
633
                fitted_df = results['fitted_df']
                fitted_df.plot(x="tau", y="fitted", logx=logx, logy=logy, ax=ax)
                if fit_model == 'biased_diff':
                    title_ = r'D={:0.3f} {:}, $\chi^2$={:0.3f}'.format(results['param']['D'], D_unit, results['redchi'])
                elif fit_model == 'PRW':
                    title_ = r'P={:0.3f} {:}, $\chi^2$={:0.3f}'.format(results['param']['P'], info['time_unit'],
                                                                       results['redchi'])
                elif fit_model == 'pure_diff':
                    title_ = r'D={:0.3f} {:}, $\chi^2$={:0.3f}'.format(results['param']['D'], D_unit, results['redchi'])
amichaut's avatar
amichaut committed
634
635
636
                ax.set_title(title_)
        ax.set_xlabel('lag time ({})'.format(info['time_unit']))
        ax.set_ylabel(r'MSD ({})'.format(D_unit))
637
        if plot_config['despine']:
amichaut's avatar
amichaut committed
638
            sns.despine(fig)
639
        fig.savefig(osp.join(plot_dir, '{}{}'.format(int(track), plot_config['format'])), dpi=plot_config['dpi'],
amichaut's avatar
amichaut committed
640
                    bbox_inches='tight')
amichaut's avatar
amichaut committed
641
642
        plt.close(fig)

amichaut's avatar
amichaut committed
643
    return msd[['tau', 'msd']]
amichaut's avatar
amichaut committed
644

amichaut's avatar
amichaut committed
645
646

def plot_param_vs_param(data_dir, x_param, y_param, df=None, hue=None, hue_order=None, set_axis_lim=None,
amichaut's avatar
amichaut committed
647
                        plot_config=None,
amichaut's avatar
amichaut committed
648
                        plot_dir=None, prefix='', suffix=''):
amichaut's avatar
amichaut committed
649
650
    """Plot a parameter of df (y_param) against another parameter (x_param). Optional: compare datasets with hue as datasets identifier."""

651
652
653
    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']
amichaut's avatar
amichaut committed
654

655
    # get df
amichaut's avatar
amichaut committed
656
    if df is None:
657
        data = tpr.get_data(data_dir)
amichaut's avatar
amichaut committed
658
        df = data['df']
amichaut's avatar
amichaut committed
659

660
    # define plotting directory
amichaut's avatar
amichaut committed
661
    if plot_dir is None:
amichaut's avatar
amichaut committed
662
663
        plot_dir = data_dir

664
665
666
667
668
669
670
671
    # make sure params are in df
    if x_param not in df.columns: 
        print("Warning: parameter {} does not exist".format(x_param))
        return -1
    if y_param not in df.columns: 
        print("Warning: parameter {} does not exist".format(y_param))
        return -1

amichaut's avatar
amichaut committed
672
673
674
    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'])
amichaut's avatar
amichaut committed
675
676

    if hue is not None:
amichaut's avatar
amichaut committed
677
        df[hue] = df[hue].astype('category')  # make sure that sns.scatterplot does not use the continuous colormap
amichaut's avatar
amichaut committed
678

679
    fig, ax = plt.subplots(1, 1, figsize=plot_config['figsize'])
amichaut's avatar
amichaut committed
680
681
    # 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)
amichaut's avatar
amichaut committed
682
683
684
    ax.set_xlabel(x_lab)
    ax.set_ylabel(y_lab)
    if set_axis_lim is not None:
amichaut's avatar
amichaut committed
685
686
687
688
689
690
691
692
693
        ax.set_xlim(set_axis_lim[0], set_axis_lim[1])
        ax.set_ylim(set_axis_lim[2], set_axis_lim[3])
    else:  # recalculate because sometimes matplotlib auto limit fails
        xlim_ = [df[x_param].min(), df[x_param].max()]
        ylim_ = [df[y_param].min(), df[y_param].max()]
        xlim_ = [xlim_[0] - 0.05 * (xlim_[1] - xlim_[0]), xlim_[1] + 0.05 * (xlim_[1] - xlim_[0])]
        ylim_ = [ylim_[0] - 0.05 * (ylim_[1] - ylim_[0]), ylim_[1] + 0.05 * (ylim_[1] - ylim_[0])]
        ax.set_xlim(xlim_[0], xlim_[1])
        ax.set_ylim(ylim_[0], ylim_[1])
amichaut's avatar
amichaut committed
694

695
    if plot_config['despine']:
amichaut's avatar
amichaut committed
696
697
        sns.despine(fig)

698
699
    if hue is not None:
        ax.legend(frameon=False)
700

701
702
    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')
amichaut's avatar
amichaut committed
703
704
705
    plt.close(fig)

    if export_data_pts:
amichaut's avatar
amichaut committed
706
707
        cols = [x_param, y_param, hue] if hue is not None else [x_param, y_param]
        fn = osp.join(plot_dir, prefix + '{}_vs_{}{}{}'.format(y_param, x_param, suffix, '.csv'))
amichaut's avatar
amichaut committed
708
709
        df[cols].to_csv(fn)

amichaut's avatar
amichaut committed
710

711
def plot_param_hist(data_dir, param, df=None, hue=None, hue_order=None, hist=True, kde=True,
amichaut's avatar
amichaut committed
712
                    plot_config=None,
amichaut's avatar
amichaut committed
713
                    plot_dir=None, prefix='', suffix=''):
amichaut's avatar
amichaut committed
714
    """Plot a parameter histogram. Optional: compare datasets with hue as datasets identifier."""
amichaut's avatar
amichaut committed
715

716
717
718
719
    plot_config = make_plot_config() if plot_config is None else plot_config
    color_list = plot_config['color_list']
    figsize = plot_config['figsize']
    export_data_pts = plot_config['export_data_pts']
amichaut's avatar
amichaut committed
720

721
    # get df
amichaut's avatar
amichaut committed
722
    if df is None:
723
        data = tpr.get_data(data_dir)
amichaut's avatar
amichaut committed
724
        df = data['df']
amichaut's avatar
amichaut committed
725

726
    # define plotting directory
amichaut's avatar
amichaut committed
727
    if plot_dir is None:
amichaut's avatar
amichaut committed
728
729
        plot_dir = data_dir

730
731
732
733
734
735
    # make sure param is in df
    if param not in df.columns: 
        print("Warning: parameter {} does not exist".format(param))
        return -1

    # make label
amichaut's avatar
amichaut committed
736
737
    info = tpr.get_info(data_dir)
    param_label = tpr.make_param_label(param, l_unit=info['length_unit'], t_unit=info['time_unit'])
amichaut's avatar
amichaut committed
738

739
740
741
    # make sure data is float and finite
    df[param] = df[param].astype(np.float)
    df = df[np.isfinite(df[param])]
amichaut's avatar
amichaut committed
742

743
    kind = "hist" if hist else "kde"
744
    g = sns.displot(data=df, x=param, hue=hue, kind=kind, kde=kde,facet_kws={'legend_out':False,'despine':plot_config['despine']})
745
746
    if hue is not None:
        sns.move_legend(g, "upper right", frameon=False, bbox_to_anchor=[0.95, 0.95],title=None)
747
    fig = g.figure
748
    ax = g.ax
amichaut's avatar
amichaut committed
749
    fig.set_size_inches(figsize[0], figsize[1])
750
    ax.set_xlabel(param_label)
amichaut's avatar
amichaut committed
751

752
753
    filename = osp.join(plot_dir, prefix + '{}_hist{}{}'.format(param, suffix, plot_config['format']))
    fig.savefig(filename, dpi=plot_config['dpi'], bbox_inches='tight')
amichaut's avatar
amichaut committed
754
755
756
    plt.close(fig)

    if export_data_pts:
amichaut's avatar
amichaut committed
757
758
        cols = [param, hue] if hue is not None else param
        fn = osp.join(plot_dir, prefix + '{}_hist{}{}'.format(param, suffix, '.csv'))
amichaut's avatar
amichaut committed
759
760
        df[cols].to_csv(fn)

amichaut's avatar
amichaut committed
761
762

def plot_param_boxplot(data_dir, df, x_param, param, order=None, hue=None, save_stat=False, hue_order=None,
763
764
                       boxplot=True, swarmplot=True, plot_config=None, plot_dir=None, prefix='', suffix='',
                       leg_lab=None):
amichaut's avatar
amichaut committed
765
    """Plot boxplot between categories (given by x_param). Sub-categories can be plotted too (given by hue). A ttest can be performed has well"""
amichaut's avatar
amichaut committed
766

767
768
769
770
    plot_config = make_plot_config() if plot_config is None else plot_config
    color_list = plot_config['color_list']
    figsize = plot_config['figsize']
    export_data_pts = plot_config['export_data_pts']
amichaut's avatar
amichaut committed
771

772
    # get df
amichaut's avatar
amichaut committed
773
    if df is None:
774
        data = tpr.get_data(data_dir)
amichaut's avatar
amichaut committed
775
        df = data['df']
amichaut's avatar
amichaut committed
776

777
    # define plotting directory
amichaut's avatar
amichaut committed
778
    if plot_dir is None:
amichaut's avatar
amichaut committed
779
780
        plot_dir = data_dir

781
782
783
784
785
786
787
788
    # make sure params are in df
    if x_param not in df.columns: 
        print("Warning: parameter {} does not exist".format(x_param))
        return -1
    if param not in df.columns: 
        print("Warning: parameter {} does not exist".format(param))
        return -1

amichaut's avatar
amichaut committed
789
790
    info = tpr.get_info(data_dir)
    param_label = tpr.make_param_label(param, l_unit=info['length_unit'], t_unit=info['time_unit'])
amichaut's avatar
amichaut committed
791

amichaut's avatar
amichaut committed
792
793
    # plotting
    fig, ax = plt.subplots(1, 1, figsize=figsize)
amichaut's avatar
amichaut committed
794
    if hue is None:
amichaut's avatar
amichaut committed
795
796
797
798
        width = 0.5
        color = color_list[0]
        palette = [(0.25, 0.25, 0.25), (0.25, 0.25, 0.25)]
        lw = 0
amichaut's avatar
amichaut committed
799
    else:
amichaut's avatar
amichaut committed
800
801
802
803
        width = 0.8
        color = None
        palette = color_list
        lw = 1
amichaut's avatar
amichaut committed
804
    if boxplot:
amichaut's avatar
amichaut committed
805
806
        sns.boxplot(data=df, x=x_param, y=param, ax=ax, order=order, width=width, hue=hue, hue_order=hue_order,
                    color=color)
amichaut's avatar
amichaut committed
807
    if swarmplot:
amichaut's avatar
amichaut committed
808
809
        sns.swarmplot(data=df, x=x_param, y=param, ax=ax, order=order, size=8, linewidth=lw, dodge=True,
                      palette=palette, hue=hue, hue_order=hue_order)
amichaut's avatar
amichaut committed
810
811
812
    if hue is not None:
        handles, labels = ax.get_legend_handles_labels()
        if leg_lab is not None:
amichaut's avatar
amichaut committed
813
            labs = leg_lab
amichaut's avatar
amichaut committed
814
        else:
amichaut's avatar
amichaut committed
815
816
817
            labs = labels[0:len(handles) // 2]
        ax.legend(handles[0:len(handles) // 2], labs, frameon=False)
    ax.ticklabel_format(axis='y', style='sci', scilimits=(-2, 3))
amichaut's avatar
amichaut committed
818
    ax.get_yaxis().get_major_formatter().set_useMathText(True)
amichaut's avatar
amichaut committed
819
820
    ax.set_ylabel(param_label)
    ax.set_xlabel(x_param)
amichaut's avatar
amichaut committed
821

822
    if plot_config['despine']:
amichaut's avatar
amichaut committed
823
824
        sns.despine(fig)

amichaut's avatar
amichaut committed
825
    filename_basis = osp.join(plot_dir, prefix + '{}_box{}'.format(param, suffix))
826
    fig.savefig(filename_basis + plot_config['format'], dpi=plot_config['dpi'], bbox_inches='tight')
amichaut's avatar
amichaut committed
827
    plt.close(fig)
amichaut's avatar
amichaut committed
828
829

    # stat
amichaut's avatar
amichaut committed
830
    if save_stat and x_param is not None:
amichaut's avatar
amichaut committed
831
        x_list = df[x_param].unique() if order is None else order
amichaut's avatar
amichaut committed
832
        if hue is not None:
amichaut's avatar
amichaut committed
833
834
835
836
            hue_list = df[hue].unique() if hue_order is None else hue_order
            ind = pd.MultiIndex.from_product([x_list, hue_list], names=[x_param, hue])
            df_mean = pd.DataFrame(index=ind, columns=['mean', 'std', 'n'])
            df_pval = pd.DataFrame(index=ind, columns=hue_list)
amichaut's avatar
amichaut committed
837
            for xp in x_list:
amichaut's avatar
amichaut committed
838
839
                subdf = df[df[x_param] == xp]
                data_dict = {}
amichaut's avatar
amichaut committed
840
                for h in hue_list:
amichaut's avatar
amichaut committed
841
842
843
844
                    sub_nonan = subdf[subdf[hue] == h][param].dropna()
                    data_dict[h] = sub_nonan.values
                    df_mean.loc[(xp, h), :] = [np.mean(data_dict[h]), np.std(data_dict[h]), data_dict[h].shape[0]]
                pairs = list(itertools.combinations(data_dict.keys(), 2))
amichaut's avatar
amichaut committed
845
                for p in pairs:
amichaut's avatar
amichaut committed
846
847
                    ttest_ = stats.ttest_ind(data_dict[p[0]], data_dict[p[1]], equal_var=False)
                    df_pval.loc[(xp, p[0]), p[1]] = ttest_.pvalue
amichaut's avatar
amichaut committed
848
849

        else:
amichaut's avatar
amichaut committed
850
851
852
            df_mean = pd.DataFrame(index=x_list, columns=['mean', 'std', 'n'])
            df_pval = pd.DataFrame(index=x_list, columns=x_list)
            data_dict = {}
amichaut's avatar
amichaut committed
853
            for xp in x_list:
amichaut's avatar
amichaut committed
854
855
856
                sub_nonan = df[df[x_param] == xp][param].dropna()
                data_dict[xp] = sub_nonan.values
                df_mean.loc[xp, :] = [np.mean(data_dict[