使用numpy.argsort时规避内存错误及结果差异的技术咨询

Hey there! Let's break down your two questions one by one:

The core issue is that your first approach uses an incorrect indexing method that doesn't match the logic of the loop version:

• np.argsort(x, axis=2) returns c_sort, a (100,100,100) array where c_sort[i,j,k] is the index of the k-th sorted element along the 2nd axis for the (i,j) slice.
• Your loop version is correct: for each (i,j) slice, it uses the sorted indices to pull values from the corresponding slice of y along axis 2.
• But y[c_sort] treats y as a flattened 1D array! Every value in c_sort is used as an index into y.flatten(), not as an index within the (i,j) slice's axis 2. For example, if c_sort[i,j,0] = 2, y[c_sort] would grab y.flatten()[2] instead of y[i,j,2]—these are totally different positions in the array.

That's why the results don't match: the first approach's indexing logic is fundamentally wrong.

If your actual array is much larger than the (100,100,100) example (like (1000,1000,1000)), the incorrect y[c_sort] isn't just wrong—it can cause MemoryError due to temporary arrays created during numpy's advanced indexing. Here are a few solid solutions:

1. Stick with the Loop (Best for Memory) #

Your loop version is already memory-friendly:

• It only processes one (i,j) slice at a time, so no huge temporary arrays are created.
• Memory usage is just the size of x, y, and f—no extra overhead.
• If you want faster speed, use numba to JIT-compile the loop (it'll run almost as fast as native numpy):

  import numba
  import numpy as np
  
  @numba.jit(nopython=True)
  def sort_y_by_x(x, y):
      f = np.zeros_like(y)
      for i in range(x.shape[0]):
          for j in range(x.shape[1]):
              f[i,j,:] = y[i,j, np.argsort(x[i,j,:])]
      return f
  
  x = np.random.random([100,100,100])
  y = np.random.random([100,100,100])
  f = sort_y_by_x(x, y)
  

2. Use Correct Numpy Advanced Indexing (Best for Speed) #

If you prefer vectorized numpy operations, you need to build the right index arrays instead of using c_sort directly:

import numpy as np

x = np.random.random([100,100,100])
y = np.random.random([100,100,100])
c_sort = np.argsort(x, axis=2)

# Create grid indices for i and j that match c_sort's shape
i_indices = np.arange(x.shape[0])[:, np.newaxis, np.newaxis]
j_indices = np.arange(x.shape[1])[np.newaxis, :, np.newaxis]

# This indexing matches the loop's logic exactly
f = y[i_indices, j_indices, c_sort]

This is vectorized and faster than pure Python loops, though it uses a bit more memory for the index arrays (they're integer arrays, so much smaller than float arrays).

3. Process in Chunks (For Extremely Large Arrays) #

If your array is so big that even x and y strain memory, split the array into chunks and process each piece:

import numpy as np

x = np.random.random([1000,1000,1000])
y = np.random.random([1000,1000,1000])
f = np.zeros_like(y)

chunk_size = 100  # Adjust based on your available memory
for start in range(0, x.shape[0], chunk_size):
    end = start + chunk_size
    x_chunk = x[start:end]
    y_chunk = y[start:end]
    c_sort_chunk = np.argsort(x_chunk, axis=2)
    
    # Create indices for the current chunk
    i_chunk = np.arange(x_chunk.shape[0])[:, np.newaxis, np.newaxis]
    j_chunk = np.arange(x_chunk.shape[1])[np.newaxis, :, np.newaxis]
    
    f[start:end] = y_chunk[i_chunk, j_chunk, c_sort_chunk]

This only loads a portion of the array into memory at a time, making it feasible for huge datasets.

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