Add support for several grids in same realm

Author: rigoudyg

dr2xml/Xwrite.py

@@ -464,8 +464,8 @@ def write_xios_file_def(filename, svars_per_table, year, dummies, skipped_vars_p
     set_config_variable("domain_defs", OrderedDict())
     # Add xml_file_definition
     xml_file_definition = DR2XMLElement(tag="file_definition")
-    _, hgrid, _, _, _ = internal_dict['grids'][internal_dict["select_grid_choice"]][context]
     files_list = determine_files_list(svars_per_table, enddate, year, debug)
+    _, hgrid, _, _, _ = internal_dict['grids'][internal_dict["select_grid_choice"]][context]["default"]
     for file_dict in files_list:
         write_xios_file_def_for_svars_list(hgrid=hgrid, xml_file_definition=xml_file_definition, dummies=dummies,
                                            skipped_vars_per_table=skipped_vars_per_table,
@@ -751,6 +751,7 @@ def get_grid_info(sv, grid, table):
     grid_choice = internal_dict["grid_choice"]
     context_index = get_config_variable("context_index")
     if grid in ["", ]:
+        internal_grids = internal_dict["grids"][grid_choice][context]
         # either native or close-to-native
         if sv.type in ["dev", ]:
             target_grid = sv.description.split('|')[1]
@@ -767,10 +768,10 @@ def get_grid_info(sv, grid, table):
                         grids_dev[sv.label][grid_choice][context]
             else:
                 grid_label, target_hgrid_id, zgrid_id, grid_resolution, grid_description = \
-                    internal_dict["grids"][grid_choice][context]
+                    internal_grids.get(sv.label, internal_grids["default"])
         else:
             grid_label, target_hgrid_id, zgrid_id, grid_resolution, grid_description = \
-                internal_dict["grids"][grid_choice][context]
+                internal_grids.get(sv.label, internal_grids["default"])
     else:
         if grid in ['cfsites', ]:
             target_hgrid_id = cfsites_domain_id

dr2xml/projects/dr2xml.json

@@ -731,14 +731,38 @@
         }
       ]
     },
+    "variables_per_grid_type": {
+      "help": "List of variables associated with a grid type",
+      "values": [
+        {
+          "type": "value",
+          "origin": "laboratory",
+          "keys": ["variables_per_grid_type"]
+        },
+        {}
+      ],
+      "target_type": "dict"
+    },
     "grids": {
       "help": "Grids : per model resolution and per context :\\n- CMIP6 qualifier (i.e. 'gn' or 'gr') for the main grid chosen (because you  may choose has main production grid a regular one, when the native grid is e.g. unstructured)\\n- Xios id for the production grid (if it is not the native grid)\\n- Xios id for the latitude axis used for zonal means (mist match latitudes for grid above)\\n- resolution of the production grid (using CMIP6 conventions)\\n- grid description",
       "fatal": true,
       "values": [
         {
-          "type": "value",
-          "origin": "laboratory",
-          "keys": ["grids"]
+          "type": "func",
+          "origin": "functions_file",
+          "keys": ["format_grids"],
+          "options": {
+            "grids": {
+              "type": "value",
+              "origin": "laboratory",
+              "keys": ["grids"]
+            },
+            "variables_per_grid_type": {
+              "type": "value",
+              "origin": "internal",
+              "keys": ["variables_per_grid_type"]
+            }
+          }
         }
       ]
     },
@@ -1354,9 +1378,10 @@
               "keys": ["excluded_dimensions"]
             }
           ],
-          "not_values": [null]
+          "not_values": [[]]
         }
-      ]
+      ],
+      "target_type": "list"
     },
     "select_excluded_expgroups": {
       "help": "Excluded experiments groups for variable selection.",

dr2xml/projects/dr2xml_func.py

@@ -7,6 +7,9 @@
 
 from __future__ import print_function, division, absolute_import, unicode_literals
 
+import copy
+from collections import defaultdict
+
 from utilities.logger import get_logger
 
 
@@ -64,4 +67,41 @@ def sort_mips(mips):
 			rep = rep | mips[grid]
 	else:
 		rep = mips
-	return sorted(list(rep))
+	return sorted(list(rep))
+
+
+def format_grids(grids, variables_per_grid_type):
+	logger = get_logger()
+	if not isinstance(grids, dict) or any([not isinstance(grids[elt], dict) for elt in grids]):
+		logger.error("Grids must be defined as a dict of resolution, realm and grid_type")
+		raise ValueError("Grids must be defined as a dict of resolution, realm and grid_type")
+	else:
+		issues = list()
+		for (resolution, dict_resolution) in grids.items():
+			for (realm, dict_realm) in dict_resolution.items():
+				if isinstance(dict_realm, list) and len(dict_realm) == 5:
+					new_dict_realm = dict(default=dict_realm)
+					grids[resolution][realm] = copy.deepcopy(new_dict_realm)
+				elif isinstance(dict_realm, dict) and \
+						all([isinstance(elt, dict) and len(elt) == 5 for elt in dict_realm.values()]) and \
+						"default" in dict_realm:
+					pass
+				else:
+					issues.append((resolution, realm))
+		if len(issues) > 0:
+			logger.error("Issues in grids definition for the following (resolution, realm) couples:\n%s" % issues)
+			raise ValueError("Issues in grids definition for the following (resolution, realm) couples:\n%s" % issues)
+		else:
+			rep = defaultdict(lambda: defaultdict(dict))
+			for (resolution, dict_resolution) in grids.items():
+				for (realm, dict_realm) in dict_resolution.items():
+					rep[resolution][realm]["default"] = copy.deepcopy(grids[resolution][realm]["default"])
+					for (grid_type, grid_type_list) in dict_realm.items():
+						grid_def = copy.deepcopy(grids[resolution][realm][grid_type])
+						if grid_type in grids[resolution][realm]:
+							for var in grid_type_list:
+								rep[resolution][realm][var] = grid_def
+						else:
+							logger.warning("Could not find grid_type %s for resolution %s and realm %s, use default." %
+							               (grid_type, resolution, realm))
+			return rep