Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Jin Zeng, Dawei Zhan, Ruchi Guo, Chaozhen Wei

TL;DR

The paper addresses efficient computation of generalized Wasserstein gradient flows with nonlinear mobilities under a structure-preserving JKO framework. It extends the transformed primal–dual (TPD) approach to a Variable-Preconditioned Transformed Primal–Dual (VPTPD) method that combines proximal splitting with semi-implicit-explicit updates, and employs diagonal preconditioners derived from a regularized Hessian to balance per-iteration cost and convergence speed. A theoretical result guarantees the existence and uniqueness of a bounded proximal solution and a bound-preserving Newton solver, complemented by an adaptive step-size strategy to handle poor Lipschitz continuity of the energy derivative. Numerical experiments across 1D–3D problems (including Saturation, Cahn–Hilliard, wetting, Keller–Segel, aggregation-drift, and fracture models) show up to 20× speedups over existing methods, highlighting the method’s robustness and broad applicability for challenging simulations.

Abstract

We propose a Variable-Preconditioned Transformed Primal-Dual (VPTPD) method for solving generalized Wasserstein gradient flows based on the structure-preserving JKO scheme. This is a nontrivial extension of the TPD method [Chen et al. (2025) SIAM J. Sci. Comput.] 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 iterative scheme that combines proximal gradient steps with explicit gradient steps to treat the nonsmooth and smooth terms respectively; (ii) variable-dependent preconditioners constructed from the Hessian of a regularized objective to balance iteration count and per-iteration cost; (iii) a proof of existence and uniqueness of bounded solutions for the generalized proximal operator with the chosen preconditioner, along with a convergent and bound-preserving Newton solver; and (iv) 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-achieving up to a 20$\times$ speedup over existing methods-thereby highlighting its broad applicability through several challenging simulations.

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

TL;DR

The paper addresses efficient computation of generalized Wasserstein gradient flows with nonlinear mobilities under a structure-preserving JKO framework. It extends the transformed primal–dual (TPD) approach to a Variable-Preconditioned Transformed Primal–Dual (VPTPD) method that combines proximal splitting with semi-implicit-explicit updates, and employs diagonal preconditioners derived from a regularized Hessian to balance per-iteration cost and convergence speed. A theoretical result guarantees the existence and uniqueness of a bounded proximal solution and a bound-preserving Newton solver, complemented by an adaptive step-size strategy to handle poor Lipschitz continuity of the energy derivative. Numerical experiments across 1D–3D problems (including Saturation, Cahn–Hilliard, wetting, Keller–Segel, aggregation-drift, and fracture models) show up to 20× speedups over existing methods, highlighting the method’s robustness and broad applicability for challenging simulations.

Abstract

We propose a Variable-Preconditioned Transformed Primal-Dual (VPTPD) method for solving generalized Wasserstein gradient flows based on the structure-preserving JKO scheme. This is a nontrivial extension of the TPD method [Chen et al. (2025) SIAM J. Sci. Comput.] 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 iterative scheme that combines proximal gradient steps with explicit gradient steps to treat the nonsmooth and smooth terms respectively; (ii) variable-dependent preconditioners constructed from the Hessian of a regularized objective to balance iteration count and per-iteration cost; (iii) a proof of existence and uniqueness of bounded solutions for the generalized proximal operator with the chosen preconditioner, along with a convergent and bound-preserving Newton solver; and (iv) 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-achieving up to a 20 speedup over existing methods-thereby highlighting its broad applicability through several challenging simulations.

Paper Structure

This paper contains 19 sections, 1 theorem, 40 equations, 10 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Fully discrete dynamic JKO scheme
  3. Transformed Primal Dual (TPD) method
  4. Review of the TPD method
  5. Comparison with the PDJKO, PrePDJKO methods
  6. Variable-Preconditioned Transformed Primal-Dual method for JKO scheme
  7. The VPTPD iterative scheme
  8. Computation of proximal operator
  9. Adaptive step-size strategy
  10. Complete algorithm of VPTPD method
  11. Numerical experiments
  12. Comparison experiments
  13. 1D Saturation equation
  14. 2D Cahn-Hilliard equation
  15. 3D wetting phenomenon of droplet
  16. ...and 4 more sections

Key Result

Theorem 4.2

For concave mobilites in the form of $M(\rho) = (\rho - \alpha)^{\xi}(\beta - \rho)^{\xi}$ and $M(\rho) = \rho^{\xi}$ with $\xi\in(0,1]$, the problem subproblem admits a unique solution $(\rho^{*},\boldsymbol{m}^{*})$ within the bounds such that $M(\rho^*)\geq 0$. The distribution of the solution $\

Figures (10)

  • Figure 1: Evolution of solutions, energy and relative mass error for the 1D Saturation equation \ref{['1Dsaturation']} for $t \in [0, 15]$. $M = 3.32,\; \Delta x = 0.02,\; \tau = 0.01, \; \lambda_0=0.05$.
  • Figure 2: Evolution of solutions, energy and relative mass error for the 2D CH equation \ref{['2DCH_random']} with randomized initial value for $t \in [0, 10]$. $N_{x} = 128, \; N_{y} = 128, \; \tau = 0.001$.
  • Figure 3: The shape evolution of 3D droplet wetting with $\beta_{w} = \pi/4$ (top) and $3\pi/4$ (bottom) for $t \in [0, 1]$ and $\Omega = [-0.5,0.5]^{2} \times [0,0.4]$. $N_{x} = 64,\; N_{y} = 64,\; N_{z} = 40,\;\tau = 0.1$.
  • Figure 4: Evolution of energy and relative mass error for 3D droplet wetting for $t \in [0, 1]$, $\beta_{w}=\pi/4$. $N_{x} = 64,\; N_{y} = 64,\; N_{z} = 40,\;\tau = 0.1$.
  • Figure 5: Evolution of solutions of the 1D Keller-Segel equation \ref{['Keller-Segel equation']} with different initial values $C = 1$ (left) and $C=15$ (right) for $t \in [0, 2]$. $N_{x} = 800,\; \tau = 0.01$.
  • ...and 5 more figures

Theorems & Definitions (3)

  • Remark 4.1: Selection of $I_u$ and $I_p$
  • Theorem 4.2
  • Proof 1