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?
Answers:
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.
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.]])
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
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.
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?
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.
Start with 2 dimensions that are different (in label at least)
Xshape(n,m)ashape(n,)bshape(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.]])
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
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.