diff --git a/src/cr/cube/crunch_cube.py b/src/cr/cube/crunch_cube.py index 7619a2a0c..fe57d2c4a 100644 --- a/src/cr/cube/crunch_cube.py +++ b/src/cr/cube/crunch_cube.py @@ -237,7 +237,7 @@ def _as_array(self, include_missing=False, get_non_selected=False, if prune and not self.has_mr: # Remove columns and rows where marginal value is 0 or np.nan - if len(self.dimensions) == 1: + if len(self.dimensions) == 1 or len(res.shape) == 1: margin = self.margin( include_transforms_for_dims=include_transforms_for_dims, axis=1, diff --git a/tests/integration/test_crunch_cube.py b/tests/integration/test_crunch_cube.py index 3d11b3603..edda18e5c 100644 --- a/tests/integration/test_crunch_cube.py +++ b/tests/integration/test_crunch_cube.py @@ -2005,3 +2005,11 @@ def test_row_unweighted_margin_when_has_means(self): ]) actual = cube.margin(axis=1, weighted=False, prune=True) np.testing.assert_array_equal(actual, expected) + + def test_ca_with_single_cat_pruning(self): + cube = CrunchCube(FIXT_CA_SINGLE_CAT) + # The last 0 of the expectation is not visible in whaam because of + # pruning, which is not the responsibility of cr.cube. + expected = np.array([79, 80, 70]) + actual = cube.as_array(weighted=False, prune=True) + np.testing.assert_almost_equal(actual, expected)