MatrixAlgebraKit.jl/src/implementations/schur.jl at 91b010dfc7d0d3bf8e7afc4c1706c93fc4e0195d · 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
69
70
71
72
# Inputs
# ------
copy_input(::typeof(schur_full), A) = copy_input(eig_full, A)
copy_input(::typeof(schur_vals), A) = copy_input(eig_vals, A)

# check input
function check_input(::typeof(schur_full!), A::AbstractMatrix, TZv, ::AbstractAlgorithm)
    m = LinearAlgebra.checksquare(A)
    T, Z, vals = TZv
    @assert T isa AbstractMatrix && Z isa AbstractMatrix && vals isa AbstractVector
    @check_size(T, (m, m))
    @check_scalar(T, A)
    @check_size(Z, (m, m))
    @check_scalar(Z, A)
    @check_size(vals, (n,))
    @check_scalar(vals, A, complex)
    return nothing
end
function check_input(::typeof(schur_vals!), A::AbstractMatrix, vals, ::AbstractAlgorithm)
    m = LinearAlgebra.checksquare(A)
    @assert vals isa AbstractVector
    @check_size(vals, (m,))
    @check_scalar(vals, A, complex)
    return nothing
end

# Outputs
# -------
function initialize_output(::typeof(schur_full!), A::AbstractMatrix, ::AbstractAlgorithm)
    n = size(A, 1) # square check will happen later
    Z = similar(A, (n, n))
    vals = similar(A, complex(eltype(A)), n)
    return (A, Z, vals)
end
function initialize_output(::typeof(schur_vals!), A::AbstractMatrix, ::AbstractAlgorithm)
    n = size(A, 1) # square check will happen later
    vals = similar(A, complex(eltype(A)), n)
    return vals
end

# Implementation
# --------------
function schur_full!(A::AbstractMatrix, TZv, alg::LAPACK_EigAlgorithm)
    check_input(schur_full!, A, TZv, alg)
    T, Z, vals = TZv
    if alg isa LAPACK_Simple
        isempty(alg.kwargs) ||
            throw(ArgumentError("LAPACK_Simple Schur (gees) does not accept any keyword arguments"))
        YALAPACK.gees!(A, Z, vals)
    else # alg isa LAPACK_Expert
        isempty(alg.kwargs) ||
            throw(ArgumentError("LAPACK_Expert Schur (geesx) does not accept any keyword arguments"))
        YALAPACK.geesx!(A, Z, vals)
    end
    T === A || copy!(T, A)
    return T, Z, vals
end

function schur_vals!(A::AbstractMatrix, vals, alg::LAPACK_EigAlgorithm)
    check_input(schur_vals!, A, vals, alg)
    Z = similar(A, eltype(A), (size(A, 1), 0))
    if alg isa LAPACK_Simple
        isempty(alg.kwargs) ||
            throw(ArgumentError("LAPACK_Simple (gees) does not accept any keyword arguments"))
        YALAPACK.gees!(A, Z, vals)
    else # alg isa LAPACK_Expert
        isempty(alg.kwargs) ||
            throw(ArgumentError("LAPACK_Expert (geesx) does not accept any keyword arguments"))
        YALAPACK.geesx!(A, Z, vals)
    end
    return vals
end