Warm-starting finite difference methods with PINNs
In this post, we’ll modify the previous example and solve a nonlinear ordinary differential equation (ODE) with physics-informed neural networks (PINNs).
We’ll take the second-order ODE describing the steady-state diffusion of heat due to radiation in a 1D rod:
\(\begin{equation} \frac{d^2 T(x)}{d x^2} = \alpha(T^4(x) - T_{\infty}^4) \end{equation}\)
where \(\alpha\) is a coefficient that combines the thermal conductivity and emissivity and \(T_{\infty}\) is the surrounding temperature.
We’ll do something fancy at the end though—we’ll pass the pre-computed PINN solution as an initial guess to a finite difference method solving a linearized version of this ODE and show that the convergence can be (two iterations 😉) faster! This shows the potential of a PINN model to warm-start finite difference methods.
Modest code modification
The great part of using PINNs to solve ODEs is that not a whole lot needs to change in our previous code to accommodate a different right-hand-side.
Except for specifying the new numeric parameters of this problem:
T_L = 300 # K
T_R = 400 # K
L = 1 # m
α = 1e-8 # 1/(K^3m^2)
T_infty = 900 # K
everything else that needs to change is how we compute the new residual! This ODE in a residual form looks like this:
\(\begin{equation} \frac{d^2 T(x)}{d x^2} - \alpha(T^4(x) - T_{\infty}^4) = 0 \end{equation}\)
Hence, we need to replace this part in our training loop:
residual = T_xx - k * (T_pred - T_infty)
with this:
residual = T_xx - α * (T_pred**4 - T_infty**4)
and we’re ready to train the new PINN model!
Here’s the PINN approximation along with a finite-difference numerical approximation using a second-order-accurate stencil:
As before, the PINN approximation is quite significantly different from the finite-difference approximation. Which begs the question what could PINNs be actually useful for… 🤔
Warm-starting finite difference methods
Let’s try to use this pre-computed PINN solution as an initial guess when we solve this ODE numerically and see if it can be helpful!
I linearized this ODE to get a set of linear equations of the form:
\(\begin{equation} \frac{T_{i-1} - 2T_i + T_{i + 1}}{\Delta x^2} = \alpha (- 3{T}_{\text{guess}, i}^4 + 4 {T}_{\text{guess}, i}^3 T_i - T_{\infty}^4) \end{equation}\)
where, in general, \({T}_{\text{guess}}(x)\) is our initial guess for the temperature distribution in the entire domain.
I solve these in an iterative process of updating the guess with the current numeric solution. My own initial guess would be to set \(T(x) = \frac{T_{\infty}}{2}\). In this case, my solver gets to the \(L_2\) norm between the previous and current guess lower than \(10^{-6}\) in 6 iterations:
10946.85015189905
2721.2344585582314
442.33582868664246
10.182994079198922
0.005245534201349177
3.8544634728714925e-09
Converged in 6 iterations.
When I use the pre-computed PINN solution as the initial guess, I converge in 4 iterations:
703.0258369833273
22.588183443400222
0.025565367593299866
1.2049626265198362e-07
Converged in 4 iterations.
Maybe PINNs are useful for something after all! 😉
Note that for this warm-starting application, you might not even need to fully converge on the residual loss when training a PINN model. A reasonably good approximation that captures the general trend well can be still way more accurate than your own best guess. In fact, I tried training for just 300 epochs instead of 5000, where the residual loss is still pretty high, and we get just the same saving of two iterations. In this case, the PINN approximation is worse:
but works just as well as an initial guess!
1137.3367124605265
39.527513439514806
0.06727458340750923
1.8620638206290557e-07
Converged in 4 iterations.