A Benamou-Brenier Proximal Splitting Method for Constrained Unbalanced Optimal Transport

Abstract: The dynamic formulation of optimal transport, also known as the Benamou-Brenier formulation, has been extended to the unbalanced case by introducing a source term in the continuity equation. When this source term is penalized based on the Fisher-Rao metric, the resulting model is referred to as the Wasserstein-Fisher-Rao (WFR) setting, and allows for the comparison between any two positive measures without the need for equalized total mass. In recent work, we introduced a constrained variant of this model, in which affine integral equality constraints are imposed along the measure path. In the present paper, we propose a further generalization of this framework, which allows for constraints that apply not just to the density path but also to the momentum and source terms, and incorporates affine inequalities in addition to equality constraints. We prove, under suitable assumptions on the constraints, the well-posedness of the resulting class of convex variational problems. The paper is then primarily devoted to developing an effective numerical pipeline that tackles the corresponding constrained optimization problem based on finite difference discretizations and parallel proximal schemes. Our proposed framework encompasses standard balanced and unbalanced optimal transport, as well as a multitude of natural and practically relevant constraints, and we highlight its versatility via several synthetic and real data examples.