add optional caching in open_zarr

Author: floriankrb

src/anemoi/datasets/data/stores.py

@@ -102,7 +102,7 @@ def __contains__(self, key):
         return key in self.store
 
 
-def open_zarr(path, dont_fail=False):
+def open_zarr(path, dont_fail=False, cache=None):
     try:
         store = path
 
@@ -117,6 +117,9 @@ def open_zarr(path, dont_fail=False):
                 store = zarr.storage.DirectoryStore(store)
             store = DebugStore(store)
 
+        if cache is not None:
+            store = zarr.LRUStoreCache(store, max_size=cache)
+
         return zarr.convenience.open(store, "r")
     except zarr.errors.PathNotFoundError:
         if not dont_fail: