Create a numpy matrix with elements as a function of indices

How can I create a numpy matrix with its elements being a function of its ind开发者_运维问答ices? For example, a multiplication table: a[i,j] = i*j

An Un-numpy and un-pythonic would be to create an array of zeros and then loop through.

There is no doubt that there is a better way to do this, without a loop.

However, even better would be to create the matrix straight-away.

A generic solution would be to use np.fromfunction()

From the doc:

numpy.fromfunction(function, shape, **kwargs)

Construct an array by executing a function over each coordinate. The resulting array therefore has a value fn(x, y, z) at coordinate (x, y, z).

The below snippet should provide the required matrix.

import numpy as np

np.fromfunction(lambda i, j: i*j, (5,5))

Output:

array([[  0.,   0.,   0.,   0.,   0.],
       [  0.,   1.,   2.,   3.,   4.],
       [  0.,   2.,   4.,   6.,   8.],
       [  0.,   3.,   6.,   9.,  12.],
       [  0.,   4.,   8.,  12.,  16.]])

The first parameter to the function is a callable which is executed for each of the coordinates. If foo is a function that you pass as the first argument, foo(i,j) will be the value at (i,j). This holds for higher dimensions too. The shape of the coordinate array can be modified using the shape parameter.

Edit:

Based on the comment on using custom functions like lambda x,y: 2*x if x > y else y/2, the following code works:

import numpy as np

def generic_f(shape, elementwise_f):
    fv = np.vectorize(elementwise_f)
    return np.fromfunction(fv, shape)


def elementwise_f(x , y):
    return 2*x if x > y else y/2

print(generic_f( (5,5), elementwise_f))

Output:

[[0.  0.5 1.  1.5 2. ]
 [2.  0.5 1.  1.5 2. ]
 [4.  4.  1.  1.5 2. ]
 [6.  6.  6.  1.5 2. ]
 [8.  8.  8.  8.  2. ]]

The user is expected to pass a scalar function that defines the elementwise operation. np.vectorize is used to vectorize the user-defined scalar function and is passed to np.fromfunction().

Here's one way to do that:

>>> indices = numpy.indices((5, 5))
>>> a = indices[0] * indices[1]
>>> a
array([[ 0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4],
       [ 0,  2,  4,  6,  8],
       [ 0,  3,  6,  9, 12],
       [ 0,  4,  8, 12, 16]])

To further explain, numpy.indices((5, 5)) generates two arrays containing the x and y indices of a 5x5 array like so:

>>> numpy.indices((5, 5))
array([[[0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2],
        [3, 3, 3, 3, 3],
        [4, 4, 4, 4, 4]],

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

When you multiply these two arrays, numpy multiplies the value of the two arrays at each position and returns the result.

For the multiplication

np.multiply.outer(np.arange(5), np.arange(5))  # a_ij = i * j

and in general

np.frompyfunc(
    lambda i, j: f(i, j), 2, 1
).outer(
    np.arange(5),
    np.arange(5),
).astype(np.float64)  # a_ij = f(i, j)

basically you create an np.ufunc via np.frompyfunc and then outer it with the indices.

Edit

Speed comparision between the different solutions.

Small matrices:

Eyy![1]: %timeit np.multiply.outer(np.arange(5), np.arange(5))
100000 loops, best of 3: 4.97 µs per loop

Eyy![2]: %timeit np.array( [ [ i*j for j in xrange(5)] for i in xrange(5)] )
100000 loops, best of 3: 5.51 µs per loop

Eyy![3]: %timeit indices = np.indices((5, 5)); indices[0] * indices[1]
100000 loops, best of 3: 16.1 µs per loop

Bigger matrices:

Eyy![4]: %timeit np.multiply.outer(np.arange(4096), np.arange(4096))
10 loops, best of 3: 62.4 ms per loop

Eyy![5]: %timeit indices = np.indices((4096, 4096)); indices[0] * indices[1]
10 loops, best of 3: 165 ms per loop

Eyy![6]: %timeit np.array( [ [ i*j for j in xrange(4096)] for i in xrange(4096)] )
1 loops, best of 3: 1.39 s per loop

I'm away from my python at the moment, but does this one work?

array( [ [ i*j for j in xrange(5)] for i in xrange(5)] )

Just wanted to add that @Senderle's response can be generalized for any function and dimension:

dims = (3,3,3) #i,j,k
ii = np.indices(dims)

You could then calculate a[i,j,k] = i*j*k as

a = np.prod(ii,axis=0)

or a[i,j,k] = (i-1)*j*k:

a = (ii[0,...]-1)*ii[1,...]*ii[2,...]

etc