Building a General PDE Solving Framework with Symbolic-Numeric Scientific Machine Learning

Wednesday April 3rd, 4:30pm, Room 2-255 at MIT Building a General PDE Solving Framework with Symbolic-Numeric Scientific Machine Learning The dream is to be able to write down a symbolic expression of a PDE and just get a solution. Not only get a solution, but get that solution efficiently. The question is, how do we effectively build a software ecosystem to achieve this goal, and what are the remaining mathematical problems that are necessary to solve in order to fill the gaps? In this talk we will discuss the Julia SciML ecosystem's approach to a general PDE solving framework. We will focus on 5 aspects: (1) a symbolic structural hierarchical classification of PDEs for shuttling high level PDE descriptions to appropriate solution approaches, (2) new time-stepping methods for accelerating the solution of semi-discretized equations and generalizing approaches from PDEs to structured ODE forms, (3) symbolic-numeric approaches for automating the transformation of semi-discretizations to simpler and more numerically stable equations before solving, and (4) new improvements to adjoint methods to decrease the memory requirements for differentiation of PDE solutions, and (5) scientific machine learning approaches to generate accelerated approximations (surrogates) to PDE simulators. We will demonstrate these methods with open source software that starts from symbolic expressions and solves industrial-scale PDEs with minimal lines of code.

1. Building a General PDE Solving Framework with Symbolic-Numeric Scientific Machine Learning Chris Rackauckas VP of Modeling and Simulation, JuliaHub Research Affiliate, MIT CSAIL

2. Connect with SciML Follow on Twitter @ChrisRackauckas Follow on LinkedIn https://www.linkedin.com/in/chrisrackauckas/ Star the GitHub repositories https://github.com/SciML Join our community chatrooms julialang.org/slack Dr. Chris Rackauckas VP of Modeling and Simulation JuliaHub Research Affiliate MIT CSAIL Director of Scientific Research Pumas-AI

3. PDEs Done Wrong! Writing Loops for a Single PDE

4. Fast Simulation: Good Algorithms + Fast Implementation (“Fast” + Implicit Euler / Newmark Beta? Get Outta Here!)

5. All good PDE methods can be expressed as lower level primitives: Symbolic transformations to standard computable forms Numerical methods which exploit structure on standard computable forms ModelingToolkit PDE Hypothesis:

6. The SciML Organization

7. , Christopher, and Qing Nie. "Confederated modular differential equation APIs for accelerated algorithm development and benchmarking." Advances in Engineering Software 132 (2019): 1-6. SciML: Open Source Tooling For Modeling and Simulation https://github.com/SciML/SciMLBenchmarks.jl Rackauckas, Christopher, and Qing Nie. "Differentialequations.jl–a performant and feature-rich ecosystem for solving differential equations in julia." Journal of Open Research Software 5.1 (2017). Rackauckas, Christopher, and Qing Nie. "Confederated modular differential equation APIs for accelerated algorithm development and benchmarking." Advances in Engineering Software 132 (2019): 1-6. Speed Stability Stochasticity Adjoints and Inference Parallelism Non-Stiff ODE: Rigid Body System 8 Stiff ODEs: HIRES Chemical Reaction Network Generally: 50x faster than SciPy 50x faster than MATLAB 100x faster than R’s deSolve Methods that outperform classic C and Fortran methods (CVODE, DASSL) Specialized situational methods (IMEX, stage parallel, etc.) Active Developer Base: ~100 contributors in a month across SciML Paid summer programs training ~10 new devs MIT Lab (~20 people) driven development

8. NonlinearSolve.jl: Robust Nonlinear Root-Finding in Julia NonlinearSolve.jl outperforms standard softwares like Sundials & MINPACK by 10-100x. Built-In Sparsity Detection & Automatic Differentiation Support enables scaling to large systems. Pal, Avik, Holtorf, Larsson, Loman, Utkarsh, Schaefer, Qu, Alan Edelman, and Chris Rackauckas. "NonlinearSolve.jl: High-Performance and Robust Solvers for Systems of Nonlinear Equations in Julia." Manuscript in Preparation. MINPACK is the slowest NonlinearSolve.jl has the fastest solvers Sundials is also very slow! NonlinearSolve.jl successfully solves novel problems where other software fail.

9. Modeling and Simulation Tools Across the Spectrum of Mathematics https://scimlbase.sciml.ai/dev/ The SciML Common Interface for Julia Equation Solvers LinearSolve.jl: Unified Linear Solver Interface   NonlinearSolve.jl: Unified Nonlinear Solver Interface   DifferentialEquations.jl: Unified Interface for all Differential Equations     Optimization.jl: Unified Optimization Interface   Integrals.jl: Unified Quadrature Interface   Unified Partial Differential Equation Interface  

10. What is Scientific Machine Learning (SciML)? Scientific Computing ↔ Machine Learning Machine Learning Neural Nets Bayesian Modeling Automatic Differentiation Scientific Computing Scientific Machine Learning Model Building Robust Solvers Control Systems Differentiable Simulators Surrogates and ROM Inverse Problems & Calibration Automatic Equation Discovery Applicable to Small Data and more …. ‹#›

11. JuliaSim Model Discovery: Autocompleting Models with SciML

12. Universal Differential Equations Predict Chemical Processes UDEs in advection-diffusion transform the learning problem to low dimensional spaces where small data is sufficient

13. Universal Differential Equations Predict Chemical Processes Santana, V. V., Costa, E., Rebello, C. M., Ribeiro, A. M., Rackauckas, C., & Nogueira, I. B. (2023). Efficient hybrid modeling and sorption model discovery for non-linear advection-diffusion-sorption systems: A systematic scientific machine learning approach. arXiv preprint arXiv:2303.13555 . Recovers equations with the same 2nd order Taylor expansion

14. PDE Discretization via ModelingToolkit Is Transformaiton

15. modular differential equation APIs for accelerated algorithm development and benchmarking." Advances in Engineering Software 132 (2019): 1-6. SciML’s Modeling Ecosystem: ModelingToolkit’s Numerical Counterpart Every improvement by every package developer feeds into one pipeline Rackauckas, Christopher, and Qing Nie. "Confederated modular differential equation APIs for accelerated algorithm development and benchmarking." Advances in Engineering Software 132 (2019): 1-6. Youtube: Differential Equations in 2021 SciML’s Common Interface: One consistent interface for all numerics Symbolic modeling for all forms Automated inverse problems and adjoints Composes across the whole package ecosystem Fully embraces generic programming Uses and embraces the work of other developers

16. Heat Equation using MethodOfLines, ModelingToolkit, DomainSets @parameters t, x @variables u ( .. ) Dt = Differential (t); Dx = Differential (x) Dxx = Differential (x) ^ 2 α = 1.1 eq = Dt ( u (t, x)) ~ α * Dxx ( u (t, x)) domain = [x ∈ Interval ( 0.0 , 10.0 ), t ∈ Interval ( 0.0 , 1.0 )] ic_bc = [ u ( 0.0 , x) ~ exp ( - (x - 4.0 ) ^ 2 ) + exp ( - (x - 6.0 ) ^ 2 ), u (t, 0.0 ) ~ 0.0 , u (t, 10.0 ) ~ 0.0 ] @named sys = PDESystem (eq, ic_bc, domain, [t, x], [ u (t, x)]) SciML/MethodOfLines.jl

17. SciML/MethodOfLines.jl # Method of lines discretization dx = 0.1 order = 2 discretization = MOLFiniteDifference ([x => dx], t, approx_order = order) # Convert the PDE problem into an ODE problem prob = discretize (sys, discretization) using OrdinaryDiffEq # Solve ODE problem sol = solve (prob, Tsit5 (), saveat = 0.1 )

18. SciML/MethodOfLines.jl solu = sol[ u (t,x)] x_grid = sol[x] t_grid = sol[t] using Plots heatmap (t_grid, x_grid, transpose (solu), …)

19. SciML/MethodOfLines.jl

20. SciML/MethodOfLines.jl anim = @animate for i in 1 : length (t_grid) plot (x_grid, solu[i, :], …) end gif (anim, "heat.gif" , fps = 8 )

21. SciML/MethodOfLines.jl

22. SciML/MethodOfLines.jl Burger’s Equation

23. @parameters x t @variables u ( .. ) Dx = Differential (x) Dxx = Dx ^ 2 Dt = Differential (t) x_min = - 1.0 x_max = 1.0 t_min = 0.0 t_max = 3.0 v = 0.01 eq = Dt ( u (t, x)) ~ - u (t, x) * Dx ( u (t, x)) + v * Dxx ( u (t,x)) bcs = [ u ( , x) ~ 0.11 * cospi ( 4 x), u (t, x_min) ~ u (t, x_max)] domains = [t ∈ Interval (t_min, t_max), x ∈ Interval (x_min, x_max)] @named pdesys = PDESystem (eq, bcs, domains, [t, x], [ u (t, x)]) Burger’s Equation SciML/MethodOfLines.jl

24. Implemented schemes Centered difference scheme to arbitrary order, for even ordered spatial derivative terms. Upwind difference scheme (hardcoded to first order) for odd ordered spatial derivatives. WENO Scheme Integral schemes Nonlinear laplacian scheme. Spherical laplacian scheme. Discrete callbacks Staggered Grid SciML/MethodOfLines.jl

25. ModelingToolkit PDEs: Method Of Lines Finite Difference

26. ModelingToolkit PDEs : Physics-Informed Neural Networks Easy and Customizable PINN PDE Solving with NeuralPDE.jl, JuliaCon 2021

27. Modeling and Simulation Tools Across the Spectrum of Mathematics https://scimlbase.sciml.ai/dev/ The SciML Common Interface for Julia Equation Solvers LinearSolve.jl: Unified Linear Solver Interface   NonlinearSolve.jl: Unified Nonlinear Solver Interface   DifferentialEquations.jl: Unified Interface for all Differential Equations     Optimization.jl: Unified Optimization Interface   Integrals.jl: Unified Quadrature Interface   Unified Partial Differential Equation Interface  

28. ModelingToolkit PDEs : Extensible PDE Interface Many extra forms have partial implementations being productionized: Finite Volume methods (FiniteVolumeMethods.jl) Spectral methods (?) Finite element methods (via Ferrite.jl) Discrete Galerkin methods (via Trixi.jl) Exterior Calculus (via Decapodes.jl) … Attempting to unify the difference PDE discretization methods into a one interface compatible with component-based modeling!

29. But a general PDE solver interface means nothing if it’s slow. How we take a general PDE discretization engine and make it as fast as a code that is hand-optimized to specific PDEs?

30. ModelingToolkit Symbolic-Numeric Optimizations

31. , Christopher, and Qing Nie. "Confederated modular differential equation APIs for accelerated algorithm development and benchmarking." Advances in Engineering Software 132 (2019): 1-6. Symbolics.jl MTK DiffEq

32. Example of Tearing Nonlinear Systems

33. Example of Tearing Nonlinear Systems It automatically reduced your 5 equation system to 1!

34. Example of Tearing Nonlinear Systems Only solves one equation numerically But can generate the other variables

35. Inline Integration: Tearing of Internal Solves Perform tearing on the internal nonlinear solve of an implicit method

36. Structural Simplify Performance: DAE vs AE DAE AE Derivative Form Analytic Computational Analytic Computational Alias Elimination 18 20 18 18 Structural Simplify 6 12 4 8

37. Structural Simplify Performance: DAE vs AE Model Input Output Eqs/States DAE AE f4 x ddy 97 Alias Elimination: 52 Structural Simplify: 27 eqs/26 states 50 11

38. Inlined Linear Solver Optimization

39. ExactODEReduction.jl

40. Improving Automated Solvers with Sparsity Detection Symbolics.jl: Automated Sparsity Analysis on Code Will be a part of the automatic solver choice!

41. What Kinds of Transformations Do You Get? DAE Index Reduction Not solvable by standard numerical solvers! Differentiate the last equation twice, do a few substitutions… Easy to solve! If you don’t know the details about why this makes a better numerical simulation, then you should be using ModelingToolkit.

42. What Kinds of Transformations Do You Get? DAE Index Reduction Let me fix that for you… Structural_simplify: Pantelides algorithm For index reduction

43. But that form has numerical drift

44. Dummy Derivative: Include the Implicit Constraints

45. Many “specialized CFD” methods are simply a standard ODE solver after Pantelides transformation of discretized Navier-Stokes Inlined Linear Solver optimizations covers many simple semi-implicit schemes Tearing of nonlinear systems on finite difference is a generalization of ghost points While this already hits good runtimes, there’s compilation issues I’ve glossed over Some immediate consequences that would take too long to prove:

46. Interested in Symbolic-Numerics? Lots of starter projects! Get in touch

47. Implicit-Explicit schemes, ADI, etc. all handled by this aspect. We have yet to find a PDE scheme that cannot be expressed by some combination of the previous features Scaling and automatic parallelism scheduling are among the current ongoing research questions Some immediate consequences that would take too long to prove:

48. sys = structural_simplify(complete_motor) JuliaSimCompiler: Accelerated ModelingToolkit 2 lines of code to turn on, enables enormous scalability improvements Solves a major scaling problem in acasual systems Conclusion: ModelingToolkit is a widely used open modeling platform, and with JuliaSim it’s also the most scalable. using JuliaSimCompiler complete_motor_ir = IRSystem(complete_motor) sys_ir = structural_simplify(complete_motor_ir) MTK JuliaSimCompiler

49. Loop Rerolling systems = @named begin sine = Sine(frequency = 10 ) source = Voltage() resistors[ 1 :n] = Resistor() capacitors[ 1 :n] = Capacitor() ground = Ground() end;

50. Loop Rerolling resistors_1₊v(t) ~ resistors_1₊p₊v(t) - resistors_1₊n₊v(t) ~ resistors_1₊p₊i(t) + resistors_1₊n₊i(t) resistors_1₊i(t) ~ resistors_1₊p₊i(t) resistors_1₊v(t) ~ resistors_1₊ R *resistors_1₊i(t) resistors_2₊v(t) ~ -resistors_2₊n₊v(t) + resistors_2₊p₊v(t) ~ resistors_2₊p₊i(t) + resistors_2₊n₊i(t) resistors_2₊i(t) ~ resistors_2₊p₊i(t) resistors_2₊v(t) ~ resistors_2₊ R *resistors_2₊i(t) Variable classes: {{resistors_1₊v, resistors_2₊v, …}, {resistors_1₊p₊v, resistors_2₊p₊v}, …} Equation classes: {0 = f₁(x, y, z) = x - (y - z), 0 = f₂(x, y) = x + y, …}

51. Loop Rerolling for var "%33" = 1 : 97 var "%34" = var "%33" - var "%32" var "%35" = var "%29" + var "%34" var "%36" = Base.getindex( var "###in 1###" , var "%35" ) var "%37" = var "%30" + var "%34" var "%38" = Base.getindex( var "###in 1###" , var "%37" ) var "%39" = var "%31" + var "%34" var "%40" = Base.getindex( var "###in 1###" , var "%39" ) var "%41" = var "%24" * var "%38" var "%42" = var "%41" + var "%36" var "%43" = var "%40" + var "%42" var "%44" = var "%31" + var "%33" var "%45" = Base.setindex!( var "###out###" , var "%43" , var "%44" ) end

52. Loop Rerolling on JuliaSimBattery (Single-Particle Model (SPM), Lithium Nickel Manganese Cobalt Oxide (NMC))

53. Exploiting Structure in ODE Solvers

54. “Refinements” to ODE Problems can fall into different classes

55. Approximate Matrix Factorization = Generalized ADI

56. Preliminary results indicate upto 7x acceleration on PDE solutions with upto 2x lesser memory consumption** Idea: Non-linear solve is typically the most time consuming step in PDE simulations. Perform non-linear solves in mixed precision with GPUs and also without affecting convergence?* 32-bit precision non-linear solve can achieve machine precision accuracy! * Kelley, C. T. "Newton's Method in Mixed Precision." SIAM Review 64.1 (2022): 191-211. **Manuscript in preparation. GPU Accelerated Mixed Precision Methods for Semi-discretized PDEs

57. Side Note: Some Special Forms

58. Solving Semilinear Parabolic PDEs (Nonlinear Black-Scholes) by Learning Structured Stochastic Differential Equations   Simplified: Transform the PDE into a Backwards SDE: a stochastic boundary value problem. The unknown is a function! Learn the unknown function via neural network . Once learned, the PDE solution is known. Solving high-dimensional partial differential equations using deep learning, 2018, PNAS, Han, Jentzen, E

59. Solving Semilinear Parabolic PDEs (Nonlinear Black-Scholes) by Learning Structured Stochastic Differential Equations Solving high-dimensional partial differential equations using deep learning, 2018, PNAS, Han, Jentzen, E

60. Solving Semilinear Parabolic PDEs (Nonlinear Black-Scholes) by Learning Structured Stochastic Differential Equations This is equivalent to training the universal SDE: Solves a 100 dimensional PDE in minutes on a laptop! As a universal SDE, we can solve this with high order (less neural network evaluations), adaptivity, etc. methods using DiffEqFlux.jl. Thus we automatically extend the deep BSDE method to better discretizations by using this interpretation!

61. Extending the Method: Learning Nonlinearities While Solving?

62. Differentiation for Inverse Problems: Issues with Adjoints

63. Advancements in Structural Identifiability Analysis of ODE control systems Novel autonomous algorithms based on differential algebra that uncover the source of non-identifiability Classic problem: impossible to deduce if we only measure and 𝑏 is unknown Also: first steps towards PDE identifiability in "Algebraic identifiability of partial differential equation models," arxiv 2024. Solution: model reparametrization. Fewer parameters. All new variables are identifiable. is not identifiable Reparametrization is already available in SciML/StructuralIdentifiability.jl : Fully automatic. Domain agnostic. Works with non-linear ODEs of dimensions around 20.

64. The adjoint equation is an ODE! How do you get z(t)? One suggestion: Reverse the ODE Timeseries is not stored, therefore O(1) in memory!   Machine Learning Neural Ordinary Differential Equations Chen, Ricky TQ, et al. "Neural ordinary differential equations." Advances in neural information processing systems. 2018.

65. Rackauckas, Christopher, and Qing Nie. "Differentialequations. jl–a performant and feature-rich ecosystem for solving differential equations in julia." Journal of Open Research Software 5.1 (2017). Rackauckas, Christopher, and Qing Nie. "Confederated modular differential equation APIs for accelerated algorithm development and benchmarking." Advances in Engineering Software 132 (2019): 1-6. “Adjoints by reversing” also is unconditionally unstable on some problems! Advection Equation: Approximating the derivative in x has two choices: forwards or backwards If you discretize in the wrong direction you get unconditional instability You need to understand the engineering principles and the numerical simulation properties of domain to make ML stable on it. Result: the method in the PyTorch library is unconditionally unstable!!!

66. Problems With Naïve Adjoint Approaches On Stiff Equations   Kim, Suyong, Weiqi Ji, Sili Deng, and Christopher Rackauckas. "Stiff neural ordinary differential equations."  Chaos  (2021). How do you get u(t) while solving backwards? 3 options! 1. 2. Store u(t) while solving forwards (dense output) 3. Checkpointing   Unstable High memory More Compute Each choices has an engineering trade-off!

67. Problems With Naïve Adjoint Approaches On Stiff Equations     Kim, Suyong, Weiqi Ji, Sili Deng, and Christopher Rackauckas. "Stiff neural ordinary differential equations."  Chaos  (2021).

68. Gives about a 10x performance improvement for adjoint calculations on PDEs “Adjoint” can mean many many different things… Ma, Yingbo, Vaibhav Dixit, Michael J. Innes, Xingjian Guo, and Chris Rackauckas. "A comparison of automatic differentiation and continuous sensitivity analysis for derivatives of differential equation solutions." In  2021 IEEE High Performance Extreme Computing Conference (HPEC) , pp. 1-9. IEEE, 2021.

69. The Test Problem

70. Julia Derivative Results All methods compute the same derivative to many digits

71. Jax Implementation of the Test Problem Jax requires a lot more code

72. Jax Derivative Results: BacksolveAdjoint BacksolveAdjoint works here, but it’s the unstable one!

73. Jax Derivative Results: RecursiveCheckpointingAdjoint and DirectAdjoint Jax’s other methods don’t converge, >60% error!

74. Memory-efficient Continuous Adjoint Sensitivity Method for Universal Ordinary Differential Equations In comparison to QuadratureAdjoint, GaussAdjoint is much more memory efficient! will work on GPUs will be compatible with SDEs Goal: compute the gradient of a loss function with respect to many parameters Sapienza et al. , "Differentiable Programming for Differential Equations: A Review" in preparation (2024). This gradient can be computed as follows Solve the differential equation to obtain u(t): Solve the adjoint differential equation to obtain λ(t): Solve the integral: Idea: GaussAdjoint computes the integral in a memory-efficient manner by accumulating the Gauss quadrature approximation of the integral over time! This circumvents the need of a continuous solution for λ(t)!

75. The PyTorch library uses a numerically unstable method The Jax library uses a non-convergent method. The Julia library is numerically stable and convergent. Calculating derivatives is hard. Conclusion:

76. SciML Surrogates for PDEs

77. ModelOrderReduction: Intrusive MOR done Nonintrusively

78. Using structural PDE information reduces the amount of data required for the surrogates to achieve 5% error by 100x! Pestourie, R., Mroueh, Y., Rackauckas, C. et al. Physics-enhanced deep surrogates for partial differential equations. Nat Mach Intell 5, 1458–1465 (2023). https://doi.org/10.1038/s42256-023-00761-y Physics-Enhanced Deep Neural Surrogates (PEDS)

79. Steady State Surrogate Jet Engine Pipeline C Binary Simulator Model Ingestion Parallel Data Generation using JuliaHub Simulator Simulator Simulator Sampled Config Mass Model Experimentation using JuliaHub Deployment

80. Surrogates which converge to a user-chosen accuracy Provisional patent filed JuliaSimSurrogates Auto-Train Fully automated training infrastructure for surrogates

81. JuliaSim Surrogates Cloud Deployment Automating the parallelization in SciML training Surrogates have multiple points of natural parallelization: Neural network training should take place on large GPU systems Surrogate data training is an embarrassingly-parallel process, generally on CPU Optimal parallelism thus requires different machines for different parts of the process, and expertise than can optimize the hardware choices automatically.

82. 2D electrochemical model of a single battery 300 equations, 5 ms solve time Battery pack: repeat the model 200 times Very long simulation time Most tools cannot scale physically-accurate models to full packs JuliaSim Batteries overview

83. Digital Echo on Live Battery Models

84. Surrogate Validation App

85. Connect with SciML Follow on Twitter @ChrisRackauckas Follow on LinkedIn https://www.linkedin.com/in/chrisrackauckas/ Star the GitHub repositories https://github.com/SciML Join our community chatrooms julialang.org/slack Dr. Chris Rackauckas VP of Modeling and Simulation JuliaHub Research Affiliate MIT CSAIL Director of Scientific Research Pumas-AI