| 1360 | assert_identical(expected, actual) |
| 1361 | |
| 1362 | def test_groupby_sum(self) -> None: |
| 1363 | array = self.da |
| 1364 | grouped = array.groupby("abc") |
| 1365 | |
| 1366 | expected_sum_all = Dataset( |
| 1367 | { |
| 1368 | "foo": Variable( |
| 1369 | ["abc"], |
| 1370 | np.array( |
| 1371 | [ |
| 1372 | self.x[:, :9].sum(), |
| 1373 | self.x[:, 10:].sum(), |
| 1374 | self.x[:, 9:10].sum(), |
| 1375 | ] |
| 1376 | ).T, |
| 1377 | ), |
| 1378 | "abc": Variable(["abc"], np.array(["a", "b", "c"])), |
| 1379 | } |
| 1380 | )["foo"] |
| 1381 | assert_allclose(expected_sum_all, grouped.reduce(np.sum, dim=...)) |
| 1382 | assert_allclose(expected_sum_all, grouped.sum(...)) |
| 1383 | |
| 1384 | expected = DataArray( |
| 1385 | [ |
| 1386 | array["y"].values[idx].sum() |
| 1387 | for idx in [slice(9), slice(10, None), slice(9, 10)] |
| 1388 | ], |
| 1389 | [["a", "b", "c"]], |
| 1390 | ["abc"], |
| 1391 | ) |
| 1392 | actual = array["y"].groupby("abc").map(np.sum) |
| 1393 | assert_allclose(expected, actual) |
| 1394 | actual = array["y"].groupby("abc").sum(...) |
| 1395 | assert_allclose(expected, actual) |
| 1396 | |
| 1397 | expected_sum_axis1 = Dataset( |
| 1398 | { |
| 1399 | "foo": ( |
| 1400 | ["x", "abc"], |
| 1401 | np.array( |
| 1402 | [ |
| 1403 | self.x[:, :9].sum(1), |
| 1404 | self.x[:, 10:].sum(1), |
| 1405 | self.x[:, 9:10].sum(1), |
| 1406 | ] |
| 1407 | ).T, |
| 1408 | ), |
| 1409 | "abc": Variable(["abc"], np.array(["a", "b", "c"])), |
| 1410 | } |
| 1411 | )["foo"] |
| 1412 | assert_allclose(expected_sum_axis1, grouped.reduce(np.sum, "y")) |
| 1413 | assert_allclose(expected_sum_axis1, grouped.sum("y")) |
| 1414 | |
| 1415 | @pytest.mark.parametrize("use_flox", [True, False]) |
| 1416 | @pytest.mark.parametrize("shuffle", [True, False]) |