routingkit-cch 0.1.4

Rust bindings for RoutingKit's Customizable Contraction Hierarchies (CCH)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

from dataclasses import dataclass
import routingkit_cch as rk
import geopandas as gpd
from tqdm import tqdm
import more_itertools
from concurrent.futures import ThreadPoolExecutor
import os
import time
import pickle


@dataclass
class Data:
    x: list[float]
    y: list[float]
    tail: list[int]
    head: list[int]
    weights: list[int]
    trips: list[list[int]]


def read_data(path: str):
    edge_df = gpd.read_file(os.path.join(path, "map/edges.shp"), ignore_geometry=True)
    node_df = gpd.read_file(os.path.join(path, "map/nodes.shp"), ignore_geometry=True)
    map_node_osmid_to_id = {
        j: i for i, j in enumerate(node_df[["osmid"]].to_numpy().flatten())
    }
    edge_df["u"] = edge_df["u"].map(map_node_osmid_to_id)
    edge_df["v"] = edge_df["v"].map(map_node_osmid_to_id)

    with open(os.path.join(path, "preprocessed_train_trips_all.pkl"), "rb") as f:
        data_train: list[tuple[tuple[str, int], list[int], tuple[int, int]]] = (
            pickle.load(f)
        )

    x: list[float] = node_df["x"].tolist()
    y: list[float] = node_df["y"].tolist()
    tail: list[int] = edge_df["u"].tolist()
    head: list[int] = edge_df["v"].tolist()
    weights = (edge_df["length"] * 1e3).astype(int).tolist()
    trips: list[list[int]] = [trip for _, trip, _ in data_train]
    return Data(x, y, tail, head, weights, trips)


def sequential(path: str = "data/beijing_data"):
    data = read_data(path)
    order = rk.compute_order_inertial(
        len(data.x),
        data.tail,
        data.head,
        data.x,
        data.y,
    )

    cch = rk.CCH(order, data.tail, data.head, False)
    metric = rk.CCHMetric(cch, data.weights)
    q = rk.CCHQuery(metric)

    start = time.time()
    pbar = tqdm(total=len(data.trips))
    for trip in data.trips:
        pbar.update(1)
        u = data.tail[trip[0]]
        v = data.head[trip[-1]]

        res = q.run(u, v)
        node_path = res.node_path
        assert node_path[0] == u and node_path[-1] == v
        assert sum(data.weights[i] for i in res.arc_path) == res.distance
        assert sum(data.weights[i] for i in trip) >= res.distance
        del res
    return time.time() - start


def parallel(path: str = "data/beijing_data", max_workers: int = 2):
    data = read_data(path)
    order = rk.compute_order_inertial(
        len(data.x),
        data.tail,
        data.head,
        data.x,
        data.y,
    )

    cch = rk.CCH(order, data.tail, data.head, False)
    metric = rk.CCHMetric(cch, data.weights)

    executor = ThreadPoolExecutor(max_workers=max_workers)

    def worker(trips: list[list[int]]):
        q = rk.CCHQuery(metric)
        for trip in trips:
            u = data.tail[trip[0]]
            v = data.head[trip[-1]]

            res = q.run(u, v)
            node_path = res.node_path
            assert node_path[0] == u and node_path[-1] == v
            assert sum(data.weights[i] for i in res.arc_path) == res.distance
            assert sum(data.weights[i] for i in trip) >= res.distance
            del res

    start = time.time()
    futures = [
        executor.submit(worker, list(chunk))
        for chunk in more_itertools.chunked(
            data.trips, len(data.trips) // max_workers + 1
        )
    ]
    for future in futures:
        future.result()
    return time.time() - start


if __name__ == "__main__":
    for city in os.listdir("data"):
        print(f"Processing city: {city}")
        path = os.path.join("data", city)
        print(f"Sequential took {sequential(path)} seconds")
        print(f"Parallel took {parallel(path)} seconds")