如何利用NumPy广播实现多数组双高斯模型的完全向量化优化？

Great question! Eliminating that final loop over wavelength_array will definitely give you a nice speed boost, especially with larger datasets. Let's refactor your function to leverage NumPy's broadcasting capabilities entirely, getting rid of all explicit loops.

Step-by-Step Explanation #

First, let's recap what your original code does: it generates a double Gaussian for each sub-array in wavelength_array and each width value. The key is to reshape all your arrays so NumPy can broadcast them across all combinations automatically, instead of looping through each sub-array.

1. Calculate Fixed Values
   Start with your separation and central wavelength values, but compute lambda0 for all sub-arrays at once instead of per iteration:

   separation = (3729.875 - 3727.092) / 2
   # Get the middle element of each sub-array in wavelength_array
   lambda0 = wavelength_array[:, (wavelength_array.shape[1] + 1) // 2]
   

2. Reshape Arrays for Broadcasting
   We need to adjust the dimensions of our arrays so NumPy knows how to align them across all sub-arrays, wavelengths, and width values:

   • Expand lambda0 (and its shifted versions) to a 3D shape (N, 1, 1) where N is the number of sub-arrays in wavelength_array. This lets it align with every wavelength in each sub-array.
   • Expand width to a 3D shape (1, 1, W) where W is the number of width values. This lets it apply to every sub-array and wavelength.
   • Expand wavelength_array to (N, M, 1) where M is the length of each sub-array, to align with the width dimension.

3. Compute the Double Gaussian in One Go
   With all arrays properly shaped, we can compute both Gaussians and sum them without any loops.

Full Vectorized Function #

Here's the refactored, fully vectorized version of your function:

import numpy as np

def Model(wavelength_array, width):
    """Returns the model
    Parameters:
    wavelength_array: Full wavelength array (2D array: [num_subarrays, num_wavelengths])
    width: width of the Gaussians (1D array of varying widths)
    """
    separation = (3729.875 - 3727.092) / 2
    
    # Get central wavelength for each sub-array
    lambda0 = wavelength_array[:, (wavelength_array.shape[1] + 1) // 2]
    
    # Reshape arrays for broadcasting
    # Shape (N, 1, 1): aligns with each wavelength in each sub-array
    mean1 = (lambda0 - separation)[:, np.newaxis, np.newaxis]
    mean2 = (lambda0 + separation)[:, np.newaxis, np.newaxis]
    
    # Shape (1, 1, W): applies to all sub-arrays and wavelengths
    width_expanded = width[np.newaxis, np.newaxis, :]
    
    # Shape (N, M, 1): aligns with each width value
    wavelengths_expanded = wavelength_array[:, :, np.newaxis]
    
    # Compute both Gaussians and sum
    gaussian1 = np.exp(-((wavelengths_expanded - mean1) / width_expanded) ** 2)
    gaussian2 = np.exp(-((wavelengths_expanded - mean2) / width_expanded) ** 2)
    
    return gaussian1 + gaussian2

Why This Works #

NumPy's broadcasting rules handle the alignment automatically:

• wavelengths_expanded (shape (N, M, 1)) broadcasts with mean1 (shape (N, 1, 1)) to produce a (N, M, 1) array of wavelength-mean differences.
• This difference array then broadcasts with width_expanded (shape (1, 1, W)) to apply every width value to every wavelength and sub-array.
• The final result is a (N, M, W) array, identical to what your original loop-based code produces, but computed entirely with optimized C-level operations instead of slow Python loops.

Testing with Your Example #

Using your sample inputs:

wavelength_array = np.array([[4500, 4502, 4504, 4506],
                             [4502, 4504, 4506, 4508]])
width = np.arange(.1, .5, .1)

# Get the vectorized result
vectorized_models = Model(wavelength_array, width)

# Compare with original loop-based result (should be identical)
def OriginalModel(wavelength_array, width):
    Gaussian = lambda x, mean, std: np.exp(-((x[:, np.newaxis] - mean)/std)**2)
    separation = (3729.875-3727.092)/2
    lambda0 = wavelength_array[:,(len(wavelength_array[0])+1)//2]
    models = []
    for i,j in enumerate(wavelength_array):
        models.append(Gaussian(j, lambda0[i]-separation, width) + \
                      Gaussian(j, lambda0[i]+separation, width))
    return np.asarray(models)

original_models = OriginalModel(wavelength_array, width)
print(np.allclose(vectorized_models, original_models))  # Prints True

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