Add support for cross product

xarray/__init__.py

@@ -18,7 +18,16 @@
 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, unify_chunks, where
+from .core.computation import (
+    apply_ufunc,
+    corr,
+    cov,
+    cross,
+    dot,
+    polyval,
+    unify_chunks,
+    where,
+)
 from .core.concat import concat
 from .core.dataarray import DataArray
 from .core.dataset import Dataset
@@ -56,6 +65,7 @@
     "dot",
     "cov",
     "corr",
+    "cross",
     "full_like",
     "infer_freq",
     "load_dataarray",

xarray/core/computation.py

@@ -41,6 +41,7 @@
     from .dataset import Dataset
 
     T_DSorDA = TypeVar("T_DSorDA", DataArray, Dataset)
+    T_DSorDAorVar = TypeVar("T_DSorDAorVar", Dataset, DataArray, Variable)
 
 _NO_FILL_VALUE = utils.ReprObject("<no-fill-value>")
 _DEFAULT_NAME = utils.ReprObject("<default-name>")
@@ -1396,6 +1397,209 @@ def _get_valid_values(da, other):
         return corr
 
 
+def cross(
+    a: "T_DSorDAorVar",
+    b: "T_DSorDAorVar",
+    dim: str,
+) -> "T_DSorDAorVar":
+    """
+    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, 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, Dataset or Variable
+        Components of the first and second vector(s).
+    dim : hashable
+        The dimension along which the cross product will be computed.
+        Must be available in both vectors.
+
+    Examples
+    --------
+    Vector cross-product with 3 dimensions.
+
+    >>> a = xr.DataArray([1, 2, 3])
+    >>> b = xr.DataArray([4, 5, 6])
+    >>> xr.cross(a, b, "dim_0")
+    <xarray.DataArray (dim_0: 3)>
+    array([-3,  6, -3])
+    Dimensions without coordinates: dim_0
+
+    Vector cross-product with 2 dimensions, returns in the perpendicular
+    direction:
+
+    >>> a = xr.DataArray([1, 2])
+    >>> b = xr.DataArray([4, 5])
+    >>> xr.cross(a, b, "dim_0")
+    <xarray.DataArray ()>
+    array(-3)
+
+    Vector cross-product with 3 dimensions but zeros at the last axis
+    yields the same results as with 2 dimensions:
+
+    >>> a = xr.DataArray([1, 2, 0])
+    >>> b = xr.DataArray([4, 5, 0])
+    >>> xr.cross(a, b, "dim_0")
+    <xarray.DataArray (dim_0: 3)>
+    array([ 0,  0, -3])
+    Dimensions without coordinates: dim_0
+
+    One vector with dimension 2.
+
+    >>> a = xr.DataArray(
+    ...     [1, 2],
+    ...     dims=["cartesian"],
+    ...     coords=dict(cartesian=(["cartesian"], ["x", "y"])),
+    ... )
+    >>> b = xr.DataArray(
+    ...     [4, 5, 6],
+    ...     dims=["cartesian"],
+    ...     coords=dict(cartesian=(["cartesian"], ["x", "y", "z"])),
+    ... )
+    >>> xr.cross(a, b, "cartesian")
+    <xarray.DataArray (cartesian: 3)>
+    array([12, -6, -3])
+    Coordinates:
+      * cartesian  (cartesian) object 'x' 'y' 'z'
+
+    One vector with dimension 2 but coords in other positions.
+
+    >>> a = xr.DataArray(
+    ...     [1, 2],
+    ...     dims=["cartesian"],
+    ...     coords=dict(cartesian=(["cartesian"], ["x", "z"])),
+    ... )
+    >>> b = xr.DataArray(
+    ...     [4, 5, 6],
+    ...     dims=["cartesian"],
+    ...     coords=dict(cartesian=(["cartesian"], ["x", "y", "z"])),
+    ... )
+    >>> xr.cross(a, b, "cartesian")
+    <xarray.DataArray (cartesian: 3)>
+    array([-10,   2,   5])
+    Coordinates:
+      * cartesian  (cartesian) object 'x' 'y' 'z'
+
+    Multiple vector cross-products. Note that the direction of the
+    cross product vector is defined by the right-hand rule.
+
+    >>> a = xr.DataArray(
+    ...     [[1, 2, 3], [4, 5, 6]],
+    ...     dims=("time", "cartesian"),
+    ...     coords=dict(
+    ...         time=(["time"], [0, 1]),
+    ...         cartesian=(["cartesian"], ["x", "y", "z"]),
+    ...     ),
+    ... )
+    >>> b = xr.DataArray(
+    ...     [[4, 5, 6], [1, 2, 3]],
+    ...     dims=("time", "cartesian"),
+    ...     coords=dict(
+    ...         time=(["time"], [0, 1]),
+    ...         cartesian=(["cartesian"], ["x", "y", "z"]),
+    ...     ),
+    ... )
+    >>> xr.cross(a, b, "cartesian")
+    <xarray.DataArray (time: 2, cartesian: 3)>
+    array([[-3,  6, -3],
+           [ 3, -6,  3]])
+    Coordinates:
+      * time       (time) int64 0 1
+      * cartesian  (cartesian) <U1 'x' 'y' 'z'
+
+    See Also
+    --------
+    numpy.cross : Corresponding numpy function
+    """
+    from .dataarray import DataArray
+    from .dataset import Dataset
+
+    all_dims: List[Hashable] = []
+    arrays: List["T_DSorDAorVar"] = [a, b]
+    for i, arr in enumerate(arrays):
+        if isinstance(arr, Dataset):
+            # Turn the dataset to a stacked dataarray to follow the
+            # normal code path. Then at the end turn it back to a
+            # dataset.
+            is_dataset = True
+            arrays[i] = arr = arr.to_stacked_array(
+                variable_dim=dim, new_dim="variable", sample_dims=arr.dims
+            ).unstack("variable")
+            if is_duck_dask_array(arr.data):
+                arrays[i] = arr = arr.chunk({dim: -1})
+        elif isinstance(arr, (DataArray, Variable)):
+            is_dataset = False
+        else:
+            raise TypeError(
+                "Only xr.DataArray, xr.Dataset and xr.Variable are supported, "
+                f"got {type(arr)}."
+            )
+
+        # TODO: Find spatial dim default by looking for unique
+        # (3 or 2)-valued dim?
+        if dim not in arr.dims:
+            raise ValueError(f"Dimension {dim} not in {arr}.")
+
+        if not 1 <= arr.sizes[dim] <= 3:
+            raise ValueError(
+                "Incompatible dimensions for cross product,\n"
+                "dimension with coords must be 1, 2 or 3."
+            )
+
+        all_dims += [d for d in arr.dims if d not in all_dims]
+
+    if arrays[0].sizes[dim] != arrays[1].sizes[dim]:
+        # Arrays have different sizes. Append zeros where the smaller
+        # array is missing a value, zeros will not affect np.cross:
+        i = 1 if arrays[0].sizes[dim] > arrays[1].sizes[dim] else 0
+        array_small, array_large = arrays[i], arrays[1 - i]
+
+        if getattr(array_large, "coords", False) and getattr(
+            array_small, "coords", False
+        ):
+            # if all([getattr(arr, "coords", False) for arr in arrays]):
+            # If the arrays have coords we know which indexes to fill
+            # with zeros:
+            arrays[i] = array_small.reindex_like(array_large, fill_value=0)
+        elif array_small.sizes[dim] == 2:
+            # If the array doesn't have coords we can only infer
+            # that it is composite values if the size is 2:
+            arrays[i] = array_small.pad({dim: (0, 1)}, constant_values=0)
+            if is_duck_dask_array(arrays[i].data):
+                arrays[i] = arrays[i].chunk({dim: -1})
+        else:
+            # Size is 1, then we do not know if the array is a constant or
+            # composite value:
+            raise ValueError(
+                "Incompatible dimensions for cross product,\n"
+                "dimension without coords must be 2 or 3."
+            )
+
+    c = apply_ufunc(
+        np.cross,
+        *arrays,
+        input_core_dims=[[dim], [dim]],
+        output_core_dims=[[dim] if arrays[0].sizes[dim] == 3 else []],
+        dask="parallelized",
+        output_dtypes=[np.result_type(*arrays)],
+    )
+    c = c.transpose(*[d for d in all_dims if d in c.dims])
+    if is_dataset:
+        c = c.stack(variable=[dim]).to_unstacked_dataset("variable")
+        c = c.expand_dims(
+            list({d: s for ds in arrays for d, s in ds.sizes.items() if s == 1})
+        )
+
+    return c
+
+
 def dot(*arrays, dims=None, **kwargs):
     """Generalized dot product for xarray objects. Like np.einsum, but
     provides a simpler interface based on array dimensions.

xarray/core/dataset.py

@@ -3866,8 +3866,8 @@ def stack(
     def to_stacked_array(
         self,
         new_dim: Hashable,
-        sample_dims: Sequence[Hashable],
-        variable_dim: str = "variable",
+        sample_dims: Mapping[Hashable, Sequence[Hashable]],
+        variable_dim: Hashable = "variable",
         name: Hashable = None,
     ) -> "DataArray":
         """Combine variables of differing dimensionality into a DataArray

xarray/tests/test_computation.py

@@ -1930,3 +1930,128 @@ 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("use_dask", [False, True])
+@pytest.mark.parametrize(
+    "a, b, ae, be, dim, axis",
+    [
+        [
+            xr.DataArray([1, 2, 3]),
+            xr.DataArray([4, 5, 6]),
+            [1, 2, 3],
+            [4, 5, 6],
+            "dim_0",
+            -1,
+        ],
+        [
+            xr.DataArray([1, 2]),
+            xr.DataArray([4, 5, 6]),
+            [1, 2],
+            [4, 5, 6],
+            "dim_0",
+            -1,
+        ],
+        [
+            xr.Variable(dims=["dim_0"], data=[1, 2, 3]),
+            xr.Variable(dims=["dim_0"], data=[4, 5, 6]),
+            [1, 2, 3],
+            [4, 5, 6],
+            "dim_0",
+            -1,
+        ],
+        [
+            xr.Variable(dims=["dim_0"], data=[1, 2]),
+            xr.Variable(dims=["dim_0"], data=[4, 5, 6]),
+            [1, 2],
+            [4, 5, 6],
+            "dim_0",
+            -1,
+        ],
+        [
+            xr.Dataset({0: ("dim_0", [1]), 1: ("dim_0", [2]), 2: ("dim_0", [3])}),
+            xr.Dataset({0: ("dim_0", [4]), 1: ("dim_0", [5]), 2: ("dim_0", [6])}),
+            [1, 2, 3],
+            [4, 5, 6],
+            "cartesian",
+            -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"], ["x", "y", "z"]),
+                    var=(["var"], [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"], ["x", "y", "z"]),
+                    var=(["var"], [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"], ["z"])),
+            ),
+            xr.DataArray(
+                np.array([4, 5, 6]),
+                dims=["cartesian"],
+                coords=dict(cartesian=(["cartesian"], ["x", "y", "z"])),
+            ),
+            [0, 0, 1],
+            [4, 5, 6],
+            "cartesian",
+            -1,
+        ],
+        [  # Test filling inbetween with coords:
+            xr.DataArray(
+                [1, 2],
+                dims=["cartesian"],
+                coords=dict(cartesian=(["cartesian"], ["x", "z"])),
+            ),
+            xr.DataArray(
+                [4, 5, 6],
+                dims=["cartesian"],
+                coords=dict(cartesian=(["cartesian"], ["x", "y", "z"])),
+            ),
+            [1, 0, 2],
+            [4, 5, 6],
+            "cartesian",
+            -1,
+        ],
+    ],
+)
+def test_cross(a, b, ae, be, dim, axis, use_dask):
+    expected = np.cross(ae, be, axis=axis)
+
+    if use_dask:
+        if not has_dask:
+            pytest.skip("test for dask.")
+        a = a.chunk()
+        b = b.chunk()
+
+    actual = xr.cross(a, b, dim=dim)
+
+    if isinstance(actual, xr.Dataset):
+        actual = (
+            actual.to_stacked_array(
+                variable_dim=dim, new_dim="variable", sample_dims=actual.dims
+            )
+            .unstack("variable")
+            .squeeze()
+        )
+
+    xr.testing.assert_duckarray_allclose(expected, actual)