compute_matrix.py

#!/usr/bin/env python3
import pandas as pd
import numpy as np
import scipy.sparse
import sys
import os


def create_matrix(times, locations):
    stations = locations['stop_name']
    rows = []
    cols = []
    data = []
    table = {}
    for i, name in enumerate(stations):
        table[name] = i
    print("table created")
    for index, row in times.iterrows():
        if index % 1000 == 0:
            print(index)
        rows.append(table[row["departure_station"]])
        cols.append(table[row["arrival_station"]])
        data.append(row['time'] + 0.0001)

    n = len(stations)
    graph = scipy.sparse.csc_matrix(scipy.sparse.coo_matrix((data,
                                    (rows, cols)), shape=(n, n), dtype=np.float64))
    return scipy.sparse.csgraph.floyd_warshall(graph, return_predecessors=True)


def parse_arguments(argv):
    parser = argparse.ArgumentParser(description='Compute distance matrix')
    parser.add_argument('times_path', type=str,
                        help="Path to the times data generated by process_data.py")
    parser.add_argument('locations_path', type=str,
                        help="Path to the locations data generated by process_data.py")
    parser.add_argument("output_path", type=str,
                        help="Path to output folder.")
    return parser.parse_args(argv)


def main(argv):
    args = parse_arguments(argv)
    times = args.times_path
    locations = args.locations_path
    df = pd.read_csv(times)
    loc_df = pd.read_csv(locations)
    dist_matrix, predecessors = create_matrix(df, loc_df)
    np.save(os.path.join(args.output_path, "dist_matrix"), dist_matrix)
    np.save(os.path.join(args.output_path, "predecessors"), predecessors)

if __name__ == "__main__":
    main(sys.argv[1:])