Inconsistent `output_tensor` dimensions in `EnsemblePlotMixin`

[dietervdb-meteo]

What happened?

The process function takes the members argument, which can be an integer, a list or None.

Although not strictly speaking a bug (tests using this function seem to pass without problem) the behaviour of the function is somewhat unexpected/incosistent. In particular, when members = None , torch slicing adds an extra (needless) dimension at the position of what previously was the ensemble dimension, making it harder to track what dimension corresponds to what data. Similarly, if members = 4 we will loose the ensemble dimension, while if members = [1,2] it remains.

I would suggest a rewriting of the code that keeps the dimension of output_tensor fixed in all cases, something along the lines of

if members is None:
  members = slice(members)
elif not isinstance(members, list):
  members = [members] 

What are the steps to reproduce the bug?

Have a look at the code or use the PlotEnsSample callback.

Version

current main

Platform (OS and architecture)

irrelevant

Relevant log output

Accompanying data

No response

Organisation

RMI

[mc4117]

Hi @dietervdb-meteo I agree that this looks wrong. My only question is how does

if members is None:
  members = slice(members)

work if members is None?
Also at least in the code I was looking at, it doesn't look like None is a valid option?

[anaprietonem]

@jakob-schloer maybe it's also useful for you to check this

[dietervdb-meteo]

@mc4117 , None is actually the default option as specified in the default config. As the comment in the config, and the docstring indicate, the desired behavior when members = None is that all members should be selected. This is strictly speaking the case in the current implementation, since tensor[...,None, :, :] does return the full tensor, but it also adds a useless extra dimension. The preferred behavior would be tensor[...,:,:,:], and since slice(None) = : I think what I proposed above would do this. But could very well be I miss something or that alternative/better solutions exist.

[mc4117]

Thanks for clarifying! Then I think the solution you propose sounds good from my perspective

[mc4117]

@dietervdb-meteo would you be able to make a small PR to main to fix this?