如何用NumPy创建坐标元组矩阵并为每个元素应用自定义函数？

Hey there! You're already off to a great start with np.meshgrid—let's break down how to turn those coordinate grids into a matrix of function results without messy loops.

Step 1: Generate Your Coordinate Grids #

First, use np.meshgrid to create 2D arrays representing the i and j coordinates for your grid. For a 3×3 matrix, here's how:

import numpy as np

# Define the range of coordinates (0 to 2 for 3x3)
x = np.arange(3)
y = np.arange(3)

# Create 2D grids for i (xx) and j (yy) coordinates
xx, yy = np.meshgrid(x, y)

Now xx holds the row indices (i-values) for every position, and yy holds the column indices (j-values). For a 3×3 grid, xx looks like:

[[0 1 2]
 [0 1 2]
 [0 1 2]]

And yy looks like:

[[0 0 0]
 [1 1 1]
 [2 2 2]]

Step 2: Apply Your Custom Function to the Coordinates #

You have two straightforward ways to feed these coordinates into myfunc(i,j) -> k:

Method 1: Vectorize Your Function (Quick & Simple)

NumPy's np.vectorize wraps your scalar function to work with arrays directly. This is perfect for small grids or simple functions:

# Define your custom function
def myfunc(i, j):
    # Example: return i squared plus j cubed
    return i**2 + j**3

# Vectorize the function to handle arrays
vec_myfunc = np.vectorize(myfunc)

# Apply to the coordinate grids—this gives your result matrix directly
result_matrix = vec_myfunc(xx, yy)

For our example myfunc, the output will be:

[[0 1 4]
 [1 2 5]
 [8 9 12]]

Method 2: Flatten Coordinates & Reshape (More Control)

If you prefer to work with coordinate tuples explicitly (or need better performance for large grids), flatten the grids into pairs of (i,j) coordinates, apply the function, then reshape back to your grid size:

# Flatten the grids and stack into (i,j) tuples
coords = np.stack((xx.flatten(), yy.flatten()), axis=1)

# Apply myfunc to each coordinate pair (using a list comprehension here)
results = np.array([myfunc(i, j) for i, j in coords])

# Reshape back to 3x3 matrix
result_matrix = results.reshape(3, 3)

This gives the exact same result as the vectorize method, but avoids the overhead of np.vectorize (which is a convenience wrapper, not true vectorization under the hood).

Why This Works Better Than Manual Traversal #

Instead of looping through each coordinate tuple individually (which is slow in Python for large grids), we leverage NumPy's array operations to handle the entire grid at once. This is the "NumPy way"—it's faster and more readable once you get the hang of it.

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