(614f) An Efficient and Differentiable Scientific Machine Learning Framework for Crystallization Modelling

Crystallization is a key separation and purification step in pharmaceutical manufacturing, in which the particle size and shape distribution (PSSD) of the product affects downstream processing and its function. Within the framework of Industry 4.0, predicting crystallization performance is essential for improving process design and control.¹ For this purpose, population balance equations (PBEs) are used to describe the evolution of the PSSD of crystals undergoing growth, dissolution, breakage, among others, in a crystallization step at the process scale.² However, numerically solving multidimensional PBEs is challenging, as a high computing power is required to achieve reasonably accurate solutions. Moreover, accurately describing the underlying physics of various fundamental phenomena within the PBE involves laborious trial-and-error efforts. To tackle these challenges, cutting-edge scientific machine learning (SciML) techniques can be leveraged to enhance both computational efficiency and physics comprehension. Unlocking the full potential of SciML techniques requires a differentiable PBE framework, i.e. it can be used to obtain exact derivatives flexibly, without reliance on computationally expensive numerical differentiation.

To this end, as a first step, we develop an efficient, differentiable framework for numerically solving PBE models, which is implemented in JAX,³ a state-of-the-art SciML Python library. This framework is GPU accelerated and uses just-in-time compilation to efficiently solve PBEs. We have used our framework to model the crystallization process in a batch system incorporating fundamental crystallization phenomena such as growth and dissolution. We benchmark our framework by comparing its speed and accuracy against different programming languages and libraries, namely: (i) C++ on CPU, (ii) CUDA C++ on GPU, and (iii) Python NumPy on CPU. Our framework significantly surpasses these benchmarks, providing a speed-up by a factor of up to 40× when compared to NumPy. Finally, as well as offering significant performance gains, our framework is differentiable, which we demonstrate the advantage of by applying our framework to inverse problems, such as parameter estimation. Going forward, we will use our framework as a foundation for incorporating novel SciML tools, facilitating the discovery of the underlying physics, and further computational acceleration.

References
(1) Cote, A.; Erdemir, D.; Girard, K. P.; Green, D. A.; Lovette, M. A.; Sirota, E.; Nere, N. K. Cryst. Growth Des. 2020, 20, 7568–7581.
(2) Ramkrishna, D.; Singh, M. R. Annu. Rev. Chem. Biomol. Eng 2014, 5, DOI: 10.1146/ annurev-chembioeng-060713-040241.
(3) JAX: High-Performance Array Computing — JAX Documentation https://jax.readthedocs.io/en/latest/index.html (accessed 12/22/2023).