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

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.