Add support for cross product

xarray/__init__.py

@@ -18,7 +18,7 @@
 from .core.alignment import align, broadcast
 from .core.combine import combine_by_coords, combine_nested
 from .core.common import ALL_DIMS, full_like, ones_like, zeros_like
-from .core.computation import apply_ufunc, corr, cov, dot, polyval, where
+from .core.computation import apply_ufunc, corr, cov, cross, dot, polyval, where
 from .core.concat import concat
 from .core.dataarray import DataArray
 from .core.dataset import Dataset
@@ -56,6 +56,7 @@
     "dot",
     "cov",
     "corr",
+    "cross",
     "full_like",
     "infer_freq",
     "load_dataarray",

xarray/core/computation.py

@@ -1527,6 +1527,137 @@ def dot(*arrays, dims=None, **kwargs):
     return result.transpose(*[d for d in all_dims if d in result.dims])
 
 
+def cross(a, b, dim=None):
+    """
+    Return the cross product of two (arrays of) vectors.
+
+    The cross product of `a` and `b` in :math:`R^3` is a vector perpendicular
+    to both `a` and `b`.  If `a` and `b` are arrays of vectors, the vectors
+    are defined by the last axis of `a` and `b` by default, and these axes
+    can have dimensions 2 or 3.  Where the dimension of either `a` or `b` is
+    2, the third component of the input vector is assumed to be zero and the
+    cross product calculated accordingly.  In cases where both input vectors
+    have dimension 2, the z-component of the cross product is returned.
+
+    Parameters
+    ----------
+    a, b : DataArray
+        something
+    dim : hashable or tuple of hashable
+        something
+
+    Examples
+    --------
+    Vector cross-product.
+
+    >>> x = xr.DataArray(np.array([1, 2, 3]))
+    >>> y = xr.DataArray(np.array([4, 5, 6]))
+    >>> xr.cross(x, y)
+    array([-3,  6, -3])
+
+    One vector with dimension 2.
+
+    >>> a = xr.DataArray(
+    ...     np.array([1, 2]),
+    ...     dims=["cartesian"],
+    ...     coords=dict(cartesian=(["cartesian"], np.array(["x", "z"]))),
+    ... )
+    >>> b = xr.DataArray(
+    ...     np.array([4, 5, 6]),
+    ...     dims=["x"],
+    ...     coords=dict(cartesian=(["cartesian"], np.array(["x", "y", "z"]))),
+    ... )
+    >>> xr.cross(a, b)
+    array([12, -6, -3])
+
+    Multiple vector cross-products. Note that the direction of the
+    cross product vector is defined by the right-hand rule.
+
+    >>> x = xr.DataArray(np.array([[1, 2, 3], [4, 5, 6]]), dims=("a", "b"))
+    >>> y = xr.DataArray(np.array([[4, 5, 6], [1, 2, 3]]), dims=("a", "b"))
+    >>> xr.cross(x, y)
+    array([[-3,  6, -3],
+           [ 3, -6,  3]])
+
+    Change the vector definition of x and y using axisa and axisb.
+
+    >>> x = xr.DataArray(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
+    >>> y = xr.DataArray(np.array([[7, 8, 9], [4, 5, 6], [1, 2, 3]]))
+    >>> np.cross(x, y)
+    array([[ -6,  12,  -6],
+           [  0,   0,   0],
+           [  6, -12,   6]])
+    >>> np.cross(x, y, axisa=0, axisb=0)
+    array([[-24,  48, -24],
+           [-30,  60, -30],
+           [-36,  72, -36]])
+
+    See Also
+    --------
+    numpy.cross : Corresponding numpy function
+    """
+    from .dataarray import DataArray
+
+    arrays = [a, b]
+    for arr in arrays:
+        if not isinstance(arr, (DataArray)):
+            raise TypeError(
+                f"Only xr.DataArray and xr.Variable are supported, got {type(arr)}."
+            )
+
+        if dim is None:
+            # TODO: Find spatial dim default by looking for unique
+            # (3 or 2)-valued dim?
+            dim = arr.dims[-1]
+        elif dim not in arr.dims:
+            raise ValueError(f"Dimension {dim} not in {arr}.")
+
+        s = arr.sizes[dim]
+        if s < 1 or s > 3:
+            raise ValueError(
+                "incompatible dimensions for cross product\n"
+                "(dimension with coords must be 1, 2 or 3)"
+            )
+
+    if a.sizes[dim] == b.sizes[dim]:
+        # Arrays have the same size, no need to do anything:
+        pass
+    else:
+        # Arrays have different sizes. Append zeros where the smaller
+        # array is missing a value, zeros will not affect np.cross:
+        ind = 1 if a.sizes[dim] > b.sizes[dim] else 0
+
+        if a.coords:
+            # If the array has coords we know which indexes to fill
+            # with zeros:
+            arrays[ind] = arrays[ind].reindex_like(arrays[1 - ind], fill_value=0)
+        elif arrays[ind].sizes[dim] > 1:
+            # If it doesn't have coords we can can only that infer that
+            # it is composite values if the size is 2.
+            arrays[ind] = arrays[ind].pad({dim: (0, 1)}, constant_values=0)
+        else:
+            # Size is 1, then we do not know if it is a constant or
+            # composite value:
+            raise ValueError(
+                "incompatible dimensions for cross product\n"
+                "(dimension without coords must be 2 or 3)"
+            )
+
+    # Figure out the output dtype:
+    output_dtype = np.cross(
+        np.empty((2, 2), dtype=arrays[0].dtype), np.empty((2, 2), dtype=arrays[1].dtype)
+    ).dtype
+
+    return apply_ufunc(
+        np.cross,
+        *arrays,
+        input_core_dims=[[dim], [dim]],
+        output_core_dims=[[dim]],
+        dask="parallelized",
+        output_dtypes=[output_dtype],
+    )
+
+
 def where(cond, x, y):
     """Return elements from `x` or `y` depending on `cond`.
 

xarray/tests/test_computation.py

@@ -1898,3 +1898,86 @@ def test_polyval(use_dask, use_datetime):
     da_pv = xr.polyval(da.x, coeffs)
 
     xr.testing.assert_allclose(da, da_pv.T)
+
+
+@pytest.mark.parametrize(
+    "a, b, ae, be, dim, axis",
+    [
+        [
+            xr.DataArray(np.array([1, 2, 3])),
+            xr.DataArray(np.array([4, 5, 6])),
+            np.array([1, 2, 3]),
+            np.array([4, 5, 6]),
+            None,
+            -1,
+        ],
+        [
+            xr.DataArray(np.array([1, 2])),
+            xr.DataArray(np.array([4, 5, 6])),
+            np.array([1, 2]),
+            np.array([4, 5, 6]),
+            None,
+            -1,
+        ],
+        [  # Test dim in the middle:
+            xr.DataArray(
+                np.arange(0, 5 * 3 * 4).reshape((5, 3, 4)),
+                dims=["time", "cartesian", "var"],
+                coords=dict(
+                    time=(["time"], np.arange(0, 5)),
+                    cartesian=(["cartesian"], np.array(["x", "y", "z"])),
+                    var=(["var"], np.array([1, 1.5, 2, 2.5])),
+                ),
+            ),
+            xr.DataArray(
+                np.arange(0, 5 * 3 * 4).reshape((5, 3, 4)) + 1,
+                dims=["time", "cartesian", "var"],
+                coords=dict(
+                    time=(["time"], np.arange(0, 5)),
+                    cartesian=(["cartesian"], np.array(["x", "y", "z"])),
+                    var=(["var"], np.array([1, 1.5, 2, 2.5])),
+                ),
+            ),
+            np.arange(0, 5 * 3 * 4).reshape((5, 3, 4)),
+            np.arange(0, 5 * 3 * 4).reshape((5, 3, 4)) + 1,
+            "cartesian",
+            1,
+        ],
+        [  # Test 1 sized arrays with coords:
+            xr.DataArray(
+                np.array([1]),
+                dims=["cartesian"],
+                coords=dict(cartesian=(["cartesian"], np.array(["z"]))),
+            ),
+            xr.DataArray(
+                np.array([4, 5, 6]),
+                dims=["cartesian"],
+                coords=dict(cartesian=(["cartesian"], np.array(["x", "y", "z"]))),
+            ),
+            np.array([0, 0, 1]),
+            np.array([4, 5, 6]),
+            None,
+            -1,
+        ],
+        [  # Test filling inbetween with coords:
+            xr.DataArray(
+                np.array([1, 2]),
+                dims=["cartesian"],
+                coords=dict(cartesian=(["cartesian"], np.array(["x", "z"]))),
+            ),
+            xr.DataArray(
+                np.array([4, 5, 6]),
+                dims=["cartesian"],
+                coords=dict(cartesian=(["cartesian"], np.array(["x", "y", "z"]))),
+            ),
+            np.array([1, 0, 2]),
+            np.array([4, 5, 6]),
+            None,
+            -1,
+        ],
+    ],
+)
+def test_cross(a, b, ae, be, dim, axis):
+    expected = np.cross(ae, be, axis=axis)
+    actual = xr.cross(a, b, dim=dim)
+    xr.testing.assert_duckarray_allclose(expected, actual)