Broadcast an operation along specific axis in python

Question

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?

Answer 1

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.

Answer 2

Start with 2 dimensions that are different (in label at least)

• X shape (n,m)
• a shape (n,)
• b shape (m,)

The ways to combine these are:

(n,m)-(n,) => (n,m)-(n,1) => (n,m)
X - a[:,None]     

(n,m)-(m,) => (n,m)-(1,m) => (n,m)
X - b[None,:]
X - b      # [None,:] is automatic, if needed.

The basic point is that when the number dimensions differ, numpy can add new dimensions at the start, but you have to be explicit about adding new dimensions at the end.

Or to combine 2 1d arrays in a outer product (difference):

(n,) - (m,) => (n,1)-(1,m) => (n,m)
a[:,None] - b[None,:]
a[:,None] - b

Without the these rules, a-b could result in a (n,m) or (m,n) or something else.

And with 2 matching length arrays:

(n,) - (n,) => (n,)
a - a

or

(n,) - (n,) => (n,1)-(1,n) => (n,n)
a[:,None]-a[None,:]

=============

To write a function that would take an axis parameter, you could use np.expand_dims:

In [220]: np.expand_dims([1,2,3],0)
Out[220]: array([[1, 2, 3]])    # like [None,:]
In [221]: np.expand_dims([1,2,3],1)
Out[221]:             # like [:,None]
array([[1],
       [2],
       [3]])

def foo(X, a, axis=0):
    return X - np.expand_dims(a, axis=axis)

to be used as:

In [223]: foo(np.eye(3),[1,2,3],axis=0)
Out[223]: 
array([[ 0., -2., -3.],
       [-1., -1., -3.],
       [-1., -2., -2.]])
In [224]: foo(np.eye(3),[1,2,3],axis=1)
Out[224]: 
array([[ 0., -1., -1.],
       [-2., -1., -2.],
       [-3., -3., -2.]])