synthetic_data.py 9.44 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-2022 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/>.                           #
##########################################################################

amichaut's avatar
amichaut committed
26
27
import os.path as osp

amichaut's avatar
amichaut committed
28
import matplotlib.pyplot as plt
29
import matplotlib as mpl
amichaut's avatar
amichaut committed
30
import numpy as np
amichaut's avatar
amichaut committed
31
import pandas as pd
32
from skimage.color import rgb2gray
amichaut's avatar
amichaut committed
33
34
35
36
from skimage.util import img_as_ubyte
import seaborn as sns
import tifffile as tifff

37
from track_analyzer import prepare as tpr
amichaut's avatar
amichaut committed
38
39

# Plotting parameters
40
41
42
43
color_list = [c['color'] for c in list(plt.rcParams['axes.prop_cycle'])] + sns.color_palette("Set1", n_colors=9,
                                                                                             desat=.5)
plot_param = {'figsize': (5, 5), 'dpi': 300, 'color_list': color_list, 'format': '.png', 'despine': True, 'logx': False,
              'logy': False, 'invert_yaxis': True, 'export_data_pts': False}
amichaut's avatar
amichaut committed
44

45
mpl.use('agg')
amichaut's avatar
amichaut committed
46

47
48
49
50
51
52
53
54
55
56
57
58
def make_diff_traj(part_index=0, grid_size=[500, 500, 500], dim=3, tmax=10, periodic=True, noise_amp=10,
                   x0=[250, 250, 250], bias=[0, 0, 0]):
    """
    Generate a trajectory with a diffusive trajectory, with a bias.
    bias gives the amplitude at each step along each dimension.
    """
    # time and index
    t = np.arange(tmax)
    index = np.ones(tmax) * part_index
    # displacement
    displacement = pd.DataFrame(np.random.randn(tmax, dim), columns=list('xyz')[0:dim])
    displacement['r2'] = 0
amichaut's avatar
amichaut committed
59
    for i in range(dim):
60
61
        displacement['r2'] += displacement[list('xyz')[i]] ** 2
    displacement['r'] = np.sqrt(displacement['r2'])
amichaut's avatar
amichaut committed
62
    for i in range(dim):
63
        displacement[list('xyz')[i]] /= displacement['r']  # normalize raw displacement
64
        displacement[list('xyz')[i]] *= noise_amp  # apply amplitude
65
66
        displacement[list('xyz')[i]] += bias[i]  # add bias
    displacement = displacement[list('xyz')[0:dim]].values
amichaut's avatar
amichaut committed
67

68
69
    # traj
    traj = np.zeros((tmax, dim))
amichaut's avatar
amichaut committed
70
    for i in range(dim):
71
        traj[:, i] = np.cumsum(displacement[:, i]) + x0[i]
amichaut's avatar
amichaut committed
72
        if periodic:
73
74
75
            traj[:, i] = np.remainder(traj[:, i], grid_size[i])
    return pd.DataFrame(np.concatenate([index[:, None], t[:, None], traj], axis=1),
                        columns=['traj', 'frame'] + list('xyz')[0:dim])
amichaut's avatar
amichaut committed
76

77
78
79
80

def make_spatial_gradient(part_num=100, grid_size=[500, 500, 500], dim=3, tmax=10, periodic=True,
                          bias_basis=[0, 0, 0],
                          diff_grad={'min': 0, 'max': 10}, bias_grad={'min': 0, 'max': 10, 'dim': 0},
81
                          grad={'step_num': 4, 'dim': 0, 'boundaries':None},
82
                          x0_range={'x': [0.1, 0.9], 'y': [0.1, 0.9], 'z': [0.1, 0.9]}, dt=1):
83
84
85
    """
    Make a spatial gradient in diffusion or bias, along a specific dimension, given by grad['dim'].
    The number of steps on the gradient is given by grad['step_num']. 
amichaut's avatar
amichaut committed
86
87
88
    The gradient can be in diffusion with diff_grad or bias_grad. min and max give the extrema of the gradient, and bias_grad['dim'] give the dimension along the gradient in bias is applied.
    An overall constant bias can be passed by bias_basis. 
    """
89
90
91
92
93
94

    df = pd.DataFrame([], columns=['traj', 'frame'] + list('xyz')[0:dim])
    df_param = pd.DataFrame([], columns=['traj', 'v', 'D'])

    diff_grad_ = np.linspace(diff_grad['min'], diff_grad['max'], grad['step_num'])
    bias_grad_ = np.linspace(bias_grad['min'], bias_grad['max'], grad['step_num'])
95
    
96
97
98
99
100
    # spatial boundaries of the regions of particles
    lims = [[x0_range['x'][0] * grid_size[0], x0_range['x'][1] * grid_size[0]],
            [x0_range['y'][0] * grid_size[1], x0_range['y'][1] * grid_size[1]],
            [x0_range['z'][0] * grid_size[2], x0_range['z'][1] * grid_size[2]]]

101

102
    part_count = 0
amichaut's avatar
amichaut committed
103
    for i in range(grad['step_num']):
104
105
106
107
108
109
110
111
        grad_increment = (lims[grad['dim']][1] - lims[grad['dim']][0]) / grad['step_num']
        lims_ = lims[:]
        lims_[grad['dim']] = [lims_[grad['dim']][0] + i * grad_increment,
                              lims_[grad['dim']][0] + (i + 1) * grad_increment]
        noise_amp = diff_grad_[i]
        bias = bias_basis[:]
        bias[bias_grad['dim']] = bias_grad_[i]
        bias_ampl = 0
amichaut's avatar
amichaut committed
112
        for k in range(dim):
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
            bias_ampl += bias[k] ** 2
        bias_ampl = np.sqrt(bias_ampl)

        for j in range(int(part_num / grad['step_num'])):
            x0 = [np.random.uniform(lims_[0][0], lims_[0][1]),
                  np.random.uniform(lims_[1][0], lims_[1][1]),
                  np.random.uniform(lims_[2][0], lims_[2][1])]

            traj = make_diff_traj(part_index=part_count, noise_amp=noise_amp, x0=x0, bias=bias, tmax=tmax,
                                  periodic=periodic, dim=dim)
            df = pd.concat([df, traj])
            v = bias_ampl / dt
            D = noise_amp ** 2 / (2. * dim * dt)
            df_param.loc[part_count, :] = [part_count, v, D]

            part_count += 1
    return df, df_param


def make_attraction_node(part_num=100, grid_size=[500, 500, 500], dim=3, tmax=10, periodic=True, noise_amp=10,
                         bias_basis=[0, 0, 0],
                         attraction_ampl=10, node=None, x0_range={'x': [0.1, 0.9], 'y': [0.1, 0.9], 'z': [0.1, 0.9]},
                         dt=1):
amichaut's avatar
amichaut committed
136
    """Make array of diffusive particles biased toward a node (or away if attraction_ampl is negative)"""
137
138
139

    df = pd.DataFrame([], columns=['traj', 'frame'] + list('xyz')[0:dim])
    df_param = pd.DataFrame([], columns=['traj', 'v', 'D'])
amichaut's avatar
amichaut committed
140
141

    if node is None:
142
143
144
145
146
147
        node = [grid_size[d] / 2 for d in range(dim)]  # by default center

    # spatial boundaries of the regions of particles
    lims = [[x0_range['x'][0] * grid_size[0], x0_range['x'][1] * grid_size[0]],
            [x0_range['y'][0] * grid_size[1], x0_range['y'][1] * grid_size[1]],
            [x0_range['z'][0] * grid_size[2], x0_range['z'][1] * grid_size[2]]]
amichaut's avatar
amichaut committed
148
149

    for i in range(part_num):
150
151
152
153
154
155
156
        x0 = [np.random.uniform(lims[0][0], lims[0][1]),
              np.random.uniform(lims[1][0], lims[1][1]),
              np.random.uniform(lims[2][0], lims[2][1])]

        # unit vector towards node
        node_vec = np.array([node[d] - x0[d] for d in range(dim)])
        sum_ = 0
amichaut's avatar
amichaut committed
157
        for d in range(dim):
158
159
160
161
162
163
164
            sum_ += node_vec[d] ** 2
        node_vec /= np.sqrt(sum_)

        bias = node_vec * attraction_ampl
        bias = bias + np.array(bias_basis)

        bias_ampl = 0
amichaut's avatar
amichaut committed
165
        for k in range(dim):
166
167
168
169
170
171
172
173
174
175
176
177
178
179
            bias_ampl += bias[k] ** 2
        bias_ampl = np.sqrt(bias_ampl)

        traj = make_diff_traj(part_index=i, noise_amp=noise_amp, x0=x0, bias=bias, tmax=tmax, periodic=periodic,
                              dim=dim)
        df = pd.concat([df, traj])
        v = bias_ampl / dt
        D = noise_amp ** 2 / (2. * dim * dt)
        df_param.loc[i, :] = [i, v, D]

    return df, df_param


def plot_synthetic_stack(df, outdir, dpi=300, grid_size=[500, 500, 500], tmax=10):
amichaut's avatar
amichaut committed
180
181
    """Plot synthetic data and save it as a grayscaled tiff stack"""

182
    stack = []  # stack to store images
183
    groups = df.groupby('frame')
184
185
    
    # plot frames
amichaut's avatar
amichaut committed
186
    for i in range(tmax):
187
        group = groups.get_group(i).reset_index(drop=True)
188
189
190

        figsize = (grid_size[0] / dpi, grid_size[1] / dpi)
        fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
amichaut's avatar
amichaut committed
191
192
        ax = fig.add_axes([0, 0, 1, 1])
        for k in range(group.shape[0]):
193
194
195
            ax.scatter(group.loc[k, 'x'], group.loc[k, 'y'], s=10)
        ax.set_xlim(0, grid_size[0])
        ax.set_ylim(0, grid_size[1])
amichaut's avatar
amichaut committed
196
197
        ax.invert_yaxis()
        ax.axis('off')
198
199
200
201
202
203
204
205
206
207
208
209
        
        fig.canvas.draw()
        fig_image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        fig_image = fig_image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        fig_image = rgb2gray(fig_image)  # convert to grayscale
        fig_image = img_as_ubyte(fig_image)  # convert to 8 bit
        stack.append(fig_image)
        plt.close(fig)

    # save
    stack = np.array(stack)
    tifff.imsave(osp.join(outdir,'stack.tiff'), stack)
amichaut's avatar
amichaut committed
210
211