🐛[BUG]: prep_data_array implicitly assumes order of coords

[swbg]

Version

main

On which installation method(s) does this occur?

No response

Describe the issue

earth2studio.data.utils.prep_data_array assumes that "time", "lead_time", and "variable" dimensions occur before other dimensions (e.g., "lat", "lon", "ensemble"). This leads to inconsistent tensors/coords. Most importantly, this prevents the "ensemble" dimension from being moved around, e.g., for broadcasting.

Example:

from earth2studio.data import NCAR_ERA5, prep_data_array
from earth2studio.utils.time import to_time_array

times = to_time_array(["2023-03-24 12:00:00"])
variables = ["t2m", "tcwv"]
ncar_era5 = NCAR_ERA5()
da = ncar_era5(times, variables)
x, coords = prep_data_array(da.transpose("lat", "lon", ...))

Result:
x.shape - torch.Size([721, 1440, 1, 2])
list(coords) - ['time', 'variable', 'lat', 'lon']