python - IndexError: index 31 is out of bounds for axis 1 with size 31

Question

I am using this code to solve a nonlinear PDE. Since I am using an implicit (backwards-difference) method I am required to solve a set of nonlinear algebraic equations. In my attempt to do this, I get the error:

    f[1:M, q] = U[1:M, q+1] - r * (U[2: M+1, q+1] + U[0:M-1, q+1]) + 2 * r * U[1:M, q+1] 
IndexError: index 31 is out of bounds for axis 1 with size 31 

Code:

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))
        #f[0, q] = 0
        #f[M, q] = 0
        f[1:M, q] = U[1:M, q+1] - r * (U[2: M+1, q+1] + U[0:M-1, q+1]) + 2 * r * U[1:M, q+1] 
                    - r * U[1:M, q+1] * U[2:M+1, q+1]**2 + r * U[1:M, q+1]**2 * U[2:M+1, q+1] - r * U[1:M, 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, 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_xlabel('t')
ax.set_ylabel('x')
ax.set_zlabel('u(x, t)')

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

The solution should look like this (which I produced from an explicit method)

Update:

Ok so I adjusted my code accordingly

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

which produces

I'm pretty sure the issue is with f. I don't think I have expressed the fact that I want:

f[0, q]  = 0
f[1:M-1, q] = 0
f[-1, q] = 0 

correctly.

Answer 1

等待大神答复