Repository: xCDAT/xcdat

Issue:
[Bug]: `xarray.open_mfdataset` drops `.encoding` attributes, need to handle with `xcdat.open_mfdataset()`
### What happened?

Related to https://github.com/xCDAT/xcdat/pull/302#discussion_r942905147.

> Nice catch! I found an open xarray issue related to xr.open_mfdataset() dropping encoding attributes.
>
> - https://github.com/pydata/xarray/issues/2436
> - https://github.com/pydata/xarray/issues/2436#issuecomment-449737841
> 
> Somebody mentions this about the issue here https://github.com/pydata/xarray/issues/2436#issuecomment-888204512:
> 
>
>>     However, there is another problem related to xarray.open_mfdataset:
>>     The encoding dictionary gets lost somewhere during the merging operation of the datasets of the respective files (https://github.com/pydata/xarray/issues/2436).
>> 
>>     This leads to problems for example with cf-xarray when trying to detect coordinates or bounds, but also leads to problems related to the time axis encoding apparently (as seen in the linked issue). I managed at least to avoid the problems for cf-xarray bounds and coordinates detection by using the decode functionality of xarray only after the datasets have been read in (leaving however the unnecessary time dimension in place ...):
>> 
>>     ds = xarray.open_mfdataset("/path/to/files/*.nc")
>>     ds = xarray.decode_cf(ds, decode_coords="all")
>
> 
> I think we might need to use the preprocess function to save the .encoding dict, or some other workaround like from one of the comments.
> 

The temporal averaging APIs require decoded time coordinates and the `ds.time.encoding` dict to be populated with the `"calendar"` attr.

### What did you expect to happen?

The `.encoding` attributes should remain persistent after running `xcdat.open_mfdataset()`. 

This does not affect `xcdat.open_dataset()`.

### Minimal Complete Verifiable Example

```python
>> import xcdat

>> ds = xcdat.open_mfdataset(
    "/p/css03/cmip5_css02/data/cmip5/output1/ICHEC/EC-EARTH/historical/mon/atmos/Amon/r1i1p1/tas/1/*.nc")

>> ds.time.encoding

{}
```


### Relevant log output

_No response_

### Anything else we need to know?

Possible Solutions
1. Store .encoding and copy it back to the final merged dataset (maybe with `preprocess kwarg`)
  - https://github.com/pydata/xarray/issues/2436#issuecomment-610008589

Related Issues

### Environment

xcdat 0.3.0

INSTALLED VERSIONS
------------------
commit: None
python: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:56:21) 
[GCC 10.3.0]
python-bits: 64
OS: Linux
OS-release: 3.10.0-1160.45.1.el7.x86_64
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.12.1
libnetcdf: 4.8.1

xarray: 2022.3.0
pandas: 1.4.3
numpy: 1.22.4
scipy: 1.9.0
netCDF4: 1.5.8
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: None
cftime: 1.6.0
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: 2022.6.1
distributed: 2022.6.1
matplotlib: 3.5.2
cartopy: None
seaborn: None
numbagg: None
fsspec: 2022.7.1
cupy: None
pint: None
sparse: 0.13.0
setuptools: 63.3.0
pip: 22.1.2
conda: None
pytest: 7.1.2
IPython: 8.4.0
sphinx: 4.5.0

Discussion:
tomvothecoder: > Since the fix is only a couple lines, I wonder if it makes sense to go with [Option 2](https://github.com/xCDAT/xcdat/pull/309/files#r943756817) and just fix this in the xcdat.open_mfdataset function? It might make sense to keep it in _postprocess_dataset() if there were instances where the encoding wasn't handled correctly during an xcdat.open_dataset, but I don't think this should 

Thanks for the input. I agree, `_postprocess_dataset()` is a reusable function so the logic for time encoding should be coupled to `open_mfdataset()` where it is needed. I just pushed a commit to follow option 2.

> I don't view this as an urgent PR – so it might also make sense to fix the broken test (which > looks like a surprising amount of work?).

I marked this as critical because it `TemporalAccessor` requires `.encoding` attrs (specifically "calendar"), which is not properly set with `xr.open_mfdataset()`. 

A possible quick fix to the test is to have two minimal, dummy `.nc` dataset with just a time axis and coordinates. It probably needs to be generated outside of xarray. I'll figure this out, and then merge this PR.
codecov-commenter: # [Codecov](https://codecov.io/gh/xCDAT/xcdat/pull/309?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) Report
> Merging [#309](https://codecov.io/gh/xCDAT/xcdat/pull/309?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) (86760e7) into [main](https://codecov.io/gh/xCDAT/xcdat/commit/4c51879a65992e08b29eabd2d31f7332aca78e61?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) (4c51879) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff            @@
##              main      #309   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           14        14           
  Lines         1170      1187   +17     
=========================================
+ Hits          1170      1187   +17     
```


| [Impacted Files](https://codecov.io/gh/xCDAT/xcdat/pull/309?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) | Coverage Δ | |
|---|---|---|
| [xcdat/dataset.py](https://codecov.io/gh/xCDAT/xcdat/pull/309/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT#diff-eGNkYXQvZGF0YXNldC5weQ==) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT)


Patch:
diff --git a/xcdat/dataset.py b/xcdat/dataset.py
index b778441..d661eb8 100644
--- a/xcdat/dataset.py
+++ b/xcdat/dataset.py
@@ -3,7 +3,7 @@ import pathlib
 from datetime import datetime
 from functools import partial
 from glob import glob
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Hashable, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import xarray as xr
@@ -14,7 +14,7 @@ from xarray.coding.times import convert_times
 
 from xcdat import bounds  # noqa: F401
 from xcdat.axis import center_times as center_times_func
-from xcdat.axis import get_axis_coord, swap_lon_axis
+from xcdat.axis import get_axis_coord, get_axis_dim, swap_lon_axis
 from xcdat.logger import setup_custom_logger
 
 logger = setup_custom_logger(__name__)
@@ -22,6 +22,16 @@ logger = setup_custom_logger(__name__)
 #: List of non-CF compliant time units.
 NON_CF_TIME_UNITS: List[str] = ["months", "years"]
 
+# Type annotation for the `paths` arg.
+Paths = Union[
+    str,
+    pathlib.Path,
+    List[str],
+    List[pathlib.Path],
+    List[List[str]],
+    List[List[pathlib.Path]],
+]
+
 
 def open_dataset(
     path: str,
@@ -87,10 +97,9 @@ def open_dataset(
     """
     if decode_times:
         cf_compliant_time: Optional[bool] = _has_cf_compliant_time(path)
+        # xCDAT attempts to decode non-CF compliant time coordinates.
         if cf_compliant_time is False:
-            # XCDAT handles decoding time values with non-CF units.
             ds = xr.open_dataset(path, decode_times=False, **kwargs)
-            # attempt to decode non-cf-compliant time axis
             ds = decode_non_cf_time(ds)
         else:
             ds = xr.open_dataset(path, decode_times=True, **kwargs)
@@ -103,14 +112,7 @@ def open_dataset(
 
 
 def open_mfdataset(
-    paths: Union[
-        str,
-        pathlib.Path,
-        List[str],
-        List[pathlib.Path],
-        List[List[str]],
-        List[List[pathlib.Path]],
-    ],
+    paths: Paths,
     data_var: Optional[str] = None,
     add_bounds: bool = True,
     decode_times: bool = True,
@@ -201,10 +203,19 @@ def open_mfdataset(
 
     .. [2] https://xarray.pydata.org/en/stable/generated/xarray.open_mfdataset.html
     """
+    # `xr.open_mfdataset()` drops the time coordinates encoding dictionary if
+    # multiple files are merged with `decode_times=True` (refer to
+    # https://github.com/pydata/xarray/issues/2436). The workaround is to store
+    # the time encoding from the first dataset as a variable, and add the time
+    # encoding back to final merged dataset.
+    time_encoding = None
+
     if decode_times:
+        time_encoding = _keep_time_encoding(paths)
+
         cf_compliant_time: Optional[bool] = _has_cf_compliant_time(paths)
-        # XCDAT handles decoding time values with non-CF units using the
-        # preprocess kwarg.
+        # xCDAT attempts to decode non-CF compliant time coordinates using the
+        # preprocess keyword arg with `xr.open_mfdataset()`.
         if cf_compliant_time is False:
             decode_times = False
             preprocess = partial(_preprocess_non_cf_dataset, callable=preprocess)
@@ -218,6 +229,12 @@ def open_mfdataset(
     )
     ds = _postprocess_dataset(ds, data_var, center_times, add_bounds, lon_orient)
 
+    if time_encoding is not None:
+        time_dim = get_axis_dim(ds, "T")
+        ds[time_dim].encoding = time_encoding
+        # Update "original_shape" to reflect the final time coordinates shape.
+        ds[time_dim].encoding["original_shape"] = ds[time_dim].shape
+
     return ds
 
 
@@ -405,16 +422,43 @@ def decode_non_cf_time(dataset: xr.Dataset) -> xr.Dataset:
     return ds
 
 
-def _has_cf_compliant_time(
-    path: Union[
-        str,
-        pathlib.Path,
-        List[str],
-        List[pathlib.Path],
-        List[List[str]],
-        List[List[pathlib.Path]],
-    ]
-) -> Optional[bool]:
+def _keep_time_encoding(paths: Paths) -> Dict[Hashable, Any]:
+    """
+    Returns the time encoding attributes from the first dataset in a list of
+    paths.
+
+    Time encoding information is critical for several xCDAT operations such as
+    temporal averaging (e.g., uses the "calendar" attr). This function is a
+    workaround to the undesired xarray behavior/quirk with
+    `xr.open_mfdataset()`, which drops the `.encoding` dict from the final
+    merged dataset (refer to https://github.com/pydata/xarray/issues/2436).
+
+    Parameters
+    ----------
+    paths: Paths
+        The paths to the dataset(s).
+
+    Returns
+    -------
+    Dict[Hashable, Any]
+        The time encoding dictionary.
+    """
+    first_path = _get_first_path(paths)
+
+    # xcdat.open_dataset() is called instead of xr.open_dataset() because
+    # we want to handle decoding non-CF compliant as well.
+    # FIXME: Remove `type: ignore` comment after properly handling the type
+    # annotations in `_get_first_path()`.
+    ds = open_dataset(first_path, decode_times=True, add_bounds=False)  # type: ignore
+    time_coord = get_axis_coord(ds, "T")
+
+    time_encoding = time_coord.encoding
+    time_encoding["source"] = paths
+
+    return time_coord.encoding
+
+
+def _has_cf_compliant_time(paths: Paths) -> Optional[bool]:
     """Checks if a dataset has time coordinates with CF compliant units.
 
     If the dataset does not contain a time dimension, None is returned.
@@ -444,19 +488,8 @@ def _has_cf_compliant_time(
     performance because it is slower to combine all files then check for CF
     compliance.
     """
-    first_file: Optional[Union[pathlib.Path, str]] = None
-
-    if isinstance(path, str) and "*" in path:
-        first_file = glob(path)[0]
-    elif isinstance(path, str) or isinstance(path, pathlib.Path):
-        first_file = path
-    elif isinstance(path, list):
-        if any(isinstance(sublist, list) for sublist in path):
-            first_file = path[0][0]  # type: ignore
-        else:
-            first_file = path[0]  # type: ignore
-
-    ds = xr.open_dataset(first_file, decode_times=False)
+    first_path = _get_first_path(paths)
+    ds = xr.open_dataset(first_path, decode_times=False)
 
     if ds.cf.dims.get("T") is None:
         return None
@@ -474,6 +507,38 @@ def _has_cf_compliant_time(
     return cf_compliant
 
 
+def _get_first_path(path: Paths) -> Optional[Union[pathlib.Path, str]]:
+    """Returns the first path from a list of paths.
+
+    Parameters
+    ----------
+    path : Paths
+        A list of paths.
+
+    Returns
+    -------
+    str
+        Returns the first path from a list of paths.
+    """
+    # FIXME: This function should throw an exception if the first file
+    # is not a supported type.
+    # FIXME: The `type: ignore` comments should be removed after properly
+    # handling the types.
+    first_file: Optional[Union[pathlib.Path, str]] = None
+
+    if isinstance(path, str) and "*" in path:
+        first_file = glob(path)[0]
+    elif isinstance(path, str) or isinstance(path, pathlib.Path):
+        first_file = path
+    elif isinstance(path, list):
+        if any(isinstance(sublist, list) for sublist in path):
+            first_file = path[0][0]  # type: ignore
+        else:
+            first_file = path[0]  # type: ignore
+
+    return first_file
+
+
 def _postprocess_dataset(
     dataset: xr.Dataset,
     data_var: Optional[str] = None,


Test patch:
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 571bb33..82d55f8 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -305,6 +305,38 @@ class TestOpenMfDataset:
         self.file_path1 = f"{dir}/file1.nc"
         self.file_path2 = f"{dir}/file2.nc"
 
+    def test_mfdataset_keeps_time_encoding_dict(self):
+        ds1 = generate_dataset(cf_compliant=True, has_bounds=True)
+        ds1.to_netcdf(self.file_path1)
+
+        # Create another dataset that extends the time coordinates by 1 value,
+        # to mimic a multifile dataset.
+        ds2 = generate_dataset(cf_compliant=True, has_bounds=True)
+        ds2 = ds2.isel(dict(time=slice(0, 1)))
+        ds2["time"].values[:] = np.array(
+            ["2002-01-16T12:00:00.000000000"],
+            dtype="datetime64[ns]",
+        )
+        ds2.to_netcdf(self.file_path2)
+
+        result = open_mfdataset([self.file_path1, self.file_path2], decode_times=True)
+        expected = ds1.merge(ds2)
+
+        assert result.identical(expected)
+
+        # We mainly care for the "source" and "original_shape" attrs (updated
+        # internally by xCDAT), and the "calendar" and "units" attrs. We don't
+        # perform equality assertion on the entire time `.encoding` dict because
+        # there might be different encoding attributes added or removed between
+        # xarray versions (e.g., "bzip2", "ztsd", "blosc", and "szip" are added
+        # in v2022.06.0), which makes that assertion fragile.
+        paths = result.time.encoding["source"]
+        assert self.file_path1 in paths[0]
+        assert self.file_path2 in paths[1]
+        assert result.time.encoding["original_shape"] == (16,)
+        assert result.time.encoding["calendar"] == "standard"
+        assert result.time.encoding["units"] == "days since 2000-01-01"
+
     def test_non_cf_compliant_time_is_not_decoded(self):
         ds1 = generate_dataset(cf_compliant=False, has_bounds=True)
         ds1.to_netcdf(self.file_path1)