Xarray based computations fail after latest code update in xarray engine

[sandorkertesz]

What happened?

When using this code:

#!/usr/bin/env python3

from earthkit import data as ekd

seas5_data = ekd.from_source("url", "https://sites.ecmwf.int/repository/earthkit-data/test-data/seas5_2m_temperature_201501-201503_europe_1deg.grib")
seas5_xr = seas5_data.to_xarray(
    time_dim_mode="forecast", add_valid_time_coord=True
).rename({"2t": "t2m"})   # .compute() #  .compute required for latest ek-d version


seas5_xr_gb = seas5_xr.groupby("valid_time.day")

print(seas5_xr_gb)

The following error is generated:

 
--- Logging error ---
Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/xr_engine/builder.py", line 261, in _raw_indexing_method
    result = r.to_array(index=field_index, array_namespace=self.array_namespace, dtype=self.dtype)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 89, in wrapped
    return func(self, *args, **_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 452, in to_array
    return self._to_array(source_to_array_func, index=index)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 439, in _to_array
    arr, current_field_shape = self._prepare_tensor_data(source_to_array_func, index=index)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 411, in _prepare_tensor_data
    arr = source_to_array_func(index=index)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/indexed.py", line 51, in to_array
    return self._as_array("to_array", empty_array_namespace=ns, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/indexed.py", line 64, in _as_array
    vals = _vals(first)
           ^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/indexed.py", line 57, in _vals
    return getattr(f, accessor)(**kwargs) if not is_property else getattr(f, accessor)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/core/field.py", line 537, in to_array
    v = v[index]
        ~^^^^^^^
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (31,) (41,) 
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/logging/__init__.py", line 1160, in emit
    msg = self.format(record)
          ^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/logging/__init__.py", line 999, in format
    return fmt.format(record)
           ^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/logging/__init__.py", line 703, in format
    record.message = record.getMessage()
                     ^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/logging/__init__.py", line 392, in getMessage
    msg = msg % self.args
          ~~~~^~~~~~~~~~~
TypeError: not all arguments converted during string formatting
Call stack:
  File "/Users/edwardcomyn-platt/Work/Git_Repositories/EARTHKIT/earthkit-transforms/dev/seasonal-computation-error.py", line 15, in <module>
    seas5_xr_gb = seas5_xr.groupby("valid_time.day")
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/util/deprecation_helpers.py", line 119, in inner
    return func(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/dataset.py", line 10188, in groupby
    return DatasetGroupBy(self, rgroupers, restore_coord_dims=restore_coord_dims)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/groupby.py", line 681, in __init__
    _ensure_1d(group=codes, obj=obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/groupby.py", line 281, in _ensure_1d
    newobj = obj.stack({stacked_dim: orig_dims})
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/util/deprecation_helpers.py", line 144, in wrapper
    return func(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/dataset.py", line 5293, in stack
    result = result._stack_once(dims, new_dim, index_cls, create_index)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/dataset.py", line 5208, in _stack_once
    stacked_var = exp_var.stack({new_dim: dims})
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/util/deprecation_helpers.py", line 144, in wrapper
    return func(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/variable.py", line 1548, in stack
    result = result._stack_once(dims, new_dim)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/variable.py", line 1511, in _stack_once
    new_data = duck_array_ops.reshape(reordered.data, new_shape)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/variable.py", line 456, in data
    duck_array = self._data.get_duck_array()
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/indexing.py", line 846, in get_duck_array
    array = self.array[self.key]
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/xr_engine/builder.py", line 235, in __getitem__
    return indexing.explicit_indexing_adapter(
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/indexing.py", line 1156, in explicit_indexing_adapter
    result = raw_indexing_method(raw_key.tuple)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/xr_engine/builder.py", line 263, in _raw_indexing_method
    LOG.exception("Error in to_array:", e)
Message: 'Error in to_array:'
Arguments: (IndexError('shape mismatch: indexing arrays could not be broadcast together with shapes (31,) (41,) '),)
Traceback (most recent call last):
  File "/Users/..../Work/Git_Repositories/EARTHKIT/earthkit-transforms/dev/seasonal-computation-error.py", line 15, in <module>
    seas5_xr_gb = seas5_xr.groupby("valid_time.day")
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/util/deprecation_helpers.py", line 119, in inner
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/dataset.py", line 10188, in groupby
    return DatasetGroupBy(self, rgroupers, restore_coord_dims=restore_coord_dims)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/groupby.py", line 681, in __init__
    _ensure_1d(group=codes, obj=obj)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/groupby.py", line 281, in _ensure_1d
    newobj = obj.stack({stacked_dim: orig_dims})
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/util/deprecation_helpers.py", line 144, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/dataset.py", line 5293, in stack
    result = result._stack_once(dims, new_dim, index_cls, create_index)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/dataset.py", line 5208, in _stack_once
    stacked_var = exp_var.stack({new_dim: dims})
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/util/deprecation_helpers.py", line 144, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/variable.py", line 1548, in stack
    result = result._stack_once(dims, new_dim)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/variable.py", line 1511, in _stack_once
    new_data = duck_array_ops.reshape(reordered.data, new_shape)
                                      ^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/variable.py", line 456, in data
    duck_array = self._data.get_duck_array()
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/indexing.py", line 846, in get_duck_array
    array = self.array[self.key]
            ~~~~~~~~~~^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/xr_engine/builder.py", line 235, in __getitem__
    return indexing.explicit_indexing_adapter(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/xarray/core/indexing.py", line 1156, in explicit_indexing_adapter
    result = raw_indexing_method(raw_key.tuple)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/xr_engine/builder.py", line 261, in _raw_indexing_method
    result = r.to_array(index=field_index, array_namespace=self.array_namespace, dtype=self.dtype)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 89, in wrapped
    return func(self, *args, **_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 452, in to_array
    return self._to_array(source_to_array_func, index=index)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 439, in _to_array
    arr, current_field_shape = self._prepare_tensor_data(source_to_array_func, index=index)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/tensor.py", line 411, in _prepare_tensor_data
    arr = source_to_array_func(index=index)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/indexed.py", line 51, in to_array
    return self._as_array("to_array", empty_array_namespace=ns, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/indexed.py", line 64, in _as_array
    vals = _vals(first)
           ^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/indexing/indexed.py", line 57, in _vals
    return getattr(f, accessor)(**kwargs) if not is_property else getattr(f, accessor)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/test/lib/python3.12/site-packages/earthkit/data/core/field.py", line 537, in to_array
    v = v[index]

What are the steps to reproduce the bug?

See above

Version

release/1.0.0rc0

Platform (OS and architecture)

any

Relevant log output

Accompanying data

No response

Organisation

No response

[pawel-wolff]

Fixed by PR #941