Refactor whitespace and formatting in atmosphere.py, calc_atmos.py, ACCESS-ESM1.6_mappings.json, and vocabulary_processors.py for improved readability

Author: rbeucher

src/access_moppy/atmosphere.py

@@ -1,5 +1,5 @@
-import warnings
 import re
+import warnings
 
 import numpy as np
 import xarray as xr
@@ -72,66 +72,72 @@ def calculate_missing_bounds_variables(self, bnds_required):
     def remove_spurious_time_dimensions(self, required_vars):
         """
         Remove spurious time dimensions from coordinate and auxiliary variables.
-        
+
         This method addresses a common issue in xarray when combining datasets:
-        spatial bounds (lat_bnds, lon_bnds) and other coordinate variables can incorrectly 
+        spatial bounds (lat_bnds, lon_bnds) and other coordinate variables can incorrectly
         gain time dimensions during multi-file dataset operations, even though they are time-invariant.
-        
+
         Why this is necessary:
-        - When using xr.open_mfdataset() with combine_coords="time", xarray 
+        - When using xr.open_mfdataset() with combine_coords="time", xarray
           conservatively assumes all coordinate-linked variables might vary with time
         - This causes spatial bounds and coordinates to be broadcasted along the time dimension
         - Results in redundant data storage and non-CF-compliant files
-        
+
         Why this is reasonable for ACCESS Models:
         - ACCESS Models use static grids throughout model runs
         - Latitude, longitude coordinates (and their bounds) are time-invariant
         - The grid definition remains constant across all timesteps
         - Only time_bnds and data variables should legitimately have a time dimension
         - This optimization is safe and improves storage efficiency
-        
+
         Args:
             required_vars (list): Variables that should keep their time dimension
         """
         # Identify all variables that have gained spurious time dimensions
         # Include bounds variables and any other coordinate variables
         problematic_vars = [
-            name for name in self.ds.variables
+            name
+            for name in self.ds.variables
             if "time" not in name  # Don't touch time_bnds or time coordinate
             and name not in required_vars  # Don't touch required data variables
-            and name in self.ds 
+            and name in self.ds
             and "time" in self.ds[name].coords
             and self.ds[name].dims != ("time",)  # Skip pure time variables
         ]
-        
+
         if problematic_vars:
             # Process all problematic variables efficiently in a single operation
             corrections = {
-                name: self.ds[name].isel(time=0).drop_vars("time") 
+                name: self.ds[name].isel(time=0).drop_vars("time")
                 for name in problematic_vars
             }
             self.ds = self.ds.assign(corrections)
 
     def select_and_process_variables(self):
-
         # Select input variables required for the CMOR variable
         required_vars = self.mapping[self.cmor_name]["model_variables"]
-        
-        required_axes, axes_rename_map = self.vocab._get_axes(self.mapping)        
-        required_bounds, bounds_rename_map = self.vocab._get_required_bounds_variables(self.mapping)
 
-        required = set(required_vars + list(axes_rename_map.keys()) + list(bounds_rename_map.keys()))
+        required_axes, axes_rename_map = self.vocab._get_axes(self.mapping)
+        required_bounds, bounds_rename_map = self.vocab._get_required_bounds_variables(
+            self.mapping
+        )
+
+        required = set(
+            required_vars
+            + list(axes_rename_map.keys())
+            + list(bounds_rename_map.keys())
+        )
         self.load_dataset(required_vars=required)
 
         # Remove spurious time dimensions from spatial bounds and coordinates
         self.remove_spurious_time_dimensions(required_vars)
 
         # Ensure time dimension is sorted
         self.sort_time_dimension()
-        
+
         ## Calculate missing bounds variables
         ##self.calculate_missing_bounds_variables(required_bounds)
-        
+
         calc = self.mapping[self.cmor_name]["calculation"]
 
         # Handle the calculation type
@@ -161,12 +167,18 @@ def select_and_process_variables(self):
             raise ValueError(f"Unsupported calculation type: {calc['type']}")
 
         # Rename axes and bounds variables
-        rename_map = {k: v for k, v in {**bounds_rename_map, **axes_rename_map}.items() if k in self.ds}
+        rename_map = {
+            k: v
+            for k, v in {**bounds_rename_map, **axes_rename_map}.items()
+            if k in self.ds
+        }
         self.ds = self.ds.rename(rename_map)
 
         # Transpose the data variable according to the CMOR dimensions
-        cmor_dims = re.sub(r'\w*level', 'lev', self.vocab.variable["dimensions"]).split()
-        
+        cmor_dims = re.sub(
+            r"\w*level", "lev", self.vocab.variable["dimensions"]
+        ).split()
+
         transpose_order = [
             self.vocab.axes[dim]["out_name"]
             for dim in cmor_dims
@@ -194,7 +206,7 @@ def update_attributes(self):
         )
 
         cmor_attrs = self.vocab.variable
-        #self._check_units(self.cmor_name, cmor_attrs.get("units"))
+        # self._check_units(self.cmor_name, cmor_attrs.get("units"))
 
         self.ds[self.cmor_name].attrs.update(
             {k: v for k, v in cmor_attrs.items() if v not in (None, "")}

src/access_moppy/derivations/calc_atmos.py

@@ -250,8 +250,11 @@
 def cli_level_to_height(ds):
     # Handle level coordinate transformation
     if "theta_level_height" in ds:
-        ds = (ds.assign_coords({"lev": ds["theta_level_height"]})
-                        .swap_dims({"model_theta_level_number": "lev"})
-                        .drop_vars(["theta_level_height", "model_theta_level_number"], errors="ignore"))
+        ds = (
+            ds.assign_coords({"lev": ds["theta_level_height"]})
+            .swap_dims({"model_theta_level_number": "lev"})
+            .drop_vars(
+                ["theta_level_height", "model_theta_level_number"], errors="ignore"
+            )
+        )
     return ds
-

src/access_moppy/mappings/ACCESS-ESM1.6_mappings.json

@@ -77,7 +77,7 @@
                 "type": "hybrid_height",
                 "coordinate_variables": {
                     "sigma_theta": "b",
-                    "surface_altitude": "orog", 
+                    "surface_altitude": "orog",
                     "theta_level_height": "lev"
                 },
             "formula": "z = lev + b*orog"
@@ -86,7 +86,7 @@
                 "type": "formula",
                 "operation": "swap_dimensions",
                 "operands": [
-                    "fld_s02i261", 
+                    "fld_s02i261",
                     {"literal": "model_level_number"},
                     {"literal": "lev"}
                 ]
@@ -109,7 +109,7 @@
                 "type": "hybrid_height",
                 "coordinate_variables": {
                     "sigma_theta": "b",
-                    "surface_altitude": "orog", 
+                    "surface_altitude": "orog",
                     "theta_level_height": "lev"
                 },
                 "formula": "z = a + b*orog"

src/access_moppy/vocabulary_processors.py

@@ -258,7 +258,6 @@ def _get_variable_suggestions(self) -> List[str]:
         return suggestions
 
     def _get_axes(self, mapping) -> Dict[str, Any]:
-
         # Resolve resource inside the module path
         coord_entry = files(self.table_dir) / "CMIP6_coordinate.json"
 
@@ -268,48 +267,54 @@ def _get_axes(self, mapping) -> Dict[str, Any]:
 
         dims = self.variable["dimensions"].split()
         vars_required = {}
-        
+
         for dim in dims:
             if dim in axes and dim not in ["alevel"]:
                 coord = axes[dim]
                 vars_required[dim] = {k: v for k, v in coord.items() if v != ""}
-        
+
         # Add z-axis coordinate variables if applicable
         if "zaxis" in mapping[self.cmor_name]:
-
             # Get z-axis type from mapping
-            zaxis_type = mapping[self.cmor_name]["zaxis"].get(
-                "type", {}
-            )
+            zaxis_type = mapping[self.cmor_name]["zaxis"].get("type", {})
 
             # Process main z-axis coordinate
             zcoord = axes.get(zaxis_type, {})["out_name"]
-            vars_required[zcoord] = {k: v for k, v in axes[zaxis_type].items() if v != ""}
-            
+            vars_required[zcoord] = {
+                k: v for k, v in axes[zaxis_type].items() if v != ""
+            }
+
             # Process z_factors
             zfactors_str = axes.get(zaxis_type, {}).get("z_factors", "")
-            
+
             zfactors = {}
             if zfactors_str:
                 parts = zfactors_str.split()
-                zfactors = {parts[i].rstrip(':'): parts[i+1] 
-                           for i in range(0, len(parts), 2) if i+1 < len(parts)}
-            
+                zfactors = {
+                    parts[i].rstrip(":"): parts[i + 1]
+                    for i in range(0, len(parts), 2)
+                    if i + 1 < len(parts)
+                }
+
             formula_entry = files(self.table_dir) / "CMIP6_formula_terms.json"
             with as_file(formula_entry) as fpath:
                 with open(fpath, "r", encoding="utf-8") as ff:
-                    formula_terms = json.load(ff)["formula_entry"]  
+                    formula_terms = json.load(ff)["formula_entry"]
 
             for factor_name, _ in zfactors.items():
                 if factor_name in formula_terms:
                     zcoord = formula_terms[factor_name]
-                    vars_required[factor_name] = {k: v for k, v in zcoord.items() if v != ""}
+                    vars_required[factor_name] = {
+                        k: v for k, v in zcoord.items() if v != ""
+                    }
 
         # Let's map the axis and formula terms to the inputs
         vars_rename_map = {}
-        extended_mapping = mapping[self.cmor_name]["dimensions"] | mapping[self.cmor_name].get("zaxis", {}).get("coordinate_variables", {})
-        inverted_extended_mapping = {v: k for k, v in extended_mapping.items()} 
-        
+        extended_mapping = mapping[self.cmor_name]["dimensions"] | mapping[
+            self.cmor_name
+        ].get("zaxis", {}).get("coordinate_variables", {})
+        inverted_extended_mapping = {v: k for k, v in extended_mapping.items()}
+
         for _, v in vars_required.items():
             input_dim = inverted_extended_mapping.get(v["out_name"])
             if input_dim:
@@ -322,21 +327,23 @@ def _get_axes(self, mapping) -> Dict[str, Any]:
     def _get_required_bounds_variables(self, mapping: Dict[str, Any]) -> tuple:
         """
         Get required bounds variables based on CMOR vocabulary axes.
-        
+
         Args:
             mapping: Variable mapping dictionary containing dimensions
-            
+
         Returns:
             tuple: (bnds_required, bounds_rename_map) where
                 - bnds_required: list of required bounds variable names
                 - bounds_rename_map: dict mapping input bounds names to output bounds names
         """
         bnds_required = {}
         bounds_rename_map = {}
-        
-        extended_mapping = mapping[self.cmor_name]["dimensions"] | mapping[self.cmor_name].get("zaxis", {}).get("coordinate_variables", {})
-        inverted_extended_mapping = {v: k for k, v in extended_mapping.items()} 
-        
+
+        extended_mapping = mapping[self.cmor_name]["dimensions"] | mapping[
+            self.cmor_name
+        ].get("zaxis", {}).get("coordinate_variables", {})
+        inverted_extended_mapping = {v: k for k, v in extended_mapping.items()}
+
         axes, _ = self._get_axes(mapping)
         for _, v in axes.items():
             if v.get("must_have_bounds") == "yes":
@@ -345,8 +352,10 @@ def _get_required_bounds_variables(self, mapping: Dict[str, Any]) -> tuple:
                 input_bounds = input_dim + "_bnds"
                 output_bounds = v["out_name"] + "_bnds"
                 bounds_rename_map[input_bounds] = output_bounds
-                bnds_required[output_bounds] = {key: val for key, val in v.items() if val != ""}
-                
+                bnds_required[output_bounds] = {
+                    key: val for key, val in v.items() if val != ""
+                }
+
         return bnds_required, bounds_rename_map
 
     def get_variant_components(self) -> Dict[str, int]: