CODEWITHSUNDEEP
numpytensorflow
Bob
Bob

Reputation: 571

Index variable range in numpy

I have a numpy zero matrix A of the shape (2, 5).

A = [[ 0.,  0.,  0.,  0.,  0.],
    [ 0.,  0.,  0.,  0.,  0.]]

I have another array seq of size 2. This is same as the first axis of A.

seq = [2, 3]

I want to create another matrix B which looks like this:

B = [[ 1.,  1.,  0.,  0.,  0.],
    [ 1.,  1.,  1.,  0.,  0.]]

B is constructed by changing the first seq[i] elements in the ith row of A with 1.

This is a toy example. A and seq can be large so efficiency is required. I would be extra thankful if someone knows how to do this in tensorflow.

Upvotes: 0

Views: 216

Answers (2)

mrry
mrry

Reputation: 126154

You can do this in TensorFlow (and with some analogous code in NumPy) as follows:

seq = [2, 3]

b = tf.expand_dims(tf.range(5), 0)   # A 1 x 5 matrix.
seq_matrix = tf.expand_dims(seq, 1)  # A 2 x 1 matrix.
b_bool = tf.greater(seq_matrix, b)   # A 2 x 5 bool matrix.

B = tf.to_int32(b_bool)              # A 2 x 5 int matrix.

Example output:

In [7]: b = tf.expand_dims(tf.range(5), 0)
        [[0 1 2 3 4]]

In [21]: b_bool = tf.greater(seq_matrix, b)
In [22]: op = sess.run(b_bool)
In [23]: print(op)
[[ True  True False False False]
 [ True  True  True False False]]

In [24]: bint = tf.to_int32(b_bool)
In [25]: op = sess.run(bint)
In [26]: print(op)
[[1 1 0 0 0]
 [1 1 1 0 0]]

Upvotes: 2

hpaulj
hpaulj

Reputation: 231385

This @mrry's solution, expressed a little differently

In [667]: [[2],[3]]>np.arange(5)
Out[667]: 
array([[ True,  True, False, False, False],
       [ True,  True,  True, False, False]], dtype=bool)
In [668]: ([[2],[3]]>np.arange(5)).astype(int)
Out[668]: 
array([[1, 1, 0, 0, 0],
       [1, 1, 1, 0, 0]])

The idea is to compare [2,3] with [0,1,2,3,4] in an 'outer' broadcasting sense. The result is boolean which can be easily changed to 0/1 integers.

Another approach would be to use cumsum (or another ufunc.accumulate function):

In [669]: A=np.zeros((2,5))
In [670]: A[range(2),[2,3]]=1
In [671]: A
Out[671]: 
array([[ 0.,  0.,  1.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  0.]])
In [672]: A.cumsum(axis=1)
Out[672]: 
array([[ 0.,  0.,  1.,  1.,  1.],
       [ 0.,  0.,  0.,  1.,  1.]])
In [673]: 1-A.cumsum(axis=1)
Out[673]: 
array([[ 1.,  1.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  0.,  0.]])

Or a variation starting with 1's:

In [681]: A=np.ones((2,5))
In [682]: A[range(2),[2,3]]=0
In [683]: A
Out[683]: 
array([[ 1.,  1.,  0.,  1.,  1.],
       [ 1.,  1.,  1.,  0.,  1.]])
In [684]: np.minimum.accumulate(A,axis=1)
Out[684]: 
array([[ 1.,  1.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  0.,  0.]])

Upvotes: 1

Related Questions