search

InkToYou / Tree-Wasserstein

A python module for fast calculation of the wasserstein distance on tree metrics implemented in C++.
BSD 3-Clause "New" or "Revised" License
1 stars 2 forks source link

Get flow for the wasserstein distance #2

Open 100rab-S opened 1 year ago

100rab-S commented 1 year ago

@InkToYou Can the implementation return the wasserstein flow (a.k.a optimal transport matrix) instead of just distance? It would be great, if that's possible. I looked into the code, but there isn't any implementation for the flow.

InkToYou commented 1 year ago

@100rab-S The current implementation gets the dual solutions for computing the wasserstein distance on a tree. So it easily returns the optimal flow on each edge but does not return the primary solution (optimal transport matrix).

I mean that it is possible to get the outputs like:

import tw
# Input
a = [0.1, 0.3, 0.6]
b = [0.3, 0.4, 0.3]
edges = [(0, 1, 1), (0, 2, 1)]

distance, flow = tw.distance(a, b, edges)

# flow: [{(0, 1, 1): -0.1, (0, 2, 1): 0.3}]