diff --git a/src/cr/cube/crunch_cube.py b/src/cr/cube/crunch_cube.py index f8a573bff..30536ab23 100644 --- a/src/cr/cube/crunch_cube.py +++ b/src/cr/cube/crunch_cube.py @@ -195,11 +195,16 @@ def _as_array(self, include_missing=False, get_non_selected=False, ) res = res + adjusted - if prune and not self.has_mr: + if prune: return self._prune(res, include_transforms_for_dims) return res + def _mr_prune(self, res): + margin = self.margin(axis=0) + ind_prune = margin == 0 + return res[~ind_prune] + def _prune(self, res, transforms=None): '''Prune the result based on margins content. @@ -207,6 +212,9 @@ def _prune(self, res, transforms=None): 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).