1
\$\begingroup\$

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
\$\endgroup\$
2
  • \$\begingroup\$ Please show sample input and output. \$\endgroup\$
    – Reinderien
    Mar 28, 2022 at 15:31
  • \$\begingroup\$ You can use the inputs and outputs given as reference. \$\endgroup\$
    – Ruan
    Mar 28, 2022 at 15:44

1 Answer 1

1
\$\begingroup\$

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
\$\endgroup\$
1
  • \$\begingroup\$ 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.. \$\endgroup\$
    – Ruan
    Mar 29, 2022 at 8:52

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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