[Paper Review] Physics-constrained neural differential equations for learning multi-ionic transport

Danyal Rehman, John H. Lienhard|arXiv (Cornell University)|Mar 7, 2023

Fuel Cells and Related Materials4 citations

TL;DR

This paper proposes a physics-constrained neural differential equation model that learns multi-ionic transport through polyamide nanopores by integrating neural ODEs with continuum physics models and experimental data. The method leverages charge conservation, ion activity coefficients, and partitioning mechanisms as inductive biases, achieving strong agreement with experimental rejection data after pre-training on simulated data and fine-tuning with noise-augmented real measurements.

ABSTRACT

Continuum models for ion transport through polyamide nanopores require solving partial differential equations (PDEs) through complex pore geometries. Resolving spatiotemporal features at this length and time-scale can make solving these equations computationally intractable. In addition, mechanistic models frequently require functional relationships between ion interaction parameters under nano-confinement, which are often too challenging to measure experimentally or know a priori. In this work, we develop the first physics-informed deep learning model to learn ion transport behaviour across polyamide nanopores. The proposed architecture leverages neural differential equations in conjunction with classical closure models as inductive biases directly encoded into the neural framework. The neural differential equations are pre-trained on simulated data from continuum models and fine-tuned on independent experimental data to learn ion rejection behaviour. Gaussian noise augmentations from experimental uncertainty estimates are also introduced into the measured data to improve model generalization. Our approach is compared to other physics-informed deep learning models and shows strong agreement with experimental measurements across all studied datasets.

Motivation & Objective

• To develop a deep learning model that accurately predicts ion transport and rejection in highly-selective polyamide nanopores without relying solely on expensive experiments.
• To address the computational intractability of solving PDEs for complex, irregular nanopore geometries under nano-confinement.
• To incorporate physical laws—such as charge conservation and ion partitioning mechanisms—directly into a neural architecture as inductive biases.
• To improve model generalization by pre-training on simulated continuum model data and fine-tuning on independent experimental data with Gaussian noise augmentations.
• To enable accurate prediction of ion selectivity across diverse ionic mixtures and concentration conditions using a single unified framework.

Proposed method

• Neural ordinary differential equations (NODEs) parameterize the continuous evolution of ion concentrations along flux, with the hidden state dynamics governed by a learnable ODE: $ \frac{d\mathbf{h}(J_v)}{dJ_v} = f_\theta(\mathbf{h}(J_v), J_v; \theta) $.
• Inductive biases are embedded via physical constraints: charge conservation $ \sum_{j=1}^{d} z_j \mathbf{h}_j(J_v) = 0 $, Pitzer-Kim activity coefficients, and three-part boundary conditions (steric, dielectric, Donnan exclusion).
• The model is pre-trained on simulated data generated by the DSPM-DE continuum model, which solves coupled PDEs for ion flux, electric potential, and concentration with under-relaxation schemes for convergence.
• Latent membrane parameters ($ r_p, \Delta x_e, \zeta_p, \chi_d $) are optimized via simulated annealing and Nelder-Mead to match experimental rejection data, forming the basis for synthetic training data.
• Fine-tuning uses real experimental data with added Gaussian noise to improve robustness and generalization across unseen concentration inputs.
• Backpropagation through the ODE solver uses the continuous adjoint method, and integration is performed via Dormand–Prince 5(4) with PyTorch and SciPy.

Experimental results

Research questions

• RQ1 Can a physics-informed neural ODE model accurately learn ion transport behavior in polyamide nanopores without requiring direct measurement of all ion interaction parameters?
• RQ2 How well can a neural ODE with embedded physical constraints generalize across diverse ionic mixtures and concentration ranges when fine-tuned on limited experimental data?
• RQ3 To what extent does incorporating experimental uncertainty via noise augmentation improve the robustness and predictive accuracy of the model?
• RQ4 Can the integration of mechanistic closure models (e.g., Pitzer-Kim, Donnan exclusion) into a neural framework outperform purely data-driven or purely physics-based models in predicting ion rejection?
• RQ5 How effective is the hybrid optimization of latent membrane parameters in enabling accurate simulation-based pre-training for downstream fine-tuning?

Key findings

• The physics-constrained neural ODE model achieves strong agreement with experimental ion rejection data across all tested datasets, demonstrating high predictive accuracy.
• Pre-training on simulated data from the DSPM-DE model enables the neural ODE to learn complex ion transport dynamics with minimal fine-tuning on real data.
• Incorporating Gaussian noise augmentations based on experimental uncertainty estimates significantly improves model generalization and robustness to input variations.
• The model successfully encodes key physical mechanisms—charge conservation, steric exclusion, dielectric effects, and Donnan exclusion—into the neural architecture, ensuring physically plausible predictions.
• Latent parameter optimization using simulated annealing and Nelder-Mead yielded consistent values: $ r_p = 0.51 \, \text{nm} $, $ \Delta x_e = 1.27 \, \mu\text{m} $, $ \zeta_p = 43.56 $, $ \chi_d = -51.23 \, \text{mol/m}^3 $, which were used to generate high-fidelity training data.
• The model outperforms other physics-informed deep learning baselines in predicting ion rejection, particularly in complex, multi-ionic systems with non-ideal behavior.