Modern, Efficient, and Differentiable Transport Equation Models using JAX: Applications to Population Balance Equations

• URL: http://arxiv.org/abs/2411.00742v1
• Date: Fri, 01 Nov 2024 17:07:41 GMT
• Title: Modern, Efficient, and Differentiable Transport Equation Models using JAX: Applications to Population Balance Equations
• Authors: Mohammed Alsubeihi, Arthur Jessop, Ben Moseley, Cláudio P. Fonte, Ashwin Kumar Rajagopalan,
• Abstract summary: Population balance equation (PBE) models have potential to automate many engineering processes. In the pharmaceutical sector, crystallization model-based design can contribute to shortening excessive drug development timelines.
• Score: 0.6990493129893112
• License:
• Abstract: Population balance equation (PBE) models have potential to automate many engineering processes with far-reaching implications. In the pharmaceutical sector, crystallization model-based design can contribute to shortening excessive drug development timelines. Even so, two major barriers, typical of most transport equations, not just PBEs, have limited this potential. Notably, the time taken to compute a solution to these models with representative accuracy is frequently limiting. Likewise, the model construction process is often tedious and wastes valuable time, owing to the reliance on human expertise to guess constituent models from empirical data. Hybrid models promise to overcome both barriers through tight integration of neural networks with physical PBE models. Towards eliminating experimental guesswork, hybrid models facilitate determining physical relationships from data, also known as 'discovering physics'. Here, we aim to prepare for planned Scientific Machine Learning (SciML) integration through a contemporary implementation of an existing PBE algorithm, one with computational efficiency and differentiability at the forefront. To accomplish this, we utilized JAX, a cutting-edge library for accelerated computing. We showcase the speed benefits of this modern take on PBE modelling by benchmarking our solver to others we prepared using older, more widespread software. Primarily among these software tools is the ubiquitous NumPy, where we show JAX achieves up to 300x relative acceleration in PBE simulations. Our solver is also fully differentiable, which we demonstrate is the only feasible option for integrating learnable data-driven models at scale. We show that differentiability can be 40x faster for optimizing larger models than conventional approaches, which represents the key to neural network integration for physics discovery in later work.

Related papers

• Text2PDE: Latent Diffusion Models for Accessible Physics Simulation [7.16525545814044]
  We introduce several methods to apply latent diffusion models to physics simulation. We show that the proposed approach is competitive with current neural PDE solvers in both accuracy and efficiency. By introducing a scalable, accurate, and usable physics simulator, we hope to bring neural PDE solvers closer to practical use.
  arXiv  Detail & Related papers  (2024-10-02T01:09:47Z)
• Training Deep Surrogate Models with Large Scale Online Learning [48.7576911714538]
  Deep learning algorithms have emerged as a viable alternative for obtaining fast solutions for PDEs. Models are usually trained on synthetic data generated by solvers, stored on disk and read back for training. It proposes an open source online training framework for deep surrogate models.
  arXiv  Detail & Related papers  (2023-06-28T12:02:27Z)
• Evaluation of Differentially Constrained Motion Models for Graph-Based Trajectory Prediction [1.1947990549568765]
  This research investigates the performance of various motion models in combination with numerical solvers for the prediction task. The study shows that simpler models, such as low-order integrator models, are preferred over more complex, e.g., kinematic models, to achieve accurate predictions.
  arXiv  Detail & Related papers  (2023-04-11T10:15:20Z)
• Neural parameter calibration for large-scale multi-agent models [0.7734726150561089]
  We present a method to retrieve accurate probability densities for parameters using neural equations. The two combined create a powerful tool that can quickly estimate densities on model parameters, even for very large systems.
  arXiv  Detail & Related papers  (2022-09-27T17:36:26Z)
• On the balance between the training time and interpretability of neural ODE for time series modelling [77.34726150561087]
  The paper shows that modern neural ODE cannot be reduced to simpler models for time-series modelling applications. The complexity of neural ODE is compared to or exceeds the conventional time-series modelling tools. We propose a new view on time-series modelling using combined neural networks and an ODE system approach.
  arXiv  Detail & Related papers  (2022-06-07T13:49:40Z)
• Real-time Neural-MPC: Deep Learning Model Predictive Control for Quadrotors and Agile Robotic Platforms [59.03426963238452]
  We present Real-time Neural MPC, a framework to efficiently integrate large, complex neural network architectures as dynamics models within a model-predictive control pipeline. We show the feasibility of our framework on real-world problems by reducing the positional tracking error by up to 82% when compared to state-of-the-art MPC approaches without neural network dynamics.
  arXiv  Detail & Related papers  (2022-03-15T09:38:15Z)
• Interfacing Finite Elements with Deep Neural Operators for Fast Multiscale Modeling of Mechanics Problems [4.280301926296439]
  In this work, we explore the idea of multiscale modeling with machine learning and employ DeepONet, a neural operator, as an efficient surrogate of the expensive solver. DeepONet is trained offline using data acquired from the fine solver for learning the underlying and possibly unknown fine-scale dynamics. We present various benchmarks to assess accuracy and speedup, and in particular we develop a coupling algorithm for a time-dependent problem.
  arXiv  Detail & Related papers  (2022-02-25T20:46:08Z)
• Mixed Effects Neural ODE: A Variational Approximation for Analyzing the Dynamics of Panel Data [50.23363975709122]
  We propose a probabilistic model called ME-NODE to incorporate (fixed + random) mixed effects for analyzing panel data. We show that our model can be derived using smooth approximations of SDEs provided by the Wong-Zakai theorem. We then derive Evidence Based Lower Bounds for ME-NODE, and develop (efficient) training algorithms.
  arXiv  Detail & Related papers  (2022-02-18T22:41:51Z)
• Closed-form Continuous-Depth Models [99.40335716948101]
  Continuous-depth neural models rely on advanced numerical differential equation solvers. We present a new family of models, termed Closed-form Continuous-depth (CfC) networks, that are simple to describe and at least one order of magnitude faster.
  arXiv  Detail & Related papers  (2021-06-25T22:08:51Z)
• Using Differentiable Programming for Flexible Statistical Modeling [0.0]
  Differentiable programming has recently received much interest as a paradigm that facilitates taking gradients of computer programs. We show how differentiable programming can enable simple gradient-based optimization of a model by automatic differentiation. This allowed us to quickly prototype a model under time pressure that outperforms simpler benchmark models.
  arXiv  Detail & Related papers  (2020-12-07T12:33:49Z)
• Hybrid modeling: Applications in real-time diagnosis [64.5040763067757]
  We outline a novel hybrid modeling approach that combines machine learning inspired models and physics-based models. We are using such models for real-time diagnosis applications.
  arXiv  Detail & Related papers  (2020-03-04T00:44:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.