run up to _step with dummy_model

ecmwf · Oct 16, 2024 · 3b89b82 · 3b89b82
1 parent 6a9dda4
commit 3b89b82
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/anemoi/utils/data_structures.py b/src/anemoi/utils/data_structures.py
@@ -229,6 +229,9 @@ def __init__(self, arrays, **kwargs):
             arrays = {i: v for i, v in enumerate(arrays)}
         self.arrays = arrays
 
+    def keys(self):
+        return self.arrays.keys()
+
     def clone(self):
         raise NotImplementedError
 
@@ -239,8 +242,8 @@ def check_array_type(self, arrays):
             assert isinstance(a, _type), (type(a), _type)
 
     def __getitem__(self, tupl):
-        assert isinstance(tupl, (int, tuple)), type(tupl)
-        if isinstance(tupl, int):
+        assert isinstance(tupl, (int, tuple, str)), type(tupl)
+        if isinstance(tupl, int) or isinstance(tupl, str):
             return self.arrays[tupl]
         assert len(tupl) == 2
         i, j = tupl