@@ -232,25 +232,71 @@ class CMIP6_Atmosphere_CMORiser(CMIP6_CMORiser):
232232 """
233233
234234 def select_and_process_variables (self ):
235- bnds_required = [
236- v ["out_name" ] + "_bnds"
237- for v in self .vocab .axes .values ()
238- if v .get ("must_have_bounds" ) == "yes"
239- ]
235+ # Find all required bounds variables
236+ bnds_required = []
237+ bounds_rename_map = {}
238+ for dim , v in self .vocab .axes .items ():
239+ if v .get ("must_have_bounds" ) == "yes" :
240+ # Find the input dimension name that maps to this output name
241+ input_dim = None
242+ for k , val in self .mapping [self .cmor_name ]["dimensions" ].items ():
243+ if val == v ["out_name" ]:
244+ input_dim = k
245+ break
246+ if input_dim is None :
247+ raise KeyError (
248+ f"Can't find input dimension mapping for output dimension '{ v ['out_name' ]}
249+ )
250+ bnds_var = input_dim + "_bnds"
251+ if bnds_var in self .ds .variables or bnds_var in self .ds .coords :
252+ bnds_required .append (bnds_var )
253+ # Store mapping for later renaming after calculation
254+ bounds_rename_map [bnds_var ] = v ["out_name" ] + "_bnds"
240255
256+ # Select input variables
241257 input_vars = self .mapping [self .cmor_name ]["model_variables" ]
242258 calc = self .mapping [self .cmor_name ]["calculation" ]
243259 self .ds = self .ds [input_vars + bnds_required ]
244260
261+ # Handle the calculation type
245262 if calc ["type" ] == "direct" :
263+ # If the calculation is direct, just rename the variable
246264 self .ds = self .ds .rename ({input_vars [0 ]: self .cmor_name })
247265 elif calc ["type" ] == "formula" :
266+ # If the calculation is a formula, evaluate it
248267 context = {var : self .ds [var ] for var in input_vars }
249268 context .update (custom_functions )
250269 self .ds [self .cmor_name ] = evaluate_expression (calc , context )
270+ # Drop the original input variables, except the CMOR variable and keep bounds
271+ self .ds = self .ds .drop_vars (
272+ [
273+ var
274+ for var in input_vars
275+ if var != self .cmor_name and var not in bnds_required
276+ ],
277+ errors = "ignore" ,
278+ )
251279 else :
252280 raise ValueError (f"Unsupported calculation type: { calc ['type' ]} )
253281
282+ # Rename dimensions according to the CMOR vocabulary
283+ dim_rename = self .mapping [self .cmor_name ]["dimensions" ]
284+ dims_to_rename = {k : v for k , v in dim_rename .items () if k in self .ds .dims }
285+ self .ds = self .ds .rename (dims_to_rename )
286+
287+ # Also rename coordinates if needed
288+ coords_to_rename = {k : v for k , v in dim_rename .items () if k in self .ds .coords }
289+ if coords_to_rename :
290+ self .ds = self .ds .rename (coords_to_rename )
291+
292+ # Rename bounds variables
293+ for bnds_var , out_bnds_name in bounds_rename_map .items ():
294+ if bnds_var in self .ds :
295+ self .ds = self .ds .rename ({bnds_var : out_bnds_name })
296+ elif bnds_var in self .ds .coords :
297+ self .ds = self .ds .rename ({bnds_var : out_bnds_name })
298+
299+ # Transpose the data variable according to the CMOR dimensions
254300 cmor_dims = self .vocab .variable ["dimensions" ].split ()
255301 transpose_order = [
256302 self .vocab .axes [dim ]["out_name" ]
@@ -261,6 +307,7 @@ def select_and_process_variables(self):
261307 for dim in self .ds [self .cmor_name ].dims :
262308 if dim not in transpose_order and self .ds [self .cmor_name ][dim ].size == 1 :
263309 self .ds [self .cmor_name ] = self .ds [self .cmor_name ].squeeze (dim )
310+
264311 self .ds [self .cmor_name ] = self .ds [self .cmor_name ].transpose (* transpose_order )
265312
266313 def update_attributes (self ):
@@ -270,8 +317,6 @@ def update_attributes(self):
270317 if v not in None , "" )
271318 }
272319
273- self .ds = self .ds .rename (self .mapping [self .cmor_name ]["dimensions" ])
274-
275320 required_coords = {
276321 v ["out_name" ] for v in self .vocab .axes .values () if "value" in v
277322 }.union ({v ["out_name" ] for v in self .vocab .axes .values ()})
0 commit comments