Optimal Transport With Order Constraints

The otoc directory contains the following:

1. otoc: optimal transport with order constraints (first-order method).
2. search_otoc_candidates: obtain set of alternative transport plan candidates given a base transport plan.
   • The base transport plan should be a regular OT transport plan (without optimal constriants).
   • Can use the POT package to compute one, see the code examples below.

The above two techniques are described in the below publication; please consider citing if you find this repository useful.

Citation

Fabian Lim, Laura Wynter, Shiau Hong Lim, "Order Constraints in Optimal Transport", 2022. https://arxiv.org/abs/2110.07275.

Usage

To install the python wheel:

pip install git+https://github.com/IBM/otoc@main#egg=otoc

Code Examples

otoc solves the OT with OC formulation with a first-order method.

fromotocimportotocimportnumpyasnpC=np.array([
[1, 5, 3],
[0, 1, 4],
])
a=np.array([.5, .5])
b=np.array([1./3, 1./3, 1./3])
# solve the optimal transport with order constraints# meta give some statisticsX, meta=otoc(
[(0,2)], # X[0,2] >= all othersa, b, C, )
# X# array([[0.16658407, 0. , 0.3333329 ],# [0.16674927, 0.3333329 , 0. ]])X, meta=otoc(
[(0,2), (1,0)], # X[1,0] >= X[0,2] >= othersa, b, C, )
# X# array([[0. , 0.16669247, 0.333301 ],# [0.333301 , 0.16664086, 0. ]])

Alternatively, calling with implementation='iterator' will return an Iterator object, that iterates once per first-order iteration.

# ... intialize a, b, C as above# returns an iterator objectalgo=otoc(
[(0,2)], a, b, C, implementation='iterator'
)
# this iterate through the first-order method until stoppingfor_inalgo:
pass# prints solution and costprint(algo.solution)
# (array([[0.16658407, 0. , 0.3333329 ],# [0.16674927, 0.3333329 , 0. ]]),# 1.4999156499311608)# prints statisticsprint(algo.projection_statistics)
# Statistics(res_primal=8.260397888704896e-05, res_dual=2.4625149137936733e-05, norm_dual=3.570708156906683)# prints iteration countprint(algo.iteration_cnt)
# 41

otoc_candidates learns multiple candidates given a base transport plan by searching OC's in a tree-search manner. This is an Iterator class that iterates as the tree nodes are searched.

# install the python optimal transport package# https://pythonot.github.io/# pip install pot==0.7.0# --- generate a, b, C as above# compute the OT base plan using the python optimal# transport packageimportot# imports POTbase_cost, base_plan=ot.lp.emd2(
a, b, C,
return_matrix=True,
)
# search candidates on top of base planfromotocimportsearch_otoc_candidatesalgo_search=search_otoc_candidates(
a, b, C, base=base_plan['G'],
thresholds=(1.,1.), # small example, so just relax thresholdsk1=10,
k2=5,
k3=1,
)
# iterate each node (which is an Iterator object of otoc)# each node corresponds to an otoc with some learnt order constraintsfornodeinalgo_search:
# run the first order iterations of otoc# until stoppingfor_innode:
pass# third best candidatealgo_search.best_solution(n=2)
# array([[0.16676627, 0. , 0.33333243],# [0.16656706, 0.33333243, 0. ]])# history of nodes exploredalgo_search._history# each gives the index, bound_value, and the otoc iterator (algo)# [History(index=[(0, 2)], bound_value=1.3333333333333333, algo=<otoc.iterative.dual_method.DualMethod object at 0x7fb6242dc9d0>, skip=False),# History(index=[(1, 1)], bound_value=1.3333333333333333, algo=<otoc.iterative.dual_method.DualMethod object at 0x7fb6242dc1d0>, skip=False),# History(index=[(1, 0)], bound_value=1.3333333333333333, algo=<otoc.iterative.dual_method.DualMethod object at 0x7fb6242dc590>, skip=False),# History(index=[(0, 0)], bound_value=1.6666666666666665, algo=<otoc.iterative.dual_method.DualMethod object at 0x7fb6242dce10>, skip=False),# History(index=[(0, 1)], bound_value=2.3333333533333334, algo=<otoc.iterative.dual_method.DualMethod object at 0x7fb6242dd210>, skip=True),# History(index=[(1, 2)], bound_value=2.0, algo=<otoc.iterative.dual_method.DualMethod object at 0x7fb6242dd490>, skip=True)]# check which explored node was the second best candidatealgo_search.best_history_index(n=2)
# 1# confirmalgo_search._history[1].algo.solution# (array([[0.16676627, 0. , 0.33333243],# [0.16656706, 0.33333243, 0. ]]),# 1.5000959909228264)

License

The code in the otoc and tests directories are released under the MIT License. The otoc has about 1000 lines of code (ignoring space and comments) and only one package dependency (numpy).