How does xarray broadcasting work with numpy functions?
<p>I am trying to do a coordinate transformation with my xarray coordinates so </p>

<p>I have a DataArray like:</p>

<pre><code>d = xr.DataArray(np.zeros((10, 10, 1)), dims=['x', 'y', 'z']
</code></pre>

<p>and am doing operations like:</p>

<pre><code>r = np.sqrt(d.x**2 + d.y**2 + d.z**2)
theta = np.arctan2(np.sqrt(d.x**2 + d.y**2), d.z)
phi = np.arctan2(d.y, d.x)
</code></pre>

<p>I get shapes of:</p>

<pre><code>In [212]: r.shape
Out[212]: (10, 10, 1)

In [214]: theta.shape
Out[214]: (10, 10)

In [216]: phi.shape
Out[216]: (10,)
</code></pre>

<p>I would like to see the <code>(10, 10, 1)</code> shape that I get for r for all of them. </p>

<p>It looks like xarray's fancy broadcasting only kicks in for basic arithmetic operations and the numpy functions are broadcasting naively. </p>

<p>Is this correct?</p>

<p>Is there an approved way to get around this? I can write my own xarray aware arctan2 function which gets the broadcasting right, but I am hoping I don't have to.</p>

<p>Unfortunately, NumPy functions generally don't do appropriate broadcasting on xarray objects. But xarray does come with wrapped versions of many of these functions, including <code>arctan2</code>, in the <a href="http://xarray.pydata.org/en/stable/api.html#universal-functions" rel="nofollow"><code>xarray.ufuncs</code> module</a>. So <code>xarray.ufuncs.arctan2</code> should do exactly what you're looking for.</p>

python numpy python-xarray