CODEWITHSUNDEEP
pythonarraysnumpy
homm
homm

Reputation: 2252

numpy: multidimensional linspace

numpy.linspace is one-dimensional linear gradient. For example: numpy.linspace(0, 3, 4):

[0, 1, 2, 3]

It is easy to imagine 2-dimensional linear gradient:

[[[0, 0], [1, 0], [2, 0], [3, 0]],
 [[0, 1], [1, 1], [2, 1], [3, 1]],
 [[0, 2], [1, 2], [2, 2], [3, 2]],
 [[0, 3], [1, 3], [2, 3], [3, 3]]]

What I need is 3 or even 4-dimensional linear gradients using numpy.

I can easily achieve it using Python code, but it is too slow.

%time arr = numpy.array([
    [
        [
            (a, b, c)
            for a in range(65)
        ]
        for b in range(65)
    ]
    for c in range(65)
])
Wall time: 177 ms

arr.shape
(65, 65, 65, 3)

Upvotes: 2

Views: 2445

Answers (1)

hpaulj
hpaulj

Reputation: 231375

There are number of tools to do this - np.meshgrid, np.indices, etc

numpy/lib/index_tricks.py has a number of 'cute' functions and classes that do things like this:

In [387]: np.mgrid[0:3:4j, 0:3:4j]
Out[387]: 
array([[[0., 0., 0., 0.],
        [1., 1., 1., 1.],
        [2., 2., 2., 2.],
        [3., 3., 3., 3.]],

       [[0., 1., 2., 3.],
        [0., 1., 2., 3.],
        [0., 1., 2., 3.],
        [0., 1., 2., 3.]]])

Swap the axes to get the layout you show:

In [391]: IJ.transpose(1,2,0)
Out[391]: 
array([[[0., 0.],
        [0., 1.],
        [0., 2.],
        [0., 3.]],

       [[1., 0.],
        [1., 1.],
        [1., 2.],
        [1., 3.]],

       [[2., 0.],
        [2., 1.],
        [2., 2.],
        [2., 3.]],

       [[3., 0.],
        [3., 1.],
        [3., 2.],
        [3., 3.]]])
In [392]: _.shape
Out[392]: (4, 4, 2)

Another one: np.array(list(np.ndindex((4,4)))).reshape(4,4,2), or np.array(list(itertools.product(range(4),range(4)))).reshape(4,4,2)

Upvotes: 3

Related Questions