diff --git a/src/cr/cube/crunch_cube.py b/src/cr/cube/crunch_cube.py index ea67d7aaf..a49229761 100644 --- a/src/cr/cube/crunch_cube.py +++ b/src/cr/cube/crunch_cube.py @@ -357,7 +357,13 @@ def _mr_margin(self, axis, weighted, adjusted): margin = table[self.ind_selected] + table[self.ind_non_selected] margin = margin[np.ix_(*self.valid_indices)] - return np.sum(margin, 1 - self.mr_dim_ind) if axis is None else margin + if axis is None: + axis = tuple([ + i for i, _ in enumerate(self.dimensions) if i != self.mr_dim_ind + ]) + return np.sum(margin, axis) + + return margin def _inflate_dim(self, array, axis): # Explicitly check if axis is tuple (which could be the case for doing @@ -433,14 +439,15 @@ def _mr_proportions(self, axis, weighted, include_transforms_for_dims=None): (table[:, 0, :, :] + table[:, 1, :, :])) # The following are normal MR x something (not CA) - if axis == 1: + if axis == 0: + res = table[:, 0, :] / (table[:, 0, :] + table[:, 1, :]) + else: num = self.as_array( include_transforms_for_dims=include_transforms_for_dims ) - den = self.margin(axis=1)[:, np.newaxis] + den = self.margin(axis=axis)[:, np.newaxis] return num / den - else: - res = table[:, 0, :] / (table[:, 0, :] + table[:, 1, :]) + elif self.mr_dim_ind == 1: num = table[self.ind_selected][np.ix_(*self.valid_indices)] non_selected = table[self.ind_non_selected][np.ix_(*self.valid_indices)] diff --git a/tests/integration/test_crunch_cube.py b/tests/integration/test_crunch_cube.py index 66a0f2f88..4f378ea0b 100644 --- a/tests/integration/test_crunch_cube.py +++ b/tests/integration/test_crunch_cube.py @@ -885,7 +885,6 @@ def test_z_scores_admit_by_gender_weighted_rows(self): actual = cube.standardized_residuals np.testing.assert_almost_equal(actual, expected) - def test_selected_crosstab_dim_names(self): cube = CrunchCube(FIXT_SELECTED_CROSSTAB_4) expected = ['Statements agreed with about Climate', 'Gender'] @@ -976,15 +975,15 @@ def test_selected_crosstab_proportions_by_cols(self): actual = cube.proportions(axis=0) np.testing.assert_almost_equal(actual, expected) - def test_selected_crosstab_proportions_total(self): + def test_selected_crosstab_proportions_by_cell(self): cube = CrunchCube(FIXT_SELECTED_CROSSTAB_4) expected = np.array([ - [.6815894, .73842499], - [.6632888, .63436033], - [.81146893, .85004262], - [.85174892, .88409986], - [.37871497, .28475016], - [.53511338, .5426064], + [0.329036700375595, 0.381950958618156], + [0.320620717695708, 0.327723062528721], + [0.392254504152701, 0.439142047148397], + [0.418103897371069, 0.450115632023491], + [0.185809278853744, 0.14504292098248], + [0.262843097025161, 0.27608279697761], ]) actual = cube.proportions() np.testing.assert_almost_equal(actual, expected) @@ -1190,9 +1189,9 @@ def test_fruit_x_pets_proportions_by_col(self): def test_pets_x_fruit_proportions_by_cell(self): cube = CrunchCube(FIXT_PETS_X_FRUIT) expected = np.array([ - [.4285714, .5384615], - [.48, .4074074], - [.5217391, .5531915], + [0.15, 0.35], + [0.15189873, 0.27848101], + [0.17142857, 0.37142857], ]) actual = cube.proportions(axis=None) np.testing.assert_almost_equal(actual, expected) @@ -1346,6 +1345,12 @@ def test_pets_array_x_pets_cell(self): actual = cube.proportions(axis=None)[0] np.testing.assert_almost_equal(actual, expected) + def test_pets_x_pets_array_margin_by_cell(self): + cube = CrunchCube(FIXT_PETS_X_PETS_ARRAY) + expected = np.array([229, 229, 229]) + actual = cube.margin() + np.testing.assert_almost_equal(actual, expected) + def test_pets_x_pets_array_by_col(self): cube = CrunchCube(FIXT_PETS_X_PETS_ARRAY) expected = np.array([ @@ -1360,23 +1365,29 @@ def test_pets_x_pets_array_by_col(self): def test_pets_x_pets_array_by_row(self): cube = CrunchCube(FIXT_PETS_X_PETS_ARRAY) expected = np.array([ - [0.44444444, 0.41176471], - [0., 1.], - [0.5, 0.47368421], + [0.55555556, 0.19444444], + [0., 0.55555556], + [0.44444444, 0.25], ]) - actual = cube.proportions(axis=2)[0] - # Since cube is 3D, col dim is 1 (instead of 0) + actual = cube.proportions(axis=1)[0] np.testing.assert_almost_equal(actual, expected) def test_pets_x_pets_array_by_cell(self): cube = CrunchCube(FIXT_PETS_X_PETS_ARRAY) expected = np.array([ - [0.44444444, 0.41176471], - [0., 1.], - [0.5, 0.47368421], + [[0.08733624, 0.06113537], + [0., 0.17467249], + [0.069869, 0.07860262]], + + [[0., 0.14847162], + [0.06550218, 0.06113537], + [0.05676856, 0.069869]], + + [[0.09606987, 0.069869], + [0.08733624, 0.07860262], + [0., 0.16593886]] ]) - actual = cube.proportions(axis=None)[0] - # Since cube is 3D, col dim is 1 (instead of 0) + actual = cube.proportions() np.testing.assert_almost_equal(actual, expected) def test_mr_x_cat_x_cat_by_row(self):