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:
> This turns up in staggered grid calculations with xgcm where it is easy to mistakenly construct very high-dimensional arrays because of automatic broadcasting.

Do you have an example that illustrates this issue?

I'm curious to see if in that specific case using a custom index for staggered grid coordinates wouldn't work as an alternative solution. The alignment rules are pretty strict (same index type + same sequence of coordinate names & dims). Not 100% sure if that applies in the xgcm case, but in theory it would raise an error regardless of the join method if you try to align objects that have broadcastable coordinates with incompatible indexes.
> I'm curious to see if in that specific case using a custom index for staggered grid coordinates wouldn't work as an alternative solution. The alignment rules are pretty strict (same index type + same sequence of coordinate names & dims). Not 100% sure if that applies in the xgcm case, but in theory it would raise an error regardless of the join method if you try to align objects that have broadcastable coordinates with incompatible indexes.

This is on my to-do list to think about really hard :sweat_smile: 

I think one problem would be that sometimes a different grid position implies a different length coordinate array, e.g. "outer" vs "center" is larger by one element.
Basically I want the following to raise an error because I've picked the wrong variable with dimension `x1` instead of dimension `x` by mistake

``` python
xr.DataArray([1], dims="x") / xr.DataArray([1], dims="x1")
```

I think we could do it with a "staggered grid index", but it seems useful in general 
Ah ok I see, thanks! Yes I agree this may be useful.
@dcherian — I'm discussing this with @etienneschalk over at #8698 

Can we clarify your proposal? Is it:
1. All dims must match — i.e. `ds1.dims == ds2.dims`, as well as each dim being the same size as its counterpart?
2. ...or something more muted, like each dim being the same size as its counterpart, and not broadcasting a dimension of size 1 to the same-named dimension of size `n`?

The first seems very strict, I haven't hit that personally, but possibly others have...
Thoughts from the meeting:

`"join"` is about indexes, and alignment and broadcasting should be orthogonal ops. 

So instead of a new `join` option, we'll add an option to control the broadcasting. So something like:
1.  `align(..., join='exact', broadcast=False)` # new kwarg `broadcast`
2. And a new global option `arithmetic_broadcast` to complement `arithmetic_join` . This will control automatic broadcasting in binary ops.
Hello @dcherian 

I adapted changes made in #8698 with the added clear separation between the `join` and `broadcast` keyword parameters.

So while the form (API) changed, I did not change the current substance, and it is still not entirely clear for me, what is expected to implement this issue successfully. I wrote a small text following the discussions with @max-sixty about the changes I currently made, and my current understanding: https://github.com/pydata/xarray/pull/8698#issuecomment-1936982530 ; is this _aligned_ with what you would expect? (pun intended... 😅) 

You can have a look at [test_broadcast_on_vs_off_global_option](https://github.com/pydata/xarray/pull/8698/files#diff-ca5c2de2fe6e9e25fbf22bd53e4976c15da74900dfb14deb7e6e87f5377230e3R3259) in the PR, testing the behaviour you described in your [comment](https://github.com/pydata/xarray/issues/6806#issuecomment-1188339194) 
@etienneschalk I sincerely apologized. I  deeply misunderstood what `align` was doing when I wrote this issue. `align` isn't actually broadcasting, it is simply checking indexes (as it should). 

Given the discussion in the meeting, I think the simplest solution is to raise an error here if `arithmetic_broadcast is False`:
https://github.com/pydata/xarray/blob/ff0d056ec9e99dece0202d8d73c1cb8c6c20b2a1/xarray/core/variable.py#L2265-L2268

This should be a significantly smaller and easier PR.

Again, apologies for leading you down the wrong path.
No problem, even if the PR is not solving what was expected, it was the opportunity for me to dig into the xarray internal logic which is valuable learning!

I can try to implement your suggestion, reusing the existing 'arithmetic_broadcast' flag 



Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index eb5dcbda36f..bc603ab61d3 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/xarray/core/options.py b/xarray/core/options.py
index 18e3484e9c4..f5614104357 100644
--- a/xarray/core/options.py
+++ b/xarray/core/options.py
@@ -34,6 +34,7 @@
     ]
 
     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 @@ class T_Options(TypedDict):
 
 
 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 315c46369bd..cac4d3049e2 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 df0899509cb..2e6e638f5b1 100644
--- a/xarray/tests/__init__.py
+++ b/xarray/tests/__init__.py
@@ -89,6 +89,13 @@ def _importorskip(
 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_dataarray.py b/xarray/tests/test_dataarray.py
index 2829fd7d49c..5ab20b2c29c 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -51,6 +51,7 @@
     requires_bottleneck,
     requires_cupy,
     requires_dask,
+    requires_dask_expr,
     requires_iris,
     requires_numexpr,
     requires_pint,
@@ -3203,6 +3204,42 @@ def test_align_str_dtype(self) -> None:
         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 @@ def test_to_dataframe_0length(self) -> None:
         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)