Research

The recently developed Deep Potential [Phys. Rev. Lett. 120, 143001, 2018] is a powerful method to represent general inter-atomic potentials using deep neural networks. The success of Deep Potential rests on the proper treatment of locality and symmetry properties of each component of the network. In this paper, we leverage its network structure to effectively represent the mapping from the atomic configuration to the electron density in Kohn-Sham density function theory (KS-DFT). By directly targeting at the self-consistent electron density, we demonstrate that the adapted network architecture, called the Deep Density, can effectively represent the electron density as the linear combination of contributions from many local clusters. The network is constructed to satisfy the translation, rotation, and permutation symmetries, and is designed to be transferable to different system sizes. We demonstrate that using a relatively small number of training snapshots, Deep Density achieves excellent performance for one-dimensional insulating and metallic systems, as well as systems with mixed insulating and metallic characters. We also demonstrate its performance for real three-dimensional systems, including small organic molecules, as well as extended systems such as water (up to 512 molecules) and aluminum (up to 256 atoms).

Joint work with Yixiao Chen, Weile Jia, Jiefu Zhang, Linfeng Zhang and Lin Lin.

The efficient treatment of long-range interactions for point clouds is a challenging problem in many scientific machine learning applications. To extract global information, one usually needs a large window size, a large number of layers, and/or a large number of channels. This can often significantly increase the computational cost. In this work, we present a novel neural network layer that directly incorporates long-range information for a point cloud. This layer, dubbed the long-range convolutional (LRC)-layer, leverages the convolutional theorem coupled with the non-uniform Fourier transform. In a nutshell, the LRC-layer mollifies the point cloud to an adequately sized regular grid, computes its Fourier transform, multiplies the result by a set of trainable Fourier multipliers, computes the inverse Fourier transform, and finally interpolates the result back to the point cloud. The resulting global all-to-all convolution operation can be performed in nearly-linear time asymptotically with respect to the number of input points. The LRC-layer is a particularly powerful tool when combined with local convolution as together they offer efficient and seamless treatment of both short and long range interactions. We showcase this framework by introducing a neural network architecture that combines LRC-layers with short-range convolutional layers to accurately learn the energy and force associated with a $N$-body potential. We also exploit the induced two-level decomposition and propose an efficient strategy to train the combined architecture with a reduced number of samples.

Joint work with Yifan Peng, Lin Lin and Lexing Ying.

We propose an end-to-end deep learning framework that comprehensively solves the inverse wave scattering problem across all length scales. Our framework consists of the newly introduced wide-band butterfly network coupled with a simple training procedure which dynamically injects noise during training. While our trained network provides competitive results in classical imaging regimes, most notably it also succeeds in the super-resolution regime where other comparable methods fail. This encompasses both (i) reconstruction of scatterers with sub-wavelength geometric features, and (ii) accurate imaging when two or more scatterers are separated by less than the classical diffraction limit. We demonstrate these properties are retained even in the presence of strong noise and extend to scatterers not previously seen in the training set. In addition, our network is straightforward to train requiring no restarts and has an online runtime that is an order of magnitude faster than optimization-based algorithms. We perform experiments with a variety of wave scattering mediums and we demonstrate that our proposed framework outperforms both classical inversion and competing network architectures that specialize in oscillatory wave scattering data.

Joint work with Matthew Li and Laurent Demanet.

We present the first fast solver for the high-frequency Helmholtz equation that scales optimally in parallel for a single right-hand side. The L-sweeps approach achieves this scalability by departing from the usual propagation pattern, in which information flows in a $180$ degree cone from interfaces in a layered decomposition. Instead, with L- sweeps, information propagates in 90◦ cones induced by a Cartesian domain decomposition (CDD). We extend the notion of accurate transmission conditions to CDDs and introduce a new sweeping strategy to efficiently track the wave fronts as they propagate through the CDD. The new approach decouples the subdomains at each wave front, so that they can be processed in parallel, resulting in better parallel scalability than previously demonstrated in the literature. The method has an overall $\mathcal{O} (N/p)\log \omega$ empirical run-time for $N = n^d$ total degrees-of-freedom in a d-dimensional problem, frequency $\omega$, and $p = \mathcal{O}(n)$ processors. We introduce the algorithm and provide a complexity analysis for our parallel implementation of the solver. We corroborate all claims in several two- and three-dimensional numerical examples involving constant, smooth, and discontinuous wave speeds.

Joint work with Matthias Taus, Russell Hewett and Laurent Demanet.

We introduce a new multiscale artificial neural network based on the structure of $\mathcal{H}$-matrices (and $\mathcal{H}^2$-matrices here) . This network generalizes the latter to the nonlinear case by introducing a local deep neural network at each spatial scale. Numerical results indicate that the network is able to efficiently approximate discrete nonlinear maps obtained from discretized nonlinear partial differential equations. In particular, We demonstrate the performance of the multiscale neural network by approximating the solution to the nonlinear Schrödinger equation, as well as the Kohn-Sham map. These mappings are highly nonlinear, and yet can be well approximated by the proposed neural network, with a relative accuracy on the order of $10^{-4}\sim 10^{-3}$. Furthermore, we do not observe overfitting even in the presence of a relatively small set of training samples.

Joint work with Yuwei Fan, Lin Lin, and Lexing Ying.

Quantum embedding theories are playing an increasingly important role in bridging different levels of approximation to the many body Schrödinger equation in physics, chemistry and materials science. In this paper, we present a linear algebra perspective of the recently developed projection based embedding theory (PET) [Manby et al, J. Chem. Theory Comput. 8, 2564, 2012], restricted to the context of Kohn-Sham density functional theory. By partitioning the global degrees of freedom into a ''system'' part and a ''bath'' part, and by choosing a proper projector from the bath, PET is an in principle exact formulation to confine the calculation to the system part only, and hence can be carried out with reduced computational cost. Viewed from the perspective of the domain decomposition method, one particularly interesting feature of PET is that it does not enforce a boundary condition explicitly, and remains applicable even when the discretized Hamiltonian matrix is dense, such as in the context of the planewave discretization. In practice, the accuracy of PET depends on the accuracy of the projector for the bath. Based on the linear algebra reformulation, we develop a first order perturbation correction to the projector from the bath to improve its accuracy. Numerical results for real chemical systems indicate that with a proper choice of reference system, the perturbatively corrected PET can be sufficiently accurate even when strong perturbation is applied to very small systems, such as the computation of the ground state energy of a SiH$_3$F molecule, using a SiH$_4$ molecule as the reference system.

Joint work with Lin Lin.

We propose a hybrid approach to solve the high-frequency Helmholtz equation with point source terms in smooth heterogeneous media. The method is based on the ray-based finite element method (ray-FEM), whose original version can not handle the singularity close to point sources accurately. This pitfall is addressed by combining the ray-FEM, which is used to compute the smooth far-field of the solution accurately, with a high- order asymptotic expansion close to the point source, which is used to properly capture the singularity of the solution in the near-field. The method requires a fixed number of grid points per wavelength to accurately represent the wave field with an asymptotic convergence rate of $\mathcal{O}(\omega^{-1/2})$, where $\omega$ is the frequency parameter in the Helmholtz equation. In addition, a fast sweeping-type preconditioner is used to solve the resulting linear system. We present numerical examples in 2D to show both accuracy and efficiency of our method as the frequency increases. In particular, we provide numerical evidence of the convergence rate, and we show empirically that the overall complexity is $\mathcal{O}(\omega^{2})$ up to a poly-logarithmic factor.

Joint work with Hongkai Zhao, Jun Fang, and Jianliang Qiang.

We present a ray-based finite element method for the high-frequency Helmholtz equation in smooth media, whose basis is learned adaptively from the medium and source. The method requires a fixed number of grid points per wavelength to represent the wave field; moreover, it achieves an asymptotic convergence rate of $\mathcal{O}(\omega^{-1/2})$, where $\omega$ is the frequency parameter in the Helmholtz equation. The local basis is motivated by the geometric optics ansatz and is composed of polynomials modulated by plane waves propagating in a few dominant ray directions. The ray directions are learned by processing a low-frequency wave field that probes the medium with the same source. Once the local ray directions are extracted, they are incorporated into the local basis to solve the high-frequency Helmholtz equation. This process can be continued to further improve the approximations for both local ray directions and high-frequency wave fields iteratively. Finally, a fast solver is developed for solving the resulting linear system with an empirical complexity $\mathcal{O}(\omega ^d)$ up to a poly-logarithmic factor. Numerical examples in 2D are presented to corroborate the claims.

Joint work with Hongkai Zhao, Jun Fang, and Jianliang Qiang.

We present a fast iterative solver for Lippmann-Schwinger equation for high frequency waves scattered by a smooth and compactly supported perturbation. The solver is based on the sparsifying preconditioner and the method of polarized traces. The algorithm has two levels: we first precondition the Lippmann-Schwinger equation using the sparsifying preconditioner that relies on solving a sparse linear system, which is then solved fast using a matrix-free variant of the method of polarized traces in a domain decomposition setting using four sweeps in different directions. The assymptotic cost of applying the preconditioner is $\mathcal{O}(N\log{N})$, where $N$ is the number of degrees of freedom. Numerical experiments in 2D indicate that the number of iterations in both levels depends weakly on the frequency.

Joint work with Hongkai Zhao.

We present a fast solver for the 3D high-frequency Helmholtz equation with heterogeneous, constant density, acoustic media. The solver is based on the nested polarized-traces method, coupled with distributed linear algebra libraries and pipelining to obtain a solver with online runtime $\mathcal{O}(\max(1,R/n)N)$ where $N = n^3$ is the total number of degrees of freedom and $R$ is the number of right-hand sides.

The solver has two levels, the outer level, in which we seek the solution in the whole domain; and the inner level, in which we seek the solution of the Helmholtz equation in a subdomain. We present two implementations of the method depending on the choice for the inner solver. We either use high-quality off-the-shelf distributed linear algebra libraries to solve the local problem at each subdomain, or we use the nested variant of the method of polarized traces, in which each layer is decomposed further in smaller subproblems that we call beams, and the local problem is solved iteratively within each layer. In order to speed up the nested approach we use distributed linear algebra to solve the problems within each beam.

Joint work with Laurent Demanet, Russell Hewett and Adrien Scheuer.

We present a solver for the 2D high-frequency Helmholtz equation in heterogeneous acoustic media, with online parallel complexity that scales optimally as $\mathcal{O}(N/L)$, where $N$ is the number of volume unknowns, and $L$ is the number of processors, as long as L grows at most like a small fractional power of N. The solver decomposes the domain into layers, and uses transmission conditions in boundary integral form to explicitly define "polarized traces", i.e., up- and down-going waves sampled at interfaces. Local direct solvers are used in each layer to precompute traces of local Green's functions in an embarrassingly parallel way (the offline part), and incomplete Green's formulas are used to propagate interface data in a sweeping fashion, as a preconditioner inside a GMRES loop (the online part). Adaptive low-rank partitioning of the integral kernels is used to speed up their application to interface data. The method uses second-order finite differences. The complexity scalings are empirical but motivated by an analysis of ranks of off-diagonal blocks of oscillatory integrals. They continue to hold in the context of standard geophysical community models such as BP and Marmousi 2, where convergence occurs in 5 to 10 GMRES iterations.

Another related work is Preconditioning the 2D Helmholtz equation with polarized traces; some extensions are presented in A short note on the nested-sweep polarized traces method for the 2D Helmholtz equation and in Nested domain decomposition with polariced traces for the 2D Helmholtz Equation. An extension to HDG discretizations can be found in A short note on a fast and high-order Hybridazable Discontinuous Galerkin solver for the 2D high-frequency Helmholtz equation.

Joint work with Laurent Demanet.

In this project we study the application of a timestepping method unconstrained by the CFL condition, for computational acoustic wave propagation in the context of seismic waveform inversion. The numerical scheme is a locally one-dimensional (LOD) variant of alternating dimension implicit (ADI) method. Its main feature is that it allows "upscaled timestepping", i.e., the time step is only restricted by the Nyquist sampling rate. The advantage over traditional explicit timestepping methods is threefold: (1) the complexity is automatically lowered in the low-frequency case, without having to coarsen the spatial mesh, (2) the size of the time step is not adversely affected by high velocity contrasts, and (3) perfectly matched layers (PML) can be much steeper, hence much thinner, without a penalty on the time step. The novelty of the project, from a numerical analysis viewpoint, is that a PML is added to the LOD scheme.

Joint work with Laurent Demanet and Russell Hewett.

In this project, we treat the extrapolation of wave-fields using limited aperture Dirichlet data on a broken line. We prove that, in uniform velocity and under geometrical assumptions, it is possible to extrapolate a harmonic wave-field several wavelengths within machine precision. This extrapolation is accurate inside a cone, whose slope we determine using an asymptotic development. We also provide bounds on the extrapolation error. The analysis relies on the properties of the prolate spheroidal wave functions (PSWF). The theoretical results are illustrated by numerical experiments, and we propose a solver for the Helmholtz equation based on this approach.