Parameter conditioned interpretable U-Net surrogate model for data-driven predictions of convection-diffusion-reaction processes

TL;DR

This paper develops a parameter-conditioned U-Net surrogate model with FiLM and coordinate encoding for efficient, accurate predictions of complex convection-diffusion-reaction processes, validated on cancer cell modeling scenarios.

Contribution

It introduces a novel parameter-conditioned U-Net architecture with FiLM and coordinate encoding for PDE surrogate modeling, demonstrating improved generalization and prediction efficiency.

Findings

The surrogate achieves low prediction error on test data.

Prediction times are significantly reduced using GPU parallelization.

Approximation difficulty depends mainly on PDE regime conditioning, not initial conditions.

Abstract

We present a combined numerical and data-driven workflow for efficient prediction of nonlinear, instationary convection-diffusion-reaction dynamics on a two-dimensional phenotypic domain, motivated by macroscopic modeling of cancer cell plasticity. A finite-difference solver, implemented in C++, is developed using second-order spatial discretizations and a step-size controlled Runge-Kutta time integrator. A mesh refinement study confirms the second-order convergence for the spatial discretizations error. Based on simulated input-output pairs and corresponding parameterizations for the diffusion, advection, and reaction mechanisms, we train a parameter-conditioned U-Net surrogate to approximate the fixed-horizon solution map. The surrogate incorporates Feature-wise Linear Modulation (FiLM) for parameter conditioning, coordinate encoding to incorporate spatial location information, and residual blocks to enable multiscale representation learning in combination with the U-Nets skip connections. The trained model achieves low prediction error on held-out test data and provides favorable prediction times due to the GPU based parallelization. Generalization is analyzed using a factorial test dataset, separating initial conditions from parameter conditioning. The results reveal that approximation difficulty varies primarily with the conditioning vector (i.e., the induced PDE regime), rather than with the initial conditions.