Merge remote-tracking branch 'origin/rel-5.1.30'

src/cr/cube/crunch_cube.py

@@ -196,52 +196,77 @@ def _as_array(self, include_missing=False, get_non_selected=False,
         res = res + adjusted
 
         if prune:
-            return self._prune(res, include_transforms_for_dims)
+            return self._prune_body(res, include_transforms_for_dims)
 
         return res
 
-    def _mr_prune(self, res):
-        if len(res.shape) > 2:
-            # Only perform pruning for 1-D MR cubes.
-            return res
-        margin = self.margin(axis=0)
-        ind_prune = margin == 0
-        return res[~ind_prune]
-
-    def _prune(self, res, transforms=None):
+    def _prune_body(self, res, transforms=None):
         '''Prune the result based on margins content.
 
         Pruning is the removal of rows or columns, whose corresponding
         marginal elements are either 0 or not defined (np.nan).
         '''
 
-        if self.has_mr:
-            return self._mr_prune(res)
-
         # Note: If the cube contains transforms (H&S), and they happen to
         # have the marginal value 0 (or NaN), they're NOT pruned. This is
         # done to achieve parity with how the front end client works (whaam).
+        row_prune = self._row_prune_indices(transforms)
+
+        if len(self.dimensions) == 1 or len(res.shape) == 1:
+            # For 1D, margin is calculated as the row margin.
+            return res[~row_prune]
+
+        # Prune columns by column margin values.
+        col_prune = self._col_prune_indices(transforms)
+        res = res[:, ~col_prune]
+
+        # Prune rows by row margin values.
+        return res[~row_prune, :]
 
-        inserted_inds = self.inserted_hs_indices()
+    def prune_indices(self, transforms=None):
+        row_indices = self._row_prune_indices(transforms)
+
+        if len(self.dimensions) == 1:
+            return [row_indices]
+
+        col_indices = self._col_prune_indices(transforms)
+
+        return [row_indices, col_indices]
+
+    def _row_prune_indices(self, transforms=None):
+        def _get_axis_for_margin_calculations():
+            # when dealing with 1D MR cubes, axis for margin should be 0
+            if len(self.dimensions) == 1 and self.has_mr:
+                return 0
+            # when dealing with 3D cubes, axis should be 3
+            elif len(self.dimensions) == 3:
+                return 2
+            return 1
+
+        inserted_inds = self.inserted_hs_indices() if transforms else []
         inserted_rows = np.array(inserted_inds[0] if len(inserted_inds) else [])
-        row_margin = self.margin(include_transforms_for_dims=transforms, axis=1)
+        axis = _get_axis_for_margin_calculations()
+        row_margin = self.margin(include_transforms_for_dims=transforms, axis=axis)
         ind_inserted = np.zeros(row_margin.shape, dtype=bool)
         if any(inserted_rows):
             ind_inserted[inserted_rows] = True
         row_prune = np.logical_or(row_margin == 0, np.isnan(row_margin))
         row_prune = np.logical_and(row_prune, ~ind_inserted)
 
-        if len(self.dimensions) == 1 or len(res.shape) == 1:
-            # For 1D, margin is calculated as the row margin.
-            return res[~row_prune]
+        return row_prune
 
-        # Prune columns by column margin values.
-        col_margin = self.margin(include_transforms_for_dims=transforms, axis=0)
+    def _col_prune_indices(self, transforms=None):
+        inserted_inds = self.inserted_hs_indices() if transforms else []
+        inserted_cols = np.array(inserted_inds[1] if len(inserted_inds) > 1 else [])
+        axis = 0 if len(self.dimensions) < 3 else 1
+        col_margin = self.margin(include_transforms_for_dims=transforms, axis=axis)
+        ind_inserted = np.zeros(col_margin.shape, dtype=bool)
+        if any(inserted_cols):
+            ind_inserted[inserted_cols] = True
         col_prune = np.logical_or(col_margin == 0, np.isnan(col_margin))
-        res = res[:, ~col_prune]
+        col_prune = np.logical_and(col_prune, ~ind_inserted)
 
-        # Prune rows by row margin values.
-        return res[~row_prune, :]
+        return col_prune
 
     @classmethod
     def _fix_valid_indices(cls, valid_indices, insertion_index, dim):
@@ -411,10 +436,9 @@ def _margin(self, axis=None, weighted=True, adjusted=False,
         array = self._inflate_dim(array, axis)
         res = np.sum(array, axis)
 
-        if prune:
-            # Remove values if 0 or np.nan
-            ind_prune = np.logical_or(res == 0, np.isnan(res))
-            res = res[~ind_prune]
+        if prune and axis is not None and type(res) is np.ndarray:
+            prune_indices = self.prune_indices(include_transforms_for_dims)
+            res = res[~prune_indices[1 - axis]]
 
         if len(res.shape) == 0:
             # Each margin needs to be iterable, even if it only has