Merge branch 'ca-x-mr-with-hs-#157406153' into rel-5.1.144

src/cr/cube/crunch_cube.py

@@ -228,6 +228,10 @@ def _prune_3d_body(self, res, transforms):
                 slice_mask = slice_mask[np.newaxis, :, np.newaxis, :]
             elif self.mr_dim_ind == (0, 2):
                 slice_mask = slice_mask[np.newaxis, :, :, np.newaxis]
+            elif self.mr_dim_ind == 1 and self.ndim == 3:
+                slice_mask = slice_mask[:, np.newaxis, :]
+            elif self.mr_dim_ind == 2 and self.ndim == 3:
+                slice_mask = slice_mask[:, :, np.newaxis]
 
             mask[i] = slice_mask
         res = np.ma.masked_array(res, mask=mask)
@@ -682,6 +686,7 @@ def _mr_props_as_0th(self, axis, table, hs_dims, prune):
     def _mr_props_as_1st(self, axis, table, hs_dims, prune):
         num = table[self.ind_selected][np.ix_(*self.valid_indices)]
         non_selected = table[self.ind_non_selected][np.ix_(*self.valid_indices)]
+
         if num.ndim >= 3:
             if axis == (1, 2) or axis is None:
                 den = np.sum(num + non_selected, 2)[:, :, None]
@@ -691,16 +696,22 @@ def _mr_props_as_1st(self, axis, table, hs_dims, prune):
                 den = num + non_selected
             res = num / den
             res = self._transform(res, hs_dims, None)
+            if prune:
+                res = np.ma.masked_array(res,
+                                         mask=self.as_array(prune=prune).mask)
             return res
+
         if axis == 0:
             den = np.sum(num, 0)
         elif axis == 1 or axis == 2 and len(num.shape) >= 3:
             den = num + non_selected
         else:
             axis = 0 if len(num.shape) < 3 else (1, 2)
             den = np.sum(num + non_selected, axis)
+
         res = num / den
         res = self._transform(res, hs_dims, None)
+
         if prune:
             res = np.ma.masked_array(res, mask=self.as_array(prune=prune).mask)
         return res
@@ -711,6 +722,8 @@ def _mr_props_as_2nd(self, axis, hs_dims):
             if axis == 1 else
             self.margin(axis=0)
         )
+        if self.ndim == 3 and axis is None:
+            margin = np.sum(margin, axis=1)[:, np.newaxis, :]
         return self.as_array(
             include_transforms_for_dims=hs_dims
         ) / margin

tests/integration/fixtures/__init__.py

@@ -122,3 +122,5 @@ def _load(cube_file):
 CAT_X_MR_X_CAT_MISSING = _load('cat-x-mr-x-cat-missing.json')
 MR_X_CA_HS = _load('mr-x-ca-with-hs.json')
 CA_X_MR_WEIGHTED_HS = _load('ca-x-mr-weighted-hs.json')
+CA_X_MR_HS = _load('ca-x-mr-with-hs.json')
+CA_CAT_X_MR_X_CA_SUBVAR_HS = _load('ca-x-mr-hs.json')

tests/integration/test_crunch_cube.py

@@ -1190,6 +1190,10 @@ def test_pets_array_x_pets_cell(self):
             [0.44836533, 0., 0.48261546],
             [0.39084967, 1., 0.47843137],
         ])
+        expected = np.array([
+            [0.24992768, 0.        , 0.26901938],  # noqa
+            [0.17298235, 0.44258027, 0.21174429],
+        ])
         actual = cube.proportions(axis=None)[0]
         np.testing.assert_almost_equal(actual, expected)
 

tests/integration/test_headers_and_subtotals.py

@@ -19,6 +19,8 @@
 from .fixtures import CA_X_CAT_HS
 from .fixtures import CAT_X_MR_WEIGHTED_HS
 from .fixtures import MR_X_CA_HS
+from .fixtures import CA_X_MR_HS
+from .fixtures import CA_CAT_X_MR_X_CA_SUBVAR_HS
 
 from cr.cube.crunch_cube import CrunchCube
 
@@ -770,3 +772,37 @@ def test_mr_x_ca_props_by_row_with_hs(self):
         ])
         actual = cube.proportions(include_transforms_for_dims=[0, 1, 2])
         np.testing.assert_almost_equal(actual, expected)
+
+    def test_ca_cat_x_mr_x_ca_subvar_as_array_pruning(self):
+        cube = CrunchCube(CA_X_MR_HS)
+        expected = np.array([
+            [False, False, True],
+            [False, False, True],
+            [False, False, True],
+            [True, True, True],
+        ])
+        actual = cube.as_array(prune=True)[0].mask
+        np.testing.assert_array_equal(actual, expected)
+
+    def test_ca_cat_x_mr_x_ca_subvar_proportions_pruning(self):
+        cube = CrunchCube(CA_X_MR_HS)
+        expected = np.array([
+            [False, False, True],
+            [False, False, True],
+            [False, False, True],
+            [True, True, True],
+        ])
+        actual = cube.proportions(prune=True)[0].mask
+        np.testing.assert_array_equal(actual, expected)
+
+    def test_ca_x_mr_counts_pruning(self):
+        cube = CrunchCube(CA_CAT_X_MR_X_CA_SUBVAR_HS)
+        actual = cube.as_array(prune=True)[0].mask
+        expected = np.array([
+            [False, False, False, True],
+            [False, False, False, True],
+            [False, False, False, True],
+            [False, False, False, True],
+            [False, False, False, True]
+        ])
+        np.testing.assert_array_equal(actual, expected)

tests/integration/test_multiple_response.py

@@ -338,9 +338,13 @@ def test_array_x_mr_by_row(self):
 
     def test_array_x_mr_by_cell(self):
         cube = CrunchCube(ARRAY_X_MR)
+        # expected = np.array([
+        #     [0.41922353375674093, 0.03471395310157275, 0.5832027484767315],
+        #     [0.5143557893611596, 1, 0.5199603338915276],
+        # ])
         expected = np.array([
-            [0.41922353375674093, 0.03471395310157275, 0.5832027484767315],
-            [0.5143557893611596, 1, 0.5199603338915276],
+            [0.23701678, 0.01962626, 0.32972586],
+            [0.223554, 0.43462911, 0.2259899],
         ])
         # Only compare the first slice (parity with whaam tests)
         actual = cube.proportions(axis=None)[0]