Merge pull request #76 from matt-long/master

sel_time method added; change behavior of time_manager

docs/source/api.rst

@@ -40,6 +40,7 @@ Utilities
 .. autosummary::
    utils.time.compute_time_var
    utils.time.uncompute_time_var
+   utils.time.sel_time
 
 
 
@@ -80,3 +81,4 @@ Utilities
 
 .. autofunction:: compute_time_var
 .. autofunction:: uncompute_time_var
+.. autofunction:: sel_time

esmlab/climatology.py

@@ -38,6 +38,7 @@ def compute_mon_climatology(dset, time_coord_name=None):
     """
 
     tm = time_manager(dset, time_coord_name)
+    dset = tm.compute_time()
     time_coord_name = tm.time_coord_name
 
     static_variables = get_static_variables(dset, time_coord_name)
@@ -114,6 +115,7 @@ def compute_mon_anomaly(dset, slice_mon_clim_time=None, time_coord_name=None):
     """
 
     tm = time_manager(dset, time_coord_name)
+    dset = tm.compute_time()
     time_coord_name = tm.time_coord_name
 
     static_variables = get_static_variables(dset, time_coord_name)
@@ -173,6 +175,7 @@ def compute_ann_mean(dset, weights=None, time_coord_name=None):
     """
 
     tm = time_manager(dset, time_coord_name)
+    dset = tm.compute_time()
     time_coord_name = tm.time_coord_name
 
     static_variables = get_static_variables(dset, time_coord_name)

esmlab/utils/time.py

@@ -17,7 +17,7 @@ class time_manager(object):
     """Class to support managing the time coordinate of datasets.
     """
 
-    def __init__(self, ds, time_coord_name=None):
+    def __init__(self, ds, time_coord_name=None, year_offset=None):
         """
         Parameters
         ----------
@@ -27,16 +27,19 @@ def __init__(self, ds, time_coord_name=None):
            Name of the time coordinate of `ds`; if not provided, the
            code will attempt to infer it.
         """
-        self._ds = ds
+        self._ds = xr.Dataset(ds)
         self.orig_time_coord_name = None
         self.orig_time_coord_decoded = None
+        self._time_computed = False
+
         if time_coord_name is None:
             self.time_coord_name = self._infer_time_coord_name()
         else:
             self.time_coord_name = time_coord_name
 
+        self.year_offset = year_offset
+
         self._infer_time_bound_var()
-        self._compute_time()
 
     @property
     def time(self):
@@ -116,7 +119,7 @@ def get_time_undecoded(self):
             self.orig_time_coord_decoded = True
         return time
 
-    def get_time_decoded(self, midpoint=True, year_offset=np.nan):
+    def get_time_decoded(self, midpoint=True):
         """Return time decoded.
         """
         # to compute a time midpoint, we need a time_bound variable
@@ -130,7 +133,7 @@ def get_time_decoded(self, midpoint=True, year_offset=np.nan):
             # if time has already been decoded and there's no year_offset,
             # just return the time as is
             if self.time.dtype == np.dtype("O"):
-                if np.isnan(year_offset):
+                if self.year_offset is None:
                     return time_values
 
                 # if we need to un-decode time to apply the year_offset,
@@ -141,9 +144,9 @@ def get_time_decoded(self, midpoint=True, year_offset=np.nan):
             else:
                 time_values = self.time
 
-        if not np.isnan(year_offset):
+        if self.year_offset is not None:
             time_values += cftime.date2num(
-                datetime(int(year_offset), 1, 1),
+                datetime(int(self.year_offset), 1, 1),
                 units=self.time_attrs["units"],
                 calendar=self.time_attrs["calendar"],
             )
@@ -156,9 +159,13 @@ def get_time_decoded(self, midpoint=True, year_offset=np.nan):
         )
         return time_out
 
-    def _compute_time(self):
+    def compute_time(self):
         """Compute the mid-point of the time bounds.
         """
+
+        if self._time_computed:
+            return self._ds
+
         self.orig_time_coord_name = "t" + uuid.uuid4().hex
 
         self._ds[self.orig_time_coord_name] = self.get_time_undecoded()
@@ -171,9 +178,15 @@ def _compute_time(self):
 
         self._ds[self.time_coord_name].values = groupby_coord.values
 
+        self._time_computed = True
+
+        return self._ds
+
     def restore_dataset(self, ds):
         """Return the original time variable.
         """
+        if not self._time_computed:
+            raise ValueError("time was not computed; cannot restore dataset")
         time_values = ds[self.orig_time_coord_name].values
         if self.orig_time_coord_decoded:
             time_values = xr.CFTimeIndex(
@@ -230,7 +243,7 @@ def time_year_to_midyeardate(ds, time_coord_name):
     return ds
 
 
-def compute_time_var(ds, midpoint=True, year_offset=np.nan):
+def compute_time_var(ds, midpoint=True, year_offset=None):
     """Compute the time coordinate of a dataset.
 
     Parameters
@@ -247,8 +260,10 @@ def compute_time_var(ds, midpoint=True, year_offset=np.nan):
     ds : `xarray.Dataset`
        The dataset with time coordinate modified.
     """
-    tm = time_manager(ds)
-    ds[tm.time_coord_name] = tm.get_time_decoded(midpoint, year_offset)
+    tm = time_manager(ds, year_offset=year_offset)
+    ds = tm.compute_time()
+    ds[tm.time_coord_name] = tm.get_time_decoded(midpoint)
+    ds = ds.drop(tm.orig_time_coord_name)
     return ds
 
 
@@ -266,5 +281,31 @@ def uncompute_time_var(ds):
        The dataset with time coordinate modified.
     """
     tm = time_manager(ds)
+    ds = tm.compute_time()
     ds[tm.time_coord_name] = tm.get_time_undecoded()
+    ds = ds.drop(tm.orig_time_coord_name)
     return ds
+
+
+def sel_time(ds, indexer_val, year_offset=None):
+    """Return dataset truncated to specified time range.
+
+    Parameters
+    ----------
+    ds : `xarray.Dataset`
+       The dataset on which to operate
+    indexer_val : scalar, slice, or array_like
+       value passed to ds.isel(**{time_coord_name: indexer_val})
+    year_offset : numeric, optional
+       Integer year by which to offset the time axis.
+
+    Returns
+    -------
+    ds : `xarray.Dataset`
+       The dataset with time coordinate truncated.
+    """
+
+    tm = time_manager(ds, year_offset=year_offset)
+    ds = tm.compute_time()
+    ds = ds.sel(**{tm.time_coord_name: indexer_val})
+    return tm.restore_dataset(ds)

tests/utils/test_time.py

@@ -43,3 +43,10 @@ def test_uncompute_time_var(dset, time_coord_name="time"):
     assert dset[time_coord_name].dtype == np.dtype("O")
     dset_with_uncomputed_time = utils.time.uncompute_time_var(dset)
     assert np.issubdtype(dset_with_uncomputed_time[time_coord_name].dtype, np.number)
+
+
+def test_sel_time(dset):
+    dset = utils.time.sel_time(
+        dset, indexer_val=slice("1850-01-01", "1850-12-31"), year_offset=1850
+    )
+    assert len(dset.time) == 12