🐛[BUG]: Dimension ordering in run.deterministic workflows on batched input data

[Corkman99]

Version

0.11.0

On which installation method(s) does this occur?

source

Describe the issue

I am attempting to use the run.deterministic workflow on prognostic models, for single time but batched input data. I have created my own datasource which simply wraps a xarray.dataset in memory. Batched input data can for example correspond to different perturbations, although I am not using the Earth2Studio perturbation API.

Using an xarray without a batch dimension works without issue for many prognostic models.

I am now considering a xarray.dataarray with dimensions and coordinates batch, time, variable, lat, lon as requested by the api. I get a ValueError that the "lead_time" dimension could not be found at the right index (index 2). Upon further inquiry, I find that fetch_data returns coords in the order batch, lead_time, time, lat, lon, which then gets passed down to the iterator and fails the handshake_dim.

I have tried to patch fetch_data and return the expected ordering of batch, time, lead_time, lat, lon, but I then get further errors from io.write. I see that while the model has generated outputs for the first step, they are all nan.

Do you have any tips on how to operate with batched inputs? Is my io handling the issue or does the batch.py module only work in junction with the perturbation api?

Thanks!

[NickGeneva]

HI @Corkman99

Interesting, yeah the expected order should be <misc>, time, lead_time, variable, lat, lon.
I think the issue is because with data sources, the expected output data array is of dim [time, variable,lat, lon]. The reason is that data sources call only has input parameters time and variables.

The addition of another dim (I would suggest using ensemble instead of batch, batch is used internally a lot so better to avoid calling it batch for debugging reasons if possible) typically happens after fetch_data. So I don't think the deterministic would work out of the box for you.

For example the ensemble workflow does this repeat of the ensemble dimension after fetch data.

https://github.com/NVIDIA/earth2studio/blob/main/earth2studio/run.py#L409

You are likely needing a custom workflow implementation.

[NickGeneva]

I can look to see what would be needed to generalize the the deterministic workflow to support extra dimensions.

In the immediate term, creating your own workflow is best referencing the deterministic and ensemble workflows to provide guidance on running the models and handling an ensemble dim respectively.