necludov / action-matching

Implementation of Action Matching for the Schrödinger equation

Action Matching for Schrödinger Equation Simulation

We demonstrate that Action Matching can learn a wide range of stochastic dynamics by applying it to the dynamics of a quantum system evolving according to the Schrödinger equation. The Schrödinger equation describes the evolution of many quantum systems, and in particular, it describes the physics of molecular systems. Here, for the ground truth dynamics, we take the dynamics of an excited state of the hydrogen atom, which is described by the following equation

$i\frac{\partial}{\partial t}\psi(x,t) = -\frac{1}{|x|}\psi(x,t) -\frac{1}{2}\nabla^2\psi(x,t).$

The function $\psi(x,t): \mathbb{R}^3\times \mathbb{R} \to \mathbb{C}$ is called a wavefunction and it completely describes the state of the quantum system. In particular, it defines the distribution of the coordinates $x$ by defining its density as $q_t(x) := |\psi(x,t)|^2$, which dynamics is defined by the dynamics of $\psi(x,t)$. Below we demonstrate the ground truth dynamics where we project the density $q_t(x)$ on three different planes.

In what follows we illustrate the histograms for the learned dynamics. Since the original distribution is in $\mathbb{R}^3$ we project the samples onto three different planes and draw 2d-histograms. The top row for every model corresponds to the ground truth dynamics (the training data), and the bottom rows corresspond to the learned models.

Action Matching (AM) results visualization

Score Matching (SM) results visualization

Sliced Score Matching (SSM) results visualization