DIMON: Learning Solution Operators of Partial Differential Equations on a Diffeomorphic Family of Domains

TL;DR

DIMON tackles the computational bottleneck of solving PDEs on multiple domains by learning a latent operator on a fixed reference domain and transporting inputs across a family of diffeomorphic domains via $\varphi_\theta$. It combines Large Deformation Diffeomorphic Mapping with neural operator networks (MIONet) to approximate a parameterized PDE operator $\mathcal{F}_0$ and then maps back to the original domain, supported by an extended universal approximation theorem. The framework is demonstrated on three problems—Laplace, reaction–diffusion, and patient-specific LV electrophysiology—achieving accurate predictions with substantial speedups over conventional solvers. DIMON offers a flexible, geometry-aware approach to fast PDE prediction on diffeomorphic domains with potential impact in engineering design and precision medicine, while acknowledging limitations such as transporting only scalar fields and the need for further convergence analysis.

Abstract

The solution of a PDE over varying initial/boundary conditions on multiple domains is needed in a wide variety of applications, but it is computationally expensive if the solution is computed de novo whenever the initial/boundary conditions of the domain change. We introduce a general operator learning framework, called DIffeomorphic Mapping Operator learNing (DIMON) to learn approximate PDE solutions over a family of domains $\{Ω_θ}_θ$, that learns the map from initial/boundary conditions and domain $Ω_θ$ to the solution of the PDE, or to specified functionals thereof. DIMON is based on transporting a given problem (initial/boundary conditions and domain $Ω_θ$) to a problem on a reference domain $Ω_{0}$, where training data from multiple problems is used to learn the map to the solution on $Ω_{0}$, which is then re-mapped to the original domain $Ω_θ$. We consider several problems to demonstrate the performance of the framework in learning both static and time-dependent PDEs on non-rigid geometries; these include solving the Laplace equation, reaction-diffusion equations, and a multiscale PDE that characterizes the electrical propagation on the left ventricle. This work paves the way toward the fast prediction of PDE solutions on a family of domains and the application of neural operators in engineering and precision medicine.

Figures (8)

• Figure 1: Schematics of DIMON for operator learning on a family of diffeomorphic domains.a: A partial differential operator $\mathcal{G}_{\theta}$ is posed on ${\Omega^{\theta_i}}$ together with functions $\mathbf{v}^{\theta_i} (i=1,\dots, N)$ as input. A map $\varphi_{\theta}$, parameterized by the shape parameter $\theta$, transports $\mathbf{v}^{0}_{\theta}$ posed on reference domain ${\Omega^0}$ onto ${\Omega^\theta}$ and identifies point correspondences across all ${\Omega^\theta}$ and ${\Omega^0}$ (indicated by red marks and blue dots). There exists a latent operator ${\mathcal{F}_{0}}$ such that $\mathcal{G}_{\theta}(\mathbf{v}^\theta) = {\mathcal{F}_{0}}(\theta, \mathbf{v}^0_\theta)\circ \varphi_{\theta}^{-1}$. b: The operator ${\mathcal{F}_{0}}$ is approximated by a neural operator $\tilde{\mathcal{F}}_{0}$: a shape parameter $\theta$ (typically estimated from data) and the transported function $\mathbf{v}^{0}_{\theta}$ are the inputs of branch networks; the coordinates of the reference domain $x$ are the inputs of the trunk network.
• Figure 2: DIMON for learning solution operator of the Laplace equation on parameterized 2D domains.a: Data generation on $N$ parametric 2D domains. Landmark points are matched across 2D domains with different meshing. Shape parameter $\tilde{\theta}\in \mathbb{R}^{p'}$ is calculated by PCA on the deformation field at the landmark points. Here, we adopt the first 15 principal components as $\tilde{\theta}$. The PDE is solved using finite elements and mapped onto the reference domain. b: Neural operator structure. Shape parameter $\tilde{\theta}$ and boundary condition $\tilde{v}$ are fed into the geometry (Geo.) and boundary condition (B.C.) branch, respectively. Coordinates in the reference domain $\mathbf{x}$ are the input of the trunk net. c: Network predictions ($u_{NN}$) and their absolute errors with numerical solutions ($|u_{FEM} - u_{NN}|$) on three testing domains. d: Distributions of mean absolute errors $\varepsilon_{abs}$ and relative $L_{2}$ error ($\varepsilon_{l2}$) on the testing dataset in log scale. The testing dataset contains $n=200$ cases. e The average of $\varepsilon_{l2}$ ($\bar{\varepsilon}_{l2}$) for the training (red) and testing dataset (blue) is plotted against the number of training cases.
• Figure 3: DIMON for learning solution operator of the Reaction-Diffusion (RD) equation on parametric 2D domains.a: Data preparation on $N$ parametric domains. Each domain is parameterized by either the deformation or the momentum field at the landmark points with $p'$ principal component coefficients. A finite element method is adopted to numerically solve the RD equation with various initial conditions (I.C.) on ${\Omega^{\theta_i}}$, set as an annulus, and map the solution to ${\Omega^0}$ at different time snapshots. b: Neural operator structure. Truncated shape parameter $\tilde{\theta}$ and boundary condition $\tilde{v}$ are fed into the geometry (Geo.) and boundary condition (B.C.) branch, respectively. Coordinates in the reference domain $\mathbf{x}$ are the input of the trunk net. c: Network predictions with true shape parameters of the affine transformation at $t = 0, 1,$ and $2$s and their absolute errors on three testing domains. d: Average relative $L_{2}$ error ($\varepsilon_{l2}$) on the testing cases based on $\tilde{\theta}$ obtained from 4 approaches: affine with true shape parameters, LDDMM with PCA on momentum and deformation, and LDDMM with true shape parameters. The averaged $\varepsilon_{l2}$ for different approaches range from $3.2\times10^{-2}$ to $7.5\times 10^{-2}$, indicating robustness of DIMON with respect to the way the diffeomorphisms are learned and approximated. The number of testing cases is $n=600$.
• Figure 4: Predicting electrical signal propagation on patient-specific left ventricles (LV). a: Data preparation for the network training. The cohort assembled consists of 1007 cardiac magnetic resonance (CMR) scans of patients from Johns Hopkins Hospital, based on which LV meshes and fiber fields are generated. Next, we simulate a paced beat for 600 ms at an endocardial site until the LV is fully repolarized and record the activation times (ATs) and repolarization times (RTs) for network training. b: Pacing locations (marked in white balls) located on the endocardium with a reduced American Heart Association (AHA) LV segments from zone 0 to 6. c: Network structure. c: Universal ventricular coordinate (UVC) calculation for registering points across geometries.
• Figure 5: Network prediction of ATs and RTs from the testing cohort. Top panel: LV shape with their predicted ATs and RTs and absolute errors. Red stars are the pacing locations. The colorbar is selected to be consistent with the CARTO clinical mapping system, with red denoting early activation and progressing to purple for late activation. Bottom left panel: Relative and absolution error distributions for ATs and RTs. Bottom right panel: Bar plots of the average relative $L_2$ errors from 10-fold cross-validation.

Theorems & Definitions (1)

• Theorem 1: Adapted from jin2022mionet