从函数创建矩阵

Question

I want to create a matrix from a function, such that the (3,3) matrix C has values equal to 1 if the row index is smaller than a given threshold k.

import numpy as np

k = 3
C = np.fromfunction(lambda i,j: 1 if i < k else 0, (3,3))

However, this piece of code throws an error

"The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()" and I do not really understand why.

Answer 1

The code for fromfunction is:

dtype = kwargs.pop('dtype', float)
args = indices(shape, dtype=dtype)
return function(*args, **kwargs)

You see it calls function just once - with the whole array of indices . It is not iterative.

In [672]: idx = np.indices((3,3))                                                    
In [673]: idx                                                                        
Out[673]: 
array([[[0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]],

       [[0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]]])

Your lambda expects scalar i,j values, not a 3d array

 lambda i,j: 1 if i < k else 0

idx<3 is a 3d boolean array. The error arises when that is use in an if context.

np.vectorize or np.frompyfunc is better if you want to apply a scalar function to a set of arrays:

In [677]: np.vectorize(lambda i,j: 1 if i < 2 else 0)(idx[0],idx[1])                 
Out[677]: 
array([[1, 1, 1],
       [1, 1, 1],
       [0, 0, 0]])

However it isn't faster than more direct iterative approaches, and way slower than functions that operation on whole arrays.

One of many whole-array approaches:

In [680]: np.where(np.arange(3)[:,None]<2, np.ones((3,3),int), np.zeros((3,3),int))  
Out[680]: 
array([[1, 1, 1],
       [1, 1, 1],
       [0, 0, 0]])

Answer 2

As suggested by @MarkSetchell you need to vectorize your function:

k = 3
f = lambda i,j: 1 if i < k else 0

C = np.fromfunction(np.vectorize(f), (3,3))

and you get:

C
array([[1, 1, 1],
       [1, 1, 1],
       [1, 1, 1]])

Answer 3

The problem is that np.fromfunction does not iterate over all elements, it only returns the indices in each dimension. You can use np.where() to apply a condition based on those indices, choosing from two alternatives depending on the condition:

import numpy as np

k = 3
np.fromfunction(lambda i, j: np.where(i < k, 1, 0), (5,3))

which gives:

array([[1, 1, 1],
       [1, 1, 1],
       [1, 1, 1],
       [0, 0, 0],
       [0, 0, 0]])

This avoids naming the lambda without things becoming too unwieldy. On my laptop, this approach was about 20 times faster than np.vectorize() .