synthetic_data.py 9.39 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
37
from skimage.util import img_as_ubyte
import seaborn as sns
import tifffile as tifff

# Plotting parameters
38
39
40
41
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
42

43
mpl.use('agg')
amichaut's avatar
amichaut committed
44

45
46
47
48
49
50
51
52
53
54
55
56
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
57
    for i in range(dim):
58
59
        displacement['r2'] += displacement[list('xyz')[i]] ** 2
    displacement['r'] = np.sqrt(displacement['r2'])
amichaut's avatar
amichaut committed
60
    for i in range(dim):
61
        displacement[list('xyz')[i]] /= displacement['r']  # normalize raw displacement
62
        displacement[list('xyz')[i]] *= noise_amp  # apply amplitude
63
64
        displacement[list('xyz')[i]] += bias[i]  # add bias
    displacement = displacement[list('xyz')[0:dim]].values
amichaut's avatar
amichaut committed
65

66
67
    # traj
    traj = np.zeros((tmax, dim))
amichaut's avatar
amichaut committed
68
    for i in range(dim):
69
        traj[:, i] = np.cumsum(displacement[:, i]) + x0[i]
amichaut's avatar
amichaut committed
70
        if periodic:
71
72
73
            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
74

75
76
77
78

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},
79
                          grad={'step_num': 4, 'dim': 0, 'boundaries':None},
80
                          x0_range={'x': [0.1, 0.9], 'y': [0.1, 0.9], 'z': [0.1, 0.9]}, dt=1):
81
82
83
    """
    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
84
85
86
    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. 
    """
87
88
89
90
91
92

    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'])
93
    
94
95
96
97
98
    # 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]]]

99

100
    part_count = 0
amichaut's avatar
amichaut committed
101
    for i in range(grad['step_num']):
102
103
104
105
106
107
108
109
        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
110
        for k in range(dim):
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
            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
134
    """Make array of diffusive particles biased toward a node (or away if attraction_ampl is negative)"""
135
136
137

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

    if node is None:
140
141
142
143
144
145
        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
146
147

    for i in range(part_num):
148
149
150
151
152
153
154
        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
155
        for d in range(dim):
156
157
158
159
160
161
162
            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
163
        for k in range(dim):
164
165
166
167
168
169
170
171
172
173
174
175
176
177
            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
178
179
    """Plot synthetic data and save it as a grayscaled tiff stack"""

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

        figsize = (grid_size[0] / dpi, grid_size[1] / dpi)
        fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)
amichaut's avatar
amichaut committed
189
190
        ax = fig.add_axes([0, 0, 1, 1])
        for k in range(group.shape[0]):
191
192
193
            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
194
195
        ax.invert_yaxis()
        ax.axis('off')
196
197
198
199
200
201
202
203
204
205
206
207
        
        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
208
209