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

import time
import numpy as np
from joblib import Parallel, delayed
from tqdm import trange
import traj_dist.distance as tdist

NUM_TRAJECTORIES = 20
TRAJS = [np.random.uniform(size=(1000, 2)) for _ in range(NUM_TRAJECTORIES)]
WARMUP_RUNS = 1


def test_func(func, num_runs):
    # Warmup
    for _ in range(WARMUP_RUNS):
        _ = func()

    # Measure time
    times = []
    for _ in trange(num_runs):
        start = time.perf_counter()
        _ = func()
        end = time.perf_counter()
        times.append(end - start)
    return np.median(times)


if __name__ == "__main__":
    def test1():
        for i in range(NUM_TRAJECTORIES):
            for j in range(i + 1, NUM_TRAJECTORIES):
                tdist.dtw(TRAJS[i], TRAJS[j], "euclidean")

    def test2():
        tdist.pdist(TRAJS, metric="dtw", type_d="euclidean")

    def test3():
        Parallel(n_jobs=-1)(
            delayed(lambda x, y: tdist.dtw(x, y, "euclidean"))(TRAJS[i], TRAJS[j])
            for i in range(NUM_TRAJECTORIES)
            for j in range(i + 1, NUM_TRAJECTORIES)
        )

    print(test_func(test1, 5))
    print(test_func(test2, 5))
    print(test_func(test3, 5))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

fn compute_pdist_parallel<T, D>(trajectories: &[T], calculator: &D) -> Vec<f64>
where
    T: CoordSequence + Sync,
    D: Distance<T>,
{
    let n = trajectories.len();

    // Create index pairs for all unique pairs (i, j) where i < j
    let pairs: Vec<(usize, usize)> = (0..n)
        .flat_map(|i| ((i + 1)..n).map(move |j| (i, j)))
        .collect();

    // Compute distances in parallel using Rayon's global thread pool
    pairs
        .into_par_iter()
        .map(|(i, j)| calculator.distance(&trajectories[i], &trajectories[j]))
        .collect()
}

 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

/// Python wrapper for the Rust DpResult struct
///
/// This class wraps the Rust DpResult and provides Python-friendly access
/// to the distance and optional matrix.
#[cfg(feature = "python-binding")]
#[gen_stub_pyclass]
#[pyclass(name = "DpResult")]
pub struct PyDpResult {
    /// The inner Rust DpResult
    pub inner: crate::distance::DpResult,
}

#[cfg(feature = "python-binding")]
#[gen_stub_pymethods]
#[pymethods]
impl PyDpResult {
    /// Pickle serialization support using __reduce__
    ///
    /// Uses bincode to serialize the entire DpResult::inner as bytes for better performance.
    /// Returns a tuple (callable, args) that pickle can use to reconstruct the object.
    fn __reduce__(&self, py: Python) -> PyResult<(Py<PyAny>, Py<PyAny>, Py<PyAny>)> {
        use pyo3::prelude::*;
        use pyo3::types::{PyBytes, PyTuple};

        // Import the module and get the helper function
        let module = py.import("traj_dist_rs")?;
        let helper_func = module.getattr("__dp_result_from_pickle")?;

        // Serialize the entire DpResult using bincode
        let serialized =
            bincode::encode_to_vec(&self.inner, bincode::config::standard()).map_err(|e| {
                pyo3::exceptions::PyRuntimeError::new_err(format!("Serialization failed: {}", e))
            })?;

        // Create args tuple containing the serialized bytes
        let bytes_py = PyBytes::new(py, &serialized);
        let args_tuple = PyTuple::new(py, [bytes_py.as_any()])?;

        // Return (helper_func, args, state) where state is None
        Ok((helper_func.unbind(), args_tuple.unbind().into(), py.None()))
    }
}

/// Helper function to create DpResult from pickle data
///
/// Deserializes the DpResult from bincode-encoded bytes.
#[cfg(feature = "python-binding")]
#[gen_stub_pyfunction]
#[pyfunction]
pub fn __dp_result_from_pickle(
    #[gen_stub(override_type(type_repr = "bytes"))] data: &[u8],
) -> PyResult<PyDpResult> {
    bincode::decode_from_slice(data, bincode::config::standard())
        .map(|(dp_result, _)| PyDpResult { inner: dp_result })
        .map_err(|e| {
            pyo3::exceptions::PyRuntimeError::new_err(format!("Deserialization failed: {}", e))
        })
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import time
import polars as pl
from traj_dist_rs import Metric, pdist


def get_trajs():
    df = pl.read_parquet("trajcl_samples.parquet")
    start_idx, end_idx = 0, 7000

    trajs = [df["seq"][idx].to_numpy() for idx in range(start_idx, end_idx)]
    return trajs


if __name__ == "__main__":
    trajs = get_trajs()
    t = time.time()
    pdist(trajs, metric=Metric.dtw(), parallel=True).shape
    print(time.time() - t)

 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

import time
import math
import multiprocessing as mp
from typing import List
import numpy as np
import polars as pl
import traj_dist.distance as tdist

# 优化 1：直接接收 List[np.ndarray]，抛弃 Pandas DataFrame
def _simi_matrix(fn, trajs: List[np.ndarray]) -> List[List[float]]:
    length = len(trajs)
    batch_size = 50
    assert length % batch_size == 0

    tasks = []
    for i in range(math.ceil(length / batch_size)):
        if i < math.ceil(length / batch_size) - 1:
            # 优化 2：把原本传 DataFrame 改为传 trajs 列表
            tasks.append((fn, trajs, list(range(batch_size * i, batch_size * (i + 1)))))
        else:
            tasks.append((fn, trajs, list(range(batch_size * i, length))))

    num_cores = int(mp.cpu_count())

    pool = mp.Pool(num_cores)
    lst_simi = pool.starmap(_simi_comp_operator, tasks)
    pool.close()

    return lst_simi

def _simi_comp_operator(fn, trajs: List[np.ndarray], sub_idx: List[int]):
    simi = []
    length = len(trajs)
    for _i in sub_idx:
        t_i = trajs[_i]  # 优化 3：直接 List 索引，O(1) 且极快，干掉 .iloc
        simi_row = []
        for _j in range(_i + 1, length):
            t_j = trajs[_j] # 同上
            simi_row.append(fn(t_i, t_j))
        simi.append(simi_row)
    return simi

def get_trajs():
    # 统一使用 polars 读取，保证数据准备阶段绝对公平
    df = pl.read_parquet("trajcl_samples.parquet")
    start_idx, end_idx = 0, 7000
    trajs = [df["seq"][idx].to_numpy() for idx in range(start_idx, end_idx)]
    return trajs

if __name__ == "__main__":
    trajs = get_trajs()
    t = time.time() # 计时开始（不包含数据读取部分）
    _simi_matrix(tdist.dtw, trajs)
    print(f"Optimized Baseline Time: {time.time() - t}s")