How do I apply some function to a Python meshgrid?

A meshgrid creates coordinate matrices from coordinate vectors, allowing you to apply functions across all combinations of input values. Python's NumPy provides efficient ways to apply functions to meshgrids using vectorization.

Basic Function Application with Lists

You can apply functions to coordinate vectors using NumPy's vectorize decorator ?

import numpy as np

@np.vectorize
def foo(a, b):
    return a + b

x = [0.0, 0.5, 1.0]
y = [0.0, 1.0, 8.0]
print("Function Output:", foo(x, y))

Function Output: [0. 1.5 9. ]

Creating and Using Meshgrids

For true meshgrid operations, use np.meshgrid() to create coordinate matrices ?

import numpy as np

def distance_function(x, y):
    return np.sqrt(x**2 + y**2)

# Create coordinate vectors
x = np.array([0, 1, 2])
y = np.array([0, 1, 2])

# Create meshgrid
X, Y = np.meshgrid(x, y)
print("X grid:")
print(X)
print("\nY grid:")
print(Y)

# Apply function to meshgrid
result = distance_function(X, Y)
print("\nDistance from origin:")
print(result)

X grid:
[[0 1 2]
 [0 1 2]
 [0 1 2]]

Y grid:
[[0 0 0]
 [1 1 1]
 [2 2 2]]

Distance from origin:
[[0.         1.         2.        ]
 [1.         1.41421356 2.23606798]
 [2.         2.23606798 2.82842712]]

Complex Function Example

Apply more complex mathematical functions to meshgrids ?

import numpy as np

def wave_function(x, y):
    return np.sin(x) * np.cos(y)

# Create finer meshgrid
x = np.linspace(0, 2*np.pi, 5)
y = np.linspace(0, 2*np.pi, 5)
X, Y = np.meshgrid(x, y)

# Apply wave function
Z = wave_function(X, Y)
print("Wave function results:")
print(np.round(Z, 3))

Wave function results:
[[ 0.     0.     0.     0.     0.   ]
 [ 0.707  0.     0.    -0.     0.707]
 [ 1.     0.    -0.    -0.     1.   ]
 [ 0.707 -0.    -0.     0.     0.707]
 [ 0.     0.     0.     0.     0.   ]]

Comparison of Methods

Conclusion

Use np.meshgrid() to create coordinate matrices, then apply functions directly using NumPy's vectorized operations. For complex functions, @np.vectorize provides a convenient decorator approach.