Generative AI models for learning flow maps of stochastic dynamical systems in bounded domains

Simulating stochastic differential equations (SDEs) in bounded domains presents significant computational challenges due to particle exit phenomena, which require the 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 hybrid diffusion model that combines a conditional diffusion model with an exit prediction neural network to capture both interior stochastic dynamics and boundary exit phenomena. Specifically, the proposed hybrid diffusion model consists of two major components: a neural network that learns exit probabilities using binary cross-entropy loss with rigorous convergence guarantees, and a conditional 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 with three test cases: a simple one-dimensional problem for theoretical verification, a two-dimensional advection-diffusion problem in a bounded domain, and a three-dimensional transport problem of interest to magnetically confined fusion plasmas.