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.

Computational Synthetic Biology Enabled through JAX: A Showcase

SBMLtoODEjax: Efficient Simulation and Optimization of Biological Network Models in JAX

A Novel Lecture Program in Bioinformatics: Interdisciplinary Cell Visualization and Modeling

CAX: Cellular Automata Accelerated in JAX

Differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics

SynJax: Structured Probability Distributions for JAX

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

Athena: Modular CAM/CAD Software for Synthetic Biology

JaxLife: An Open-Ended Agentic Simulator

Enabling AI in synthetic biology through Construction File specification

Math Agents: Computational Infrastructure, Mathematical Embedding, and Genomics

evosax: JAX-based Evolution Strategies

Integrating interactive computational modeling in biology curricula

Using Jupyter for Reproducible Scientific Workflows

Slax: A Composable JAX Library for Rapid and Flexible Prototyping of Spiking Neural Networks

$\texttt{synax}$: A Differentiable and GPU-accelerated Synchrotron Simulation Package

Bio Simulators: a web UI for biological simulation

Toward AI-Driven Digital Organism: Multiscale Foundation Models for Predicting, Simulating and Programming Biology at All Levels

Computer-Aided Whole-Cell Design: Taking a Holistic Approach by Integrating Synthetic With Systems Biology

Proof-of-concept: Using ChatGPT to Translate and Modernize an Earth System Model from Fortran to Python/JAX

Human Comprehensible Active Learning of Genome-Scale Metabolic Networks