Rapid prototyping of differentiable simulators

Differentiable simulators can seamlessly integrate with modern data-driven and optimization techniques to accelerate the pace of scientific discovery and engineering design. Not only do they produce simulation outputs, but can also compute gradients of these outputs with respect to the inputs. This is crucial for  optimization, sensitivity analysis, and uncertainty quantification, as it enables use of  gradient-based optimization and machine learning to efficiently explore the parameter space and make informed decisions.

What do we do?

We research how to construct computational ecosystems for rapid prototyping and differentiable simulators by developing several exemplar ecosystems: MRST and OPM Flow for flow in porous media, BattMo for electrochemical devices and processes, and Jutul for multiphysics applications.

This entails:

1. Developing suitable software abstractions and modularity to hide specific details (i.e., by using a general unstructured format for all grids) and separating different contexts such as physical models, discretizations and discrete operators (see figure below), time-stepping and nonlinear solver, and linearization, assembly, and solution of linear systems.
2. Using automatic differentiation (AD) to compute derivatives, gradients, and Jacobians. The idea is simple: you implement the discretized model equations in residual form, and the AD library computes the necessary derivatives to linerize the resulting system of nonlinear algebraic equations to obtain the Jacobian matrix required by the nonlinear solver. This is particularly advantageous if you wish to quickly introduce new functional dependencies or couple multiple submodels.
3. Defining computational graphs to trace and manage functional dependencies reduces redundant computations and enables efficient use of compute cache and lazy evaluation. This, paired with automatic differentiation, facilitates use of generic functional interfaces, simplifying the integration of new parameter and variable dependencies. These concepts have been realized in slighlty different ways in MRST (and Jutul) and in BattMo.
4. Using  adjoint methods to compute gradients of cost functions relative to model parameters by solving adjoint equations derived from the original PDE model to determine how changes in model input impact the cost function. This backward propagation through the computational graph is akin to backpropagation in machine learning, but adjoint methods are broader in scope and applicable to general PDE-based models in scientific and engineering.

Automatic differentiation

We have researched different AD approaches:

• The data organization of the AD library in MRST revolves around variables and typically treats one variable/intermediate quantity at the time, updating its values and derivatives across all cells in the domain.  This enables vectorization and gives a more natural representation of physical quantities and equations in the code, mirroring the mathematical formulation of the PDEs. However, inevitably also results in redundant computations and excessive memory use. To counteract, we have researched more memory-optimized backends written in compiled languages.
• The AD library in OPM Flow, in contrast, is organized around grid cells, and loosely speaking updates all values/derivatives associated with a cell before moving on to the next.  This offers better memory access  and cache utilization at the expense of a somewhat less intuitive code.
• Jutul features a highly optimized automatic differentiation library that uses static, hard-coded stencils to ensure fast assembly of linearized systems. The software can also utilize AD functionalities from the Julia language.

Read more

• KA Lie. Differentiable simulators: Implications for software development and applications to reduced-order and data-driven modelling. SIAM Conference on Mathematical & Computational Issues in the Geosciences (GS23). June 19 - 22, 2023. Radisson Blu Royal Hotel, Bergen, Norway.