use `link_iter` to link positions into trajectories for colloids

link.py

#!/usr/bin/env python3
"""
This script demonstrated the way to use low-level `link_iter` api
    from `trackpy` to link two successive frames.
"""
import numpy as np
import trackpy as tp
from itertools import product
import matplotlib.pyplot as plt
from matplotlib import cm

frame_0 = np.array(list(product(np.arange(4), repeat=2)))
movement = np.random.random(frame_0.shape) - 0.5  # ~ U(-0.5, 0.5)
frame_1 = frame_0 + movement

# linker is a generator, real calculate is done in its iteration
linker = tp.link_iter([frame_0, frame_1], search_range=1.2)
identities = []
for result in linker:  # real calculation performed
    identities.append(result[1])

# plot particles id-by-id
plt.title("Same Color = Same ID")
for i in set(np.concatenate(identities)):
    color = cm.rainbow(np.random.random())
    p1 = frame_0[identities[0] == i].ravel()
    p2 = frame_1[identities[1] == i].ravel()
    if (len(p1) > 0):
        plt.scatter(*p1, marker='o', color=color, ec='k')
    if (len(p2) > 0):
        plt.scatter(*p2, marker='s', color=color, ec='k')
    if (len(p1) > 0) and (len(p2) > 0):
        plt.plot(*np.vstack((p1, p2)).T, color='k')
plt.show()