We Keep Coding

numpy nansum across first index

I have an example 2 x 2 x 2 array:
np.array([[[ 1, 2],
[ 3, 4]],
[[ 5, 6],
[ 7 , 8]]])
I want the nansum of the array across the first index as follows:
Sum all values in:
[[ 1, 2],
[ 3, 4]]
and
[[ 5, 6],
[ 7 , 8]]
The sum of the first array would be 10 and the second would be 26
i.e.
array([10, 26])

I think you are looking for this
a = np.array([[[ 1, 2],
[ 3, 4]],
[[ 5, 6],
[ 7 , 8]]])
np.nansum(a,axis=(1,2))
# array([10, 26])
because you want to sum on axis 1 and 2 only, and get one number per axis 0

Select from a 3-dimensional array with a 2-dimensional array

I have two arrays:
a: a 3-dimensional source array (N x M x 2)
b: a 2-dimensional index array (N x M) containing 0 and 1s.
I want to use the indices in b to select the corresponding elements of a in its third dimension. The resulting array should have the dimensions N x M. Here is the example as code:
import numpy as np
a = np.array( # dims: 3x3x2
[[[ 0, 1],
[ 2, 3],
[ 4, 5]],
[[ 6, 7],
[ 8, 9],
[10, 11]],
[[12, 13],
[14, 15],
[16, 17]]]
)
b = np.array( # dims: 3x3
[[1, 1, 1],
[1, 1, 1],
[1, 1, 1]]
)
# select the elements in a according to b
# to achieve this result:
desired = np.array(
[[ 1, 3, 5],
[ 7, 9, 11],
[13, 15, 17]]
)
At first, I thought this must have a simple solution but I could not find one at all. Since I would like to port it to tensorflow, I would appreciate if somebody knows a numpy-type solution for this.
Edit: The third dimension of a might contain more than two elements. Hence, b might also contain indices different from 0 and 1 - it is not a boolean mask.

We can use np.where for this:
np.where(b, a[:, :, 1], a[:, :, 0])
Output:
array([[ 1, 3, 5],
[ 7, 9, 11],
[13, 15, 17]])

As #jdehesa sugests, we can use np.ogrid to obtain the indices for the first two axes:
ax0, ax1 = np.ogrid[:b.shape[0], :b.shape[1]]
And then we can use b to directly index along the last axis. Note that ax0 and ax1 will be broadcast to the shape of b:
desired = a[ax0, ax1 ,b]
print(desired)
array([[ 1, 3, 5],
[ 7, 9, 11],
[13, 15, 17]])

I added some solutions for tensorflow.
import tensorflow as tf
a = tf.constant([[[ 0, 1],[ 2, 3],[ 4, 5]],
[[ 6, 7],[ 8, 9],[10, 11]],
[[12, 13],[14, 15],[16, 17]]],dtype=tf.float32)
b = tf.constant([[1, 1, 1],[1, 1, 1],[1, 1, 1]],dtype=tf.int32)
# 1. use tf.gather_nd
colum,row = tf.meshgrid(tf.range(a.shape[0]),tf.range(a.shape[1]))
idx = tf.stack([row, colum, b], axis=-1) # Thanks for #jdehesa's suggestion
result1 = tf.gather_nd(a,idx)
# 2. use tf.reduce_sum
mask = tf.one_hot(b,depth=a.shape[-1],dtype=tf.float32)
result2 = tf.reduce_sum(a*mask,axis=-1)
# 3. use tf.boolean_mask
mask = tf.one_hot(b,depth=a.shape[-1],dtype=tf.float32)
result3 = tf.reshape(tf.boolean_mask(a,mask),b.shape)
with tf.Session() as sess:
print('method 1: \n',sess.run(result1))
print('method 2: \n',sess.run(result2))
print('method 3: \n',sess.run(result3))
method 1:
[[ 1. 3. 5.]
[ 7. 9. 11.]
[13. 15. 17.]]
method 2:
[[ 1. 3. 5.]
[ 7. 9. 11.]
[13. 15. 17.]]
method 3:
[[ 1. 3. 5.]
[ 7. 9. 11.]
[13. 15. 17.]]

You can use np.take_along_axis:
import numpy as np
a = np.array(
[[[ 0, 1],
[ 2, 3],
[ 4, 5]],
[[ 6, 7],
[ 8, 9],
[10, 11]],
[[12, 13],
[14, 15],
[16, 17]]])
b = np.array(
[[1, 1, 1],
[1, 1, 1],
[1, 1, 1]])
print(np.take_along_axis(a, b[..., np.newaxis], axis=-1)[..., 0])
# [[ 1 3 5]
# [ 7 9 11]
# [13 15 17]]

Numpy assignment like 'numpy.take'

Is it possible to assign to a numpy array along the lines of how the take functionality works?
E.g. if I have a an array a, a list of indices inds, and a desired axis, I can use take as follows:
import numpy as np
a = np.arange(12).reshape((3, -1))
inds = np.array([1, 2])
print(np.take(a, inds, axis=1))
[[ 1 2]
[ 5 6]
[ 9 10]]
This is extremely useful when the indices / axis needed may change at runtime.
However, numpy does not let you do this:
np.take(a, inds, axis=1) = 0
print(a)
It looks like there is some limited (1-D) support for this via numpy.put, but I was wondering if there was a cleaner way to do this?

In [222]: a = np.arange(12).reshape((3, -1))
...: inds = np.array([1, 2])
...:
In [223]: np.take(a, inds, axis=1)
Out[223]:
array([[ 1, 2],
[ 5, 6],
[ 9, 10]])
In [225]: a[:,inds]
Out[225]:
array([[ 1, 2],
[ 5, 6],
[ 9, 10]])
construct an indexing tuple
In [226]: idx=[slice(None)]*a.ndim
In [227]: axis=1
In [228]: idx[axis]=inds
In [229]: a[tuple(idx)]
Out[229]:
array([[ 1, 2],
[ 5, 6],
[ 9, 10]])
In [230]: a[tuple(idx)] = 0
In [231]: a
Out[231]:
array([[ 0, 0, 0, 3],
[ 4, 0, 0, 7],
[ 8, 0, 0, 11]])
Or for a[inds,:]:
In [232]: idx=[slice(None)]*a.ndim
In [233]: idx[0]=inds
In [234]: a[tuple(idx)]
Out[234]:
array([[ 4, 0, 0, 7],
[ 8, 0, 0, 11]])
In [235]: a[tuple(idx)]=1
In [236]: a
Out[236]:
array([[0, 0, 0, 3],
[1, 1, 1, 1],
[1, 1, 1, 1]])
PP's suggestion:
def put_at(inds, axis=-1, slc=(slice(None),)):
return (axis<0)*(Ellipsis,) + axis*slc + (inds,) + (-1-axis)*slc
To be used as in a[put_at(ind_list,axis=axis)]
I've seen both styles on numpy functions. This looks like one used for extend_dims, mine was used in apply_along/over_axis.
earlier thoughts
In a recent take question I/we figured out that it was equivalent to arr.flat[ind] for some some raveled index. I'll have to look that up.
There is an np.put that is equivalent to assignment to the flat:
Signature: np.put(a, ind, v, mode='raise')
Docstring:
Replaces specified elements of an array with given values.
The indexing works on the flattened target array. `put` is roughly
equivalent to:
a.flat[ind] = v
Its docs also mention place and putmask (and copyto).
numpy multidimensional indexing and the function 'take'
I commented take (without axis) is equivalent to:
lut.flat[np.ravel_multi_index(arr.T, lut.shape)].T
with ravel:
In [257]: a = np.arange(12).reshape((3, -1))
In [258]: IJ=np.ix_(np.arange(a.shape[0]), inds)
In [259]: np.ravel_multi_index(IJ, a.shape)
Out[259]:
array([[ 1, 2],
[ 5, 6],
[ 9, 10]], dtype=int32)
In [260]: np.take(a,np.ravel_multi_index(IJ, a.shape))
Out[260]:
array([[ 1, 2],
[ 5, 6],
[ 9, 10]])
In [261]: a.flat[np.ravel_multi_index(IJ, a.shape)] = 100
In [262]: a
Out[262]:
array([[ 0, 100, 100, 3],
[ 4, 100, 100, 7],
[ 8, 100, 100, 11]])
and to use put:
In [264]: np.put(a, np.ravel_multi_index(IJ, a.shape), np.arange(1,7))
In [265]: a
Out[265]:
array([[ 0, 1, 2, 3],
[ 4, 3, 4, 7],
[ 8, 5, 6, 11]])
Use of ravel is unecessary in this case but might useful in others.

I have given an example for use of
numpy.take in 2 dimensions. Perhaps you can adapt that to your problem

You can juste use indexing in this way :
a[:,[1,2]]=0

convert separate 1D np.arrays into a list of 2D np.arrays

I'm trying to convert three 1D arrays into a list of 2D arrays. I've managed to do this by creating an empty ndarray and populating it line by line. Could someone show me a more elegant approach?
import numpy as np
import pandas as pd
one=np.arange(1,4,1)
two=np.arange(10,40,10)
three=np.arange(100,400,100)
df=pd.DataFrame({'col1':one,'col2':two,'col3':three})
desired_output=[np.array([[1.,10.],[1.,100.]]),np.array([[2.,20.],[2.,200.]]),np.array([[3.,30.],[3.,300.]])]
current, inelegant approach that works:
output=[]
for i in range(len(df)):
temp=np.zeros(shape=(2,2))
temp[0][0]=df.iloc[i,0]
temp[0][1]=df.iloc[i,1]
temp[1][0]=df.iloc[i,0]
temp[1][1]=df.iloc[i,2]
output.append(temp)

so first of all you can get array from df values by simply doing the following
In [61]:
arr = df.values
arr
Out[61]:
array([[ 1, 10, 100],
[ 2, 20, 200],
[ 3, 30, 300]])
then add the first column in the array again
In [73]:
arr_mod = np.hstack((arr , arr[: , 0][:, np.newaxis]))
arr_mod
Out[73]:
array([[ 1, 10, 100, 1],
[ 2, 20, 200, 2],
[ 3, 30, 300, 3]])
swap the column you've just added with the last column in the array
In [74]:
arr_mod[: , [2 , 3]] = arr_mod [: , [3 , 2]]
arr_mod
Out[74]:
array([[ 1, 10, 1, 100],
[ 2, 20, 2, 200],
[ 3, 30, 3, 300]])
then convert this 2d array to 3d array and convert it to list
In [78]:
list(arr_mod.reshape( -1, 2 , 2))
Out[78]:
[array([[ 1, 10],
[ 1, 100]]), array([[ 2, 20],
[ 2, 200]]), array([[ 3, 30],
[ 3, 300]])]

Here's one approach using np.column_stack and np.vsplit -
arr2D = np.column_stack((df['col1'],df['col2'],df['col1'],df['col3']))
out_list = np.vsplit(arr2D.reshape(-1,2),arr2D.shape[0])
Basically, we use np.column_stack to stack column-1 with column-2 and then again column-1 with column-3 to give us a 2D NumPy array arr2D of shape N x 4. Next, we reshape arr2D to a 2*N X 2 array and split along the rows with np.vsplit to give us the expected list of 2D arrays.
Sample run -
>>> df
col1 col2 col3
0 1 10 100
1 2 20 200
2 3 30 300
3 4 40 400
4 5 50 500
5 6 60 600
>>> arr2D = np.column_stack((df['col1'],df['col2'],df['col1'],df['col3']))
>>> out_list = np.vsplit(arr2D.reshape(-1,2),arr2D.shape[0])
>>> print out_list
[array([[ 1, 10],
[ 1, 100]]), array([[ 2, 20],
[ 2, 200]]), array([[ 3, 30],
[ 3, 300]]), array([[ 4, 40],
[ 4, 400]]), array([[ 5, 50],
[ 5, 500]]), array([[ 6, 60],
[ 6, 600]])]

Python: efficient operations on numpy arrays

Say I have a numpy array x:
x = array([[ 3, 2, 1],
[ 3, 25, 34],
[ 33, 333, 3],
[ 43, 32, 2]])
I want to carry out the following operations without explicitly writing a for loop i.e. say a method which uses automatic in built looping;
1) Replace the 2nd column by a column of all 1 i.e.
x = array([[ 3, 1, 1],
[ 3, 1, 34],
[ 33, 1, 3],
[ 43, 1, 2]])
2) In the original array , replace 3rd column with the product of 2nd and 3rd i.e.
x = array([[ 3, 2, 1*2],
[ 3, 25, 34*25],
[ 33, 333, 3*333],
[ 43, 32, 2*32]])
3) Finally, I would like to replace the 2nd column in the original array based on a condition i.e.
x[1] = 0 if x[0] > 5 else 4
i.e. the array now looks like:
x = array([[ 3, 4, 1],
[ 3, 4, 34],
[ 33, 0, 3],
[ 43, 0, 2]])
Any suggestions ?
Thanks !

The documentation on numpy is well worth reading as this is fairly basic stuff...
x[:,1] = 1
x[:,2] *= x[:,1]
x[:,1] = np.where( x[:,0] > 5, 0, 4 )