3

I have a regular three dimensional Cartesian Grid with numpy.linspace(a1,an,na), numpy.linspace(b1,bn,nb), numpy.linspace(c1,cn,nc) along their respective dimensions.

I now want to evaluate a function at all points $(a_i,b_j,c_k)$ on this grid. The function f(list_of_points,*args) takes a list of points and returns the list of scalar function values for each point.

Are there functions that a) display a Cartesian grid as a list, such that I can easily use it for this function, and then b) allow easy reformatting of the returned list of scalars into an array of dimensions (na,nb,nc)?

ertl
  • 133
  • 1
  • 4
  • It's easy to do this looping through each point. However, this takes fixed computational cost at each evaluation of f, which is what I want to avoid by writing the grid as a list, and calling f only once on that list. – ertl May 23 '20 at 14:25

1 Answers1

3

You might be looking for np.meshgrid(). It's easiest to see how it works by example:

>>> import numpy as np

# Start with three different 1-D arrays, possibly with different sizes.
>>> a = np.random.random(10)
>>> b = np.random.random(15)
>>> c = np.random.random(20)

# Use np.meshgrid() to create 3-D arrays by joining the three 1-D arrays.
>>> A, B, C = np.meshgrid(a, b, c, indexing="ij")
>>> print(A.shape, B.shape, C.shape)
(10, 15, 20) (10, 15, 20) (10, 15, 20)

# It's easy to use these arrays for computations like x_ijk = a_i + b_j * c_k.
>>> X = A + B * C

# Let's do this with loops to check that we get the same result.
>>> Y = np.empty((10, 15, 20))
>>> for i in range(a.size):                    # i = 0, ..., 9
...     for j in range(b.size):                # j = 0, ..., 14
...         for k in range(c.size):            # k = 0, ..., 19
...             Y[i,j,k] = a[i] + b[j] * c[k]  # y_ijk = a_i + b_j * c_k
... 
>>> np.allclose(X,Y)
True

This lets you compute the values of any scalar-valued function $f: (a,b,c) \mapsto \mathbb{R}$ over the desired grid in an efficient, readable way; just remember to use NumPy functions where applicable (e.g., np.sin() instead of math.sin()).

For example, if $f(a,b,c) = \sin(a)b^2 + e^{c}$, start by defining your mesh arrays a, b, and c, use np.meshgrid() to get A, B, and C, and set F = np.sin(A) * B**2 + np.exp(C).

vanPelt2
  • 165
  • 1
  • 5