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

Population balance equation (PBE) models have the 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, especially in applications where real-time performance is critical. 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 scientific machine learning (SciML) models promise to overcome both barriers through the tight integration of neural networks with physical PBE models. Toward eliminating experimental guesswork, hybrid models facilitate determining physical relationships from data, also known as “discovering physics”. Toward improving efficiency, hybrid models can learn to accelerate numerical models using previous simulations and real-world data. In this study, we aim to prepare for planned SciML integration through a contemporary implementation of an existing PBE algorithm, one with computational efficiency and differentiability (learnability) 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 modeling 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 300× relative acceleration in PBE simulations. This is a significant step toward speeding up slower simulations, even before SciML integration. 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 40× faster for optimizing larger models than conventional approaches, which represents the key to neural network integration for further computational processing improvement and physics discovery in later work.