Inverse design and optimization of stochastic particle dynamics in complex flows

Understanding and controlling stochastic particle dynamics carries great significance for a wide array of industrial and biophysical processes. Despite their significance, the precise control of such dynamics has remained elusive. The recent emergence of efficient automatically differentiable simulation techniques suggests gradient-based optimization methods as promising approaches to inverse problems in complex physical systems. In this work, we expand upon a previously developed fully-differentiable numerical framework that combines a computational fluid dynamics solver (JAX-CFD) with molecular/Brownian Dynamics (JAX-MD). We account for fluid flow in complex geometries using immersed boundary methods, while explicitly representing the stochastic equations of motion governing suspended particles, which can range from the Brownian motion of colloids to the run-and-tumble motion of bacteria. Our implementation can be used for both optimization problems and data-driven model parametrization. We demonstrate the effectiveness of our approach on a diverse range of novel problems, including optimizing geometry to mitigate upstream migration of bacteria in biomedical devices and fine-tuning fluid properties to obtain desired flow characteristics through heterogeneous porous media.