Imperial logo

Scalable SciML Lab

Efficient and differentiable population balance modelling with JAX

—

by

Ben Moseley

Overview

Population balance equations (PBEs) are used to model the evolution of populations of particles over time, such as in crystallisation processes, chemical reactors, and biological cell growth. Solving these equations are crucial in sectors like pharmaceuticals, where they allow us to optimise manufacturing processes involving crystallisation and shorten drug development timelines.

However, traditional PBE solvers face significant challenges: they are computationally intensive, often requiring substantial time to achieve accurate solutions, and they often rely heavily on empirical laws which may lead to inaccurate physics in the model.

In this collaborative project with the Purification And Separation Technology Laboratory at the University of Manchester, our goal is to enable faster manufacturing design cycles by designing more efficient PBE solvers which can improve their accuracy by learning physics from experiments.

Research highlights

JAX-based PBE solver

To enhance the computational efficiency and flexibility of PBE modelling, we re-implemented a traditional, general-purpose finite volume PBE solver using JAX, a high-performance numerical computing library. The reimplementation leverages JAX’s just-in-time compilation, automatic vectorisation, and GPU acceleration to significantly accelerate the traditional solver (Fig. 1).

Fig. 1: High-level flowchart of our JAX-based PBE solver with just-in-time compilation and GPU acceleration.
Fig. 2: Inserting a neural network into our solver to represent the growth rate of particles, instead of using a pre-defined empirical law.

Furthermore, by utilising JAX’s automatic differentiation, we ensured that our solver is fully differentiable. This paves the way for training hybrid workflows where parts of the PBE model and solver are learned. For example, instead of using a pre-determined empirical law to model the growth rate of particles, a neural network can be inserted into the solver to represent the growth rate, and trained by matching the solver output to experimental concentration values (Fig. 2).

Applications and impact

Our JAX PBE solver achieves up to a 300x speedup compared to traditional NumPy-based and C++ solvers. Furthermore, its differentiability enables a 40x speedup when used for large gradient-based optimisation tasks.

Our long term goal is to enable real-time manufacturing and engineering design cycles through efficient numerical modelling and data-driven discovery of PBE models.

Future directions

Our future work includes:

  • Improving the accuracy of our solver by learning empirical laws from experimental data.
  • Accelerating the solver further by designing multi-GPU implementations.
  • Solving practical engineering tasks by working with domain scientists and industry partners.

Contact / get involved

Please contact us or Dr. Ashwin Rajagopalan if you are interested in contributing / learning more about this project.

Publications

Mohammed Alsubeihi, Arthur Jessop, Ben Moseley, Cláudio P. Fonte, Ashwin Kumar Rajagopalan. “Modern, Efficient, and Differentiable Transport Equation Models using JAX: Applications to Population Balance Equations.” Industrial & Engineering Chemistry Research (2025).

Team & collaborators

  • Mohammed Alsubeihi: University of Manchester.
  • Arthur Jessop: University of Manchester.
  • Ben Moseley: Imperial College London.
  • Cláudio Fonte: University of Manchester.
  • Ashwin Rajagopalan: University of Manchester.

This is a collaborative project with the Purification And Separation Technology Laboratory at the University of Manchester.

More projects