Internal refactor of XArray, with a new CoordXArray subtype

This allows us to simplify our internal model for XArray (it always cached
internally as a base ndarray) and supports some previously tricky aspects
involving pandas.Index objects. Noteably:
1. The dtype of arrays stored as pandas.Index objects can now be faithfully
   saved and restored. Doing math with XArray objects always yields objects
   with the right dtype, so `ds['latitude'] + 1` has dtype=float, not
   dtype=object.
2. It's no longer necessary to load index data into memory upon creating a new
   Dataset. Instead, the index data can be loaded on demand.
3. `var.data` is always an ndarray. `var.index` is always a pandas.Index.

Related issues: #17, #39, #40.

Author: shoyer

src/xray/__init__.py

@@ -1,4 +1,4 @@
-from .xarray import XArray, broadcast_xarrays
+from .xarray import as_xarray, XArray, CoordXArray, broadcast_xarrays
 from .dataset import Dataset, open_dataset
 from .dataset_array import DatasetArray, align
 from .utils import (orthogonal_indexer, decode_cf_datetime, encode_cf_datetime,

src/xray/common.py

@@ -15,25 +15,6 @@ def func(self, dimension=cls._reduce_dimension_default,
 
 
 class AbstractArray(ImplementsReduce):
-    @property
-    def dtype(self):
-        return self._data.dtype
-
-    @property
-    def shape(self):
-        return self._data.shape
-
-    @property
-    def size(self):
-        return self._data.size
-
-    @property
-    def ndim(self):
-        return self._data.ndim
-
-    def __len__(self):
-        return len(self._data)
-
     def __nonzero__(self):
         return bool(self.data)
 

src/xray/conventions.py

@@ -260,20 +260,6 @@ def encode_cf_variable(array):
             data, encoding.pop('units', None), encoding.pop('calendar', None))
         attributes['units'] = units
         attributes['calendar'] = calendar
-    elif data.dtype == np.dtype('O'):
-        # Unfortunately, pandas.Index arrays often have dtype=object even if
-        # they were created from an array with a sensible datatype (e.g.,
-        # pandas.Float64Index always has dtype=object for some reason). Because
-        # we allow for doing math with coordinates, these object arrays can
-        # propagate onward to other variables, which is why we don't only apply
-        # this check to XArrays with data that is a pandas.Index.
-        # Accordingly, we convert object arrays to the type of their first
-        # variable.
-        dtype = np.array(data.reshape(-1)[0]).dtype
-        # N.B. the "astype" call below will fail if data cannot be cast to the
-        # type of its first element (which is probably the only sensible thing
-        # to do).
-        data = np.asarray(data).astype(dtype)
 
     def get_to(source, dest, k):
         v = source.get(k)

src/xray/dataset.py

@@ -76,10 +76,10 @@ def _get_virtual_variable(self, key):
             if ref_var in self._datetimeindices():
                 if suffix == 'season':
                     # seasons = np.array(['DJF', 'MAM', 'JJA', 'SON'])
-                    month = self[ref_var].data.month
+                    month = self[ref_var].index.month
                     data = (month // 3) % 4 + 1
                 else:
-                    data = getattr(self[ref_var].data, suffix)
+                    data = getattr(self[ref_var].index, suffix)
                 return xarray.XArray(self[ref_var].dimensions, data)
         raise KeyError('virtual variable %r not found' % key)
 
@@ -130,14 +130,15 @@ def __init__(self, variables=None, attributes=None, decode_cf=False):
 
     def _as_variable(self, name, var, decode_cf=False):
         if isinstance(var, DatasetArray):
-            var = var.array
-        if not isinstance(var, xarray.XArray):
+            var = xarray.as_xarray(var)
+        elif not isinstance(var, xarray.XArray):
             try:
                 var = xarray.XArray(*var)
             except TypeError:
                 raise TypeError('Dataset variables must be of type '
                                 'DatasetArray or XArray, or a sequence of the '
-                                'form (dimensions, data[, attributes])')
+                                'form (dimensions, data[, attributes, '
+                                'encoding])')
         # this will unmask and rescale the data as well as convert
         # time variables to datetime indices.
         if decode_cf:
@@ -147,9 +148,7 @@ def _as_variable(self, name, var, decode_cf=False):
             if var.ndim != 1:
                 raise ValueError('a coordinate variable must be defined with '
                                  '1-dimensional data')
-            # create a new XArray object on which to modify the data
-            var = xarray.XArray(var.dimensions, pd.Index(var.data),
-                                var.attributes, encoding=var.encoding)
+            var = var.to_coord()
         return var
 
     def set_variables(self, variables, decode_cf=False):
@@ -487,7 +486,7 @@ def labeled_by(self, **indexers):
         Dataset.indexed_by
         Array.indexed_by
         """
-        return self.indexed_by(**remap_loc_indexers(self.variables, indexers))
+        return self.indexed_by(**remap_loc_indexers(self, indexers))
 
     def renamed(self, name_dict):
         """Returns a new object with renamed variables and dimensions.

src/xray/dataset_array.py

@@ -66,28 +66,47 @@ def __init__(self, dataset, focus):
         self.focus = focus
 
     @property
-    def array(self):
+    def variable(self):
         return self.dataset.variables[self.focus]
-    @array.setter
-    def array(self, value):
+    @variable.setter
+    def variable(self, value):
         self.dataset[self.focus] = value
 
-    # _data is necessary for AbstractArray
     @property
-    def _data(self):
-        return self.array._data
+    def dtype(self):
+        return self.variable.dtype
+
+    @property
+    def shape(self):
+        return self.variable.shape
+
+    @property
+    def size(self):
+        return self.variable.size
+
+    @property
+    def ndim(self):
+        return self.variable.ndim
+
+    def __len__(self):
+        return len(self.variable)
 
     @property
     def data(self):
-        """The array's data as a numpy.ndarray"""
-        return self.array.data
+        """The variables's data as a numpy.ndarray"""
+        return self.variable.data
     @data.setter
     def data(self, value):
-        self.array.data = value
+        self.variable.data = value
+
+    @property
+    def index(self):
+        """The variable's data as a pandas.Index"""
+        return self.variable.index
 
     @property
     def dimensions(self):
-        return self.array.dimensions
+        return self.variable.dimensions
 
     def _key_to_indexers(self, key):
         return OrderedDict(
@@ -107,7 +126,7 @@ def __setitem__(self, key, value):
             self.dataset[key] = value
         else:
             # orthogonal array indexing
-            self.array[key] = value
+            self.variable[key] = value
 
     def __delitem__(self, key):
         del self.dataset[key]
@@ -127,11 +146,11 @@ def __iter__(self):
 
     @property
     def attributes(self):
-        return self.array.attributes
+        return self.variable.attributes
 
     @property
     def encoding(self):
-        return self.array.encoding
+        return self.variable.encoding
 
     @property
     def variables(self):
@@ -178,7 +197,7 @@ def indexed_by(self, **indexers):
         if self.focus not in ds:
             # always keep focus variable in the dataset, even if it was
             # unselected because indexing made it a scaler
-            ds[self.focus] = self.array.indexed_by(**indexers)
+            ds[self.focus] = self.variable.indexed_by(**indexers)
         return type(self)(ds, self.focus)
 
     def labeled_by(self, **indexers):
@@ -236,7 +255,7 @@ def refocus(self, new_var, name=None):
         If `new_var` is a dataset array, its contents will be merged in.
         """
         if not hasattr(new_var, 'dimensions'):
-            new_var = type(self.array)(self.array.dimensions, new_var)
+            new_var = type(self.variable)(self.variable.dimensions, new_var)
         if self.focus not in self.dimensions:
             # only unselect the focus from the dataset if it isn't a coordinate
             # variable
@@ -301,7 +320,7 @@ def transpose(self, *dimensions):
         numpy.transpose
         Array.transpose
         """
-        return self.refocus(self.array.transpose(*dimensions), self.focus)
+        return self.refocus(self.variable.transpose(*dimensions), self.focus)
 
     def squeeze(self, dimension=None):
         """Return a new DatasetArray object with squeezed data.
@@ -361,7 +380,7 @@ def reduce(self, func, dimension=None, axis=None, **kwargs):
             DatasetArray with this object's array replaced with an array with
             summarized data and the indicated dimension(s) removed.
         """
-        var = self.array.reduce(func, dimension, axis, **kwargs)
+        var = self.variable.reduce(func, dimension, axis, **kwargs)
         drop = set(self.dimensions) - set(var.dimensions)
         # For now, take an aggressive strategy of removing all variables
         # associated with any dropped dimensions
@@ -495,13 +514,13 @@ def to_series(self):
         return pd.Series(self.data.reshape(-1), index=index, name=self.focus)
 
     def __array_wrap__(self, obj, context=None):
-        return self.refocus(self.array.__array_wrap__(obj, context))
+        return self.refocus(self.variable.__array_wrap__(obj, context))
 
     @staticmethod
     def _unary_op(f):
         @functools.wraps(f)
         def func(self, *args, **kwargs):
-            return self.refocus(f(self.array, *args, **kwargs),
+            return self.refocus(f(self.variable, *args, **kwargs),
                                 self.focus + '_' + f.__name__)
         return func
 
@@ -523,12 +542,12 @@ def func(self, other):
             ds = self.unselected()
             if hasattr(other, 'unselected'):
                 ds.merge(other.unselected(), inplace=True)
-            other_array = getattr(other, 'array', other)
+            other_array = getattr(other, 'variable', other)
             other_focus = getattr(other, 'focus', 'other')
             focus = self.focus + '_' + f.__name__ + '_' + other_focus
-            ds[focus] = (f(self.array, other_array)
+            ds[focus] = (f(self.variable, other_array)
                          if not reflexive
-                         else f(other_array, self.array))
+                         else f(other_array, self.variable))
             return type(self)(ds, focus)
         return func
 
@@ -537,8 +556,8 @@ def _inplace_binary_op(f):
         @functools.wraps(f)
         def func(self, other):
             self._check_coordinates_compat(other)
-            other_array = getattr(other, 'array', other)
-            self.array = f(self.array, other_array)
+            other_array = getattr(other, 'variable', other)
+            self.variable = f(self.variable, other_array)
             if hasattr(other, 'unselected'):
                 self.dataset.merge(other.unselected(), inplace=True)
             return self
@@ -555,8 +574,9 @@ def align(array1, array2):
     # TODO: automatically align when doing math with arrays, or better yet
     # calculate the union of the indices and fill in the mis-aligned data with
     # NaN.
-    overlapping_coords = {k: (array1.coordinates[k].data
-                              & array2.coordinates[k].data)
+    # TODO: generalize this function to any number of arguments
+    overlapping_coords = {k: (array1.coordinates[k].index
+                              & array2.coordinates[k].index)
                           for k in array1.coordinates
                           if k in array2.coordinates}
     return tuple(ar.labeled_by(**overlapping_coords)

src/xray/utils.py

@@ -6,6 +6,8 @@
 import numpy as np
 import pandas as pd
 
+import xarray
+
 
 def expanded_indexer(key, ndim):
     """Given a key for indexing an ndarray, return an equivalent key which is a
@@ -86,12 +88,12 @@ def all_full_slices(key_index):
 
 
 def remap_loc_indexers(indices, indexers):
-    """Given mappings of indices and label based indexers, return equivalent
-    location based indexers.
+    """Given mappings of XArray indices and label based indexers, return
+    equivalent location based indexers.
     """
     new_indexers = OrderedDict()
     for dim, loc in indexers.iteritems():
-        index = indices[dim].data
+        index = indices[dim].index
         if isinstance(loc, slice):
             indexer = index.slice_indexer(loc.start, loc.stop, loc.step)
         else:
@@ -201,11 +203,12 @@ def encode_cf_datetime(dates, units=None, calendar=None):
             and np.issubdtype(dates.dtype, np.datetime64)):
         # for now, don't bother doing any trickery like decode_cf_datetime to
         # convert dates to numbers faster
-        dates = dates.astype(datetime)
+        # TODO: don't use pandas.DatetimeIndex to do the conversion
+        dates = pd.Index(dates.reshape(-1)).to_pydatetime().reshape(dates.shape)
 
     if hasattr(dates, 'ndim') and dates.ndim == 0:
-        # unpack dates because date2num doesn't like 0-dimensional arguments
-        dates = dates[()]
+        # date2num doesn't like 0-dimensional arguments
+        dates = dates.item()
 
     num = nc4.date2num(dates, units, calendar)
     return (num, units, calendar)
@@ -235,33 +238,40 @@ def xarray_equal(v1, v2, rtol=1e-05, atol=1e-08):
     This function is necessary because `v1 == v2` for XArrays and DatasetArrays
     does element-wise comparisions (like numpy.ndarrays).
     """
+    v1, v2 = map(xarray.as_xarray, [v1, v2])
     if (v1.dimensions == v2.dimensions
-        and dict_equal(v1.attributes, v2.attributes)):
-        try:
+            and dict_equal(v1.attributes, v2.attributes)):
+        if v1._data is v2._data:
             # if _data is identical, skip checking arrays by value
-            if v1._data is v2._data:
-                return True
-        except AttributeError:
-            # _data is not part of the public interface, so it's okay if its
-            # missing
-            pass
-
-        def is_floating(arr):
-            return np.issubdtype(arr.dtype, float)
-
-        data1 = v1.data
-        data2 = v2.data
-        if hasattr(data1, 'equals'):
-            # handle pandas.Index objects
-            return data1.equals(data2)
-        elif is_floating(data1) or is_floating(data2):
-            return allclose_or_equiv(data1, data2, rtol=rtol, atol=atol)
+            return True
         else:
-            return np.array_equal(data1, data2)
+            def is_floating(arr):
+                return np.issubdtype(arr.dtype, float)
+
+            data1 = v1.data
+            data2 = v2.data
+            if is_floating(data1) or is_floating(data2):
+                return allclose_or_equiv(data1, data2, rtol=rtol, atol=atol)
+            else:
+                return np.array_equal(data1, data2)
     else:
         return False
 
 
+def safe_cast_to_index(array):
+    """Given an array, safely cast it to a pandas.Index
+
+    Unlike pandas.Index, if the array has dtype=object or dtype=timedelta64,
+    this function will not attempt to do automatic type conversion but will
+    always return an index with dtype=object.
+    """
+    kwargs = {}
+    if isinstance(array, np.ndarray):
+        if array.dtype == object or array.dtype == np.timedelta64:
+            kwargs['dtype'] = object
+    return pd.Index(array, **kwargs)
+
+
 def update_safety_check(first_dict, second_dict, compat=operator.eq):
     """Check the safety of updating one dictionary with another.