Add units and description to output netcdf files

aospy/calc.py

@@ -524,7 +524,7 @@ def region_calcs(self, arr, func):
                         **{reg.name + '_pressure': coord}
                     )
             reg_dat.update(**{reg.name: data_out})
-        return OrderedDict(sorted(reg_dat.items(), key=lambda t: t[0]))
+        return xr.Dataset(reg_dat)
 
     def _apply_all_time_reductions(self, full_ts, monthly_ts, eddy_ts):
         """Apply all requested time reductions to the data."""
@@ -587,6 +587,7 @@ def compute(self, write_to_tar=True):
         reduced = self._apply_all_time_reductions(full, monthly, eddy)
         logging.info("Writing desired gridded outputs to disk.")
         for dtype_time, data in reduced.items():
+            self._add_units_and_description_attrs(data)
             self.save(data, dtype_time, dtype_out_vert=self.dtype_out_vert,
                       save_files=True, write_to_tar=write_to_tar)
         return self
@@ -601,8 +602,6 @@ def _save_files(self, data, dtype_out_time):
                 reg_data = xr.open_dataset(path)
             except (EOFError, RuntimeError, IOError):
                 reg_data = xr.Dataset()
-            # Add the new data to the dictionary or Dataset.
-            # Same method works for both.
             reg_data.update(data)
             data_out = reg_data
         else:
@@ -767,3 +766,20 @@ def load(self, dtype_out_time, dtype_out_vert=False, region=False,
         if plot_units:
             data = self.var.to_plot_units(data, dtype_vert=dtype_out_vert)
         return data
+
+    def _add_units_and_description_attrs(self, data):
+        if isinstance(data, xr.DataArray):
+            self._add_units_and_description_attrs_da(data)
+        else:
+            for name, da in data.data_vars.items():
+                self._add_units_and_description_attrs_da(da)
+
+    def _add_units_and_description_attrs_da(self, data):
+        units = self.var.units
+        if self.dtype_out_vert == 'vert_int':
+            if units != '':
+                units = '(vertical integral of {0}): {0} kg m^-2)'.format(units)
+            else:
+                units = '(vertical integral of quantity with unspecified units)'
+        data.attrs['units'] = units
+        data.attrs['description'] = self.var.description

aospy/test/test_calc_basic.py

@@ -5,6 +5,8 @@
 import shutil
 import unittest
 
+import xarray as xr
+
 from aospy.calc import Calc, CalcInterface
 from .data.objects.examples import (
     example_proj, example_model, example_run, condensation_rain,
@@ -37,6 +39,10 @@ def test_annual_mean(self):
         calc.compute()
         assert isfile(calc.path_out['av'])
         assert isfile(calc.path_tar_out)
+        data = xr.open_mfdataset(calc.path_out['av'], decode_times=False)
+        for name, da in data.data_vars.items():
+            assert 'units' in da.attrs
+            assert 'description' in da.attrs
 
     def test_annual_ts(self):
         calc_int = CalcInterface(intvl_out='ann',
@@ -92,6 +98,10 @@ def test_simple_reg_av(self):
         calc.compute()
         assert isfile(calc.path_out['reg.av'])
         assert isfile(calc.path_tar_out)
+        data = xr.open_mfdataset(calc.path_out['reg.av'], decode_times=False)
+        for name, da in data.data_vars.items():
+            assert 'units' in da.attrs
+            assert 'description' in da.attrs
 
     def test_simple_reg_ts(self):
         calc_int = CalcInterface(intvl_out='ann',

docs/whats-new.rst

@@ -28,6 +28,9 @@ Corrected link to documentation badge on repository main page (:pull:213). By Da
 Enhancements
 ~~~~~~~~~~~~
 
+- Add units and description from ``Var`` objects to output netcdf
+  files (closes :issue:`201` via :pull:`232`). By `Micah Kim
+  <https://github.com/micahkim23>`_.
 - Remove potentially confusing attributes from example netcdf files.
   (closes :issue:`214` via :pull:`216`). By `Micah Kim
   <https://github.com/micahkim23>`_.