MatrixAlgebraKit.jl/src/common/view.jl at 349800ea54a538467c470a740cc346509fce5ffb · QuantumKitHub/MatrixAlgebraKit.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
# diagind: provided by LinearAlgebra.jl
@doc """
    diagview(D)

Return a view of the diagonal elements of a matrix `D`.

See also [`diagonal`](@ref).
""" diagview

diagview(D::Diagonal) = D.diag
diagview(D::AbstractMatrix) = view(D, diagind(D))

@doc """
    diagonal(v)

Construct a diagonal matrix view for the given diagonal vector.

See also [`diagview`](@ref).
""" diagonal

diagonal(v::AbstractVector) = Diagonal(v)

"""
    map_diagonal!(f, dst, src...)

Map the scalar function `f` over all elements of the diagonal of `src...`, returning
a diagonal result.

See also [`map_diagonal!`](@ref).
"""
map_diagonal(f, src, srcs...) = diagonal(f.(diagview(src), map(diagview, srcs)...))

"""
    map_diagonal!(f, dst, src...)

Map the scalar function `f` over all elements of the diagonal of `src...`,
into the diagonal elements of destination `dst`.

See also [`map_diagonal`](@ref).
"""
map_diagonal!(f, dst, src, srcs...) = (diagview(dst) .= f.(diagview(src), map(diagview, srcs)...); dst)

# triangularind
function lowertriangularind(A::AbstractMatrix)
    Base.require_one_based_indexing(A)
    m, n = size(A)
    I = Vector{Int}(undef, div(m * (m - 1), 2) + m * (n - m))
    offset = 0
    for j in 1:n
        r = (j + 1):m
        I[offset .- j .+ r] = (j - 1) * m .+ r
        offset += length(r)
    end
    return I
end

function uppertriangularind(A::AbstractMatrix)
    Base.require_one_based_indexing(A)
    m, n = size(A)
    I = Vector{Int}(undef, div(m * (m - 1), 2) + m * (n - m))
    offset = 0
    for i in 1:m
        r = (i + 1):n
        I[offset .- i .+ r] = i .+ m .* (r .- 1)
        offset += length(r)
    end
    return I
end