Made the get_2d_aggregate helper function general
@@ -8,6 +8,12 @@ try:
except:
dask = None
+try:
+ import xarray as xr
+except:
+ xr = None
+
+
def toarray(v, index_value=False):
"""
Interface helper function to turn dask Arrays into numpy arrays as
@@ -48,17 +54,20 @@ def reduce_fn(x):
def get_2d_aggregate(obj):
"""
- Generates a 2D aggregate by inserting NaNs at all cross-product
- locations that do not already have a value assigned, creating
- values that can be shaped into a 2D array.
+ Generates a categorical 2D aggregate by inserting NaNs at all
+ cross-product locations that do not already have a value assigned.
+ Returns a 2D gridded Dataset object.
"""
if obj.interface.gridded:
return obj
elif obj.ndims > 2:
raise Exception("Cannot aggregate more than two dimensions")
- xdim, ydim = obj.dimensions(label=True)[:2]
- d1keys = obj.dimension_values(0, False)
- d2keys = obj.dimension_values(1, False)
+
+ dims = obj.dimensions(label=True)
+ xdim, ydim = dims[:2]
+ nvdims = len(dims) - 2
+ d1keys = obj.dimension_values(xdim, False)
+ d2keys = obj.dimension_values(ydim, False)
is_sorted = np.array_equal(np.sort(d1keys), d1keys)
if is_sorted:
@@ -70,16 +79,20 @@ def get_2d_aggregate(obj):
if is_sorted:
d1keys, d2keys = np.sort(d1keys), np.sort(d2keys)
- coords = [(d1, d2) + (np.NaN,)*len(obj.vdims) for d2 in d2keys for d1 in d1keys]
+ coords = [(d1, d2) + (np.NaN,)*nvdims for d2 in d2keys for d1 in d1keys]
+
dtype = 'dataframe' if pd else 'dictionary'
dense_data = Dataset(coords, kdims=obj.kdims, vdims=obj.vdims, datatype=[dtype])
concat_data = obj.interface.concatenate([dense_data, Dataset(obj)], datatype=dtype)
- agg = concat_data.aggregate(obj.kdims, reduce_fn)
+ agg = concat_data.reindex([xdim, ydim]).aggregate([xdim, ydim], reduce_fn)
shape = (len(d2keys), len(d1keys))
grid_data = {xdim: d1keys, ydim: d2keys}
- for vdim in agg.vdims:
+
+ for vdim in dims[2:]:
data = agg.dimension_values(vdim).reshape(shape)
data = np.ma.array(data, mask=np.logical_not(np.isfinite(data)))
- grid_data[vdim.name] = data
- return obj.clone(grid_data)
+ grid_data[vdim] = data
+
+ grid_type = 'xarray' if xr else 'grid'
+ return agg.clone(grid_data, datatype=[grid_type])
