🐛[BUG]: AIFS cannot be used with CDS datasource out-of-the-box

[mike-scchen]

Version

0.10.0

On which installation method(s) does this occur?

uv

Describe the issue

After managing to build the flash-attention (PAIN!), the test code failed with:
ERROR | earth2studio.data.cds:_build_requests:256 - variable id lsm not found in CDS lexicon
This "lsm" could be land-sea mask.

[NickGeneva]

Thanks @mike-scchen
Right now AIFS models only work with the IFS data sources (on main).
Will look at getting CDS working with them.

[NickGeneva]

Merged PR for AIFS ENS to run from CDS: #594

PR Updating AIFS will be under review this week as well.

[chdamianos]

@NickGeneva
There might be a bug with #594
I've installed Earth2Studio from commit 6e75f55 and I get an error using AIFS-ENS with the CDS datasource

This is the code

from pathlib import Path

import earth2studio.run as run
from earth2studio.data import CDS
from earth2studio.io import NetCDF4Backend
from earth2studio.models.px import AIFSENS

package = AIFSENS.load_default_package()
model = AIFSENS.load_model(package)
data = CDS()
io = NetCDF4Backend(file_name=Path("/tmp/output.nc"))

nsteps = 1
io = run.deterministic(["2024-01-01"], nsteps, model, data, io)

The error is

ValueError                                Traceback (most recent call last)
Cell In[1], line 16
     11 io = NetCDF4Backend(
     12     file_name=Path("/tmp/output.nc"),
     13 )
     15 nsteps = 1
---> 16 io = run.deterministic(["2024-01-01"], nsteps, model, data, io)
 
File /opt/conda/envs/baseml/lib/python3.12/site-packages/earth2studio/run.py:141, in deterministic(time, nsteps, prognostic, data, io, output_coords, device)
    138 for step, (x, coords) in enumerate(model):
    139     # Subselect domain/variables as indicated in output_coords
    140     x, coords = map_coords(x, coords, output_coords)
--> 141     io.write(*split_coords(x, coords))
    142     pbar.update(1)
    143     if step == nsteps:
 
File /opt/conda/envs/baseml/lib/python3.12/site-packages/earth2studio/io/netcdf4.py:276, in NetCDF4Backend.write(self, x, coords, array_name)
    274 for xi, name in zip(x, array_name):
    275     if name not in self.root.variables:
--> 276         self.add_array(adjusted_coords, array_name, data=xi)
    278     else:
    279         # Get indices as list of arrays and set torch tensor
    280         self.root[name][
    281             tuple(
    282                 [
   (...)    286             )
    287         ] = xi.to("cpu").numpy()
 
File /opt/conda/envs/baseml/lib/python3.12/site-packages/earth2studio/io/netcdf4.py:191, in NetCDF4Backend.add_array(self, coords, array_name, data)
    188     data = [None] * len(array_name)
    190 if not (len(data) == len(array_name)):
--> 191     raise ValueError(
    192         f"The number of input tensors and array names must be the same but got {len(data)} and {len(array_name)}."
    193     )
    195 adjusted_coords, mapping = convert_multidim_to_singledim(coords)
    197 for c, v in adjusted_coords.items():
 
ValueError: The number of input tensors and array names must be the same but got 1 and 104.

[NickGeneva]

Thanks @chdamianos
Looking at this now, I can reproduce this.

Kinda of a dual bug, one in aifs and actually theres another in the IO. Regardless will get AIFSENS issue fixed today which should make things work

[NickGeneva]

@chdamianos main branch should be running now for AIFS ENS. Let me know if you see any issues still. Thanks!

[chdamianos]

Thanks @NickGeneva!
I tried the code again and now it runs fine