joshnguyen99 / partialot

Implementation of gradient-based partial optimal transport solvers (AAAI 2024)

Partial Optimal Transport

This repository contains implementation of algorithms for solving the partial optimal transport plan between two discrete distributions. Check out our AAAI 2024 oral paper below.

Anh Duc Nguyen, Tuan Dung Nguyen, Quang Minh Nguyen, Hoang H. Nguyen, Lam M. Nguyen, and Kim-Chuan Toh. “On Partial Optimal Transport: Revising the Infeasibility of Sinkhorn and Efficient Gradient Methods”. Proceedings of the AAAI Conference on Artificial Intelligence 38.8 (2024), 8090–8098. URL: https://ojs.aaai.org/index.php/AAAI/article/view/28648.

(Partial) Optimal Transport

Here are two algorithms for finding $\varepsilon$-approximate solutions (cf. Definition 1) presented in our paper:

• Adaptive Primal-Dual Accelerated Gradient Descent (APDAGD), with time complexity $\widetilde{\mathcal{O}}(n^{2.5} / \varepsilon)$; and
• Dual Extrapolation (DE), with time complexity $\widetilde{\mathcal{O}}(n^{2} / \varepsilon)$.

Create a Python Environment

First, create an environment.

conda create --name partialot python=3.9.12

Then install the required packages.

pip3 install -r requirements.txt

We will be using Gurobi as the linear program solver in the backend via cvxpy. You can download this solver here; the license is available for free for academic purposes. If you prefer another solver, modify it in ot_solvers/lp.py and pot_solvers/lp.py.