Generative AI Models for Learning Flow Maps of Stochastic Dynamical Systems in Bounded Domains

TL;DR

This paper introduces a hybrid machine learning approach combining a neural network and a diffusion model to accurately simulate stochastic differential equations within bounded domains, effectively capturing interior dynamics and boundary exit phenomena.

Contribution

A novel unified method that integrates a neural network for exit prediction with a diffusion model for interior dynamics, specifically addressing boundary exit challenges in bounded domain SDE simulations.

Findings

Accurately models particle exit probabilities with convergence guarantees.

Effectively simulates stochastic dynamics in bounded domains across multiple test cases.

Demonstrates applicability to complex three-dimensional physical systems.

Abstract

Simulating stochastic differential equations (SDEs) in bounded domains, presents significant computational challenges due to particle exit phenomena, which requires accurate modeling of interior stochastic dynamics and boundary interactions. Despite the success of machine learning-based methods in learning SDEs, existing learning methods are not applicable to SDEs in bounded domains because they cannot accurately capture the particle exit dynamics. We present a unified hybrid data-driven approach that combines a conditional diffusion model with an exit prediction neural network to capture both interior stochastic dynamics and boundary exit phenomena. Our ML model consists of two major components: a neural network that learns exit probabilities using binary cross-entropy loss with rigorous convergence guarantees, and a training-free diffusion model that generates state transitions for non-exiting particles using closed-form score functions. The two components are integrated through a probabilistic sampling algorithm that determines particle exit at each time step and generates appropriate state transitions. The performance of the proposed approach is demonstrated via three test cases: a one-dimensional simplified problem for theoretical verification, a two-dimensional advection-diffusion problem in a bounded domain, and a three-dimensional problem of interest to magnetically confined fusion plasmas.