Calculate cell migration trajectories from cell coordinates of all frames

calc_trajectories.py

# Calculate trajectories using data from find_cells-headless.fiji.py

import csv
import sys


def norm_sq(a, b):
    return (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2


def filter_by_length(cs, l):
    return filter(lambda a: len(a) >= l, cs)


# Connecting points in two adjacent frames.
# Pick nearest neighbors in the frames.
# Return pairs of two IDs
def connect_adjacent(fr1, fr2):
    if len(fr1) == 0 or len(fr2) == 0:
        return []
    dist_th = 100
    ps = []
    for p2 in fr2:
        nearest = min(fr1, key=lambda p1: norm_sq(p1, p2))
        if norm_sq(nearest, p2) < dist_th ** 2:
            ps.append([nearest[2], p2[2]])
    return ps


# Trace connecting pairs towards the previous frame recursively.
def trace_connection_backwards(graph, used, id, path):
    used[id] = True
    if id in graph:
        # Tail recursion
        return trace_connection_backwards(graph, used, graph[id], [id] + path)
    else:
        return [id] + path


# Find paths of points by connecting pairs of two adjacent frames.
def mk_path_from_adjacent_pairs(css):
    graph = {}
    for c in reduce(lambda a, b: a + b, css):
        graph[c[0]] = c[1]
    used = {}
    paths = []
    for cs in css:
        for c in cs:
            id = c[0]
            if not id in used:
                paths.append(trace_connection_backwards(graph, used, id, []))
    return paths


def write_files():
    with open(out_path, 'wb') as f:
        writer = csv.writer(f)
        writer.writerow(['id', 'from', 'to', 'from_frame', 'to_frame'])
        for i, row in enumerate(reduce(lambda x, y: x + y, cs)):
            writer.writerow([str(i)] + map(str, row))

    with open(out_path2, 'wb') as f:
        writer = csv.writer(f)
        writer.writerow(['id', 'from', 'to', 'from_frame', 'to_frame'])
        count = 0
        for path in paths:
            for i in range(len(path) - 1):
                count += 1
                fr = path[i + 1]
                to = path[i]
                writer.writerow(map(str, [count, fr, to, frame_table[fr], frame_table[to]]))

    with open(out_path3, 'wb') as f:
        import json
        f.write(json.dumps(paths))


def plot_migration(paths_to_draw):
    import matplotlib.pyplot as plt
    for path in paths_to_draw:
        last = path[-1]

        def subt(p):
            a = vs_by_id[p]
            b = vs_by_id[last]
            return [a[0] - b[0], a[1] - b[1]]

        # Take the only first len_thres points in order to make it possible to compare trajectories.
        ps = list(reversed(map(subt, path)))[0:len_thres]
        xs = map(lambda a: a[0], ps)
        ys = map(lambda a: a[1], ps)
        plt.plot(xs, ys)
        plt.scatter(xs[-1], ys[-1], c='blue', marker='o', s=40, edgecolors='none')
        plt.xlim([-800, 800])
        plt.ylim([-800, 800])
        ax = plt.axes()
        ax.set_aspect('equal')
        ax.set_xticks([-800, 0, 800], minor=False)
        ax.xaxis.grid(True, which='major')
        ax.set_yticks([-800, 0, 800], minor=False)
        ax.yaxis.grid(True, which='major')
        plt.gca().invert_yaxis()
        plt.gca().set_title('Trajectories for %d frames' % len_thres)

    plt.savefig(out_path4)
    plt.show()


def main():
    global i, ks, cs, len_thres, frame_table, vs, vs_by_id
    global paths, out_path, out_path2, out_path3, out_path4

    # Dict vs holds an array of coords and particle id ([[x,y,id]]) with slice number as a key.
    vs = {}

    # Dict vs_by_id holds coords and particle id ([x,y,id]) with id as a key.
    vs_by_id = {}

    # Dict frame_table holds frame number with particle ID as a key.
    frame_table = {}

    # Define file paths for input and output files.
    in_path = sys.argv[1] if len(sys.argv) >= 2 else 'coords.csv'
    out_path = in_path.replace('_coords.csv', '_connections.csv')
    out_path2 = in_path.replace('_coords.csv', '_connections_v2.csv')
    out_path3 = in_path.replace('_coords.csv', '_paths.json')
    out_path4 = in_path.replace('_coords.csv', '_trajectories.png')

    # len_thres is threhold for drawing trajectories.
    len_thres = int(sys.argv[2]) if len(sys.argv) >= 3 else 50

    print('Input: ' + in_path)
    print('Output: ' + out_path)
    print('Threshold: ' + str(len_thres))

    # Load particle coordinates and slice numbers.
    with open(in_path, 'rU') as f:
        reader = csv.reader(f, delimiter=',')
        reader.next()
        for i, row in enumerate(reader):
            sl = int(row[4])
            if sl not in vs:
                vs[sl] = []
            vs[sl].append(map(float, row[2:4]) + [int(row[0])])
            vs_by_id[int(row[0])] = map(float, row[2:4])
            frame_table[int(row[0])] = sl

    # Make pairs of particles in adjacent frames.
    ks = range(1, max(vs.keys()) + 1)
    cs = []
    for i, k in enumerate(ks):
        v = vs[k] if k in vs else []
        if i >= 1:
            v_prev = vs[ks[i - 1]] if ks[i - 1] in vs else []
            cs.append(map(lambda a: a + [ks[i - 1], ks[i]], connect_adjacent(v_prev, v)))
        else:
            cs.append([])

    # Make paths from pairs of particles.
    cs2 = mk_path_from_adjacent_pairs(cs)
    paths = filter_by_length(cs2, 5)

    write_files()

    # Long paths are chosen for drawing cell migration trajectory.
    paths_long = filter_by_length(paths, len_thres)
    plot_migration(paths_long)


main()