如何优化NumPy中paint_rect函数的数组赋值操作？

Got it, let's dive into speeding up your paint_rect function—since it's getting called thousands of times in your genetic algorithm, even small per-call optimizations will compound into big gains. Let's break down the low-hanging fruit and more impactful changes you can make:

1. Inline get_coords to Eliminate Function Call Overhead #

Function calls in Python have non-trivial overhead when repeated thousands of times. Since get_coords is a tiny helper, moving its logic directly into paint_rect cuts out that overhead entirely.

Before:

def get_coords(w, h, x0, y0, x1, y1):
    real_x0 = int(w * min(x0, x1))
    real_y0 = int(h * min(y0, y1))
    real_x1 = int(w * max(x0, x1))
    real_y1 = int(h * max(y0, y1))
    return real_x0, real_y0, real_x1, real_y1

def paint_rect(tensor, rect, w, h, opaque=True):
    [x0, y0, x1, y1, r, g, b, a, depth] = rect
    real_x0, real_y0, real_x1, real_y1 = get_coords(w, h, x0, y0, x1, y1)
    # ... rest of code

After:

def paint_rect(tensor, rect, w, h, opaque=True):
    # Note: You mentioned rects have 8 elements—double-check the 9-element unpack here!
    x0, y0, x1, y1, r, g, b, a, depth = rect
    # Inline get_coords logic directly
    real_x0 = int(w * min(x0, x1))
    real_y0 = int(h * min(y0, y1))
    real_x1 = int(w * max(x0, x1))
    real_y1 = int(h * max(y0, y1))
    # ... rest of code

2. Reuse Tensor Slices to Avoid Redundant Computation #

Right now, you're recalculating the same slice (real_y0:real_y1, real_x0:real_x1) four times. Storing this slice as a variable cuts down on repeated indexing work, and you can even batch RGB assignments to reduce attribute lookups:

Optimized Slice Reuse:

def paint_rect(tensor, rect, w, h, opaque=True):
    x0, y0, x1, y1, r, g, b, a, depth = rect
    real_x0 = int(w * min(x0, x1))
    real_y0 = int(h * min(y0, y1))
    real_x1 = int(w * max(x0, x1))
    real_y1 = int(h * max(y0, y1))
    
    # Store the target slice once
    target_slice = tensor[:, real_y0:real_y1, real_x0:real_x1]
    # Assign RGB channels in one go
    target_slice[:3] = [r, g, b]
    # Fix opaque parameter usage (your original code ignored it!)
    target_slice[3] = 1 if opaque else a

3. Skip Unnecessary Variable Unpacking (For Ultra-Tight Loops) #

Unpacking the rect array into individual variables has a small cost. For extreme cases, you can index directly into rect instead to save a tiny bit of overhead:

def paint_rect(tensor, rect, w, h, opaque=True):
    real_x0 = int(w * min(rect[0], rect[2]))
    real_y0 = int(h * min(rect[1], rect[3]))
    real_x1 = int(w * max(rect[0], rect[2]))
    real_y1 = int(h * max(rect[1], rect[3]))
    
    target_slice = tensor[:, real_y0:real_y1, real_x0:real_x1]
    target_slice[:3] = [rect[4], rect[5], rect[6]]
    target_slice[3] = 1 if opaque else rect[7]
    # Remove unused `depth` variable if it's not needed!

4. Batch Process Rectangles (Biggest Potential Gain) #

The biggest win will come from moving away from per-rectangle function calls entirely. If your genetic algorithm processes batches of rectangles, rewrite the logic to handle all rects at once using vectorized operations (works with PyTorch/Numpy/TensorFlow):

Example with Numpy:

def paint_rects_batch(tensor, rects, w, h, opaque=True):
    rects = np.array(rects)
    # Calculate all real coordinates in parallel
    x0, y0, x1, y1 = rects[:,0], rects[:,1], rects[:,2], rects[:,3]
    real_x0 = (w * np.minimum(x0, x1)).astype(int)
    real_y0 = (h * np.minimum(y0, y1)).astype(int)
    real_x1 = (w * np.maximum(x0, x1)).astype(int)
    real_y1 = (h * np.maximum(y0, y1)).astype(int)
    
    # Iterate over batch (still way faster than per-function calls)
    for i in range(len(rects)):
        rx0, ry0, rx1, ry1 = real_x0[i], real_y0[i], real_x1[i], real_y1[i]
        r, g, b = rects[i,4], rects[i,5], rects[i,6]
        a = 1 if opaque else rects[i,7]
        tensor[:, ry0:ry1, rx0:rx1] = [r, g, b, a]

If your use case allows, you can even eliminate the loop entirely with advanced indexing—this will give you the maximum speedup for large batches.

5. Final Quick Wins #

• Profile first: Use cProfile to confirm which parts of the function are actually slow. Sometimes bottlenecks aren't where you expect them to be.
• Contiguous tensors: If you're using PyTorch, ensure your tensor is contiguous with tensor.contiguous()—slice assignments are faster on contiguous memory.
• Fix unused variables: You unpack a depth variable but never use it—remove it to avoid unnecessary work and potential errors.

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