Broadcast an operation along specific axis in python
<p>In python, suppose I have a square (numpy) matrix <strong>X</strong>, of size <em>n x n</em> and I have a numpy vector <strong>a</strong> of size <em>n</em>. </p>

<p>Very simply, I want to perform a broadcasting subtraction of <strong>X - a</strong>, 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. </p>

<p>How can I specify the axis?</p>

<p>Thanks. </p>

<p>Let's generate arrays with random elems</p>

<p>Inputs :</p>

<pre><code>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])
</code></pre>

<p><strong>I. Subtraction along <code>axis=1</code></strong></p>

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

<pre><code>In [64]: X[0] - a
Out[64]: array([ 0.31044099,  0.48637503,  0.23901049, -0.00251074])
</code></pre>

<p>Going deeper there, what's happening there is : </p>

<pre><code>X[0,0] - a[0], X[0,1] - a[1], X[0,2] - a[2] , X[0,3] - a[3]
</code></pre>

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

<pre><code>X :  n x n
a :      n
</code></pre>

<p>So, we simply do <code>X-a</code> to get all subtractions :</p>

<pre><code>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]])
</code></pre>

<p>And, finally see if we have <code>X[0] - a</code> obtained earlier is here.</p>

<p><strong>Important Note :</strong>  Thing to be noted here is that <code>a</code> 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 <code>axis=1</code>, elems of <code>a</code> would be broadcasted along the <code>axis=0</code>.</p>

<p><strong>II. Subtraction along <code>axis=0</code></strong></p>

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

<pre><code>In [67]: X[:,0]-a
Out[67]: array([ 0.31044099,  0.56985738,  0.10774888,  0.01992247])
</code></pre>

<p>Going deeper there, what's happening there is : </p>

<pre><code>X[0,0] - a[0], X[1,0] - a[1], X[2,0] - a[2] , X[3,0] - a[3]
</code></pre>

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

<pre><code>X          :  n x n
a[:,None]  :  n x 1
</code></pre>

<p>So, we do <code>X-a[:,None]</code> to get all subtractions :</p>

<pre><code>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]])
</code></pre>

<p>And, finally see if we have <code>X[:,0] - a</code> obtained earlier is here.</p>

python arrays numpy matrix numpy-broadcasting