Generating data from meshgrid data (Numpy) - python

I'd like to ask how to generate corresponding values from a meshgrid. I have a function "foo" that takes one 1D array with the length of 2, and returns some real number.
import numpy as np
def foo(X):
#this function takes a vector, e.g., np.array([2,3]), and returns a real number.
return sum(X)**np.sin( sum(X) );
x = np.arange(-2, 1, 1) # points in the x axis
y = np.arange( 3, 8, 1) # points in the y axis
X, Y = np.meshgrid(x, y) # X, Y : grid
I generate X and Y grids using meshgrid.
Then, how can I generate corresponding Z values using "foo" function, in order to plot them in 3D, e.g., plotting using plot_surface function with X,Y,Z values?
Here the question is how to generate Z values, which has the same shape to X and Y, using "foo" function. Since my "foo" function only takes an 1D array, I do not know how I can uses this function with X and Y to generate corresponding Z values.

Stack your two numpy arrays in "depth" using np.dstack, and then modify your foo function, so that it operates on only the last axis of your stacked array. This is easily done using np.sum with parameter axis=-1, instead of using the builtin sum:
import numpy as np
def foo(xy):
return np.sum(xy, axis=-1) ** np.sin(np.sum(xy, axis=-1))
x = np.arange(-2, 1, 1) # points in the x axis
y = np.arange( 3, 8, 1) # points in the y axis
X, Y = np.meshgrid(x, y) # X, Y : grid
XY = np.dstack((X, Y))
And now, you should get:
>>> XY.shape
(5, 3, 2)
>>> foo(XY)
array([[ 1. , 1.87813065, 1.1677002 ],
[ 1.87813065, 1.1677002 , 0.35023496],
[ 1.1677002 , 0.35023496, 0.2136686 ],
[ 0.35023496, 0.2136686 , 0.60613935],
[ 0.2136686 , 0.60613935, 3.59102217]])
If you want to achieve the same effect, but without modifying foo, then you can use np.apply_along_axis, which should do exactly what you need:
>>> np.apply_along_axis(foo, -1, XY)
array([[ 1. , 1.87813065, 1.1677002 ],
[ 1.87813065, 1.1677002 , 0.35023496],
[ 1.1677002 , 0.35023496, 0.2136686 ],
[ 0.35023496, 0.2136686 , 0.60613935],
[ 0.2136686 , 0.60613935, 3.59102217]])

Related

Looking up index of value in numpy 3D arrays

import numpy as np
# The 3D arrays have the axis: Z, X, Y
arr_keys = np.random.rand(20, 5, 5)
arr_vals = np.random.rand(20, 5, 5)
arr_idx = np.random.rand(5, 5)
For each grid cell in arr_idx, I want to look up the Z-position of the value closest to it in arr_keys (but with the same X, Y location) and return the value at the corresponding position in arr_vals array. Is there a way to do this without using nested for loops?
So, if the value at X=0, Y=0 for arr_idx is 0.5, I want to find the number closest to it at X=0, Y=0, Z ranges from 0 to 10
in arr_keys, and then I want to use the Z position of that number (lets call it Z_prime) to find the value in arr_vals (Z_prime, X=0, Y=0)
This is the type of problem for which np.take_along_axis was created:
# shape (20, 5, 5)
diff = np.abs(arr_idx - arr_keys)
# argmin(..., keepdims=True) doesn't exist yet - this emulates it
# shape (1, 5, 5)
inds = np.expand_dims(np.argmin(diff, axis=0), axis=0)
# shape (1, 5, 5)
res = np.take_along_axis(arr_vals, inds, axis=0)
# shape (5, 5)
res = res.squeeze(axis=0)
I think #xnx's answer is pretty good. Mine is longer but I'll post it anyway ;).
Also, a note: NumPy is made to handle large multi-dimensional arrays efficiently by vectorizing the operations. So I'd suggest avoiding for loops as much as possible. Whatever the task you're looking for, there is (usually) a way to do it while avoiding loops.
arr_keys = np.split(arr_keys, 20)
arr_keys = np.stack(arr_keys, axis=-1)[0]
arr_vals = np.split(arr_vals, 20)
arr_vals = np.stack(arr_vals, axis=-1)[0]
arr_idx = np.expand_dims(arr_idx, axis=-1)
difference = np.abs(arr_keys - arr_idx)
minimum = np.argmin(difference, axis=-1)
result = np.take_along_axis(arr_vals, np.expand_dims(minimum, axis=-1), axis=-1)
result = np.squeeze(result, axis=-1)
I think this might work: roll the axes into the correct orientation, find the index of the value of the (absolute) minimum for each of the 5x5 X,Y values and take the corresponding Z-values from arr_vals:
idx = np.argmin(np.abs(np.rollaxis(arr_keys,0,3) - arr_idx[:,:,None]), axis=2)
i,j = np.ogrid[:5,:5]
arr_vals[idx[i,j],i,j]
To test this, try the (3,2,2) case:
In [15]: arr_keys
Out[15]:
array([[[ 0.19681533, 0.26897784],
[ 0.60469711, 0.09273087]],
[[ 0.04961604, 0.3460404 ],
[ 0.88406912, 0.41284309]],
[[ 0.46298201, 0.33809574],
[ 0.99604152, 0.4836324 ]]])
In [16]: arr_vals
Out[16]:
array([[[ 0.88865681, 0.88287688],
[ 0.3128103 , 0.24188022]],
[[ 0.23947227, 0.57913325],
[ 0.85768064, 0.91701097]],
[[ 0.78105669, 0.84144339],
[ 0.81071981, 0.69217687]]])
In [17]: arr_idx
Out[17]:
array([[[ 0.31352609],
[ 0.75462329]],
[[ 0.44445286],
[ 0.97086161]]])
gives:
array([[ 0.88865681, 0.57913325],
[ 0.3128103 , 0.69217687]])
A little verbose than the already posted solution but easier to understand.
import numpy as np
# The 3D arrays have the axis: Z, X, Y
arr_keys = np.random.rand(20, 5, 5)
arr_vals = np.random.rand(20, 5, 5)
arr_idx = np.random.rand(5, 5)
arr_idx = arr_idx[np.newaxis, :, :]
dist = np.abs(arr_idx - arr_keys)
dist_ind = np.argmin(dist, axis=0)
x = np.arange(0, 5, 1)
y = np.arange(0, 5, 1)
xx, yy = np.meshgrid(x, y)
res = arr_vals[dist_ind, yy, xx]

Making a multidimensional list of vectors

I am quite new to Python so bear with me. I am writing a program to calculate some physical quantity, let's call it A. A is a function of several variables, let's call them x, y, z. So I have three nested loops to calculate A for the values of x, y, z that I am interested in.
for x in xs:
for y in ys:
for z in zs:
A[x, y, z] = function_calculating_value(x,y,z)
Now, the problem is that A[x,y,z] is two-dimensional array containing both the mean value and the variance so that A[x,y,z] = [mean, variance]. From other languages I am used to initializing A using function similar to np.zeros(). How do I do that here? What is the easiest way to achieve what I want, and how do I access the mean and variance easily for a given (x,y,z)?
(the end goal is to be able to plot the mean with the variance as error bars, so if there is an even more elegant way of doing this, I appreciate that as well)
thanks in advance!
You can create and manipulate your multi-dimensional array with numpy
# Generate a random 4d array that has nx = 3, ny = 3, and nz = 3, with each 3D point having 2 values
mdarray = np.random.random( size = (3,3,3,2) )
# The overall shape of the 4d array
mdarray
Out[66]:
array([[[[ 0.80091246, 0.28476668],
[ 0.94264747, 0.27247111],
[ 0.64503087, 0.13722768]],
[[ 0.21371798, 0.41006764],
[ 0.79783723, 0.02537987],
[ 0.80658387, 0.43464532]],
[[ 0.04566927, 0.74836831],
[ 0.8280196 , 0.90288647],
[ 0.59271082, 0.65910184]]],
[[[ 0.82533798, 0.29075978],
[ 0.76496127, 0.1308289 ],
[ 0.22767752, 0.01865939]],
[[ 0.76849458, 0.7934015 ],
[ 0.93313128, 0.88436557],
[ 0.06897508, 0.00307739]],
[[ 0.15975812, 0.00792386],
[ 0.40292818, 0.21209199],
[ 0.48805502, 0.71974702]]],
[[[ 0.66522525, 0.49797465],
[ 0.29369336, 0.68743839],
[ 0.46411967, 0.69547356]],
[[ 0.50339875, 0.66423777],
[ 0.80520751, 0.88115054],
[ 0.08296022, 0.69467829]],
[[ 0.76572574, 0.45332754],
[ 0.87982243, 0.15773385],
[ 0.5762041 , 0.91268172]]]])
# Both values for this specific sample at x = 0, y = 1 and z = 2
mdarray[0,1,2]
Out[67]: array([ 0.80658387, 0.43464532])
mdarray[0,1,2,0] # mean only at the same point
Out[68]: 0.8065838666297338
mdarray[0,1,2,1] # variance only at the same point
Out[69]: 0.43464532443865489
You can also get only the means or the variance values separately by slicing the array:
mean = mdarray[:,:,:,0]
variance = mdarray[:,:,:,1]
mean
Out[74]:
array([[[ 0.80091246, 0.94264747, 0.64503087],
[ 0.21371798, 0.79783723, 0.80658387],
[ 0.04566927, 0.8280196 , 0.59271082]],
[[ 0.82533798, 0.76496127, 0.22767752],
[ 0.76849458, 0.93313128, 0.06897508],
[ 0.15975812, 0.40292818, 0.48805502]],
[[ 0.66522525, 0.29369336, 0.46411967],
[ 0.50339875, 0.80520751, 0.08296022],
[ 0.76572574, 0.87982243, 0.5762041 ]]])
I'm still unsure how I would have preferred to plot this data, will think about this a bit and update this answer.

Reshaping numpy array without using two for loops

I have two numpy arrays
import numpy as np
x = np.linspace(1e10, 1e12, num=50) # 50 values
y = np.linspace(1e5, 1e7, num=50) # 50 values
x.shape # output is (50,)
y.shape # output is (50,)
I would like to create a function which returns an array shaped (50,50) such that the first x value x0 is evaluated for all y values, etc.
The current function I am using is fairly complicated, so let's use an easier example. Let's say the function is
def func(x,y):
return x**2 + y**2
How do I shape this to be a (50,50) array? At the moment, it will output 50 values. Would you use a for loop inside an array?
Something like:
np.array([[func(x,y) for i in x] for j in y)
but without using two for loops. This takes forever to run.
EDIT: It has been requested I share my "complicated" function. Here it goes:
There is a data vector which is a 1D numpy array of 4000 measurements. There is also a "normalized_matrix", which is shaped (4000,4000)---it is nothing special, just a matrix with entry values of integers between 0 and 1, e.g. 0.5567878. These are the two "given" inputs.
My function returns the matrix multiplication product of transpose(datavector) * matrix * datavector, which is a single value.
Now, as you can see in the code, I have initialized two arrays, x and y, which pass through a series of "x parameters" and "y parameters". That is, what does func(x,y) return for value x1 and value y1, i.e. func(x1,y1)?
The shape of matrix1 is (50, 4000, 4000). The shape of matrix2 is (50, 4000, 4000). Ditto for total_matrix.
normalized_matrix is shape (4000,4000) and id_mat is shaped (4000,4000).
normalized_matrix
print normalized_matrix.shape #output (4000,4000)
data_vector = datarr
print datarr.shape #output (4000,)
def func(x, y):
matrix1 = x [:, None, None] * normalized_matrix[None, :, :]
matrix2 = y[:, None, None] * id_mat[None, :, :]
total_matrix = matrix1 + matrix2
# transpose(datavector) * matrix * datavector
# by matrix multiplication, equals single value
return np.array([ np.dot(datarr.T, np.dot(total_matrix, datarr) ) ])
If I try to use np.meshgrid(), that is, if I try
x = np.linspace(1e10, 1e12, num=50) # 50 values
y = np.linspace(1e5, 1e7, num=50) # 50 values
X, Y = np.meshgrid(x,y)
z = func(X, Y)
I get the following value error: ValueError: operands could not be broadcast together with shapes (50,1,1,50) (1,4000,4000).
reshape in numpy as different meaning. When you start with a (100,) and change it to (5,20) or (10,10) 2d arrays, that is 'reshape. There is anumpy` function to do that.
You want to take 2 1d array, and use those to generate a 2d array from a function. This is like taking an outer product of the 2, passing all combinations of their values through your function.
Some sort of double loop is one way of doing this, whether it is with an explicit loop, or list comprehension. But speeding this up depends on that function.
For at x**2+y**2 example, it can be 'vectorized' quite easily:
In [40]: x=np.linspace(1e10,1e12,num=10)
In [45]: y=np.linspace(1e5,1e7,num=5)
In [46]: z = x[:,None]**2 + y[None,:]**2
In [47]: z.shape
Out[47]: (10, 5)
This takes advantage of numpy broadcasting. With the None, x is reshaped to (10,1) and y to (1,5), and the + takes an outer sum.
X,Y=np.meshgrid(x,y,indexing='ij') produces two (10,5) arrays that can be used the same way. Look at is doc for other parameters.
So if your more complex function can be written in a way that takes 2d arrays like this, it is easy to 'vectorize'.
But if that function must take 2 scalars, and return another scalar, then you are stuck with some sort of double loop.
A list comprehension form of the double loop is:
np.array([[x1**2+y1**2 for y1 in y] for x1 in x])
Another is:
z=np.empty((10,5))
for i in range(10):
for j in range(5):
z[i,j] = x[i]**2 + y[j]**2
This double loop can be sped up somewhat by using np.vectorize. This takes a user defined function, and returns one that can take broadcastable arrays:
In [65]: vprod=np.vectorize(lambda x,y: x**2+y**2)
In [66]: vprod(x[:,None],y[None,:]).shape
Out[66]: (10, 5)
Test that I've done in the past show that vectorize can improve on the list comprehension route by something like 20%, but the improvement is nothing like writing your function to work with 2d arrays in the first place.
By the way, this sort of 'vectorization' question has been asked many times on SO numpy. Beyond these broad examples, we can't help you without knowning more about that more complicated function. As long as it is a black box that takes scalars, the best we can help you with is np.vectorize. And you still need to understand broadcasting (with or without meshgrid help).
I think there is a better way, it is right on the tip of my tongue, but as an interim measure:
You are operating on 1x2 windows of a meshgrid. You can use as_strided from numpy.lib.stride_tricks to rearrange the meshgrid into two-element windows, then apply your function to the resultant array. I like to use a generic nd solution, sliding_windows (http://www.johnvinyard.com/blog/?p=268) (Not mine) to transform the array.
import numpy as np
a = np.array([1,2,3])
b = np.array([.1, .2, .3])
z= np.array(np.meshgrid(a,b))
def foo((x,y)):
return x+y
>>> z.shape
(2, 3, 3)
>>> t = sliding_window(z, (2,1,1))
>>> t
array([[ 1. , 0.1],
[ 2. , 0.1],
[ 3. , 0.1],
[ 1. , 0.2],
[ 2. , 0.2],
[ 3. , 0.2],
[ 1. , 0.3],
[ 2. , 0.3],
[ 3. , 0.3]])
>>> v = np.apply_along_axis(foo, 1, t)
>>> v
array([ 1.1, 2.1, 3.1, 1.2, 2.2, 3.2, 1.3, 2.3, 3.3])
>>> v.reshape((len(a), len(b)))
array([[ 1.1, 2.1, 3.1],
[ 1.2, 2.2, 3.2],
[ 1.3, 2.3, 3.3]])
>>>
This should be somewhat faster.
You may need to modify your function's argument signature.
If the link to the johnvinyard.com blog breaks, I've posted the the sliding_window implementation in other SO answers - https://stackoverflow.com/a/22749434/2823755
Search around and you'll find many other tricky as_strided solutions.
In response to your edited question:
normalized_matrix
print normalized_matrix.shape #output (4000,4000)
data_vector = datarr
print datarr.shape #output (4000,)
def func(x, y):
matrix1 = x [:, None, None] * normalized_matrix[None, :, :]
matrix2 = y[:, None, None] * id_mat[None, :, :]
total_matrix = matrix1 + matrix2
# transpose(datavector) * matrix * datavector
# by matrix multiplication, equals single value
# return np.array([ np.dot(datarr.T, np.dot(total_matrix, datarr))])
return np.einsum('j,ijk,k->i',datarr,total_matrix,datarr)
Since datarr is shape (4000,), transpose does nothing. I believe you want the result of the 2 dots to be shape (50,). I'm suggesting using einsum. But it can be done with tensordot, or I think even np.dot(np.dot(total_matrix, datarr),datarr). Test the expression with smaller arrays, focusing on getting the shapes right.
x = np.linspace(1e10, 1e12, num=50) # 50 values
y = np.linspace(1e5, 1e7, num=50) # 50 values
z = func(x,y)
# X, Y = np.meshgrid(x,y)
# z = func(X, Y)
X,Y is wrong. func takes x and y that are 1d. Notice how you expand the dimensions with [:, None, None]. Also you aren't creating a 2d array from an outer combination of x and y. None of your arrays in func is (50,50) or (50,50,...). The higher dimensions are provided by nomalied_matrix and id_mat.
When showing us the ValueError you should also indicate where in your code that occurred. Otherwise we have to guess, or recreate the code ourselves.
In fact when I run my edited func(X,Y), I get this error:
----> 2 matrix1 = x [:, None, None] * normalized_matrix[None, :, :]
3 matrix2 = y[:, None, None] * id_mat[None, :, :]
4 total_matrix = matrix1 + matrix2
5 # transpose(datavector) * matrix * datavector
ValueError: operands could not be broadcast together with shapes (50,1,1,50) (1,400,400)
See, the error occurs right at the start. normalized_matrix is expanded to (1,400,400) [I'm using smaller examples]. The (50,50) X is expanded to (50,1,1,50). x expands to (50,1,1), which broadcasts just fine.
To address the edit and the broadcasting error in the edit:
Inside your function you are adding dimensions to arrays to try to get them to broadcast.
matrix1 = x [:, None, None] * normalized_matrix[None, :, :]
This expression looks like you want to broadcast a 1d array with a 2d array.
The results of your meshgrid are two 2d arrays:
X,Y = np.meshgrid(x,y)
>>> X.shape, Y.shape
((50, 50), (50, 50))
>>>
When you try to use X in in your broadcasting expression the dimensions don't line up, that is what causes the ValueError - refer to the General Broadcasting Rules:
>>> x1 = X[:, np.newaxis, np.newaxis]
>>> nm = normalized_matrix[np.newaxis, :, :]
>>> x1.shape
(50, 1, 1, 50)
>>> nm.shape
(1, 4000, 4000)
>>>
You're on the right track with your list comprehension, you just need to add in an extra level of iteration:
np.array([[func(i,j) for i in x] for j in y])

linear interpolation in numpy

I have 2 numpy arrays
X = [[2 3 6], [7 2 9], [7 1 4]]
a = [0 0.0005413307 0.0010949014 0.0015468832 0.0027740823 0.0033288284]
b = [0 0.0050251256 0.0100502513 0.0150753769 0.0201005025 0.0251256281]
new = []
for z in range(3):
new.append(interp1d(a, z[0], b, 'linear'))
I am getting error as :
if xi is not None and shape[axis] != len(xi):
TypeError: tuple indices must be integers, not str
I need to find the linear interpolation of the same. How can I find that?
I have values X with respect to time a but I want to find interpolation for time b.
Linear interpolation will give me 3 points as in X for every a[i] and b[i] ?
You put the arguments in wrong order. Flowing is the help message of interp1d, check it out:
interp1d(x, y, kind='linear', axis=-1, copy=True, bounds_error=True,fill_value=np.nan)
Interpolate a 1-D function.
x and y are arrays of values used to approximate some function f:
y = f(x) .
This class returns a function whose call method uses interpolation
to find the value of new points.
interp1d is a function whose return value is a new function. This new function can then be called with values in the given interpolation range:
from scipy.interpolate import interp1d
x1 = [ 0., 0.04007922, 0.04723573, 0.05440107, 0.06178645, 0.06837938]
x2 = [ 0., 0.00502513, 0.01005025, 0.01507538, 0.0201005, 0.02512563]
f = interp1d(x1, x2)
f([0.0, 0.01, 0.02, 0.03, 0.068])
#array([ 0. , 0.0012538 , 0.0025076 , 0.0037614 , 0.02483647])

Numpy meshgrid in 3D

Numpy's meshgrid is very useful for converting two vectors to a coordinate grid. What is the easiest way to extend this to three dimensions? So given three vectors x, y, and z, construct 3x3D arrays (instead of 2x2D arrays) which can be used as coordinates.
Numpy (as of 1.8 I think) now supports higher that 2D generation of position grids with meshgrid. One important addition which really helped me is the ability to chose the indexing order (either xy or ij for Cartesian or matrix indexing respectively), which I verified with the following example:
import numpy as np
x_ = np.linspace(0., 1., 10)
y_ = np.linspace(1., 2., 20)
z_ = np.linspace(3., 4., 30)
x, y, z = np.meshgrid(x_, y_, z_, indexing='ij')
assert np.all(x[:,0,0] == x_)
assert np.all(y[0,:,0] == y_)
assert np.all(z[0,0,:] == z_)
Here is the source code of meshgrid:
def meshgrid(x,y):
"""
Return coordinate matrices from two coordinate vectors.
Parameters
----------
x, y : ndarray
Two 1-D arrays representing the x and y coordinates of a grid.
Returns
-------
X, Y : ndarray
For vectors `x`, `y` with lengths ``Nx=len(x)`` and ``Ny=len(y)``,
return `X`, `Y` where `X` and `Y` are ``(Ny, Nx)`` shaped arrays
with the elements of `x` and y repeated to fill the matrix along
the first dimension for `x`, the second for `y`.
See Also
--------
index_tricks.mgrid : Construct a multi-dimensional "meshgrid"
using indexing notation.
index_tricks.ogrid : Construct an open multi-dimensional "meshgrid"
using indexing notation.
Examples
--------
>>> X, Y = np.meshgrid([1,2,3], [4,5,6,7])
>>> X
array([[1, 2, 3],
[1, 2, 3],
[1, 2, 3],
[1, 2, 3]])
>>> Y
array([[4, 4, 4],
[5, 5, 5],
[6, 6, 6],
[7, 7, 7]])
`meshgrid` is very useful to evaluate functions on a grid.
>>> x = np.arange(-5, 5, 0.1)
>>> y = np.arange(-5, 5, 0.1)
>>> xx, yy = np.meshgrid(x, y)
>>> z = np.sin(xx**2+yy**2)/(xx**2+yy**2)
"""
x = asarray(x)
y = asarray(y)
numRows, numCols = len(y), len(x) # yes, reversed
x = x.reshape(1,numCols)
X = x.repeat(numRows, axis=0)
y = y.reshape(numRows,1)
Y = y.repeat(numCols, axis=1)
return X, Y
It is fairly simple to understand. I extended the pattern to an arbitrary number of dimensions, but this code is by no means optimized (and not thoroughly error-checked either), but you get what you pay for. Hope it helps:
def meshgrid2(*arrs):
arrs = tuple(reversed(arrs)) #edit
lens = map(len, arrs)
dim = len(arrs)
sz = 1
for s in lens:
sz*=s
ans = []
for i, arr in enumerate(arrs):
slc = [1]*dim
slc[i] = lens[i]
arr2 = asarray(arr).reshape(slc)
for j, sz in enumerate(lens):
if j!=i:
arr2 = arr2.repeat(sz, axis=j)
ans.append(arr2)
return tuple(ans)
Can you show us how you are using np.meshgrid? There is a very good chance that you really don't need meshgrid because numpy broadcasting can do the same thing without generating a repetitive array.
For example,
import numpy as np
x=np.arange(2)
y=np.arange(3)
[X,Y] = np.meshgrid(x,y)
S=X+Y
print(S.shape)
# (3, 2)
# Note that meshgrid associates y with the 0-axis, and x with the 1-axis.
print(S)
# [[0 1]
# [1 2]
# [2 3]]
s=np.empty((3,2))
print(s.shape)
# (3, 2)
# x.shape is (2,).
# y.shape is (3,).
# x's shape is broadcasted to (3,2)
# y varies along the 0-axis, so to get its shape broadcasted, we first upgrade it to
# have shape (3,1), using np.newaxis. Arrays of shape (3,1) can be broadcasted to
# arrays of shape (3,2).
s=x+y[:,np.newaxis]
print(s)
# [[0 1]
# [1 2]
# [2 3]]
The point is that S=X+Y can and should be replaced by s=x+y[:,np.newaxis] because
the latter does not require (possibly large) repetitive arrays to be formed. It also generalizes to higher dimensions (more axes) easily. You just add np.newaxis where needed to effect broadcasting as necessary.
See http://www.scipy.org/EricsBroadcastingDoc for more on numpy broadcasting.
i think what you want is
X, Y, Z = numpy.mgrid[-10:10:100j, -10:10:100j, -10:10:100j]
for example.
Here is a multidimensional version of meshgrid that I wrote:
def ndmesh(*args):
args = map(np.asarray,args)
return np.broadcast_arrays(*[x[(slice(None),)+(None,)*i] for i, x in enumerate(args)])
Note that the returned arrays are views of the original array data, so changing the original arrays will affect the coordinate arrays.
Instead of writing a new function, numpy.ix_ should do what you want.
Here is an example from the documentation:
>>> ixgrid = np.ix_([0,1], [2,4])
>>> ixgrid
(array([[0],
[1]]), array([[2, 4]]))
>>> ixgrid[0].shape, ixgrid[1].shape
((2, 1), (1, 2))'
You can achieve that by changing the order:
import numpy as np
xx = np.array([1,2,3,4])
yy = np.array([5,6,7])
zz = np.array([9,10])
y, z, x = np.meshgrid(yy, zz, xx)

Categories

Resources