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?

Brahmanism

You keep a as it is for axis=1. Add a new axis with a[:,None] for axis=0. It's all about pushing the elems along that axis. – Hulky

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.

Hulky

The problem with this approach is that it doesn't really work well when the number of axis of X is unknown at the beginning. Also, if X has, say, 7 dimensions, writing X-a[:,None,None,None,None,None,None] looks just so wrong. – Whenas

also you can use np.newaxis instead of None, for more readability. – Catalan

@Whenas To add n dimensions, you can use a[(..., *([None] * n))]. To explain what's happening here: Tuple used for indexing; Ellipsis ... used to expand over all existing dimensions; Multiplying [None] to get a list with n times None; Unpacking the list of Nones with a leading *. Btw. I'd prefer np.newaxis over None – Pterodactyl

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.]])

Gaddy

Answer 3

You need to make the axis of a that you want to subtract b along the last axis of a, and then move the axis back to its original position afterwards.

a = np.random.uniform(-1, 1, (4, 4))
b = np.random.uniform(-1, 1, 4)

axis = 0
result = np.moveaxis(np.moveaxis(a, axis, -1) - b, -1, axis)

we can create a function which does this for any binary arithmetic operation, and which works with arrays that have an arbitrary number of dimensions:

def apply_along_axis(ufunc, a, b, axis):
    fromaxis = tuple(np.atleast1d(axis))
    toaxis = tuple(np.r_[-len(fromaxis):0])
    return np.moveaxis(ufunc(np.moveaxis(a, fromaxis, toaxis), b), toaxis, fromaxis)

An example of using this function for subtraction across an arbitrary axis:

result = apply_along_axis(np.subtract, a, b, 0)

NB: Not to be confused with numpy.apply_along_axis

Millrun

Answer 4

supose the size n is 4, you can create a broadcast like this:

>>> v=np.arange(4)
>>> v
array([0, 1, 2, 3])
>>> np.broadcast_to(v, (v.shape[0], v.shape[0]))
array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])
>>> np.transpose(np.broadcast_to(v, (v.shape[0], v.shape[0])))
array([[0, 0, 0, 0],
       [1, 1, 1, 1],
       [2, 2, 2, 2],
       [3, 3, 3, 3]])

then you can do your operation.

The broadcast is just a view of the original vector, no memory allocating.

from numpy documents:

Returns: broadcast: array; A readonly view on the original array with the given shape. It is typically not contiguous. Furthermore, more than one element of a broadcasted array may refer to a single memory location.

Walkout