如何利用多进程加速cov_p函数中Numpy数组的填充效率？

Looking at your code, the biggest bottleneck is definitely the nested loops in cov_p where each (i,j) pair triggers a call to R_final—and inside that, four separate infinite integrals via quad. Since each (i,j) calculation is completely independent of the others, this is a perfect candidate for parallelization with multiprocessing (we use multiprocessing instead of multithreading here because CPU-bound tasks like numerical integrals are blocked by Python's GIL).

Let's break down the solution into actionable steps:

Step 1: Optimize Single-Threaded Code First #

Before jumping into parallelization, let's eliminate redundant computations in your integral functions—this will give us a speedup even before adding parallelism.

Notice that S_xr and S_xc compute the same h_x value twice. We can merge these into a single integrand function that returns both real and imaginary parts, cutting down on repeated matrix operations:

import numpy as np
from scipy.integrate import quad

# Keep L0 as a constant (no need to redefine it every time)
L_0 = np.eye(2)

def compute_R_x(Phi, l_mat, s_f, t, i, j):
    def integrand(omega):
        # Compute h_x once, not twice
        h_x = np.real(np.linalg.multi_dot(
            [Phi, np.linalg.inv(-1j*omega*np.linalg.inv(l_mat) + np.eye(l_mat.shape[0])), Phi.T]
        ))
        term = np.exp(1j*omega*t) * np.linalg.multi_dot([np.conjugate(h_x), s_f, h_x.T])
        return np.real(term[i][j]), np.imag(term[i][j])
    
    # Compute real and imaginary integrals in one pass (well, two calls, but no redundant h_x)
    real_val, _ = quad(lambda omega: integrand(omega)[0], -np.inf, np.inf, limit=10000)
    imag_val, _ = quad(lambda omega: integrand(omega)[1], -np.inf, np.inf, limit=10000)
    return real_val + 1j*imag_val

def R_final(Phi, l_mat, s_f, t):
    return np.array([[compute_R_x(Phi, l_mat, s_f, t, i, j) for j in range(2)] for i in range(2)])

Also, replace the I function with a direct i == j check—function calls add unnecessary overhead for such a simple operation:

# Remove the old I(m,p) function entirely

Step 2: Parallelize the Nested Loops with ProcessPoolExecutor #

We'll wrap the computation for each (i,j) block in a worker function, then use concurrent.futures.ProcessPoolExecutor to run these workers in parallel across multiple CPU cores.

Worker Function #

This function takes a single (i,j) pair and all necessary parameters, computes the 2x2 block for that position, and returns the indices along with the block:

def compute_cov_block(i, j, Phi, l_mat, s_f, sigma_11, sigma_22, tstep):
    # Compute the R_final term
    r_block = np.linalg.multi_dot([L_0, R_final(Phi, l_mat, s_f, (i-j)*tstep), L_0.T])
    # Add the sigma identity term (only when i == j)
    sigma_block = np.array([[sigma_11, 0], [0, sigma_22]], dtype=complex) if i == j else np.zeros((2,2), dtype=complex)
    return i, j, r_block + sigma_block

Modified cov_p Function #

Instead of nested loops, we generate all (i,j) pairs, submit them to the process pool, then assemble the results into the final array:

from concurrent.futures import ProcessPoolExecutor

def cov_p(Phi, l_mat, s_f, sigma_11, sigma_22, N_p):
    # Initialize the empty answer array
    ans = np.zeros([N_p, N_p, 2, 2], dtype=complex)
    
    # Generate all (i,j) pairs to compute
    all_pairs = [(i, j) for i in range(N_p) for j in range(N_p)]
    
    # Use a process pool to compute blocks in parallel
    # By default, this uses all available CPU cores
    with ProcessPoolExecutor() as executor:
        # Submit all tasks to the pool
        futures = [
            executor.submit(
                compute_cov_block, i, j, Phi, l_mat, s_f, sigma_11, sigma_22, tstep
            ) for i, j in all_pairs
        ]
        
        # Collect results and fill the answer array
        for future in futures:
            i, j, block = future.result()
            ans[i, j] = block
    
    # Reshape as before
    return ans.transpose(0, 2, 1, 3).reshape(2*N_p, -1)

Step 3: Critical Note for Windows Users #

If you're running this on Windows, you must wrap your test code in an if __name__ == '__main__' guard to avoid infinite process spawning:

if __name__ == '__main__':
    # Test parameters
    phi = np.array([[-0.0529255 +0.00662948j, -0.0529255 -0.00662948j, -0.03050694-0.00190298j, -0.03050694+0.00190298j], [-0.04149906+0.00171591j, -0.04149906-0.00171591j, 0.01974404-0.00194719j, 0.01974404+0.00194719j]])
    lamb_mat = np.array([[-1.00390867 +6.28783994j, 0. +0.j , 0. +0.j , 0. +0.j ], [ 0. +0.j , -1.00390867 -6.28783994j, 0. +0.j , 0. +0.j ], [ 0. +0.j , 0. +0.j , -0.25859133+12.09860357j, 0. +0.j ], [ 0. +0.j , 0. +0.j , 0. +0.j , -0.25859133-12.09860357j]])
    S_f = np.array([[100,0],[0,100]])
    tstep = 0.1
    sigma_11 = 0.3
    sigma_22 = 0.4
    
    # Run the test
    result = cov_p(phi, lamb_mat, S_f, sigma_11, sigma_22, 10)
    print(f"Result shape: {result.shape}")

Advanced Optimizations (For Even More Speed) #

1. Exploit Symmetry #

If your R_final(t) has symmetry properties (e.g., R_final(-t) = np.conj(R_final(t)).T), you can compute only half the (i,j) pairs and mirror the results. For example:

• Compute all pairs where i >= j
• For i < j, set ans[i,j] = np.conj(ans[j,i]).T

This cuts the number of computations in half—huge win for larger N_p.

2. Use Shared Memory for Large Arrays #

If Phi or l_mat are very large, passing them to each worker process involves copying data, which adds overhead. You can use multiprocessing.Array to create shared memory arrays that all processes can access without copying:

from multiprocessing import Array

def create_shared_numpy_array(arr):
    # Create a shared memory array matching the input array's dtype and shape
    shared_arr = Array(arr.dtype.char, arr.size)
    # Wrap it in a numpy array
    np_shared = np.frombuffer(shared_arr.get_obj(), dtype=arr.dtype).reshape(arr.shape)
    # Copy the original data into the shared array
    np_shared[:] = arr[:]
    return np_shared

Then in your main code, create shared versions of large arrays and pass those to the worker instead of the original arrays.

内容的提问来源于stack exchange，提问作者Shakir Rather