A scalable matrix-free spectral element approach for unsteady PDE constrained optimization using PETSc/TAO

Highlights

• Solving partial differential constrained optimization requires multiple expertise.

• Adjoints effective for solving partial differential constrained optimization.

• Math libraries ideal for solving partial differential constrained optimization.

Abstract

We provide a new approach for the efficient matrix-free application of the transpose of the Jacobian for the spectral element method for the adjoint-based solution of partial differential equation (PDE) constrained optimization. This results in optimizations of nonlinear PDEs using explicit integrators where the integration of the adjoint problem is not more expensive than the forward simulation. Solving PDE constrained optimization problems entails combining expertise from multiple areas, including simulation, computation of derivatives, and optimization. The Portable, Extensible Toolkit for Scientific computation (PETSc) together with its companion package, the Toolkit for Advanced Optimization (TAO), is an integrated numerical software library that contains an algorithmic/software stack for solving linear systems, nonlinear systems, ordinary differential equations, differential algebraic equations, and large-scale optimization problems and, as such, is an ideal tool for performing PDE-constrained optimization. This paper describes an efficient approach in which the software stack provided by PETSc/TAO can be used for large-scale nonlinear time-dependent problems. Time integration can involve a range of high-order methods, both implicit and explicit. The PDE-constrained optimization algorithm used is gradient-based and seamlessly integrated with the simulation of the physical problem.

Introduction

Fitting numerical or experimental observations to determine parameters, identifying boundary conditions that satisfy certain observations, optimizing an objective function of a simulation solution, accelerating simulations through their long transients, such as the spin-up problem [1], and many more computations fall within the field of partial differential equation (PDE)-constrained optimization and inverse problems. Despite their widespread utility, the solution of such problems is plagued by bottlenecks ranging from mathematical issues to high computational costs, high input–output costs, and especially software complexity. Several codes, such as JuMP [2], and Python libraries, such as the dolfin-adjoint project, [3], address these issues; however, they exhibit limitations for large-scale problems.

A unique contribution of this paper is the efficient, scalable application of the transpose of the Jacobian matrix-vector product for spectral elements at a cost not much more than the forward application of the nonlinear operator using tensor contractions. The matrix-free approach avoids the cost and storage required by forming the Jacobian explicitly. Our implementation of the optimization process leverages the extensive range of time integrators available in the Portable, Extensible Toolkit for Scientific in (PETSc) [4] and the optimization algorithms readily available in the Toolkit for Advanced Optimization (TAO) [5]. In particular, we utilize adjoints of numerical integrators as in [6].

Time-dependent PDE-constrained optimization problems can be posed in two ways: either by fully discretizing, in both space and time, the entire Karush–Kuhn–Tucker (KKT) system (that is, solving for all unknowns at all time steps simultaneously, sometimes called the all-at-once approach) [7] or by decoupling the direct problem and its adjoint [8], [9], [10], [11], [12] and solving the optimization problem via a forward/backward time-stepping loop, with appropriate initial and boundary conditions depending on the objective function. This work focuses on the latter, in particular on discrete adjoint approaches, targeting a minimal computational cost implementation of the adjoint equation. Optimization algorithms that involve time integration are computationally expensive, for example, for an optimization that requires $m$ iterations, one needs to perform $2m$ integrations each of $N$ time steps, where the backward integration can be more expensive than the forward integration. Employing all possible computational accelerations is crucial since any speedup accumulates significantly with the number of time steps and optimization iterations.

The core strategy presented in this study is to solve the inverse (or PDE-constrained optimization) problem in a purely nonlinear manner. Considerable work has previously relied on linearized adjoints; that is, the model is first linearized, and the adjoints are computed for that linearized problem [13]. Here we showcase the solution of the nonlinear time-dependent viscous Burgers equation utilizing PETSc with TAO.

In previous work [14], we addressed the often-overlooked data intensity and computational intensity aspects of PDE-constrained optimization. The data intensity for nonlinear problems stems from the nature of the backward-in-time component, which depends at every time instance on its forward counterpart. A robust treatment for balancing the compute time versus trajectory storage has been implemented in PETSc, using the checkpointing algorithm revolve [15] as done by Schnaen et al. [14]. Previous work for ordinary differential equations (ODEs) and differential algebraic equations (DAEs) can be found in the work of Zhang and Sandu [16].

It is well known to the PDE-constrained optimization community that continuous derivatives of the adjoint equation differ in nature from their discrete counterpart and may provide different gradients that affect the convergence of the optimization algorithm [17], [18]. We do not consider this issue in this paper; however, we note that the SUNDIALS package [19] provides continuous adjoint capabilities [20].

The paper is organized as follows. The time-dependent PDE-constrained optimization problem is stated in an operator fashion that allows for operators including diffusion, advection, and nonlinear ones such as the $u·\nabla u$ operator in Burgers equation. Then the main aspects of a spectral element method are presented, including a new approach for efficient matrix-free application of the transpose Jacobian matrix-vector product. This is followed by a discussion of the discrete adjoint approach. We then present an overview of the PETSc ODE/DAE integrators and their adjoint integrations, followed by a description of TAO and its gradient-based solvers. The results section focuses on the complete solution process for performing PDE-constrained optimization using the spectral element method and PETSc/TAO. We include a scalability study with over 2 billion unknowns and 130,000 cores. The final section briefly summarizes our conclusions and potential future work.