Original Stack Overflow Question: https://stackoverflow.com/questions/65683788/indexerror-index-31-is-out-of-bounds-for-axis-1-with-size-31?noredirect=1#comment116218335_65683788

PDE: u_t = u_xx + u(u_x)^2

The backward difference was used the time derivative meanwhile a central difference was used for the first space derivative.

If you follow the link you can see that I cannot get solutions that make sense and I cannot figure out how to. I'm adamant it is to do with the function f at the finite-difference algorithm stage. Currently, I understand that f is a vector and hence only generating the initial condition. But when I try something like f[1:M, q] it doesn't work.


import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import math
from scipy.sparse import spdiags
from numpy.linalg import inv

M = 30
dx = 1 / M
r = 0.25

tmax = 0.2
dt = r * dx**2
N = math.floor(tmax / dt)

x = np.linspace(0, 1, M + 1)
t = np.linspace(0, tmax, N + 1)

U = np.zeros((M + 1, N + 1))                           # Initial array for solution u(x, t)

U[:, 0] = 40 * x**2 * (1 - x) / 3                      # Initial condition 
U[0, :] = 0                                            # Boundary condition at x = 0
U[-1, :] = 0                                           # Boundary condition at x = 1 

for q in range(1, N):
    U[:, q+1] = U[:, q]
    for p in range(1, 10):
        f = np.zeros((M+1, M+1))
        for k in range(2, M):
            f = U[k, q+1] - r * (U[k+1, q+1] + U[k-1, q+1]) + 2 * r * U[k, q+1] \
              - r * U[k, q+1] * U[k+1, q+1]**2 + r * U[k+1, q+1]**2 * U[k, q+1] - r * U[k, q+1]**3

        updiag = -r - 2 * r * U[p, q + 1] * U[p+1, q + 1] + 2 * r * U[p, q + 1]**2
        diag = 1 + 2 * r + 3 * r * U[p, q + 1]**2 - 4 * r * U[p+1, q + 1] * U[p+1, q+1] - 2 * r * U[p, q + 1]**2
        lowdiag = r
        Jdata = np.array([31 * [diag], 31 * [lowdiag], 31 * [updiag]])
        Diags = [0, -1, 1]
        J = spdiags(Jdata, Diags, 31, 31).toarray()
        d = inv(J) * f
        u = U[:, q+1]
        u = u + d
        U[1, q + 1] = 0
        U[M, q + 1] = 0

T, X = np.meshgrid(t, x)

fig = plt.figure()
ax = fig.gca(projection='3d')

surf = ax.plot_surface(T, X, U, cmap=cm.coolwarm)

ax.set_zlabel('u(x, t)')

plt.savefig('BDImplSol.pdf', bbox_inches='tight')

Edit Discretisation formula:

$$ U_i^{n + 1} = U_i^n + r \left(\left(U_{i+1}^{n+1} - 2 U_i^{n+1} + U_{i-1}^{n+1} \right) + U_i^{n + 1} \left(U_{i + 1}^{n + 1} - U_i^{n+1} \right)^2 \right) $$ $$ r = \frac{dt}{dx^2} $$


Let's reconstruct this from first principles:

Defining the ODE system

In the method-of-lines discretization you solve an ODE system $\dot U=F(U)$, $U=(U_0,U_1,...,U_{M+1})$, $U_k(t)=u(x_k,t)$, and similarly $F=(F_0,F_1...,F_{M+1})$. Because of the boundary conditions $$ u(x, 0) = 40 · x^2 · (1 - x) / 3 \\ u(0, t) = u(1, t) = 0 $$ $U_0=U_{M+1}=0$ and consequently also $F_0=F_{M+1}=0$. For the remaining functions the space discretization was chosen to be $$ F_i(U)= \left(\left(U_{i+1} - 2 U_i + U_{i-1} \right) + U_i \left(U_{i + 1} - U_i \right)^2 \right) $$ Note that the implementation below actually returns $Δt·F(U)$ relative to the function above. To compensate this, the effective step size changes from $Δt$ to $1$ and does thus not occur in the code at the places where one would expect it.

def F(U): 
    dU = np.zeros(M+1)
    dU[1:-1] = r*((U[2:]+U[:-2]-2*U[1:-1]) + U[1:-1]*(U[2:]-U[1:-1])**2)
    return dU

Reference solution using odeint

With this function one can compute the solution with any standard ODE solver, such as

Uref = odeint(lambda U,t:F(U), U[:,0],t/dt).T

Backward Euler integration

The backwards Euler step is $$ U^{j+1}=U^j+Δt·F(U^{j+1}) $$ If one has an approximate value for $U^{j+1}$ one can compute an increment from the linearization $$ U^{j+1}+ΔU^{j+1}=U^j+Δt·F(U^{j+1})+Δt·F'(U^{j+1})ΔU^{j+1}\\ J·ΔU^{j+1}=(I-dt·F'(U^{j+1}))·ΔU^{j+1}=U^j-U^{j+1}+dt·F(U^{j+1}) $$

Composition of the Jacobian

$J$ now is, by the construction of $F$, a tri-diagonal matrix with non-zero entries $J_{0,0}=J{M+1,M+1}=1$ and for $1\le i\le M$ $\newcommand{\pd}[2]{\frac{\partial#1}{\partial#2}}$ \begin{align} J_{i,i-1}=-Δt·\pd{F_i(U^{j+1})}{U_{i-1}^{j+1}} &=-r·(1) \\ J_{i,i}=1-Δt·\pd{F_i(U^{j+1})}{U_{i}^{j+1}} &=1-r·(-2+(U_{i+1}^{j+1}-U_{i}^{j+1})^2-2U_{i}^{j+1}(U_{i+1}^{j+1}-U_{i}^{j+1})) \\&=1+r·(2-(U_{i+1}^{j+1}-U_{i}^{j+1})(U_{i+1}^{j+1}-3U_{i}^{j+1})) \\ J_{i,i+1}=-Δt·\pd{F_i(U^{j+1})}{U_{i+1}^{j+1}} &=-r·(1+2U_{i}^{j+1}(U_{i+1}^{j+1}-U_{i}^{j+1})) \end{align}

from scipy.sparse import diags

def Jac(U, gamma=1):
    rg = r*gamma
    Jm = -rg*np.ones(M, float)
    J0 = np.ones(M+1, float)
    Jp = -rg*np.ones(M, float)
    J0[1:-1] += rg*(2 - (U[2:]-U[1:-1])*(U[2:]-3*U[1:-1]))
    Jp[1:] -= rg*2*U[1:-1]*(U[2:]-U[1:-1])
    return diags([Jm,J0,Jp],[-1,0,1], format="csc")

Integrator time loop

Now combine everything in the time loop. Use the explicit Euler method as predictor and compute the Jacobian matrix only once for the simplified Newton method (this converges in 5 iterations, Newton is only one iteration less, to reach the accuracy of the discretization only requires 1 or 2 iterations). Using the same system matrix has the advantage that it only needs to be factorized once. This has now the simple form

from scipy.sparse.linalg import splu

for q in range(N):
    U[:, q+1] = U[:, q] + F(U[:, q])
    J = Jac(U[:, q+1])
    J_lu = splu(J)
    for p in range(5):
        f = U[:, q+1] - U[:, q] - F(U[:, q+1])
        U[:, q+1] -= J_lu.solve(f)

As this converges within nice numbers, the plot of the solution has the expected shape

enter image description here

Addendum: Crank-Nicolson

C-N is now the implicit trapezoidal method $$ U^{j+1}=U^j+\frac{Δt}2·(F(U^j)+F(U^{j+1})) $$ The linearization here is $$\begin{align} U^{j+1}+ΔU^{j+1}&=U^j+\frac{Δt}2\Bigl(F(U^j)+F(U^{j+1})+F'(U^{j+1})ΔU^{j+1}\Bigr)\\~\\ J·ΔU^{j+1}&=(I-γΔt·F'(U^{j+1}))·ΔU^{j+1}=U^j+γΔt·\Bigl(F(U^j)+F(U^{j+1})\Bigr)-U^{j+1} \end{align}$$ with $γ=\frac12$. Combine the terms in $U^j$ as they stay constant over the inner loop.

from scipy.sparse.linalg import spsolve, splu

for q in range(N):
    U[:, q+1] = U[:, q] + F(U[:, q])
    J = Jac(U[:, q+1], gamma=0.5)
    J_lu = splu(J)
    R = U[:, q] + 0.5*F(U[:, q+1])
    for p in range(4):
        #J = Jac(U[:, q+1])
        f = U[:, q+1] - 0.5*F(U[:, q+1]) - R
        U[:, q+1] -= J_lu.solve(f)
        print(q,p,max(abs(f)), max(abs(U[:,q+1])))

This works fine as well.

  • $\begingroup$ Wow I never thought to do it like this, thanks! Do you know how you could do it without using a predictor-corrector method? $\endgroup$ – AlphaArgonian Jan 16 at 13:45
  • $\begingroup$ You have a system of cubic equations in the new vector. There is no (sensible) way around the iterative numerical solution. If you call that Newton's method (with a sensible initial guess) or predictor-corrector scheme does not make a difference. Using BDF or Adam-Bashford formulas or similar you can get a higher order predictor, perhaps reducing the number of necessary corrector iterations. // One can solve polynomial systems algebraically, with horrendous time and space complexities, this is certainly not a short-cut. $\endgroup$ – Lutz Lehmann Jan 16 at 13:50
  • $\begingroup$ You could of course use any explicit ODE solver method. The problem is that the dissipation term is smoothness-reducing, using the inversion in implicit solvers reverts that to a smoothness-enhancing behavior. Then there is operator-splitting where you only solve the linear dissipation term with an implicit method or matrix exponential and the non-linear term with an explicit method. But since that leaves derivations on the explicit side, this might not be overarchingly effective. $\endgroup$ – Lutz Lehmann Jan 16 at 13:55
  • $\begingroup$ I've tried modifying this for the Crank-Nicolson scheme and I get a broadcast error for the ``` dU[1:-1] ``` stage. $\endgroup$ – AlphaArgonian Jan 16 at 16:18
  • $\begingroup$ C-N is, on the m-o-l level of abstraction, nothing more than the implicit trapezoidal method. I added the equations and code for it. $\endgroup$ – Lutz Lehmann Jan 16 at 16:41

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.