如何加快 python 中的嵌套 for 循环

Question

I want to calculate the cM between two different windows along a chromosome. My code has three nested loops. For sample, I use random number stand for the recombination map.

import random

windnr = 54800
w, h   = windnr, windnr
recmatrix = [[0 for x in range(w)] for y in range(h)]

#Generate 54800 random numbers between 10 and 30
rec_map = random.sample(range(0, 30), 54800)

for i in range(windnr):
    for j in range(windnr):
        recmatrix[i][j] = 0.25 * rec_map[i] #mean distance within own window
        if i > j:
            recmatrix[i][j] = recmatrix[i][j] + 0.5 * rec_map[j] #+ mean rdistance final window
            for k in range(i-1,j,-1):
                recmatrix[i][j] = recmatrix[i][j] + rec_map[k] #add all windows between i and j
        if i < j:
            recmatrix[i][j] = recmatrix[i][j] + 0.5 * rec_map[j] #+ mean distance final window
            for k in range(i+1,j):
                recmatrix[i][j] = recmatrix[i][j] + rec_map[k] #add all windows between i and j
        #j += 1
    if i % 10 == 0:
        print("window {}".format(i))
    #i += 1

The calculation costs a lot of time. I have to calculate almost 7 days for my data.
Can I speed up the nested for loop within 10 hours? How can I increase the performance?

Although the 2D array has 3 billion items (~96 GB when being floats), I would rule out hard disk swapping issues, since the server which does the computation has 200 GB of RAM.

Answer 1

Using Numpy will make your application much faster. It's written in C/C++, so it does not suffer from slow loops in Python .

---

I'm doing my tests on an old Intel Xeon X5550 with 2 sockets, 8 cores and 96 GB of triple channel RAM. I don't have much experience with Numpy, so bear with me, if below code is not optimal.

Array initialization

Already the initialization is much faster:

recmatrix = [[0 for x in range(w)] for y in range(h)]

needs 24 GB of RAM (integers) and takes 3:28 minutes on my PC. Whereas

recmatrix = np.zeros((windnr, windnr), dtype=np.int)

is finished after 50 ms. But since you need floats anyway, start with floats from the beginning:

recmatrix = np.zeros((windnr, windnr), dtype=np.float)

Random samples

The code

#Generate 54800 random numbers between 10 and 30
rec_map = random.sample(range(0, 30), 54800)

did not work for me, so I replaced it and increased k for more stable measurements

rec_map = random.choices(range(0, 30), k=5480000)

which runs in 2.5 seconds. The numpy replacement

rec_map = np.random.choice(np.arange(0, 30), size=5480000)

is done in 0.1 seconds.

The loop

The loop will need most work, since you'll avoid Python loops in Numpy whenever possible.

For example, if you have an array and want to multiply all elements by 2, you would not write a loop but simply multiply the whole array:

import numpy as np

single = np.random.choice(np.arange(0, 10), size=100)
doubled = single * 2
print(single, "\r\n", doubled)

I don't fully understand what the code does, but let's apply that strategy on the first part of the loop. The original is

for i in range(windnr):
    for j in range(windnr):
        recmatrix[i][j] = 0.25 * rec_map[i] #mean distance within own window

and it takes 18.5 seconds with a reduced windnr = 5480 . The numpy equivalent should be

column = 0.25 * rec_map_np
recmatrix = np.repeat(column, windnr)

and is done within 0.25 seconds. Also note: since we're assigning the variable here, we don't need the zero initialization at all.

For the if i>j: and if i<j: parts, I see that the first line is identical

recmatrix[i][j] = recmatrix[i][j] + 0.5 * rec_map[j]

That means, this calculation is applied to all elements except the ones on the diagonal. You can use a mask for that:

mask = np.ones((windnr, windnr), dtype=bool)
np.fill_diagonal(mask, False)
rec_map_2d = np.repeat(0.5 * rec_map_np, windnr-1)
recmatrix[mask] += rec_map_2d

This took only 1:20 minutes for all 54800 elements, but reached my RAM limit at 93 GB.

Answer 2

Usually in python looping always take much time. So if possible then in your case then use map this will save a lot of time for you. Where you are using a iter(list) so it will be good for this script.

example:

def func():
    your code
nu = (1, 2, 3, 4) 
output =  map(func, nu)
print(output)