2
\$\begingroup\$

I want to perform some calculations on a large dataset. The code can be found below, where I want to calculate the values for 'results_nr' over a large loop (1000 x 910) values. Can you help me out with optimizing this code? The code is already running for 3500 minutes so far... When picking a smaller loop (100 x 91), the code gives the result I'm aiming for, but the resulting dataset needs to be more detailed.

I also posted this question to stackoverflow (https://stackoverflow.com/questions/72849421/how-do-i-optimize-this-integration-loop).

from scipy import integrate
import numpy as np

#Construct 8a unit cell

a8_1 = np.array([0.25, 0.25, 0.25])
a8_final = []
for i in np.arange(0,2):
    for j in np.arange(0,2):
        for k in np.arange(0,2):
            lst1 = [a8_1[0]+ 0.5*k, a8_1[1] + 0.5*j, a8_1[2] + 0.5*i]
            a8_final.append(lst1) 

#Construct 24d unit cell

u = -0.0314 

d24_1 = np.array([u, 0, 0.25])
d24_2 = np.array([1-u,0.5,0.25])
d24_1_tot = np.array([np.roll(d24_1,0),np.roll(d24_1,1),np.roll(d24_1,2)])
d24_2_tot = np.array([np.roll(d24_2,0),np.roll(d24_2,1),np.roll(d24_2,2)])

d24_first = np.concatenate((d24_1_tot,d24_2_tot))
d24_minus = 1-d24_first
d24_sec = np.concatenate((d24_first,d24_minus))
d24_BC = 0.5-d24_sec
d24_final = np.concatenate((d24_sec,d24_BC))

mx = 5
pos = []
                
for i in np.arange(-mx,mx+1):
    for j in np.arange(-mx,mx+1):
        for k in np.arange(-mx,mx+1):
            for item1 in np.arange(0,len(a8_final)):
                lst2 = [a8_final[item1][0]+k,a8_final[item1][1]+j,a8_final[item1][2]+i]
                pos.append(lst2) 
            for item2 in np.arange(0,len(d24_final)):
                lst3 = [d24_final[item2][0]+k,d24_final[item2][1]+j,d24_final[item2][2]+i]
                pos.append(lst3)    
                
print(f'The length of the list is correct: {len(pos) == 32 * (((mx*2)+1)**3)}')

new_pos = []
for elem in pos:
    if elem not in new_pos:
        new_pos.append(elem)
pos = new_pos

a = b = c = 10.604 

alpha = beta = gamma = 90 #cubic structure

def neighbour_list(delta_x,delta_y,delta_z):
    r2 = (a**2 * delta_x**2) + (b**2 * delta_y**2) + (c**2 * delta_z**2) + (2*b*c*(math.cos(alpha*(math.pi/180)))*delta_y*delta_z) + (2*c*a*(math.cos(beta*(math.pi/180)))*delta_z*delta_x) + (2*a*b*(math.cos(gamma*(math.pi/180)))*delta_x*delta_y) 
    return np.sqrt(r2)

threshold_distance = 40 

a8_distances = []
d24_distances = []

#Neighbour list 8a sites

for items1 in list(range(0,len(pos))):
    neighbour_lst = neighbour_list(pos[items1][0]-a8_final[0][0],pos[items1][1]-a8_final[0][1],pos[items1][2]-a8_final[0][2])
    neighbour_lst2 = neighbour_list(pos[items1][0]-d24_final[0][0],pos[items1][1]-d24_final[0][1],pos[items1][2]-d24_final[0][2])
    a8_distances.append(neighbour_lst)
    d24_distances.append(neighbour_lst2)
    

a8_distances_selected = [x for x in a8_distances if x < threshold_distance and x > 0] 
a8_distances_selected = np.round(a8_distances_selected,decimals=9)

(unique1,counts1) = np.unique(a8_distances_selected,return_counts=True)
a8_neighbourlist = np.asarray((unique1,counts1)).T
    
d24_distances_selected = [x for x in d24_distances if x < threshold_distance and x > 0] 
d24_distances_selected = np.round(d24_distances_selected,decimals=9)

(unique2,counts2) = np.unique(d24_distances_selected,return_counts=True)
d24_neighbourlist = np.asarray((unique2,counts2)).T

d24_neighbourlist0 = d24_neighbourlist[:,0].tolist()
a8_neighbourlist0 = a8_neighbourlist[:,0].tolist()
d24_neighbourlist1 = d24_neighbourlist[:,1].tolist()
a8_neighbourlist1 = a8_neighbourlist[:,1].tolist()

d24_neighbourlist0_nm = []
a8_neighbourlist0_nm = []

for item in list(range(len(d24_neighbourlist0))):
    d24_neighbourlist0_nm.append(d24_neighbourlist0[item]/10)
    
for item in list(range(len(a8_neighbourlist0))):
    a8_neighbourlist0_nm.append(a8_neighbourlist0[item]/10)

kr = 3.02607165e+00
knr = 8.06622179e+12
dE = 925.013
Ea = 2.09542112e+04
k_B = 0.69503476
C_x = 1.85784648

def cross_relaxation4(t):
    elem1=1
    elem2=1
    
    for item in list(range(len(d24_neighbourlist))):
        elem1 = elem1*((1-(doping/100) + (doping/100)*np.exp((-C_x*t)/((d24_neighbourlist0_nm[item])**6)))**(d24_neighbourlist1[item]))

    for items in list(range(len(a8_neighbourlist))):    
        elem2 = elem2*((1-(doping/100) + (doping/100)*np.exp((-C_x*t)/((a8_neighbourlist0_nm[items])**6)))**(a8_neighbourlist1[items]))
    
    return ((0.75*elem1) + (0.25*elem2))*np.exp(-(kr+knr*np.exp(-Ea/((k_B)*T)))*t)


doping_handmade = list(np.arange(0.1,10.01,0.01))
temperature_K_handmade = list(np.arange(300,1200.1,0.1))

results_nr = []

for doping in doping_handmade:
    for T in temperature_K_handmade:
        val,err = integrate.quad(cross_relaxation4,0,np.inf)
        results_nr.append(val*kr)
\$\endgroup\$
7
  • \$\begingroup\$ It's on-topic here, but I'm not confident that it would be on-topic for SO. \$\endgroup\$
    – Reinderien
    Jul 3, 2022 at 20:47
  • \$\begingroup\$ @Reinderien yes, they are lists of respectively 1190 and 457 values long, but they are calculated in cells above the respective cell (Jupyter Notebook), so that shouldn't impact the speed of the process \$\endgroup\$
    – thimoooh
    Jul 4, 2022 at 7:14
  • 1
    \$\begingroup\$ @Reinderien it is all added now! \$\endgroup\$
    – thimoooh
    Jul 4, 2022 at 13:23
  • \$\begingroup\$ What kind of value does results_nr actually contain? i.e. what does the integral represent? \$\endgroup\$
    – Reinderien
    Jul 5, 2022 at 11:29
  • \$\begingroup\$ results_nr is a list of values representing the radiative decay efficiency, which should decrease with temperature and have an optimum in the doping concentration. Has to do with a phenomenon in physics. \$\endgroup\$
    – thimoooh
    Jul 6, 2022 at 9:20

1 Answer 1

2
\$\begingroup\$

There are two broad modes of failure in this script.

The first is often seen in Jupyter code: there are nearly no functions, thus a vast swamp of namespace pollution and no opportunity for the garbage collector to run. Add functions.

The second is a thorough misunderstanding of how Numpy is supposed to work; almost nothing is vectorised properly, the worst offender being quad(). It should be a fatal, laptop-incinerating syntax error when newcomers to Numpy use a for-loop. There is not a single loop required or justified in your script.

About quad: you must not consider cross_relaxation to be a scalar-valued integrand. You must consider it to be a vector-valued integrand, with dimensions for your doping and temperature. With this in mind, quad in a loop is doomed to be slow, and quad_vec will do better. Give it a workers=-1 to parallelise the computation. Windows compatibility requires that the script call freeze_support() in a __main__ guard.

Stop using list comprehensions, stop calling the list() constructor, and stop calling .tolist().

Aside from the vectorisation:

The length of the list is correct should probably be converted to an assert.

Many of your variable names need love. 0 and 1 suffixes for your distance arrays should actually just be distance and count.

Consider graphing your output.

The good news is: it's all quite recoverable!

I have not benchmarked the difference between the original and vectorised implementations because I haven't the patience, but we can be confident that it's several orders of magnitude.

Suggested

I've somewhat reduced the doping and temperature dimensions for test purposes.

from multiprocessing import freeze_support

from matplotlib import pyplot as plt
from matplotlib.colors import LogNorm
from numpy.lib.stride_tricks import sliding_window_view
from scipy.integrate import quad_vec
import numpy as np


kr = 3.02607165


def make_a8() -> np.ndarray:
    # Construct 8a unit cell
    # Numpy-vectorisation of what, in base Python, would be a call to itertools.product
    return 0.25 + 0.5 * np.indices((2, 2, 2)).T.reshape((8, 3))


def make_d24(u: float = -0.0314) -> np.ndarray:
    """Construct 24d unit cell"""

    # Construct both d24_1 and d24_2. Tile to double width for the purposes of the sliding window.
    d24_12 = np.tile(np.array((
        (    u, 0.0, 0.25),
        (1 - u, 0.5, 0.25),
    )), 2)
    # Apply the sliding window and reverse along the slid axis so that results are identical to OP.
    d24_12_tot = sliding_window_view(d24_12, window_shape=3, axis=1)[:, -1:-4:-1, :]

    # Do the 1-x and 0.5-x on dimensions 0 and 1 of a broadcast array before flattening to 2D
    d24_5d = np.broadcast_to(d24_12_tot, (2, 2, 2, 3, 3)).copy()
    d24_5d[:, 1, ...] = 1 - d24_5d[:, 1, ...]
    d24_5d[1, ...] = 0.5 - d24_5d[1, ...]
    return d24_5d.reshape((-1, 3))


def make_pos(
    a8: np.ndarray,
    d24: np.ndarray,
    mx: int = 5,
) -> np.ndarray:
    x = slice(-mx, mx+1)
    i, j, k, _ = np.ogrid[x, x, x, :1]
    m = i.size

    a8_d24 = np.concatenate((a8, d24))  # 32 x 3
    # Dimensions:  11 x 11 x 11 x (8 + 24) x 3
    pos = np.broadcast_to(
        a8_d24, (m, m, m, *a8_d24.shape),
    ).copy()

    pos[..., 0] += k
    pos[..., 1] += j
    pos[..., 2] += i
    return pos.reshape((-1, 3))


def neighbour_list(
    pos: np.ndarray,
    origin: np.ndarray,
    a: float = 10.604,  # also b, c
    alpha: float = 90,  # (also beta, gamma) cubic structure
    threshold_distance: float = 40,
    normalize: float = 10,
) -> tuple[np.ndarray, np.ndarray]:
    delta_x, delta_y, delta_z = (pos - origin[0, :]).T

    distances = a * np.sqrt(
        delta_x**2 + delta_y**2 + delta_z**2
        + np.cos(np.radians(alpha)) * (
              delta_y*delta_z + delta_z*delta_x + delta_x*delta_y
          )
    )

    distances_selected = distances[
        (distances > 0) & (distances < threshold_distance)
    ].round(decimals=9)

    unique, counts = np.unique(distances_selected, return_counts=True)
    unique /= normalize
    return unique, counts


def make_fac(distances: np.ndarray, C_x: float = 1.85784648) -> np.ndarray:
    return -C_x / distances[np.newaxis, :]**6


def make_doping(doping: np.ndarray) -> np.ndarray:
    return doping[:, np.newaxis] / 100


def make_temp_base(
    temperature_K: np.ndarray,
    knr: float = 8.06622179e+12,
    Ea: float = 2.09542112e+04,
    k_B: float = 0.69503476,
) -> np.ndarray:
    return np.exp(
        -(
            kr + knr * np.exp(-Ea / k_B / temperature_K[np.newaxis, :])
        )
    )


def calculate_elem(
    t: float,
    doping: np.ndarray,
    distance_fac: np.ndarray,
    counts: np.ndarray,
) -> np.ndarray:
    factors = (
        1 + doping*(np.exp(distance_fac*t) - 1)
    )**counts[None, :]
    return np.prod(factors, axis=1, keepdims=True)


def cross_relaxation4(
    t: float,
    doping: np.ndarray, temp_base: np.ndarray,
    d24_fac: np.ndarray, d24_counts: np.ndarray,
    a8_fac: np.ndarray, a8_counts: np.ndarray,
) -> np.ndarray:
    elem1 = calculate_elem(t, doping, d24_fac, d24_counts)
    elem2 = calculate_elem(t, doping, a8_fac, a8_counts)
    return (0.75*elem1 + 0.25*elem2) * temp_base**t


def graph(integral: np.ndarray, doping: np.ndarray, temperature_K: np.ndarray) -> None:
    fig, ax = plt.subplots()

    x = temperature_K
    y = doping
    z = integral
    powers = 10.**np.arange(-2, 0)
    levels = (
        np.array((1, 2, 5))[np.newaxis, :] * powers[:, np.newaxis]
    ).flatten()

    contour = ax.contour(
        x, y, z, levels=levels, norm=LogNorm(),
    )
    ax.clabel(contour)
    ax.set_title('Radiative decay efficiency')
    ax.set_xlabel('Temperature (°K)')
    ax.set_ylabel('Doping (%)')

    plt.show()


def main() -> None:
    print('Initialising...')
    a8 = make_a8()
    d24 = make_d24()
    pos = make_pos(a8, d24)
    a8_distances, a8_counts = neighbour_list(pos, origin=a8)
    d24_distances, d24_counts = neighbour_list(pos, origin=d24)

    doping = np.arange(0.1, 10.01, 0.1)
    temperature_K = np.arange(300, 1200.1, 10)

    print('Integrating...')
    integral, err = quad_vec(
        cross_relaxation4, a=0, b=np.inf, workers=-1,
        args=(
            make_doping(doping), make_temp_base(temperature_K),
            make_fac(d24_distances), d24_counts,
            make_fac(a8_distances), a8_counts,
        ),
    )
    integral *= kr

    print(integral)
    graph(integral, doping, temperature_K)


if __name__ == '__main__':
    freeze_support()
    main()

Output

$ time python 277855.py
[[0.874762   0.874762   0.874762   ... 0.03018313 0.02958226 0.02899388]
 [0.86336565 0.86336565 0.86336565 ... 0.03011315 0.02951437 0.02892802]
 [0.85215839 0.85215839 0.85215839 ... 0.03004337 0.02944668 0.02886235]
 ...
 [0.00837203 0.00837203 0.00837203 ... 0.00535962 0.00532487 0.00529006]
 [0.00835594 0.00835594 0.00835594 ... 0.00535277 0.00531811 0.00528338]
 [0.0083399  0.0083399  0.0083399  ... 0.00534594 0.00531136 0.00527671]]

real    0m13.105s
user    0m18.619s
sys 0m16.677s

graph

\$\endgroup\$
1
  • 2
    \$\begingroup\$ Wow, amazing! I will have a look at everything you mentioned and try to learn from it. Thanks! \$\endgroup\$
    – thimoooh
    Jul 5, 2022 at 8:23

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.