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 json
import sys
import time

import graph_tool
from graph_tool import topology
import numpy as np

import util_eddies


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


t0 = time.perf_counter()
timings_file = open("timings.txt", "w")
if len(sys.argv) != 2:
    sys.exit("Required argument: input-graph")
g = graph_tool.load_graph(sys.argv[1])
print("Input graph:")
print("Number of vertices:", g.num_vertices())
print("Number of edges:", g.num_edges())
print("Internal properties:")
g.list_properties()
t1 = time.perf_counter()
timings_file.write(f"Loading done in {t1 - t0:.0f} s\n")
t0 = t1

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.ep.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.")

t1 = time.perf_counter()
timings_file.write(f"closest_succ defined in {t1 - t0:.0f} s\n")
t0 = t1
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.

in_deg_prop = g.degree_property_map("in")
out_deg_prop = g.degree_property_map("out")

# Loop to define traj_prop:
for n in topology.topological_sort(g):
    # First, we check whether we have a phantom pattern. Phantom
    # eddies may appear in altimetric data because of
    # interpolation. In a phantom pattern, we have splitting
    # immediately followed by merging. We do not want to let phantom
    # patterns interrupt trajectories.

    if in_deg_prop[n] == 2:
        pred = g.get_in_neighbors(n)

        if (
            in_deg_prop[pred[0]]
            == 1
            == in_deg_prop[pred[1]]
            == out_deg_prop[pred[0]]
            == out_deg_prop[pred[1]]
        ):
            n1 = g.get_in_neighbors(pred[0])[0]
            n2 = g.get_in_neighbors(pred[1])[0]

            if n1 == n2 and out_deg_prop[n1] == 2:
                # We have a splitting event at n1 immediately followed
                # by a merging at n. Check the time interval between
                # splitting and merging. The date of merging is the
                # date of the first eddy in segment n and the date of
                # splitting is the date of the last eddy in segment
                # n1.
                date_merging = util_eddies.node_to_date_eddy(
                    g.vp.name[n], g.gp.e_overestim, only_date=True
                )
                date_splitting = util_eddies.node_to_date_eddy(
                    g.vp.inst_eddies[n1][-1], g.gp.e_overestim, only_date=True
                )

                if date_merging - date_splitting <= 6:
                    # We have a phantom pattern. Get the vertex on the
                    # side of the shortest path:
                    if (
                        g.ep.cost_function[(n1, pred[0])]
                        + g.ep.cost_function[(pred[0], n)]
                        < g.ep.cost_function[(n1, pred[1])]
                        + g.ep.cost_function[(pred[1], n)]
                    ):
                        v_short = pred[0]
                    else:
                        v_short = pred[1]

                    if traj_prop[v_short] != traj_prop[n1]:
                        # Correct the trajectories that were assigned
                        # according to closest successor, that is:
                        # swap the trajectories of pred[0] and
                        # pred[1]:
                        traj_prop[pred[0]], traj_prop[pred[1]] = (
                            traj_prop[pred[1]],
                            traj_prop[pred[0]],
                        )
                        (
                            traj_vert_ind[traj_prop[pred[0]]][-1],
                            traj_vert_ind[traj_prop[pred[1]]][-1],
                        ) = (
                            traj_vert_ind[traj_prop[pred[1]]][-1],
                            traj_vert_ind[traj_prop[pred[0]]][-1],
                        )

                    traj_prop[n] = traj_prop[n1]
                    traj_vert_ind[traj_prop[n1]].append(n)

    if traj_prop[n] == -1:
        if in_deg_prop[n] >= 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.ep.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 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

t1 = time.perf_counter()
timings_file.write(f"traj_prop defined in {t1 - t0:.0f} s\n")
t0 = t1
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)

t1 = time.perf_counter()
timings_file.write(f"traj_segm and expanded_traj defined in {t1 - t0:.0f} s\n")
t0 = t1

# Save:

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

print("Created traj_segm.json.")

with open("expanded_traj.json", "w") as outfile:
    json.dump(
        {
            "orientation": g.gp.orientation,
            "e_overestim": g.gp.e_overestim,
            "traj": expanded_traj,
        },
        outfile,
        separators=(",", ":"),
    )

print("Created expanded_traj.json.")
t1 = time.perf_counter()
timings_file.write(f"Saving done in {t1 - t0:.0f} s\n")
print("Done")