Variable-preconditioned transformed primal-dual method for generalized Wasserstein Gradient Flows

Abstract: We propose a Variable-Preconditioned Transformed Primal-Dual (VPTPD) method for solving generalized Wasserstein gradient flows via a structure-preserving JKO scheme. This is a nontrivial extension of the TPD method [Chen et al. (2023) arXiv:2312.12355] incorporating proximal splitting techniques to address the challenges arising from the nonsmoothness of the objective function. Our key contributions include: (i) a semi-implicit-explicit iteration that combines proximal steps for the nonsmooth part with explicit gradient steps for the smooth part, and variable preconditioners constructed from the Hessian of a regularized objective to balance iteration count and per-iteration cost; (ii) a proof of existence and uniqueness of bounded solutions for the resulting generalized proximal operator, along with a convergent and bound-preserving Newton solver; and (iii) an adaptive step-size strategy to improve robustness and accelerate convergence under poor Lipschitz conditions of the energy derivative. Comprehensive numerical experiments spanning from 1D to 3D settings demonstrate that our method achieves superior computational efficiency compared to existing methods, highlighting its broad applicability through several challenging simulations.