Function test_diag_indices

numpy/lib/tests/test_index_tricks.py:505–530

()

def test_diag_indices():
    di = diag_indices(4)
    a = np.array([[1, 2, 3, 4],
                  [5, 6, 7, 8],
                  [9, 10, 11, 12],
                  [13, 14, 15, 16]])
    a[di] = 100
    assert_array_equal(
        a, np.array([[100, 2, 3, 4],
                     [5, 100, 7, 8],
                     [9, 10, 100, 12],
                     [13, 14, 15, 100]])
        )

    # Now, we create indices to manipulate a 3-d array:
    d3 = diag_indices(2, 3)

    # And use it to set the diagonal of a zeros array to 1:
    a = np.zeros((2, 2, 2), int)
    a[d3] = 1
    assert_array_equal(
        a, np.array([[[1, 0],
                      [0, 0]],
                     [[0, 0],
                      [0, 1]]])
        )


class TestDiagIndicesFrom: