From 7b3c4cc75e1e9fe5a1de9e1d7fcf433d40454810 Mon Sep 17 00:00:00 2001
From: Anderson Banihirwe <axbanihirwe@ualr.edu>
Date: Thu, 10 Feb 2022 17:56:00 -0700
Subject: [PATCH] Collapse time dimension

---
 xwrf/accessors.py   | 13 +++++++++++--
 xwrf/postprocess.py | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/xwrf/accessors.py b/xwrf/accessors.py
index 71408010..8fd110c4 100644
--- a/xwrf/accessors.py
+++ b/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)
 
diff --git a/xwrf/postprocess.py b/xwrf/postprocess.py
index a20666e8..2efa68c0 100644
--- a/xwrf/postprocess.py
+++ b/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