Commit of dask/xarray issue to new branch

Author: dwest77a

XarrayActive/active_chunk.py

@@ -32,8 +32,6 @@ def _set_active_options(self, chunks={}, chunk_limits=True):
         self._active_chunks = chunks
         self._chunk_limits = chunk_limits
 
-
-
 # Holds all Active routines.
 class ActiveChunk:
 
@@ -43,23 +41,26 @@ def _post_process_data(self, data):
         # Perform any post-processing steps on the data here
         return data
 
-    def _standard_mean(self, axes=None, skipna=None, **kwargs):
+    def _standard_sum(self, axes=None, skipna=None, **kwargs):
         """
         Standard Mean routine matches the normal routine for dask, required at this
         stage if Active mean not available.
         """
-        size = 1
-        for i in axes:
-            size *= self.shape[i]
 
         arr = np.array(self)
         if skipna:
-            total = np.nanmean(arr, axis=axes, **kwargs) *size
+            total = np.nansum(arr, axis=axes, **kwargs)
         else:
-            total = np.mean(arr, axis=axes, **kwargs) *size
-        return {'n': self._numel(axes=axes), 'total': total}
+            total = np.sum(arr, axis=axes, **kwargs)
+        return total
+    
+    def _standard_max(self, axes=None, skipna=None, **kwargs):
+        return np.max(self, axis=axes)
 
-    def _numel(self, axes=None):
+    def _standard_min(self, axes=None, skipna=None, **kwargs):
+        return np.min(self, axis=axes)
+
+    def _numel(self, method, axes=None):
         if not axes:
             return self.size
 
@@ -72,7 +73,7 @@ def _numel(self, axes=None):
 
         return np.full(newshape, size)
 
-    def active_mean(self, axis=None, skipna=None, **kwargs):
+    def active_method(self, method, axis=None, skipna=None, **kwargs):
         """
         Use PyActiveStorage package functionality to perform mean of this Fragment.
 
@@ -83,47 +84,66 @@ def active_mean(self, axis=None, skipna=None, **kwargs):
         :returns:       A ``duck array`` (numpy-like) with the reduced array or scalar value, 
                         as specified by the axes parameter.
         """
+
+        standard_methods = {
+            'mean': self._standard_sum,
+            'sum' : self._standard_sum,
+            'max' : self._standard_max,
+            'min' : self._standard_min
+        }
+        ret = None
+        n = self._numel(method, axes=axis)
+
         try:
             from activestorage.active import Active
         except ImportError:
             # Unable to import Active package. Default to using normal mean.
             print("ActiveWarning: Unable to import active module - defaulting to standard method.")
-            return self._standard_mean(axes=axis, skipna=skipna, **kwargs)
-            
-        active = Active(self.filename, self.address)
-        active.method = "mean"
-        extent = tuple(self.get_extent())
-        data   = active[extent]
-
-        if axis == None:
-            axis = tuple([i for i in range(self.ndim)])
-
-        n = self._numel(axes=axis)
-
-        if len(axis) == self.ndim:
-
-            t = self._post_process_data(data) * n
+            ret = {
+                'n': n,
+                'total': standard_methods[method](axes=axis, skipna=skipna, **kwargs)
+            }
 
-            r = {
+        if not ret:
+            
+            active = Active(self.filename, self.address)
+            active.method = method
+            extent = tuple(self.get_extent())
+
+            if axis == None:
+                axis = tuple([i for i in range(self.ndim)])
+
+            n = self._numel(method, axes=axis)
+
+            if len(axis) == self.ndim:
+                data   = active[extent]
+                t = self._post_process_data(data) * n
+
+                ret = {
+                    'n': n,
+                    'total': t
+                }
+
+        if not ret:
+            # Experimental Recursive requesting to get each 1D column along the axes being requested.
+            range_recursives = []
+            for dim in range(self.ndim):
+                if dim not in axis:
+                    range_recursives.append(range(extent[dim].start, extent[dim].stop))
+                else:
+                    range_recursives.append(extent[dim])
+            results = np.array(self._get_elements(active, range_recursives, hyperslab=[]))
+
+            t = self._post_process_data(results) * n
+            ret = {
                 'n': n,
                 'total': t
             }
-            return r
-
-        # Experimental Recursive requesting to get each 1D column along the axes being requested.
-        range_recursives = []
-        for dim in range(self.ndim):
-            if dim not in axis:
-                range_recursives.append(range(extent[dim].start, extent[dim].stop))
-            else:
-                range_recursives.append(extent[dim])
-        results = np.array(self._get_elements(active, range_recursives, hyperslab=[]))
-
-        t = self._post_process_data(results) * n
-        return {
-            'n': n,
-            'total': t
-        }
+
+        if method == 'mean':
+            return ret
+        else:
+            return ret['total']/ret['n']
 
     def _get_elements(self, active, recursives, hyperslab=[]):
         dimarray = []
@@ -144,26 +164,3 @@ def _get_elements(self, active, recursives, hyperslab=[]):
                 dimarray.append(self._get_elements(active, recursives[1:], hyperslab=newslab))
 
         return dimarray
-
-def _determine_chunk_space(chunks, shape, dims, chunk_limits=True):
-    
-    if not chunks:
-        return None
-
-    chunk_space = [1 for i in range(len(shape))]
-
-    max_chunks = np.prod(shape)
-    if chunk_limits:
-        max_chunks = int(max_chunks/ 2e6)
-
-    for x, d in enumerate(dims):
-        if d not in chunks:
-            continue
-
-        chunks_in_dim = chunks[d]
-        if chunks_in_dim > max_chunks:
-            chunks_in_dim = max_chunks
-
-        chunk_space[x] = int(shape[x]/chunks[d])
-
-    return tuple(chunk_space)

XarrayActive/active_dask.py

@@ -1,14 +1,43 @@
 import dask.array as da
 from dask.array.reductions import mean_agg, mean_combine
+from dask.utils import deepmap
+from dask.array.core import _concatenate2
+import numpy as np
 
-
-def block_active_mean(arr, *args, **kwargs):
-    if hasattr(arr,'active_mean'):
-        return arr.active_mean(*args, **kwargs)
+def partition_mean(arr, *args, **kwargs):
+    return partition_method(arr, 'mean', *args, **kwargs)
+    
+def partition_max(arr, *args, **kwargs):
+    return partition_method(arr, 'max', *args, **kwargs)
+    
+def partition_min(arr, *args, **kwargs):
+    return partition_method(arr, 'min', *args, **kwargs)
+    
+def partition_sum(arr, *args, **kwargs):
+    return partition_method(arr, 'sum', *args, **kwargs)
+    
+def partition_method(arr, method, *args, **kwargs):
+    if hasattr(arr,'active_method'):
+        return arr.active_method(method,*args, **kwargs)
     else:
         # Here's where barebones Xarray might fall over - may need a non-CFA custom class.
+        x=input('stopped')
         raise NotImplementedError
 
+def general_combine(pairs, axis=None):
+    if not isinstance(pairs, list):
+        pairs = [pairs]
+    return _concatenate2(pairs, axes=axis)
+
+def max_agg(pairs, axis=None, **kwargs):
+    return general_combine(pairs, axis=axis).max(axis=axis, **kwargs)
+
+def min_agg(pairs, axis=None, **kwargs):
+    return general_combine(pairs, axis=axis).min(axis=axis, **kwargs)
+
+def sum_agg(pairs, axis=None, **kwargs):
+    return general_combine(pairs, axis=axis).sum(axis=axis, **kwargs)
+
 class DaskActiveArray(da.Array):
 
     description = "Dask Array Wrapper enabling the use of Active Storage."
@@ -21,17 +50,22 @@ def copy(self):
         """
         Create a new DaskActiveArray instance with all the same parameters as the current instance.
         """
-        return DaskActiveArray(self.dask, self.name, self.chunks, meta=self)
+        copy_arr = DaskActiveArray(self.dask, self.name, self.chunks, meta=self)
+        return copy_arr
 
-    def __getitem__(self, index):
-        """
-        Perform indexing for this ActiveArray. May need to overwrite further if it turns out
-        the indexing is performed **after** the dask `getter` method (i.e if retrieval and indexing
-        are separate items on the dask graph). If this is the case, will need another `from_delayed`
-        and `concatenation` method as used in ``active_mean``.
-        """
-        arr = super().__getitem__(index)
-        return DaskActiveArray(arr.dask, arr.name, arr.chunks, meta=arr)
+    #def __getitem__(self, index):
+    #    """
+    #    Perform indexing for this ActiveArray. May need to overwrite further if it turns out
+    #    the indexing is performed **after** the dask `getter` method (i.e if retrieval and indexing
+    #    are separate items on the dask graph). If this is the case, will need another `from_delayed`
+    #    and `concatenation` method as used in ``active_mean``.
+    #    """
+    #    copy = self.copy()
+
+    #    return copy._getitem(index)
+
+    #def _getitem(self, index):
+    #    return super().__getitem__(index)
 
     def active_mean(self, axis=None, skipna=None):
         """
@@ -48,12 +82,87 @@ def active_mean(self, axis=None, skipna=None):
         """
 
         newarr = da.reduction(
-            self,
-            block_active_mean,
+            self.copy(),
+            partition_mean,
             mean_agg,
             combine=mean_combine,
             axis=axis,
             dtype=self.dtype,
         )
 
         return newarr
+
+    def active_max(self, axis=None, skipna=None):
+        """
+        Perform ``dask delayed`` active mean for each ``dask block`` which corresponds to a single ``chunk``.
+        Combines the results of the dask delayed ``active_max`` operations on each block into a single dask Array,
+        which is then mapped to a new DaskActiveArray object.
+
+        :param axis:        (int) The index of the axis on which to perform the active max.
+
+        :param skipna:      (bool) Skip NaN values when calculating the max.
+
+        :returns:       A new ``DaskActiveArray`` object which has been reduced along the specified axes using
+                        the concatenations of active_means from each chunk.
+        """
+
+        newarr = da.reduction(
+            self.copy(),
+            partition_max,
+            max_agg,
+            combine=max_agg,
+            axis=axis,
+            dtype=self.dtype,
+        )
+
+        return newarr
+    
+    def active_min(self, axis=None, skipna=None):
+        """
+        Perform ``dask delayed`` active mean for each ``dask block`` which corresponds to a single ``chunk``.
+        Combines the results of the dask delayed ``active_min`` operations on each block into a single dask Array,
+        which is then mapped to a new DaskActiveArray object.
+
+        :param axis:        (int) The index of the axis on which to perform the active min.
+
+        :param skipna:      (bool) Skip NaN values when calculating the min.
+
+        :returns:       A new ``DaskActiveArray`` object which has been reduced along the specified axes using
+                        the concatenations of active_means from each chunk.
+        """
+
+        newarr = da.reduction(
+            self.copy(),
+            partition_min,
+            min_agg,
+            combine=min_agg,
+            axis=axis,
+            dtype=self.dtype,
+        )
+
+        return newarr
+    
+    def active_sum(self, axis=None, skipna=None):
+        """
+        Perform ``dask delayed`` active mean for each ``dask block`` which corresponds to a single ``chunk``.
+        Combines the results of the dask delayed ``active_sum`` operations on each block into a single dask Array,
+        which is then mapped to a new DaskActiveArray object.
+
+        :param axis:        (int) The index of the axis on which to perform the active sum.
+
+        :param skipna:      (bool) Skip NaN values when calculating the sum.
+
+        :returns:       A new ``DaskActiveArray`` object which has been reduced along the specified axes using
+                        the concatenations of active_means from each chunk.
+        """
+
+        newarr = da.reduction(
+            self.copy(),
+            partition_sum,
+            sum_agg,
+            combine=sum_agg,
+            axis=axis,
+            dtype=self.dtype,
+        )
+
+        return newarr