Hello all,

I've recently been introduced to Python and Numpy, and am still a beginner in applying it for numerical methods. I've been performing simple 1D diffusion computations. I suppose my question is more about applying python to differential methods. I'm asking it here because maybe it takes some diff eq background to understand my problem.

My overall question : Pretty beginner question, I have a picture of a specific example case below.

What are the benefits/cons of using built-in numpy functions vs using for loops in numerical computational methods (CFD & FEA), is it absolutely necessary to write explicit for loops for time stepping sometimes? (no pun intended)

I attached a link to the short excercise I've been practicing below. Thanks for your time!

Discretization Scheme



* Discretized 1D diffusion eqn, FWD time, Centered Space

My specific code example: How could you get rid of the for loop below ?

for loop

The link for the excercise :



1 Answer 1


numpy arrays and methods are incredibly helpful. They are usually optimized and much faster than looping in python. Always look for a way to use an existing numpy method for your application.

np.roll() will allow you to shift and then you just add.

I learned to use convolve() from comments on How to np.roll() faster?. I haven't checked if this is faster or not, but it may depend on the number of dimensions. See this answer for a 2D relaxation of the Laplace equation (electrostatics, a different problem)

For this kind of relaxation you'll need a bounding box, so the boolean do_me is False on the boundary.

I know that for Jacobi relaxation solutions to the Laplace equation, there are two speed-up methods. I don't know if they can be extended to solving the Heat Diffusion equation, but I'm sure something can be done:

  • Multigrids; solve on a coarse (fast) grid, then interpolate to a fine grid and iterate a little longer

  • Successive Overrelaxation; "overshoot" the averaging by changing the values by an amount larger than 100% of what you would have with simple averaging. Try values like 110% to 150%, it can go unstable if they are too large.

enter image description here

import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import convolve

T0             = np.zeros(50, dtype=float)
T0[17:25]      = 1.

T1             = T0.copy()  # method 1  np.roll()
T2             = T0.copy()  # method 2  convolve()

do_me          = np.ones_like(T0, dtype=bool)
do_me[[0, -1]] = False    #  keep the boundaries of your bounding box fixed

a              = 0.01

hold_1         = [T0.copy()]
for i in range(10001):
    Ttemp      = T1 + a*(np.roll(T1, +1) + np.roll(T1, -1) - 2*T1)
    T1[do_me]  = Ttemp[do_me]
    if not i%1000:

hold_2         = [T0.copy()]
kernel         = np.array([a, (1 - 2.*a), a])

for i in range(10001):
    Ttemp      = convolve(T2, kernel)
    T2[do_me]  = Ttemp[do_me]
    if not i%1000:

if True:
    plt.subplot(1, 2, 1)
    for thing in hold_1:

    plt.subplot(1, 2, 2)
    for thing in hold_2:


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.