fix update to get the right constant fields

ecmwf · Dec 2, 2024 · eb11bdc · eb11bdc
1 parent def5b0c
commit eb11bdc
Showing 1 changed file with 27 additions and 12 deletions.
diff --git a/src/anemoi/registry/commands/update.py b/src/anemoi/registry/commands/update.py
@@ -147,8 +147,22 @@ def entry_set_value(path, value):
             if not args.update or not args.force:
                 return
 
+        # Remove stuff added by prepml
+        for k in [
+            "build_dataset",
+            "config_format_version",
+            "config_path",
+            "dataset_status",
+            "ecflow",
+            "metadata",
+            "platform",
+            "reading_chunks",
+            "upload",
+        ]:
+            recipe.pop(k, None)
+
         if "recipe" not in entry.record["metadata"] or args.force:
-            LOG.info("%s, setting `constant_fields` 🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", name)
+            LOG.info("%s, setting `recipe` 🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", name)
             if args.dry_run:
                 LOG.info("Would set recipe %s", name)
             else:
@@ -157,6 +171,18 @@ def entry_set_value(path, value):
                 entry_set_value("/metadata/recipe", recipe)
                 entry_set_value("/metadata/updated", updated + 1)
 
+        computed_constant_fields = sorted(open_dataset(name).computed_constant_fields())
+        constant_fields = entry.record["metadata"].get("constant_fields", [])
+        if computed_constant_fields != constant_fields:
+            LOG.info("%s, setting `constant_fields`", name)
+            if args.dry_run:
+                LOG.info("Would set constant_fields %s", name)
+            else:
+                LOG.info("Setting constant_fields %s", name)
+                entry_set_value("/metadata/constant_fields", computed_constant_fields)
+                entry_set_value("/metadata/updated", updated + 1)
+                entry.record["metadata"]["constant_fields"] = computed_constant_fields
+
         if "constant_fields" in entry.record["metadata"] and "variables_metadata" in entry.record["metadata"]:
             LOG.info("%s, checking `variables_metadata` and `constant_fields`", name)
             constants = entry.record["metadata"]["constant_fields"]
@@ -210,17 +236,6 @@ def entry_set_value(path, value):
                 finally:
                     shutil.rmtree(dir)
 
-        if "constant_fields" not in entry.record["metadata"] or args.force:
-            LOG.info("%s, setting `constant_fields` 🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", name)
-            ds = open_dataset(name)
-            constant_fields = ds.computed_constant_fields()
-            LOG.info("%s", constant_fields)
-            if args.debug:
-                with open(f"{name}.constant_fields.json", "w") as f:
-                    print(json.dumps(constant_fields, indent=2), file=f)
-            entry_set_value("/metadata/constant_fields", constant_fields)
-            entry_set_value("/metadata/updated", updated + 1)
-
     def zarr_file_from_catalogue(self, path, args):
         import zarr