Solving a 3D heat diffusion PDE

I am trying to solve a heat diffusion type PDE using a finite difference method.

I would like to preface that I have seriously simplified the code. Just so that anyone who tries to help me, doesn't get caught up in the coefficients of the PDE.

I am trying to move away from using the 4x nested for loops, and instead would like to use a slicing or shifting approach, where I simultaneously calculate as many points in the 3D grid ijk at each value for n with as few possible steps. The execution time of the stacked for loops can get quite long if I, J, K, N are set in the 100's or 1000's.

Also, I'm only using 1 timestep for illustration purpose. If I were to use more than 1, then I would have to set the new boundary values, which isn't important for the purpose of optimising the code.

The code that I need help optimising is as follows:

import numpy as np

T = 1
I = 10
J = 10
K = 10
N = 1

T0 = 0

dt = (T - T0) / N

f_u = np.zeros((I + 1, J + 1, K + 1, N + 1))

f_u[:, :, :, N] = 1

for n in range(N, 0, -1):
for i in range(1, I):
for j in range(1, J):
for k in range(1, K):
f_u[i, j, k, n - 1] = f_u[i, j, k, n] + dt * ((i + 1) * f_u[i + 1, j + 1, k, n]\
+ 0.5 * f_u[i + 1, j, k, n]\
+ (i + 1) * 0.5 * f_u[i + 1, j - 1, k, n]\
+ f_u[i, j, k, n]\
+ (k + 1) * f_u[i, j, k + 1, n]\
+ (k - 1) * f_u[i, j, k - 1, n]\

+ (i - 1) * f_u[i - 1, j + 1, k, n]\
+  0.5 * f_u[i - 1, j, k, n]\
+ (i - 1) * 0.5 * f_u[i - 1, j - 1, k, n]\

+ f_u[i, j + 1, k, n]\
+ f_u[i, j - 1, k, n])

print(f_u[:,:,:,0].sum())
>> 21870.0

• Please show sample input and output. Commented Mar 28, 2022 at 15:31
• You can use the inputs and outputs given as reference.
– Ruan
Commented Mar 28, 2022 at 15:44

For your kind of data it's very important that you use dtype=int.

All of your dimensions that aren't n should be de-looped. Any time that you write an index, it turns into a slice:

• i - 1 -> :-2
• i -> 1:-1
• i + 1 -> 2:

Same for j and k.

PEP8 needs to be sacrificed for clarity. Attempt to align your terms.

Avoid using \; instead enclose your expression in parens ().

Expand your regression test from a single sum(), and add asserts.

Suggested

import numpy as np

T = 1
T0 = 0
N = 1
dt = (T - T0) / N

I = 10
J = 10
K = 10
i = np.arange(1, I)
j = np.arange(1, J)
k = np.arange(1, K)

f_u = np.zeros((I + 1, J + 1, K + 1, N + 1), dtype=int)
f_u[:, :, :, N] = 1

for n in range(N, 0, -1):
f_u[1:-1, 1:-1, 1:-1, n - 1] = (
f_u[1:-1, 1:-1, 1:-1, n]
+ dt * (
+ (i + 1)       * f_u[2:  , 1:-1, 1:-1, n]
+           0.5 * f_u[2:  , 1:-1, 1:-1, n]
+ (i + 1) * 0.5 * f_u[2:  ,  :-2, 1:-1, n]
+                 f_u[1:-1, 1:-1, 1:-1, n]
+ (k + 1)       * f_u[1:-1, 1:-1, 2:  , n]
+ (k - 1)       * f_u[1:-1, 1:-1,  :-2, n]
+ (i - 1)       * f_u[ :-2, 1:-1, 1:-1, n]
+           0.5 * f_u[ :-2, 1:-1, 1:-1, n]
+ (i - 1) * 0.5 * f_u[ :-2,  :-2, 1:-1, n]
+                 f_u[1:-1, 2:  , 1:-1, n]
+                 f_u[1:-1,  :-2, 1:-1, n]
)
)

assert f_u.shape == (11, 11, 11, 2)
assert f_u.sum() == 23201
assert np.isclose(f_u.mean(), 8.715627347858753, rtol=0, atol=1e-12)
assert f_u.min() == 0
assert f_u.max() == 50

• Hi.. Thanks for your help. The solutions holds when N=1, i.e. we get the same results on all elements of f_u. But when I change it to N=2, then the solution doesn't hold. I'm trying to figure out why..
– Ruan
Commented Mar 29, 2022 at 8:52