A parameterized Wasserstein Hamiltonian flow approach for solving the Schr\"odinger equation

TL;DR

This paper introduces a novel Wasserstein Hamiltonian flow method using neural networks to efficiently solve the time-dependent Schrödinger equation, especially in high-dimensional settings.

Contribution

It reformulates the Schrödinger equation as a Hamiltonian system in Wasserstein space and employs neural ODEs for scalable, parameterized solutions.

Findings

Effective in high-dimensional problems

Demonstrates promising numerical results

Provides a new generative modeling perspective

Abstract

In this paper, we propose a new method to compute the solution of time-dependent Schr\"odinger equation (TDSE). Using push-forward maps and Wasserstein Hamiltonian flow, we reformulate the TDSE as a Hamiltonian system in terms of push-forward maps. The new formulation can be viewed as a generative model in the Wasserstein space, which is a manifold of probability density functions. Then we parameterize the push-forward maps by reduce-order models such as neural networks. This induces a new metric in the parameter space by pulling back the Wasserstein metric on density manifold, which further results in a system of ordinary differential equations (ODEs) for the parameters of the reduce-order model. Leveraging the computational techniques from deep learning, such as Neural ODE, we design an algorithm to solve the TDSE in the parameterized push-forward map space, which provides an alternative approach with the potential to scale up to high-dimensional problems. Several numerical examples are presented to demonstrate the performance of this algorithm.