Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Generate a n-dimensional array of coordinates in numpy

Let's say I have a function f which can take coordinates as parameter and returns an integer (f(x) in this case). The coordinates can be multidimensional and are in the form of a list. My goal is to fill a numpy array with all values between two coordinates. I've tried to make a list of all possible indices and use it as input for the vectorized function.

Here is my code for 2 dimensional coordinates:

import itertools
import numpy


def index_array(lower_corner, upper_corner):
     x_range = range(lower_corner[0], upper_corner[0])
     y_range = range(lower_corner[1], upper_corner[1])
     return numpy.array(list(itertools.product(x_range, y_range)))


print(index_array([2, -2], [5, 3]))

This will return the index list like expected:

[[ 2 -2]
 [ 2 -1]
 [ 2  0]
 [ 2  1]
 [ 2  2]
 [ 3 -2]
 [ 3 -1]
 [ 3  0]
 [ 3  1]
 [ 3  2]
 [ 4 -2]
 [ 4 -1]
 [ 4  0]
 [ 4  1]
 [ 4  2]]

And here is my attempt for n dimensions:

import itertools
import numpy


def f(x):
    # dummy function
    return x + 5


def index_array(lower_corner, upper_corner):
    # returns all indices between two n-dimensional points
    range_list = []
    for n in range(len(lower_corner)):
        range_list.append(range(lower_corner[n], upper_corner[n]))
    return numpy.array(list(itertools.product(*range_list)))


lower_corner = numpy.array([2, -2])
upper_corner = numpy.array([5, 3])
indices = index_array(lower_corner, upper_corner)
vect_func = numpy.vectorize(f)
results = vect_func(indices)
print(results)

While this works it's quite slow and needs huge amounts of memory. Is it possible to write this in a more efficient way? I could think about using numpy.meshgrid but I don't know how I would use it.

like image 462
Gnarflord Avatar asked Jul 03 '16 13:07

Gnarflord


People also ask

What is N-dimensional array type in NumPy?

An ndarray is a (usually fixed-size) multidimensional container of items of the same type and size. The number of dimensions and items in an array is defined by its shape , which is a tuple of N non-negative integers that specify the sizes of each dimension.

How do I create a NumPy 3 dimensional array?

Importing the NumPy package enables us to use the array function in python. To create a three-dimensional array, we pass the object representing x by y by z in python, where x is the nested lists in the object, y is the nested lists inside the x nested lists, and z is the values inside each y nested list.


1 Answers

Indeed np.meshgrid would be one way to do it with some stacking, as shown below -

def ndim_grid(start,stop):
    # Set number of dimensions
    ndims = len(start)

    # List of ranges across all dimensions
    L = [np.arange(start[i],stop[i]) for i in range(ndims)]

    # Finally use meshgrid to form all combinations corresponding to all 
    # dimensions and stack them as M x ndims array
    return np.hstack((np.meshgrid(*L))).swapaxes(0,1).reshape(ndims,-1).T

Sample run

1) 2D Case :

In [97]: ndim_grid([2, -2],[5, 3])
Out[97]: 
array([[ 2, -2],
       [ 2, -1],
       [ 2,  0],
       [ 2,  1],
       [ 2,  2],
       [ 3, -2],
       [ 3, -1],
       [ 3,  0],
       [ 3,  1],
       [ 3,  2],
       [ 4, -2],
       [ 4, -1],
       [ 4,  0],
       [ 4,  1],
       [ 4,  2]])

2) 3D Case :

In [98]: ndim_grid([2, -2, 4],[5, 3, 6])
Out[98]: 
array([[ 2, -2,  4],
       [ 2, -2,  5],
       [ 2, -1,  4],
       [ 2, -1,  5],
       [ 2,  0,  4],
       [ 2,  0,  5],
       [ 2,  1,  4],
       [ 2,  1,  5],
       [ 2,  2,  4],
       [ 2,  2,  5],
       [ 3, -2,  4],
       [ 3, -2,  5],
       [ 3, -1,  4],
       [ 3, -1,  5],
       [ 3,  0,  4],
       [ 3,  0,  5],
       [ 3,  1,  4],
       [ 3,  1,  5],
       [ 3,  2,  4],
       [ 3,  2,  5],
       [ 4, -2,  4],
       [ 4, -2,  5],
       [ 4, -1,  4],
       [ 4, -1,  5],
       [ 4,  0,  4],
       [ 4,  0,  5],
       [ 4,  1,  4],
       [ 4,  1,  5],
       [ 4,  2,  4],
       [ 4,  2,  5]])
like image 101
Divakar Avatar answered Oct 24 '22 14:10

Divakar