()
| 435 | |
| 436 | |
| 437 | def test_apply_input_core_dimension() -> None: |
| 438 | def first_element(obj, dim): |
| 439 | def func(x): |
| 440 | return x[..., 0] |
| 441 | |
| 442 | return apply_ufunc(func, obj, input_core_dims=[[dim]]) |
| 443 | |
| 444 | array = np.array([[1, 2], [3, 4]]) |
| 445 | variable = xr.Variable(["x", "y"], array) |
| 446 | data_array = xr.DataArray(variable, {"x": ["a", "b"], "y": [-1, -2]}) |
| 447 | dataset = xr.Dataset({"data": data_array}) |
| 448 | |
| 449 | expected_variable_x = xr.Variable(["y"], [1, 2]) |
| 450 | expected_data_array_x = xr.DataArray(expected_variable_x, {"y": [-1, -2]}) |
| 451 | expected_dataset_x = xr.Dataset({"data": expected_data_array_x}) |
| 452 | |
| 453 | expected_variable_y = xr.Variable(["x"], [1, 3]) |
| 454 | expected_data_array_y = xr.DataArray(expected_variable_y, {"x": ["a", "b"]}) |
| 455 | expected_dataset_y = xr.Dataset({"data": expected_data_array_y}) |
| 456 | |
| 457 | assert_identical(expected_variable_x, first_element(variable, "x")) |
| 458 | assert_identical(expected_variable_y, first_element(variable, "y")) |
| 459 | |
| 460 | assert_identical(expected_data_array_x, first_element(data_array, "x")) |
| 461 | assert_identical(expected_data_array_y, first_element(data_array, "y")) |
| 462 | |
| 463 | assert_identical(expected_dataset_x, first_element(dataset, "x")) |
| 464 | assert_identical(expected_dataset_y, first_element(dataset, "y")) |
| 465 | |
| 466 | assert_identical(expected_data_array_x, first_element(data_array.groupby("y"), "x")) |
| 467 | assert_identical(expected_dataset_x, first_element(dataset.groupby("y"), "x")) |
| 468 | |
| 469 | def multiply(*args): |
| 470 | val = args[0] |
| 471 | for arg in args[1:]: |
| 472 | val = val * arg |
| 473 | return val |
| 474 | |
| 475 | # regression test for GH:2341 |
| 476 | with pytest.raises(ValueError): |
| 477 | apply_ufunc( |
| 478 | multiply, |
| 479 | data_array, |
| 480 | data_array["y"].values, |
| 481 | input_core_dims=[["y"]], |
| 482 | output_core_dims=[["y"]], |
| 483 | ) |
| 484 | expected = xr.DataArray( |
| 485 | multiply(data_array, data_array["y"]), dims=["x", "y"], coords=data_array.coords |
| 486 | ) |
| 487 | actual = apply_ufunc( |
| 488 | multiply, |
| 489 | data_array, |
| 490 | data_array["y"].values, |
| 491 | input_core_dims=[["y"], []], |
| 492 | output_core_dims=[["y"]], |
| 493 | ) |
| 494 | assert_identical(expected, actual) |
nothing calls this directly
no test coverage detected
searching dependent graphs…