PaperlayerView exact revision on arXiv

Source-linked AI summary

Differentiable Hybrid Modelling for Learning and Optimising Chemical Transport Processes from Experimental Data

Arthur Jessop, Mohammed Alsubeihi, Ben Moseley, Ashwin Kumar Rajagopalan

arXiv:2609.04011v1cs.CEcs.LGcs.NAmath.NA

TL;DR

Transport models can be limited by fixed constitutive laws, uncertain initial conditions, and black-box surrogates confined to their training tasks. This paper combines learnable components with a differentiable population balance solver to learn laws and initial conditions and optimize process conditions. The authors report that assimilating initial conditions recovers a trajectory loss comparable to the hybrid model and demonstrate temperature-profile optimization on a held-out experiment.

  • Problem

    Hand-picked constitutive laws and assumed-known initial conditions can limit transport models, while crystallization measurements may not capture particle shape in multiple dimensions.

    Source: §1 · pp. 2–3

  • Method

    The framework embeds learnable components in a differentiable mechanistic population balance solver to learn kinetic laws and assimilate initial conditions from experimental data.

    Source: §1 · p. 3, §2.1 · p. 5, §2.2 · p. 6

  • Results

    Assimilating initial conditions for the empirical law yields trajectory loss 0.195 against 0.176 for the hybrid, while optimizing the held-out E5 temperature profile gives an MAE of 0.05°C.

    Source: §3 · p. 12

  • Takeaways & Limitations

    Changing the optimized free variables lets the same solver support law recovery, initial-condition diagnosis, and operating-condition design, including a profile-optimization task not posed during training.

    Source: §3 · pp. 11–12

    The demonstration covers one material system and thirteen experiments, and the product-aspect-ratio optimization is simulated and requires experimental validation.

    Source: §2.3.2 · p. 10, §3 · p. 13

Abstract

from arXiv · show

Reliable transport models are essential when modelling and optimising many chemical engineering processes, yet, most models assume hand-picked constitutive laws which may not reflect reality, and often assume initial conditions are known exactly. Both restrictions can significantly bias model predictions and lead to systematic error when used in predictive and control settings. Black-box neural surrogate alternatives for modelling can better match real example data, but are confined to the task they were trained on and cannot be interrogated for physical consistency. Here we introduce a general-purpose differentiable hybrid modelling framework for transport processes, specifically for the case of population balance equations. Our framework integrates a JAX finite volume population balance solver with learnable neural network components which are trained to both discover constitutive laws and fit initial conditions from real experimental data, allowing us to better model real experimental transport systems. Furthermore, we use our framework for process optimisation, using its differentiability to allow us to direct optimising experimental settings for quantities of interest. This work highlights the huge potential of such differentiable hybrid modelling frameworks for learning and optimising any given chemical separation which involves mass, energy, and/or momentum transport.

1 Introduction

Transport models can be limited by hand-picked constitutive laws, uncertain initial conditions, and constraints in available measurements. The paper introduces a differentiable hybrid framework that jointly learns transport laws and initial conditions and supports process optimization.

1 Introduction

  • Hand-picked constitutive laws may miss underlying phenomena, while treating experimentally measured initial conditions as exact can bias model development.The passages note that dynamic-system solutions can be sensitive to initial conditions.Source: §1 · p. 2
  • Population balance models track particle size and shape distributions, but system-specific kinetics and limited measurement dimensionality make them difficult to formulate.The paper uses crystallization as a demonstration system and notes that many monitoring tools reduce particle shapes to one dimension.Source: §1 · pp. 2–3
  • Black-box surrogates can fit data but may lack physical constraints and remain confined to their training task, motivating a unified alternative.The introduction identifies a lack of a unified framework spanning forward simulation, model inversion, experimental-condition optimization, and process-outcome targeting.Source: §1 · p. 3
  • The JAX-based differentiable framework embeds learnable components in mechanistic transport equations and supports model building, simulation, and process design.Automatic differentiation propagates gradients through the solver to the optimized inputs.Source: §1 · p. 3, Figure 1 · p. 4
  • Jointly learning the rate law and initial conditions reduces validation trajectory loss sevenfold versus the conventional workflow, without additional experiments.The conventional workflow fits a hand-picked rate law to measured initial conditions.Source: §1 · p. 4

2 Results

The framework jointly learns dissolution kinetics and calibrates initial conditions within a differentiable population-balance solver, then uses differentiability to optimize operating conditions. On the glutamic-acid system, calibration improves the trajectory fit, and simulations identify operating profiles for trajectory matching and aspect-ratio reduction.

2.1 System & Framework

  • The solver tracks particle size and shape distributions over time, with swappable kinetic terms, and the demonstration models L-glutamic acid dissolution in water.The observed quantities include particle length, width, and solute concentration.Source: §2.1 · p. 5

2.2 Model Building & Data Assimilation

  • Joint training recovered dissolution rates that vanish at S = 1 and have negligible temperature dependence across 20 random-seed models.The learned supersaturation response is nonlinear; its ensemble spread widens where input density is low.Source: §2.2 · p. 7
  • Calibrated initial conditions shifted only marginally; for E7, mean length and width shifted 24.1 µm and 1.4 µm, while distribution mass was preserved.Assimilation adjusts distribution shape, with renormalization to measured seed mass.Source: §2.2 · p. 7, Figure 2 · p. 8
  • Trajectory loss was 1.264 for the conventional pipeline and 0.176 for the full differentiable pipeline with learned kinetics and assimilated initial conditions.The comparison uses experiments included in hybrid-model training, so it is a workflow comparison rather than a true validation comparison; no additional experiments were required.Source: §2.2 · pp. 6–7, Figure 2 · p. 8

2.3 End-to-End Process Optimisation

  • The differentiable solver supports operating-condition optimization with a frozen hybrid model for trajectory matching and alternative size-dependent kinetics for product-specification optimization.Both tasks apply temperature and supersaturation constraints; product optimization also requires final product volume at least as large as seed volume.Source: §2.3 · pp. 8–9

2.3.1 Trajectory Optimisation

  • T(t) MAE fell from 0.10°C to 0.05°C against a smoothed initial guess when initial conditions were assimilated for held-out experiment E5.The recovered profile remained within the Smin bound and closely matched the experimental profile.Source: §2.3.1 · p. 9

2.3.2 Product Aspect Ratio Optimisation

  • Aspect ratio fell from 6.50 to 4.88, a 25% reduction, as the simulated profile dissolved the seeds before growth restored product volume to 100%.This product-optimization result is simulated and requires experimental validation.Source: §2.3.2 · p. 10, Figure 3 · p. 11

3 Discussion

The same differentiable PBE solver supports model building, initial-condition assimilation, and operating-condition optimisation by changing the variables being optimised. The discussion shows that better trajectory fit alone does not guarantee physically consistent kinetics, while product-design results remain hypotheses requiring experimental validation.

3 Discussion

  • The same solver supports forward simulation, physical-law recovery, initial-condition assimilation, and operating-condition optimisation by changing the optimised variables.The framework keeps the solver and governing equations fixed across these tasks.Source: §3 · p. 11
  • A hybrid model lowered validation loss to 0.812 versus 1.264 yet learned nonzero dissolution at equilibrium and inverted temperature dependence, persisting across 20 initialisations.The authors attribute this inconsistency to fitting uncertain data as ground truth with a flexible neural network embedded in an initial-condition-sensitive PDE.Source: §3 · pp. 11–12
  • Assimilating initial conditions recovered physically consistent behaviour and comparable trajectory losses: 0.195 for the empirical law and 0.176 for the hybrid.For this dataset, the authors identify measured initial conditions, rather than functional form, as the limiting factor.Source: §3 · p. 12
  • Optimising temperature for held-out experiment E5 recovered a profile close to the observed one, with an MAE of 0.05°C, although that profile was not an optimisation target.The hybrid was frozen and asked to design an operating profile for a task absent from its training.Source: §3 · p. 12
  • For a product specification, the framework proposed an operating profile that preferentially redistributes mass toward crystal width through the size dependence of the growth law.This is a simulated result that requires experimental validation; the demonstration covers one material system and thirteen experiments, with a size-independent learned rate.Source: §3 · pp. 12–13
  • The framework requires no additional data over empirical modelling, preserves governing conservation laws through the embedded PBE solver, and can extend to other PBE kinetics and transport equations.The authors argue existing experimental datasets may contain information conventional fitting workflows do not extract.Source: §3 · p. 13

4.1 Experimental Data & Split

The study uses a thirteen-experiment benchmark dataset and holds out two validation experiments for process optimisation. The assimilation approach calibrates initial conditions individually using the remaining experiments.

4.1 Experimental Data & Split

  • The benchmark comprises 13 experiments, E1–E13, generated using a procedure detailed by Bӧtschi et al..The dataset is partly adapted from the same authors’ 2018 work, and DISCO enabled real-time monitoring of crystal size and shape.Source: §4.1 · p. 13
  • Experiments E5 and E12 were reserved from the validation set for process optimisation, while the standard approach used the Bӧtschi et al. split.For initial-condition assimilation, all experiments except E5 and E12 were used to train, with each initial condition calibrated individually.Source: §4.1 · p. 13
  • The time-series data show operating condition T(t) and predictors c, L1, and L2, with experiments labelled by number, seed mass, and data split.Source: §4.1 · p. 13

4.2 Modelling

The modelling setup couples a population balance equation for the crystal population with a liquid-phase mass balance. It solves the system using a fully discrete finite-volume method and forward Euler time stepping.

4.2 Modelling

  • The batch crystallisation model assumes no flow into or out of the crystallizer during each experiment.Source: §4.2 · p. 14
  • The population balance equation tracks crystal number density over characteristic lengths, with convection rate U describing crystal growth or dissolution.The system state Y includes information such as temperature and concentration.Source: §4.2 · p. 14
  • The crystal-population PBE is coupled to a mass balance for the liquid phase, where c is solute concentration and crystal volume contributes to the balance.The crystal density is 1.59 × 10^-12 g µm^-3 and the shape factor accounting for noncuboidal crystals is π/4.Source: §4.2 · p. 14
  • Initial and boundary conditions specify the initial crystal distribution f0, also called the seed distribution, and initial concentration c0.The fully discrete solver uses a finite-volume method for space and forward Euler with a CFL condition for time.Source: §4.2 · p. 14

4.3 Empirical Kinetic Models

The empirical kinetic models vary the convection term while keeping the solver and PDE fixed. The section introduces M1 as a predictive baseline and describes a combined size-dependent growth–dissolution model for aspect-ratio optimisation.

4.3 Empirical Kinetic Models

  • The system-specific convection term U is modelled as a function of temperature T and relative supersaturation S, while the underlying solver and PDE remain unchanged.The section also introduces system-specific solubility c*(T).Source: §4.3 · pp. 14–15

4.3.1 M1 - predictive performance baseline

  • M1 is the two-parameter, size-independent linear dissolution model used as the predictive-performance baseline with published parameters.M4 is a four-parameter, size-dependent nonlinear model; both were proposed by Bӧtschi et al..Source: §4.3 · p. 15, §4.3.1 · p. 15

4.3.2 Size-dependent growth and dissolution model

  • The aspect-ratio optimisation uses a combined model that applies M4 dissolution for S < 1 and a size-dependent growth model for S > 1.The growth model from is size-dependent only in its G1 growth rate.Source: §4.3.2 · p. 15

4.4 Hybrid Dissolution Rate

The hybrid dissolution model replaces the population balance equation’s convection term with a neural network while retaining the finite volume solver. Its inputs and outputs, constraints, and rate scaling are specified to model dissolution.

4.4 Hybrid Dissolution Rate

  • The model replaces the PBE convection term with a neural network while leaving the finite volume solver intact, preserving its number and mass balance constraints.The network describes the dissolution rate.Source: §4.4 · p. 16
  • The multilayer perceptron maps two state variables to two dissolution rates and clips its output to zero when S ≥ 1.The main-text configuration has three hidden layers of 20 neurons; smaller tested networks showed similar predictive performance.Source: §4.4 · p. 16
  • Inputs S and T are standardised using means and standard deviations calculated from observed experimental-data values.The standardised inputs are passed to the network.Source: §4.4 · p. 16
  • A negative softplus output ensures predicted rates are negative, while reference rates scale predictions because dissolution rates are not directly measurable.The clip sets rates to zero at S ≥ 1, but the network does not structurally force rates to approach zero as S approaches 1.Source: §4.4 · p. 16

4.5 Hybrid Training Procedure

The training procedure fits the hybrid model to time-resolved concentration and particle-size data. It compares optimizing dissolution parameters alone with also assimilating initial conditions.

4.5 Hybrid Training Procedure

  • Training uses three time-resolved predictors: dissolved-molecule concentration and average particle length and width.The loss equally weights the predictors and standardises experimental data and model predictions per predictor before calculating mean squared error.Source: §4.5 · pp. 16–17
  • Case A optimises dissolution-model parameters, while Case B also assimilates the initial particle-size distribution and concentration as free variables.The cases minimise the same loss over their respective free variables.Source: §4.5 · p. 17
  • Adam minimises the loss with a learning rate of 10^-3, and Case B applies ReLU updates to keep the particle-size distribution nonnegative.Case A trains for up to 500 epochs with validation-based early stopping; Case B uses training loss for early stopping and has no data split.Source: §4.5 · p. 17

4.6 Process Optimisation

Process optimisation varies temperature over time under equipment and kinetic-model constraints. The paper applies this setup to trajectory matching and product aspect-ratio optimisation.

Parametrisation

  • Temperature is represented by 300 linearly spaced values over the experiment, and some tasks also optimise the initial particle-size distribution and concentration.The free variables are denoted by ϕ.Source: Parametrisation · p. 17

Bounds and penalties

  • Temperature is hard-bounded to the experimental range of 20–41°C, while soft penalties constrain its rate of change and supersaturation.The vessel limit is |dT/dt| ≤ 0.05 °C min^-1; supersaturation is kept above Smin and, where a growth law is present, below Smax.Source: Bounds and penalties · p. 17

Optimisation procedure

  • Adam optimises the temperature profile for 500–1000 iterations, with constraint-penalty weights ramped during optimisation.For trajectory tracking, the task loss is mean squared error between target and simulated length and width trajectories.Source: Optimisation procedure · p. 18

4.6.1 Synthetic Trajectory Optimisation

  • Synthetic trajectory optimisation fits temperature profiles to particle-size targets generated from synthetic stepped-ramp or sigmoid reference profiles.The task uses the E5 initial condition and repeats optimisation 20 times with random shifts to the initial profile’s endpoints to test robustness to initial-guess noise.Source: §4.6.1 · p. 18

4.6.2 Experimental Trajectory Optimisation

  • Experimental trajectory optimisation fits smoothed measured particle-size trajectories, including initial conditions as free variables because they were unseen during training.Temperature starts from a smoothed version of the recorded profile.Source: §4.6.2 · p. 18

4.6.3 Product Aspect Ratio Optimisation

  • The product-aspect-ratio task minimises the product length-to-width ratio subject to a soft constraint that product volume is at least seed volume.The temperature profile is initialised as a linear ramp from dissolution into growth.Source: §4.6.3 · pp. 18–19

4.7 Computational Benchmarking

The computational benchmark compares the MATLAB and JAX solvers for forward simulation and model training. The comparison uses the same machine, although only JAX uses the GPU.

4.7 Computational Benchmarking

  • The benchmark compares MATLAB and JAX implementations using M1 and M4 kinetics, and includes JAX hybrid-dissolution training with and without initial-condition assimilation.M1 and M4 represent two and four optimised parameters, respectively.Source: §4.7 · p. 19

Forward simulation.

  • Forward-simulation time is averaged across the dataset’s 13 experiments.The comparison reports mean time for a single forward simulation.Source: Forward simulation · p. 19

Model Training.

  • For JAX training, the benchmark reports wall-clock time to early-stopping convergence and time for a fixed 500-epoch budget.MATLAB uses fmincon with numerically estimated gradients and its own convergence criterion; parameter counts and times are reported in Extended Figure 7.Source: Model Training · p. 19

Hardware.

  • Both solvers run on the same machine, but JAX can use its GPU while the MATLAB code cannot.The machine has an AMD Ryzen 3900 12-core CPU and NVIDIA RTX 4090 GPU; MATLAB is mex-compiled.Source: Hardware · p. 19

Extended Figures

The extended figures show initial-condition assimilation, improved validation fits, and experimental trajectory optimisation, alongside computational performance comparisons for the differentiable workflow.

Extended Figures

  • Assimilated particle size and shape distributions are similar in shape to the original measurements, while the two held-out experiments have no assimilated distributions.The held-out experiments were not included in hybrid model training.Source: Figure 4 · p. 25
  • Initial-condition assimilation improves validation fits for both empirical and hybrid models, with the largest gain for the hybrid across network architectures.The comparison covers four validation experiments and all three target metrics.Source: Figure 5 · p. 26
  • 0.10 °C to 0.05 °C: freeing the initial condition for assimilation reduces temperature-profile MAE and tracks experimental trajectories more closely than fixed-IC optimisation or the initial guess.On held-out experiment E5, supersaturation remains within the Smin bound in both optimisation cases.Source: Figure 6 · p. 27
  • Approximately 13× faster forward simulation: JAX takes 0.06 s versus 0.76–0.82 s in MATLAB, while adding 4.8 × 10^5 IC parameters changes JAX convergence time only marginally.For the NN 3 × 20 model, convergence time changes from 291 s to 282 s with joint IC assimilation.Source: Figure 7 · p. 28
Loading 2609.04011v1…