Merge pull request #138 from Crunch-io/enable-pruning-in-min-base-size

Enable pruning in min base size

README.md

@@ -112,6 +112,9 @@ The detailed description can be found
 
 ## Changes
 
+#### 1.9.8
+- Enable pruning for min-base-size masks
+
 #### 1.9.7
 - Implement Min Base Size suppression masks
 

src/cr/cube/__init__.py

@@ -2,4 +2,4 @@
 
 """Initialization module for crunch-cube package."""
 
-__version__ = "1.9.7"
+__version__ = "1.9.8"

src/cr/cube/cube_slice.py

@@ -278,7 +278,7 @@ def margin(
 
         return self._extract_slice_result_from_cube(margin)
 
-    def min_base_size_mask(self, size, hs_dims=None):
+    def min_base_size_mask(self, size, hs_dims=None, prune=False):
         """Returns MinBaseSizeMask object with correct row, col and table masks.
 
         The returned object stores the necessary information about the base size, as
@@ -292,7 +292,7 @@ def min_base_size_mask(self, size, hs_dims=None):
         >>> cube_slice.min_base_size_mask(50).column_mask
         >>> cube_slice.min_base_size_mask(22).table_mask
         """
-        return MinBaseSizeMask(self, size, hs_dims)
+        return MinBaseSizeMask(self, size, hs_dims=hs_dims, prune=prune)
 
     @lazyproperty
     def mr_dim_ind(self):

src/cr/cube/min_base_size_mask.py

@@ -5,7 +5,7 @@
 from __future__ import division
 import numpy as np
 
-from cr.cube.util import lazyproperty
+from cr.cube.util import lazyproperty, compress_pruned
 from cr.cube.enum import DIMENSION_TYPE as DT
 
 
@@ -18,16 +18,22 @@ class MinBaseSizeMask:
     marginal values and the shape of the underlying slice.
     """
 
-    def __init__(self, slice_, size, hs_dims=None):
+    def __init__(self, slice_, size, hs_dims=None, prune=False):
         self._slice = slice_
         self._size = size
         self._hs_dims = hs_dims
+        self._prune = prune
 
     @lazyproperty
     def column_mask(self):
         """ndarray, True where column margin <= min_base_size, same shape as slice."""
-        margin = self._slice.margin(
-            axis=0, weighted=False, include_transforms_for_dims=self._hs_dims
+        margin = compress_pruned(
+            self._slice.margin(
+                axis=0,
+                weighted=False,
+                include_transforms_for_dims=self._hs_dims,
+                prune=self._prune,
+            )
         )
         mask = margin < self._size
 
@@ -43,8 +49,13 @@ def column_mask(self):
     @lazyproperty
     def row_mask(self):
         """ndarray, True where row margin <= min_base_size, same shape as slice."""
-        margin = self._slice.margin(
-            axis=1, weighted=False, include_transforms_for_dims=self._hs_dims
+        margin = compress_pruned(
+            self._slice.margin(
+                axis=1,
+                weighted=False,
+                include_transforms_for_dims=self._hs_dims,
+                prune=self._prune,
+            )
         )
         mask = margin < self._size
 
@@ -60,7 +71,14 @@ def row_mask(self):
     @lazyproperty
     def table_mask(self):
         """ndarray, True where table margin <= min_base_size, same shape as slice."""
-        margin = self._slice.margin(axis=None, weighted=False)
+        margin = compress_pruned(
+            self._slice.margin(
+                axis=None,
+                weighted=False,
+                include_transforms_for_dims=self._hs_dims,
+                prune=self._prune,
+            )
+        )
         mask = margin < self._size
 
         if margin.shape == self._shape:
@@ -74,13 +92,13 @@ def table_mask(self):
 
     @lazyproperty
     def _shape(self):
-        shape = self._slice.get_shape(hs_dims=self._hs_dims)
+        shape = self._slice.get_shape(hs_dims=self._hs_dims, prune=self._prune)
 
         if len(shape) != self._slice.ndim:
             # TODO: This is an ugly hack that needs to happen due to the fact that we
             # purge dimensions with the count of 1, when getting the slice shape. This
             # will be addressed in a PR (already on the way) that strives to abandon
             # the ad-hoc purging of 1-element dimensions altogether.
-            shape = (shape[0], 1)
+            shape = (shape[0], 1) if shape else (1,)
 
         return shape