()
| 495 | |
| 496 | |
| 497 | def test_apply_output_core_dimension() -> None: |
| 498 | def stack_negative(obj): |
| 499 | def func(x): |
| 500 | return np.stack([x, -x], axis=-1) |
| 501 | |
| 502 | result = apply_ufunc(func, obj, output_core_dims=[["sign"]]) |
| 503 | if isinstance(result, xr.Dataset | xr.DataArray): |
| 504 | result.coords["sign"] = [1, -1] |
| 505 | return result |
| 506 | |
| 507 | array = np.array([[1, 2], [3, 4]]) |
| 508 | variable = xr.Variable(["x", "y"], array) |
| 509 | data_array = xr.DataArray(variable, {"x": ["a", "b"], "y": [-1, -2]}) |
| 510 | dataset = xr.Dataset({"data": data_array}) |
| 511 | |
| 512 | stacked_array = np.array([[[1, -1], [2, -2]], [[3, -3], [4, -4]]]) |
| 513 | stacked_variable = xr.Variable(["x", "y", "sign"], stacked_array) |
| 514 | stacked_coords = {"x": ["a", "b"], "y": [-1, -2], "sign": [1, -1]} |
| 515 | stacked_data_array = xr.DataArray(stacked_variable, stacked_coords) |
| 516 | stacked_dataset = xr.Dataset({"data": stacked_data_array}) |
| 517 | |
| 518 | assert_identical(stacked_array, stack_negative(array)) |
| 519 | assert_identical(stacked_variable, stack_negative(variable)) |
| 520 | assert_identical(stacked_data_array, stack_negative(data_array)) |
| 521 | assert_identical(stacked_dataset, stack_negative(dataset)) |
| 522 | assert_identical(stacked_data_array, stack_negative(data_array.groupby("x"))) |
| 523 | assert_identical(stacked_dataset, stack_negative(dataset.groupby("x"))) |
| 524 | |
| 525 | def original_and_stack_negative(obj): |
| 526 | def func(x): |
| 527 | return (x, np.stack([x, -x], axis=-1)) |
| 528 | |
| 529 | result = apply_ufunc(func, obj, output_core_dims=[[], ["sign"]]) |
| 530 | if isinstance(result[1], xr.Dataset | xr.DataArray): |
| 531 | result[1].coords["sign"] = [1, -1] |
| 532 | return result |
| 533 | |
| 534 | out0, out1 = original_and_stack_negative(array) |
| 535 | assert_identical(array, out0) |
| 536 | assert_identical(stacked_array, out1) |
| 537 | |
| 538 | out0, out1 = original_and_stack_negative(variable) |
| 539 | assert_identical(variable, out0) |
| 540 | assert_identical(stacked_variable, out1) |
| 541 | |
| 542 | out0, out1 = original_and_stack_negative(data_array) |
| 543 | assert_identical(data_array, out0) |
| 544 | assert_identical(stacked_data_array, out1) |
| 545 | |
| 546 | out0, out1 = original_and_stack_negative(dataset) |
| 547 | assert_identical(dataset, out0) |
| 548 | assert_identical(stacked_dataset, out1) |
| 549 | |
| 550 | out0, out1 = original_and_stack_negative(data_array.groupby("x")) |
| 551 | assert_identical(data_array, out0) |
| 552 | assert_identical(stacked_data_array, out1) |
| 553 | |
| 554 | out0, out1 = original_and_stack_negative(dataset.groupby("x")) |
nothing calls this directly
no test coverage detected
searching dependent graphs…