Collapse time dimension

xwrf/accessors.py

@@ -2,7 +2,12 @@
 
 import xarray as xr
 
-from .postprocess import _decode_times, _modify_attrs_to_cf, _remove_units_from_bool_arrays
+from .postprocess import (
+    _collapse_time_dim,
+    _decode_times,
+    _modify_attrs_to_cf,
+    _remove_units_from_bool_arrays,
+)
 
 
 class WRFAccessor:
@@ -34,7 +39,11 @@ def postprocess(self, decode_times=True) -> xr.Dataset:
         xarray.Dataset
             The postprocessed dataset.
         """
-        ds = self.xarray_obj.pipe(_modify_attrs_to_cf).pipe(_remove_units_from_bool_arrays)
+        ds = (
+            self.xarray_obj.pipe(_modify_attrs_to_cf)
+            .pipe(_remove_units_from_bool_arrays)
+            .pipe(_collapse_time_dim)
+        )
         if decode_times:
             ds = ds.pipe(_decode_times)
 

xwrf/postprocess.py

@@ -35,3 +35,21 @@ def _modify_attrs_to_cf(ds: xr.Dataset) -> xr.Dataset:
     for variable in vars_to_update:
         ds[variable].attrs.update(config.get(f'cf_attribute_map.{variable}'))
     return ds
+
+
+def _collapse_time_dim(ds: xr.Dataset) -> xr.Dataset:
+
+    # This "time dimension collapsing" assumption is wrong with moving nests
+    # and should be applied to static, nested domains.
+    lat_lon_coords = set(config.get('latitude_coords') + config.get('longitude_coords'))
+    coords = set(ds.variables).intersection(lat_lon_coords)
+    ds = ds.set_coords(coords)
+
+    for coord in ds.coords:
+        if coord in lat_lon_coords and ds[coord].ndim == 3:
+            attrs, encoding = ds[coord].attrs, ds[coord].encoding
+            ds = ds.assign_coords({coord: (ds[coord].dims[1:], ds[coord].data[0, :, :])})
+            ds[coord].attrs = attrs
+            ds[coord].encoding = encoding
+
+    return ds