CODEWITHSUNDEEP
pythonarraysnumpy
Raccoon 987
Raccoon 987

Reputation: 57

Transformation of the 3d numpy array

I have 3d array and I need to set to zero its right part. For each 2d slice (n, :, :) of the array the index of the column should be taken from vector b. This index defines separating point - the left and right parts, as shown in the figure below.

enter image description here

a_before = [[[ 1  2  3  4]
             [ 5  6  7  8]
             [ 9 10 11 12]
             [13 14 15 16]]

            [[17 18 19 20]
             [21 22 23 24]
             [25 26 27 28]
             [29 30 31 32]]

            [[33 34 35 36]
             [37 38 39 40]
             [41 42 43 44]
             [45 46 47 48]]]

a_before.shape = (3, 4, 4)
b = (2, 3, 1)

a_after_1 = [[[ 1  2  0  0]
              [ 5  6  0  0]
              [ 9 10 0 0]
              [13 14 0 0]]

             [[17 18 19 0]
              [21 22 23 0]
              [25 26 27 0]
              [29 30 31 0]]

             [[33 0 0 0]
              [37 0 0 0]
              [41 0 0 0]
              [45 0 0 0]]]

After this, for each 2d slice (n, :, :) I have to take index of the column from c vector and multiply by the corresponding value taken from the vector d.

c = (1, 2, 0)
d = (50, 100, 150)
a_after_2 = [[[ 1  100  0  0]
              [ 5  300  0  0]
              [ 9 500 0 0]
              [13 700 0 0]]

             [[17 18 1900 0]
              [21 22 2300 0]
              [25 26 2700 0]
              [29 30 3100 0]]

             [[4950 0 0 0]
              [5550 0 0 0]
              [6150 0 0 0]
              [6750 0 0 0]]]

I did it but my version looks ugly. Maybe someone can help me.

P.S. I would like to avoid for loops and use only numpy methods.

Thank You.

Upvotes: 0

Views: 362

Answers (2)

hpaulj
hpaulj

Reputation: 231325

Here's a version without loops.

In [232]: A = np.arange(1,49).reshape(3,4,4)
In [233]: b = np.array([2,3,1])
In [234]: d = np.array([50,100,150])
In [235]: I,J = np.nonzero(b[:,None]<=np.arange(4))
In [236]: A[I,:,J]=0
In [237]: A[np.arange(3),:,b-1] *= d[:,None]
In [238]: A
Out[238]: 
array([[[   1,  100,    0,    0],
        [   5,  300,    0,    0],
        [   9,  500,    0,    0],
        [  13,  700,    0,    0]],

       [[  17,   18, 1900,    0],
        [  21,   22, 2300,    0],
        [  25,   26, 2700,    0],
        [  29,   30, 3100,    0]],

       [[4950,    0,    0,    0],
        [5550,    0,    0,    0],
        [6150,    0,    0,    0],
        [6750,    0,    0,    0]]])

Before I developed this, I wrote an iterative version. It helped me visualize the problem.

In [240]: Ac = np.arange(1,49).reshape(3,4,4)
In [241]: 
In [241]: for i,v in enumerate(b):
     ...:     Ac[i,:,v:]=0
     ...: 
In [242]: for i,(bi,di) in enumerate(zip(b,d)):
     ...:     Ac[i,:,bi-1]*=di

It may be easier to understand, and in that sense, less ugly!

The fact that your A has middle dimension that is "just-going-along" for the ride, complicates "vectorizing" the problem.

With a (3,4) 2d array, the solution is just:

In [251]: Ab = Ac[:,0,:]
In [252]: Ab[b[:,None]<=np.arange(4)]=0
In [253]: Ab[np.arange(3),b-1]*=d

Upvotes: 2

pakpe
pakpe

Reputation: 5479

Here it is:

import numpy as np

a = np.arange(1,49).reshape(3,4,4)
b = np.array([2,3,1])
c = np.array([1,2,0])
d = np.array([50,100,150])

for i in range(len(b)):
    a[i,:,b[i]:] = 0

for i,j in enumerate(c):
    a[i,:,j] = a[i,:,j]* d[i]

print(a)

#
[[[   1  100    0    0]
  [   5  300    0    0]
  [   9  500    0    0]
  [  13  700    0    0]]

 [[  17   18 1900    0]
  [  21   22 2300    0]
  [  25   26 2700    0]
  [  29   30 3100    0]]

 [[4950    0    0    0]
  [5550    0    0    0]
  [6150    0    0    0]
  [6750    0    0    0]]]

Upvotes: 0

Related Questions