()
| 421 | |
| 422 | |
| 423 | def test_stack(): |
| 424 | # non-iterable input |
| 425 | assert_raises(TypeError, stack, 1) |
| 426 | |
| 427 | # 0d input |
| 428 | for input_ in [(1, 2, 3), |
| 429 | [np.int32(1), np.int32(2), np.int32(3)], |
| 430 | [np.array(1), np.array(2), np.array(3)]]: |
| 431 | assert_array_equal(stack(input_), [1, 2, 3]) |
| 432 | # 1d input examples |
| 433 | a = np.array([1, 2, 3]) |
| 434 | b = np.array([4, 5, 6]) |
| 435 | r1 = array([[1, 2, 3], [4, 5, 6]]) |
| 436 | assert_array_equal(np.stack((a, b)), r1) |
| 437 | assert_array_equal(np.stack((a, b), axis=1), r1.T) |
| 438 | # all input types |
| 439 | assert_array_equal(np.stack(list([a, b])), r1) |
| 440 | assert_array_equal(np.stack(array([a, b])), r1) |
| 441 | # all shapes for 1d input |
| 442 | arrays = [np.random.randn(3) for _ in range(10)] |
| 443 | axes = [0, 1, -1, -2] |
| 444 | expected_shapes = [(10, 3), (3, 10), (3, 10), (10, 3)] |
| 445 | for axis, expected_shape in zip(axes, expected_shapes): |
| 446 | assert_equal(np.stack(arrays, axis).shape, expected_shape) |
| 447 | assert_raises_regex(np.AxisError, 'out of bounds', stack, arrays, axis=2) |
| 448 | assert_raises_regex(np.AxisError, 'out of bounds', stack, arrays, axis=-3) |
| 449 | # all shapes for 2d input |
| 450 | arrays = [np.random.randn(3, 4) for _ in range(10)] |
| 451 | axes = [0, 1, 2, -1, -2, -3] |
| 452 | expected_shapes = [(10, 3, 4), (3, 10, 4), (3, 4, 10), |
| 453 | (3, 4, 10), (3, 10, 4), (10, 3, 4)] |
| 454 | for axis, expected_shape in zip(axes, expected_shapes): |
| 455 | assert_equal(np.stack(arrays, axis).shape, expected_shape) |
| 456 | # empty arrays |
| 457 | assert_(stack([[], [], []]).shape == (3, 0)) |
| 458 | assert_(stack([[], [], []], axis=1).shape == (0, 3)) |
| 459 | # out |
| 460 | out = np.zeros_like(r1) |
| 461 | np.stack((a, b), out=out) |
| 462 | assert_array_equal(out, r1) |
| 463 | # edge cases |
| 464 | assert_raises_regex(ValueError, 'need at least one array', stack, []) |
| 465 | assert_raises_regex(ValueError, 'must have the same shape', |
| 466 | stack, [1, np.arange(3)]) |
| 467 | assert_raises_regex(ValueError, 'must have the same shape', |
| 468 | stack, [np.arange(3), 1]) |
| 469 | assert_raises_regex(ValueError, 'must have the same shape', |
| 470 | stack, [np.arange(3), 1], axis=1) |
| 471 | assert_raises_regex(ValueError, 'must have the same shape', |
| 472 | stack, [np.zeros((3, 3)), np.zeros(3)], axis=1) |
| 473 | assert_raises_regex(ValueError, 'must have the same shape', |
| 474 | stack, [np.arange(2), np.arange(3)]) |
| 475 | |
| 476 | # do not accept generators |
| 477 | with pytest.raises(TypeError, match="arrays to stack must be"): |
| 478 | stack((x for x in range(3))) |
| 479 | |
| 480 | #casting and dtype test |
nothing calls this directly
no test coverage detected