trajectories.py

#!/usr/bin/env python3

"""This script loads the graph of segments, with cost values, and
identifies trajectories.

Input: the graph of segments, with cost values, produced by Graph-tool
(in gt or graphml format). Output: trajectories as lists of segments
and as lists of instantaneous eddies, in JSon files.

By construction, the identifying numbers of segments and of
instantaneous eddies in each trajectory are in ascending order.

"""

import graph_tool
from graph_tool import topology
import json
import sys

def new_traj(ind_traj, traj_prop, n, traj_vert_ind):
    """Assign new trajectory to vertex index n."""

    ind_traj += 1
    traj_prop[n] = ind_traj
    traj_vert_ind.append([n])
    return ind_traj

if len(sys.argv) != 2: sys.exit("Required argument: input-graph")
g = graph_tool.load_graph(sys.argv[1])

if "cost_function" in g.edge_properties:
    closest_succ = g.new_vertex_property('int')
    # Index of the closest direct successor of a node: tail of the
    # out-edge with smallest cost. - 1 means there is no successor.

    # Set closest_succ for all nodes:
    for n in g.iter_vertices():
        all_cost = g.get_out_edges(n,
                                   eprops = [g.edge_properties.cost_function])
        # numpy array with dtype float64

        if all_cost.size == 0:
            # n has no successor
            closest_succ[n] = - 1
        else:
            i_min = all_cost[:, 2].argmin()
            closest_succ[n] = all_cost[i_min, 1]
else:
    # This happens if there are only isolated segments.
    print("No cost values in the input file.")

traj_prop = g.new_vertex_property('int', val = - 1)
# Trajectory index to which a segment belongs.  - 1 means the segment
# does not belong yet to any trajectory. The trajectory index is also
# the index in lists traj_vert_ind, traj_segm and expanded_traj.

ind_traj = - 1 # initialize index of trajectory

traj_vert_ind = []
# An element of traj_vert_ind is a trajectory, represented by a list
# of vertex indices.

for n in topology.topological_sort(g):
    if g.vertex(n).in_degree() >= 1:
        # Find the index of the closest direct predecessor of a node:
        # head of the in-edge with smallest cost.

        all_cost = g.get_in_edges(n, eprops = [g.edge_properties.cost_function])
        # numpy array with dtype float64

        i_min = all_cost[:, 2].argmin()
        closest_pred = all_cost[i_min, 0]

        if closest_succ[closest_pred] == n:
            # Assign to n the trajectory of closest_pred. This means
            # updating ind_traj, traj_prop and
            # traj_vert_ind. closest_pred is already in a trajectory,
            # traj_prop[closest_pred]] != - 1. We extend the
            # trajectory of closest_pred forward.
            traj_prop[n] = traj_prop[closest_pred]
            traj_vert_ind[traj_prop[n]].append(n)
        else:
            ind_traj = new_traj(ind_traj, traj_prop, n, traj_vert_ind)
    else:
        ind_traj = new_traj(ind_traj, traj_prop, n, traj_vert_ind)

    # traj_prop[n] != - 1

print("Number of trajectories: ", len(traj_vert_ind))

# We have trajectories as lists of vertex indices. Now construct
# corresponding lists of segments and lists of instantaneous eddies:

traj_segm = []
# An element of traj_segm is a trajectory, represented by a list of
# segments.

expanded_traj = []
# An element of expanded_traj is a trajectory, represented by a list
# of instantaneous eddies.

for t in traj_vert_ind:
    # t is a trajectory as a list of vertex indices.
    list_segm = [] # list of segments in trajectory t
    list_eddies_traj = [] # list of eddies in trajectory t

    for n in t:
        list_segm.append(g.vp.name[n])

        # List of eddies in segment n:
        list_eddies_segm = list(g.vertex_properties.inst_eddies[n])

        list_eddies_traj.extend(list_eddies_segm)

    traj_segm.append(list_segm)
    expanded_traj.append(list_eddies_traj)

with open("traj_segm.json", "w") as outfile:
    json.dump(traj_segm, outfile, indent = 4)

print("Created traj_segm.json.")

with open("expanded_traj.json", "w") as outfile:
    json.dump(expanded_traj, outfile, indent = 4)

print("Created expanded_traj.json.")
print('Done')