In this presentation, we introduce a fully differentiable non-Newtonian fluid solver developed using JAX, capable of accommodating arbitrary geometries and flow scenarios without being constrained to simplified setups or predetermined constitutive forms. Leveraging differentiability, our solver facilitates end-to-end training of machine-learning models directly within the simulation pipeline, allowing efficient, data-driven parameterization of complex rheological behavior.
We specifically present a tensorial basis neural network designed to represent non-Newtonian stress tensors, enabling extraction of physically meaningful bulk properties such as viscosity and viscoelastic moduli. The inherent interpretability of our machine-learning framework allows for direct comparison with established constitutive models using the Bayesian Information Criterion (BIC), thereby determining the most appropriate constitutive relations. Sparse experimental data or bulk measurements can subsequently be used to refine these model parameters, significantly improving predictive generalizability.
We demonstrate the efficacy of our approach in flow scenarios where traditional constitutive models frequently fail due to complex deformation patterns and intricate boundary effects, such as contracting channels and porous media. Thus, our results not only illustrate improved accuracy and robustness relative to conventional methods but also yield novel mechanistic insights crucial for the design of industrial and biophysical processes. Ultimately, this work lays a theoretical, data-driven groundwork for advanced digital rheometry, capable of comprehensively characterizing non-Newtonian and viscoelastic fluid behavior under diverse and realistic conditions.