Learning to Solve Optimization Problems Constrained with Partial Differential Equations

TL;DR

The paper introduces PDE-OP, a learning-based framework that addresses PDE-constrained optimization by coupling a time-discrete neural operator for PDE dynamics with a proxy optimizer for decision variables. The two networks, a surrogate controller and a discrete-time neural operator, are trained in a self-supervised, primal–dual manner to ensure feasibility and fidelity to the PDE system. Empirical results on voltage control, heat, and Burgers’ equations show that PDE-OP achieves decision quality comparable to Direct methods, adjoint-based approaches, and MPC, while delivering several orders of magnitude faster inference. This end-to-end differentiable approach enables near real-time PDE-constrained decision-making and broad applicability across scientific and engineering domains where PDE dynamics govern optimization tasks.

Abstract

Partial differential equation (PDE)-constrained optimization arises in many scientific and engineering domains, such as energy systems, fluid dynamics and material design. In these problems, the decision variables (e.g., control inputs or design parameters) are tightly coupled with the PDE state variables, and the feasible set is implicitly defined by the governing PDE constraints. This coupling makes the problems computationally demanding, as it requires handling high dimensional discretization and dynamic constraints. To address these challenges, this paper introduces a learning-based framework that integrates a dynamic predictor with an optimization surrogate. The dynamic predictor, a novel time-discrete Neural Operator (Lu et al.), efficiently approximate system trajectories governed by PDE dynamics, while the optimization surrogate leverages proxy optimizer techniques (Kotary et al.) to approximate the associated optimal decisions. This dual-network design enables real-time approximation of optimal strategies while explicitly capturing the coupling between decisions and PDE dynamics. We validate the proposed approach on benchmark PDE-constrained optimization tasks inlacing Burgers' equation, heat equation and voltage regulation, and demonstrate that it achieves solution quality comparable to classical control-based algorithms, such as the Direct Method and Model Predictive Control (MPC), while providing up to four orders of magnitude improvement in computational speed.

Figures (11)

• Figure 1: Temperature regulation process: control actions ${\textcolor{blue}{u}}(t,x)$ steer the system from the initial state ${\textcolor{brown}{y}}(t_0,x)$ towards a target state ${\textcolor{brown}{y}}(T,x)$.
• Figure 2: At each time instant $t_k$, PDE-OP uses a dual network architecture consisting of a proxy optimization model $\mathcal{U}_\omega$ to estimate the basis functions weights $\bm{{c}}(t_k)$ and construct the control action $\hat{{{\textcolor{blue}{u}}}}(t_k,x)$, and a neural-DE model $\mathcal{Y}_\theta$ to estimate the state-variables $\hat{{{\textcolor{brown}{y}}}}(t_{k+1},x)$, with the objective function $\mathcal{J}(\hat{{{\textcolor{blue}{u}}}}(t,x), \hat{{{\textcolor{brown}{y}}}}(t,x))$ capturing the overall loss.
• Figure 3: Voltage Control Optimization; comparison of PDE-OP and classical methods for three target profiles: $y_{\text{target}}(x)=1+0.2\sin(6x)$ (left), $y_{\text{target}}(x)=x+0.5$ (center), and $y_{\text{target}}(x)=1$ (right)
• Figure 4: Voltage control optimization; comparison between the solution estimate (center) of PDE-OP's dynamic predictor $\cal Y_\theta$ at test time with the solution obtained with a numerical algorithm (left) and related absolute error (right).
• Figure 5: Optimal Control of the Heat Equation; comparison of PDE-OP and classical methods for three target profiles: $y_{\text{target}}(x) = 0.6 +0.3(\sin(2x))$, $y_{\text{target}}(x) = x+0.5$, and $y_{\text{target}}(x) = 1$.