Repository: pydata/xarray

Issue:
New alignment option: "exact" without broadcasting OR Turn off automatic broadcasting
### Is your feature request related to a problem?

If we have two objects with dims `x` and `x1`, then `xr.align(..., join="exact")` will pass because these dimensions are broadcastable.

I'd like a stricter option (`join="strict"`?) that disallows broadcasting.


### Describe the solution you'd like

```python
xr.align(
    xr.DataArray([1], dims="x"),
    xr.DataArray([1], dims="x1"),
    join="strict",
)
```
would raise an error.

It'd be nice to have this as a built-in option so we can use
``` python
with xr.set_options(arithmetic_join="strict"):
    ...
```


### Describe alternatives you've considered

An alternative would be to allow control over automatic broadcasting through the `set_options` context manager., but that seems like it would be more complicated to implement.

### Additional context

This turns up in staggered grid calculations with xgcm where it is easy to mistakenly construct very high-dimensional arrays because of automatic broadcasting.

Discussion:
etienneschalk: cc @dcherian

Patch:
diff --git a/.github/workflows/pypi-release.yaml b/.github/workflows/pypi-release.yaml
index f55caa9b..354a8b59 100644
--- a/.github/workflows/pypi-release.yaml
+++ b/.github/workflows/pypi-release.yaml
@@ -88,7 +88,7 @@ jobs:
           path: dist
       - name: Publish package to TestPyPI
         if: github.event_name == 'push'
-        uses: pypa/gh-action-pypi-publish@v1.8.12
+        uses: pypa/gh-action-pypi-publish@v1.8.14
         with:
           repository_url: https://test.pypi.org/legacy/
           verbose: true
@@ -111,6 +111,6 @@ jobs:
           name: releases
           path: dist
       - name: Publish package to PyPI
-        uses: pypa/gh-action-pypi-publish@v1.8.12
+        uses: pypa/gh-action-pypi-publish@v1.8.14
         with:
           verbose: true
diff --git a/ci/install-upstream-wheels.sh b/ci/install-upstream-wheels.sh
index b1db58b7..68e841b4 100755
--- a/ci/install-upstream-wheels.sh
+++ b/ci/install-upstream-wheels.sh
@@ -4,10 +4,12 @@
 micromamba install "cython>=0.29.20" py-cpuinfo
 # temporarily (?) remove numbagg and numba
 micromamba remove -y numba numbagg
+# temporarily remove numexpr
+micromamba remove -y numexpr
 # temporarily remove backends
-micromamba remove -y cf_units h5py hdf5 netcdf4
+micromamba remove -y cf_units hdf5 h5py netcdf4
 # forcibly remove packages to avoid artifacts
-conda uninstall -y --force \
+micromamba remove -y --force \
     numpy \
     scipy \
     pandas \
@@ -30,8 +32,17 @@ python -m pip install \
     scipy \
     matplotlib \
     pandas
+# for some reason pandas depends on pyarrow already.
+# Remove once a `pyarrow` version compiled with `numpy>=2.0` is on `conda-forge`
+python -m pip install \
+    -i https://pypi.fury.io/arrow-nightlies/ \
+    --prefer-binary \
+    --no-deps \
+    --pre \
+    --upgrade \
+    pyarrow
 # without build isolation for packages compiling against numpy
-# TODO: remove once there are `numpy>=2.0` builds for numcodecs and cftime
+# TODO: remove once there are `numpy>=2.0` builds for these
 python -m pip install \
     --no-deps \
     --upgrade \
diff --git a/ci/requirements/all-but-dask.yml b/ci/requirements/all-but-dask.yml
index c16c174f..2f47643c 100644
--- a/ci/requirements/all-but-dask.yml
+++ b/ci/requirements/all-but-dask.yml
@@ -5,6 +5,7 @@ channels:
 dependencies:
   - black
   - aiobotocore
+  - array-api-strict
   - boto3
   - bottleneck
   - cartopy
diff --git a/ci/requirements/environment-3.12.yml b/ci/requirements/environment-3.12.yml
index 805f71c0..99b260a6 100644
--- a/ci/requirements/environment-3.12.yml
+++ b/ci/requirements/environment-3.12.yml
@@ -4,6 +4,7 @@ channels:
   - nodefaults
 dependencies:
   - aiobotocore
+  - array-api-strict
   - boto3
   - bottleneck
   - cartopy
diff --git a/ci/requirements/environment-windows-3.12.yml b/ci/requirements/environment-windows-3.12.yml
index a1c78c64..5b3034b7 100644
--- a/ci/requirements/environment-windows-3.12.yml
+++ b/ci/requirements/environment-windows-3.12.yml
@@ -2,6 +2,7 @@ name: xarray-tests
 channels:
   - conda-forge
 dependencies:
+  - array-api-strict
   - boto3
   - bottleneck
   - cartopy
diff --git a/ci/requirements/environment-windows.yml b/ci/requirements/environment-windows.yml
index 1fd2cf54..cc361bac 100644
--- a/ci/requirements/environment-windows.yml
+++ b/ci/requirements/environment-windows.yml
@@ -2,6 +2,7 @@ name: xarray-tests
 channels:
   - conda-forge
 dependencies:
+  - array-api-strict
   - boto3
   - bottleneck
   - cartopy
diff --git a/ci/requirements/environment.yml b/ci/requirements/environment.yml
index c96f3fd7..d3dbc088 100644
--- a/ci/requirements/environment.yml
+++ b/ci/requirements/environment.yml
@@ -4,6 +4,7 @@ channels:
   - nodefaults
 dependencies:
   - aiobotocore
+  - array-api-strict
   - boto3
   - bottleneck
   - cartopy
diff --git a/ci/requirements/min-all-deps.yml b/ci/requirements/min-all-deps.yml
index 775c98b8..d2965fb3 100644
--- a/ci/requirements/min-all-deps.yml
+++ b/ci/requirements/min-all-deps.yml
@@ -8,6 +8,7 @@ dependencies:
   # When upgrading python, numpy, or pandas, must also change
   # doc/user-guide/installing.rst, doc/user-guide/plotting.rst and setup.py.
   - python=3.9
+  - array-api-strict=1.0  # dependency for testing the array api compat
   - boto3=1.24
   - bottleneck=1.3
   - cartopy=0.21
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index eb5dcbda..bc603ab6 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -23,6 +23,9 @@ v2024.03.0 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Do not broadcast in arithmetic operations when global option ``arithmetic_broadcast=False``
+  (:issue:`6806`, :pull:`8784`).
+  By `Etienne Schalk <https://github.com/etienneschalk>`_ and `Deepak Cherian <https://github.com/dcherian>`_.
 - Add the ``.oindex`` property to Explicitly Indexed Arrays for orthogonal indexing functionality. (:issue:`8238`, :pull:`8750`)
   By `Anderson Banihirwe <https://github.com/andersy005>`_.
 
diff --git a/pyproject.toml b/pyproject.toml
index 4b5f6b31..d2a5c6b8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -138,6 +138,7 @@ module = [
   "toolz.*",
   "zarr.*",
   "numpy.exceptions.*", # remove once support for `numpy<2.0` has been dropped
+  "array_api_strict.*",
 ]
 
 # Gradually we want to add more modules to this list, ratcheting up our total
diff --git a/xarray/core/options.py b/xarray/core/options.py
index 18e3484e..f5614104 100644
--- a/xarray/core/options.py
+++ b/xarray/core/options.py
@@ -34,6 +34,7 @@ if TYPE_CHECKING:
     ]
 
     class T_Options(TypedDict):
+        arithmetic_broadcast: bool
         arithmetic_join: Literal["inner", "outer", "left", "right", "exact"]
         cmap_divergent: str | Colormap
         cmap_sequential: str | Colormap
@@ -59,6 +60,7 @@ if TYPE_CHECKING:
 
 
 OPTIONS: T_Options = {
+    "arithmetic_broadcast": True,
     "arithmetic_join": "inner",
     "cmap_divergent": "RdBu_r",
     "cmap_sequential": "viridis",
@@ -92,6 +94,7 @@ def _positive_integer(value: int) -> bool:
 
 
 _VALIDATORS = {
+    "arithmetic_broadcast": lambda value: isinstance(value, bool),
     "arithmetic_join": _JOIN_OPTIONS.__contains__,
     "display_max_rows": _positive_integer,
     "display_values_threshold": _positive_integer,
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index 315c4636..cac4d304 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -2871,6 +2871,16 @@ def broadcast_variables(*variables: Variable) -> tuple[Variable, ...]:
 
 
 def _broadcast_compat_data(self, other):
+    if not OPTIONS["arithmetic_broadcast"]:
+        if (isinstance(other, Variable) and self.dims != other.dims) or (
+            is_duck_array(other) and self.ndim != other.ndim
+        ):
+            raise ValueError(
+                "Broadcasting is necessary but automatic broadcasting is disabled via "
+                "global option `'arithmetic_broadcast'`. "
+                "Use `xr.set_options(arithmetic_broadcast=True)` to enable automatic broadcasting."
+            )
+
     if all(hasattr(other, attr) for attr in ["dims", "data", "shape", "encoding"]):
         # `other` satisfies the necessary Variable API for broadcast_variables
         new_self, new_other = _broadcast_compat_variables(self, other)


Test patch:
diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py
index df089950..2e6e638f 100644
--- a/xarray/tests/__init__.py
+++ b/xarray/tests/__init__.py
@@ -89,6 +89,13 @@ with warnings.catch_warnings():
 has_pynio, requires_pynio = _importorskip("Nio")
 has_cftime, requires_cftime = _importorskip("cftime")
 has_dask, requires_dask = _importorskip("dask")
+with warnings.catch_warnings():
+    warnings.filterwarnings(
+        "ignore",
+        message="The current Dask DataFrame implementation is deprecated.",
+        category=DeprecationWarning,
+    )
+    has_dask_expr, requires_dask_expr = _importorskip("dask_expr")
 has_bottleneck, requires_bottleneck = _importorskip("bottleneck")
 has_rasterio, requires_rasterio = _importorskip("rasterio")
 has_zarr, requires_zarr = _importorskip("zarr")
diff --git a/xarray/tests/test_array_api.py b/xarray/tests/test_array_api.py
index 03dcfd9b..a5ffb37a 100644
--- a/xarray/tests/test_array_api.py
+++ b/xarray/tests/test_array_api.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import warnings
-
 import pytest
 
 import xarray as xr
@@ -9,10 +7,19 @@ from xarray.testing import assert_equal
 
 np = pytest.importorskip("numpy", minversion="1.22")
 
-with warnings.catch_warnings():
-    warnings.simplefilter("ignore")
-    import numpy.array_api as xp  # isort:skip
-    from numpy.array_api._array_object import Array  # isort:skip
+try:
+    import warnings
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+
+        import numpy.array_api as xp
+        from numpy.array_api._array_object import Array
+except ImportError:
+    # for `numpy>=2.0`
+    xp = pytest.importorskip("array_api_strict")
+
+    from array_api_strict._array_object import Array  # type: ignore[no-redef]
 
 
 @pytest.fixture
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index 2829fd7d..5ab20b2c 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -51,6 +51,7 @@ from xarray.tests import (
     requires_bottleneck,
     requires_cupy,
     requires_dask,
+    requires_dask_expr,
     requires_iris,
     requires_numexpr,
     requires_pint,
@@ -3203,6 +3204,42 @@ class TestDataArray:
         assert_identical(expected_b, actual_b)
         assert expected_b.x.dtype == actual_b.x.dtype
 
+    def test_broadcast_on_vs_off_global_option_different_dims(self) -> None:
+        xda_1 = xr.DataArray([1], dims="x1")
+        xda_2 = xr.DataArray([1], dims="x2")
+
+        with xr.set_options(arithmetic_broadcast=True):
+            expected_xda = xr.DataArray([[1.0]], dims=("x1", "x2"))
+            actual_xda = xda_1 / xda_2
+            assert_identical(actual_xda, expected_xda)
+
+        with xr.set_options(arithmetic_broadcast=False):
+            with pytest.raises(
+                ValueError,
+                match=re.escape(
+                    "Broadcasting is necessary but automatic broadcasting is disabled via "
+                    "global option `'arithmetic_broadcast'`. "
+                    "Use `xr.set_options(arithmetic_broadcast=True)` to enable automatic broadcasting."
+                ),
+            ):
+                xda_1 / xda_2
+
+    @pytest.mark.parametrize("arithmetic_broadcast", [True, False])
+    def test_broadcast_on_vs_off_global_option_same_dims(
+        self, arithmetic_broadcast: bool
+    ) -> None:
+        # Ensure that no error is raised when arithmetic broadcasting is disabled,
+        # when broadcasting is not needed. The two DataArrays have the same
+        # dimensions of the same size.
+        xda_1 = xr.DataArray([1], dims="x")
+        xda_2 = xr.DataArray([1], dims="x")
+        expected_xda = xr.DataArray([2.0], dims=("x",))
+
+        with xr.set_options(arithmetic_broadcast=arithmetic_broadcast):
+            assert_identical(xda_1 + xda_2, expected_xda)
+            assert_identical(xda_1 + np.array([1.0]), expected_xda)
+            assert_identical(np.array([1.0]) + xda_1, expected_xda)
+
     def test_broadcast_arrays(self) -> None:
         x = DataArray([1, 2], coords=[("a", [-1, -2])], name="x")
         y = DataArray([1, 2], coords=[("b", [3, 4])], name="y")
@@ -3381,6 +3418,7 @@ class TestDataArray:
         assert len(actual) == 0
         assert_array_equal(actual.index.names, list("ABC"))
 
+    @requires_dask_expr
     @requires_dask
     def test_to_dask_dataframe(self) -> None:
         arr_np = np.arange(3 * 4).reshape(3, 4)