Calculate the MSD

msd.py

#!/usr/bin/env python3
"""
This script demonstrated the way to link positions into trajectories
    and then calculate the MSD
"""
import numpy as np
import trackpy as tp
from itertools import product
import matplotlib.pyplot as plt
import pandas as pd


class TrackpyLinker():
    """
    Linking positions into trajectories using Trackpy
    Works with 2D and 3D data. High dimensional data not tested.
    (no expiment available)
    """
    def __init__(self, max_movement, memory=0, max_subnet_size=40, **kwargs):
        self.max_movement = max_movement
        self.memory = memory
        self.max_subnet_size = max_subnet_size
        self.kwargs = kwargs

    @staticmethod
    def _check_input(positions, time_points, labels):
        """
        Make sure the input is proper and sequence in time_points are ordered
        """
        assert len(positions) == len(time_points), "Lengths are not consistent"
        if not isinstance(labels, type(None)):
            assert len(positions) == len(labels), "Lengths are not consistent"
            for p, l in zip(positions, labels):
                assert len(p) == len(l), "Labels and positions are not matched"
        time_points = np.array(time_points)
        order_indice = time_points.argsort()
        ordered_time = time_points[order_indice]
        positions = list(positions)
        positions.sort(key=lambda x: order_indice.tolist())
        return positions, ordered_time, labels

    def __get_trajectory(self, value, link_result, positions, time_points, labels):
        with_label = False
        if isinstance(labels, type(None)):
            traj = [[], []]
        else:
            traj = [[], [], []]
            with_label = True
        for frame in link_result:
            frame_index, link_labels = frame
            if value in link_labels:
                number_index = link_labels.index(value)
                traj[0].append(time_points[frame_index])
                traj[1].append(positions[frame_index][number_index])
                if with_label:
                    current_label = labels[frame_index][link_labels.index(value)]
                    traj[2].append(current_label)
        traj[0] = np.array(traj[0])
        traj[1] = np.array(traj[1])
        return traj
    
    def __get_trajectories(self, link_result, positions, time_points, labels):
        trajectories = []
        total_labels = []
        for frame in link_result:
            frame_index, link_labels = frame
            total_labels.append(link_labels)
        for value in set(np.hstack(total_labels)):
            traj = self.__get_trajectory(value, link_result, positions, time_points, labels)
            trajectories.append(traj)
        return trajectories
    
    def link(self, positions, time_points=None, labels=None):
        """
        Args:
            positions (np.ndarray): shape (time, num, dim)
            time_points (np.ndarray): shape (time, ), time_points may not be continues
            labels (np.ndarray): if given, the result will have a 'label' attribute
                                 which specifies the label values in different frames
                                 [(frame_index, [labels, ... ]), ...], shape (time, num)
        """
        if isinstance(time_points, type(None)):
            time_points = np.arange(len(positions))
        pos, time, labels = self._check_input(positions, time_points, labels)
        tp.linking.Linker.MAX_SUB_NET_SIZE = self.max_subnet_size
        link_result = tp.link_iter(pos, search_range=self.max_movement, memory=self.memory, **self.kwargs)
        return self.__get_trajectories(list(link_result), pos, time, labels)


"""
Generate Simulated Data (Brownian Movement in 3D)
"""
frame = np.array(list(product(np.arange(5), repeat=3)))  # 3D data, shape (125, 3)

frames = [frame]
for _ in range(100):
    movement = np.random.random(frame.shape) - 0.5  # ~ U(-0.5, 0.5)
    movement /= 50  # ~ U(-0.01, 0.01)
    frame = frame + movement
    frames.append(frame)

linker = TrackpyLinker(max_movement=0.4, memory=0)

trajectories = linker.link(frames)

msd_vals = []
for traj in trajectories:
    time, positions = traj
    table = pd.DataFrame({  # create Pandas table for trackpy
        'frame': time,
        'x': positions[:, 0],
        'y': positions[:, 1],
        'z': positions[:, 2],
    })
    msd = tp.motion.msd(
            table,
            mpp=1,  # microns per pixel, need experimental data
            fps=1,  # frame per second, need experimental data
            max_lagtime=50,
            pos_columns=['x', 'y', 'z']
        )
    msd_vals.append(msd[['msd']].to_numpy())
msd_mean = np.mean(msd_vals, axis=0)

for msd_single in msd_vals:
    plt.plot(msd_single, lw=1, alpha=0.2, color='k')
plt.plot(msd_mean, lw=2, color='teal', marker='o', mfc='w')
plt.xscale('log')
plt.yscale('log')
plt.xlabel("Lag Time")
plt.ylabel("MSD")
plt.tight_layout()
plt.show()