Relax logic for comparing grid attributes between Run and Model objs (#…

…199)

(Fixes #187) Previously an error would be raised if a grid attribute found in a Run
object did not exist in the Run's corresponding Model.  This allows
aospy to use grid attributes found only in Run objects.

aospy/calc.py

@@ -275,7 +275,7 @@ def _add_grid_attributes(self, ds):
             ds_coord_name = set(names_ext).intersection(set(ds.coords) |
                                                         set(ds.data_vars))
             model_attr = getattr(self.model, name_int, None)
-            if ds_coord_name:
+            if ds_coord_name and (model_attr is not None):
                 # Force coords to have desired name.
                 ds = ds.rename({list(ds_coord_name)[0]: name_int})
                 ds = ds.set_coords(name_int)

aospy/data_loader.py

@@ -5,7 +5,8 @@
 import numpy as np
 import xarray as xr
 
-from .internal_names import ETA_STR, GRID_ATTRS, TIME_STR, TIME_BOUNDS_STR
+from .internal_names import (ETA_STR, GRID_ATTRS, TIME_STR, TIME_BOUNDS_STR,
+                             TIME_VAR_STRS)
 from .utils import times, io
 
 
@@ -71,10 +72,10 @@ def grid_attrs_to_aospy_names(data):
             # Prevents headaches when subsequently sub-setting.
             if name_int in data.dims and not data[name_int].coords:
                 data = data.assign_coords(**{name_int: data[name_int]})
-    return data
+    return set_grid_attrs_as_coords(data, set_time_vars=False)
 
 
-def set_grid_attrs_as_coords(ds):
+def set_grid_attrs_as_coords(ds, set_time_vars=True):
     """Set available grid attributes as coordinates in a given Dataset.
 
     Grid attributes are assumed to have their internal aospy names. Grid
@@ -91,7 +92,11 @@ def set_grid_attrs_as_coords(ds):
     Dataset
         Dataset with grid attributes set as coordinates
     """
-    int_names = GRID_ATTRS.keys()
+    if set_time_vars:
+        int_names = GRID_ATTRS.keys()
+    else:
+        int_names = (set(GRID_ATTRS.keys()) -
+                     set(TIME_VAR_STRS))
     grid_attrs_in_ds = set(int_names).intersection(set(ds.coords) |
                                                    set(ds.data_vars))
     ds.set_coords(grid_attrs_in_ds, inplace=True)

aospy/test/test_data_loader.py

@@ -98,15 +98,27 @@ def test_rename_grid_attrs_copy_attrs(self):
         ds = grid_attrs_to_aospy_names(ds_orig)
         self.assertEqual(ds[LAT_STR].attrs, orig_attrs)
 
-    def test_set_grid_attrs_as_coords(self):
+    def test_set_grid_attrs_as_coords_all(self):
         ds = grid_attrs_to_aospy_names(self.ds)
         sfc_area = ds[self.var_name].isel(**{TIME_STR: 0}).drop(TIME_STR)
         ds[SFC_AREA_STR] = sfc_area
 
-        assert SFC_AREA_STR not in ds[self.var_name]
+        assert SFC_AREA_STR not in ds.coords
 
         ds = set_grid_attrs_as_coords(ds)
-        assert SFC_AREA_STR in ds[self.var_name]
+        assert SFC_AREA_STR in ds.coords
+        assert TIME_BOUNDS_STR in ds.coords
+
+    def test_set_grid_attrs_as_coords_no_times(self):
+        ds = grid_attrs_to_aospy_names(self.ds)
+        sfc_area = ds[self.var_name].isel(**{TIME_STR: 0}).drop(TIME_STR)
+        ds[SFC_AREA_STR] = sfc_area
+
+        assert SFC_AREA_STR not in ds.coords
+
+        ds = set_grid_attrs_as_coords(ds, set_time_vars=False)
+        assert SFC_AREA_STR in ds.coords
+        assert TIME_BOUNDS_STR not in ds.coords
 
     def test_sel_var(self):
         time = np.array([0, 31, 59]) + 15

docs/whats-new.rst

@@ -77,7 +77,12 @@ Bug Fixes
   Hill <https://github.com/spencerahill>`_.
 - Always write output to a tar file in serial to prevent empty header file
   errors (fixes :issue:`75` via :pull:`197`).  By `Spencer Clark
-  <https://github.com/spencerkclark>`_
+  <https://github.com/spencerkclark>`_.
+- Allow ``aospy`` to use grid attributes that are only defined in ``Run``
+  objects. Previously if a grid attribute were defined only in a ``Run``
+  object and not also in the Run's corresponding ``Model``, an error would
+  be raised (fixes :issue:`187` via :pull:`199`).  By `Spencer Clark
+  <https://github.com/spencerkclark>`_.
 - When input data for a calculation has a time bounds array, overwrite
   its time array with the average of the start and end times for each
   timestep.  Prevents bug wherein time arrays equal to either the