[FIXED] Broadcast an operation along specific axis in python

Issue

In python, suppose I have a square numpy matrix X, of size n x n and I have a numpy vector a of size n.

Very simply, I want to perform a broadcasting subtraction of X - a, but I want to be able to specify along which dimension, so that I can specify for the subtraction to be either along axis 0 or axis 1.

How can I specify the axis?

Solution

Let's generate arrays with random elems

Inputs :

In [62]: X
Out[62]: 
array([[ 0.32322974,  0.50491961,  0.40854442,  0.36908488],
       [ 0.58840196,  0.1696713 ,  0.75428203,  0.01445901],
       [ 0.27728281,  0.33722084,  0.64187916,  0.51361972],
       [ 0.39151808,  0.6883594 ,  0.93848072,  0.48946276]])

In [63]: a
Out[63]: array([ 0.01278876,  0.01854458,  0.16953393,  0.37159562])

I. Subtraction along axis=1

Let's do the subtraction along axis=1, i.e. we want to subtract a from the first row of X, the second row of X and so on. For ease of inspecting correctness, let's just use the first row of X :

In [64]: X[0] - a
Out[64]: array([ 0.31044099,  0.48637503,  0.23901049, -0.00251074])

Going deeper there, what's happening there is :

X[0,0] - a[0], X[0,1] - a[1], X[0,2] - a[2] , X[0,3] - a[3]

So, we are matching the second axis of X with the first axis of a. Since, X is 2D and a is 1D, both are already aligned :

X :  n x n
a :      n

So, we simply do X-a to get all subtractions :

In [65]: X-a
Out[65]: 
array([[ 0.31044099,  0.48637503,  0.23901049, -0.00251074],
       [ 0.5756132 ,  0.15112672,  0.5847481 , -0.3571366 ],
       [ 0.26449405,  0.31867625,  0.47234523,  0.1420241 ],
       [ 0.37872932,  0.66981482,  0.76894679,  0.11786714]])

And, finally see if we have X[0] - a obtained earlier is here.

Important Note : Thing to be noted here is that a elems would be along one axis and along that subtraction would be done and the broadcasting would happen along the other axis. So, in this case, even though subtraction is happening along axis=1, elems of a would be broadcasted along the axis=0.

II. Subtraction along axis=0

Similarly, let's do the subtraction along axis=0, i.e. we want to subtract a from the first col of X, the second col of X and so on. For ease of inspecting correctness, let's just use the first col of X :

In [67]: X[:,0]-a
Out[67]: array([ 0.31044099,  0.56985738,  0.10774888,  0.01992247])

Going deeper there, what's happening there is :

X[0,0] - a[0], X[1,0] - a[1], X[2,0] - a[2] , X[3,0] - a[3]

So, we are matching the first axis of X with the first axis of a. Since, X is 2D and a is 1D, we need to extend a to 2D and keep all elems along its first axis with a[:,None] :

X          :  n x n
a[:,None]  :  n x 1

So, we do X-a[:,None] to get all subtractions :

In [68]: X-a[:,None]
Out[68]: 
array([[ 0.31044099,  0.49213085,  0.39575566,  0.35629612],
       [ 0.56985738,  0.15112672,  0.73573745, -0.00408557],
       [ 0.10774888,  0.16768691,  0.47234523,  0.34408579],
       [ 0.01992247,  0.31676379,  0.5668851 ,  0.11786714]])

And, finally see if we have X[:,0] - a obtained earlier is here.

Answered By - Divakar