fix: copy_metadata

apps/hiopy/_zarr_utils.py

@@ -4,7 +4,10 @@ import zarr
 def get_var_group(v):
     store = zarr.open(v.store)
     parent_group_path = "/".join(v.path.split("/")[:-1])
-    return store[parent_group_path]
+    if parent_group_path == "":
+        return store
+    else:
+        return store[parent_group_path]
 
 
 def get_time_axis(v):
@@ -18,3 +21,9 @@ def get_time_axis(v):
         return time_axis
     else:
         return None
+
+def get_var_parent_group(v):
+    var_group = get_var_group(v)
+    parent_var_path = var_group.attrs["hiopy::parent"]
+    parent_group = zarr.open(v.store)[parent_var_path]
+    return parent_group

apps/hiopy/worker.py

@@ -13,7 +13,7 @@ from coyote import (
 from ._data_handler import DataHandler
 from ._distribute_work import distribute_work
 from ._grids import def_grid, grid_id
-from ._zarr_utils import get_time_axis, get_var_group
+from ._zarr_utils import get_var_group, get_var_parent_group
 from .loco import LocoServer
 
 from argparse import ArgumentParser
@@ -145,8 +145,9 @@ def main():
                 src_comp, src_grid = v.attrs["hiopy::yac_source"]
             else:
                 assert "hiopy::parent" in var_group.attrs, f"No source for field {v.name} specified"
-                parent_var_name = var_group.attrs["hiopy::parent"] + "/" + v.name.split("/")[-1]
-                source_var = zarr.open(store=v.store)[parent_var_name]
+                parent_group = get_var_parent_group(v)
+                source_var = parent_group[v.basename]
+                src_name = source_var.name
                 source_var_gid = grid_id(source_var)
                 src_comp = src_grid = f"{args.process_group}_{source_var_gid}"
             time_method = v.attrs.get("hiopy::time_method", "point")
@@ -185,14 +186,13 @@ def main():
             )
 
     def get_source_triple(v):
-        group = get_var_group(v)
-        src_field = v.attrs.get("hiopy::src_name", v.basename)
-        if "hiopy::parent" in group.attrs:
-            parent_var_path = group.attrs["hiopy::parent"] + "/" + src_field
-            source_var = zarr.Group(store=v.store)[parent_var_path]
-            return get_source_triple(source_var)
+        var_group = get_var_group(v)
+        if "hiopy::parent" in var_group.attrs:
+            pgroup = get_var_parent_group(v)
+            return get_source_triple(pgroup[v.basename])
         elif "hiopy::yac_source" in v.attrs:
             src_comp, src_grid = v.attrs["hiopy::yac_source"]
+            src_field = v.attrs.get("hiopy::src_name", v.basename)
             return src_comp, src_grid, src_field
         else:
             raise RuntimeError("Invalid attributes: " + str(dict(v.attrs)))