SEABED: A JAX-powered python package for SEquential Analysis and Bayesian Experimental Design
POSTER
Abstract
Parameter estimation and experimental design are two critical steps in the scientific process and to accelerate scientific discovery it is key to automate and optimize these processes as much as possible. One of the leading ways to approach these tasks is through Bayesian inference and Bayesian experimental design, but large-scale adaptation of these methods have been hindered due to a lack of compatible tools that can be adapted across scientific domains and scaled to large computational resources. To address this, we have developed SEABED (SEquential Analysis and Bayesian Experimental Design), which implements sequential Monte Carlo, or particle-filtering, methods to perform Bayesian inference and experimental design in a black-box and application-agnostic manner. Because the software leverages JAX, functions and methods are differentiable and computations can be executed and scaled across devices like CPUs, GPUs, and TPUs. Moreover, SEABED can be straightforwardly combined with other powerful libraries in the JAX ecosystem for optimization, machine learning, and physics simulations. Finally, we present two example applications to emphasize the flexibility and utility of the software package.
* This work was supported primarily by Laboratory Directed Research and Development (LDRD) funding from Argonne National Laboratory.
Presenters
-
Paul M Kairys
Argonne National Laboratory
Authors
-
Paul M Kairys
Argonne National Laboratory
-
F. Joseph F Heremans
Argonne National Laboratory, Argonne National Lab, Argonne, University of Chicago