Problem when selecting some members in a dataset and not using an int or slice as index

[vincentchabot]

Describe the bug

When trying to use anemoi-training with an ensemble and taking only a few member from this ensemble, I get an error coming from anemoi.dataset (ensemble.py).

I join a tiny code to reproduce the problem I get outside a training.

** Version number **

I am using the following versions of anemoi dataset : 0.5.16

To Reproduce

from anemoi.datasets.data import open_dataset
import anemoi.datasets 
print("Anemoi dataset version used :",anemoi.datasets.__version__)
data_path = "/home/mlx/ai-ml/datasets/aifs-ea-an-enda-0001-mars-o96-1979-2022-6h-v6-recentered-on-oper.zarr"

# Opening the dataset specifying we want members 1 and 2
ds = open_dataset(data_path, numbers=[1,2], start= None, end= 2020, frequency= '6h', drop= [])
# No specification (=> every members required)
ds2 = open_dataset(data_path, start= None, end= 2020, frequency= '6h', drop= [])

# Getting the first item (selection on date)

res_full = ds2[1]
res_sel = ds[1]
print("Simple selection")
print(f"Shape full {res_full.shape}, Shape if selection {res_sel.shape}")

# Getting several items (e.g date) at the same time 
res_full = ds2[1:3]
print("Several items selection")
print(f"Shape full {res_full.shape}")
res_sel = ds[1:3]
print(f"Shape if selection {res_sel.shape}")

# Getting several items (e.g date) at the same time and specifying for others  
print("Adding specification on other dimensions")
res_full = ds2[1:3,:,:,:]
print(f"Shape full {res_full.shape}")
res_sel = ds[1:3,:,:,:]
print(f"Shape if selection {res_sel.shape}")

Steps to reproduce the behavior:

1. Copy this code (and change the datapath if you are not on ATOS)
2. Run it
3. See error

Anemoi dataset version used : 0.5.16
Simple selection
Shape full (101, 10, 40320), Shape if selection (101, 2, 40320)
Several items selection
Shape full (2, 101, 10, 40320)
Shape if selection (2, 101, 2, 40320)
Adding specification on other dimensions
Shape full (2, 101, 10, 40320)
Traceback (most recent call last):
  File "/etc/ecmwf/nfs/dh1_home_b/fra5361/anemoi-mf/TL_AD/bug_anemoi_dataset.py", line 29, in <module>
    res_sel = ds[1:3,:,:,:]
              ~~^^^^^^^^^^^
  File "/home/fra5361/perm/venv/bug_anemoidataset/lib/python3.11/site-packages/anemoi/datasets/data/ensemble.py", line 64, in __getitem__
    result = result[:, :, self.mask, :]
             ~~~~~~^^^^^^^^^^^^^^^^^^^^
IndexError: boolean index did not match indexed array along axis 2; size of axis is 2 but size of corresponding boolean axis is 10

URL to sample input data
Provide a URL to a sample input data, or attach a file to that report if it is small enough.

Expected behavior
A clear and concise description of what you expected to happen.

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
Add any other context about the problem here.

[b8raoult]

Fixed in #247