database-mqdt/angular.jl at main · pairinteraction/database-mqdt · 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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
using PythonCall
using MQDT

function calc_exp_qn(basisstate::BasisState, qn::String)
    state = get_state(basisstate)
    value = state.calc_exp_qn(qn)
    return pyconvert(Float64, value)
end

function calc_exp_qn(basisarray::BasisArray, qn::String)
    exp_list = Vector{Float64}(undef, size(basisarray))
    for (i, state) in enumerate(basisarray.states)
        exp_list[i] = calc_exp_qn(state, qn)
    end
    return exp_list
end

function calc_std_qn(basisstate::BasisState, qn::String)
    state = get_state(basisstate)
    value = state.calc_std_qn(qn)
    return pyconvert(Float64, value)
end

function calc_std_qn(basisarray::BasisArray, qn::String)
    std_list = Vector{Float64}(undef, size(basisarray))
    for (i, state) in enumerate(basisarray.states)
        std_list[i] = calc_std_qn(state, qn)
    end
    return std_list
end

function get_state(basisstate::BasisState)
    rydstate = pyimport("rydstate")
    channels = basisstate.channels.i
    kets = Any[]
    coeffs = Float64[]
    coefficients = basisstate.coefficients[findall(basisstate.model.core)]
    for (i, qn) in enumerate(channels)
        # TODO how to handle channels with not well defined Jc or Fc
        if isa(qn, jjQuantumNumbers) && isnan(qn.Jc)
            continue
        end
        if isa(qn, fjQuantumNumbers) && (isnan(qn.Jc) || isnan(qn.Fc))
            continue
        end
        push!(kets, get_ket(qn, basisstate.species))
        push!(coeffs, coefficients[i])
    end
    state = rydstate.angular.AngularState(coeffs, kets; warn_if_not_normalized = false)
    return state
end

function get_ket(qn::lsQuantumNumbers, species::Symbol)
    rydstate = pyimport("rydstate")
    ket = rydstate.angular.AngularKetLS(;
        s_c = qn.sc,
        s_tot = qn.S,
        l_c = qn.lc,
        l_r = qn.lr,
        l_tot = qn.L,
        j_tot = qn.J,
        f_tot = qn.F,
        species = String(species),
    )
    return ket
end

function get_ket(qn::jjQuantumNumbers, species::Symbol)
    rydstate = pyimport("rydstate")
    ket = rydstate.angular.AngularKetJJ(;
        s_c = qn.sc,
        l_c = qn.lc,
        l_r = qn.lr,
        j_c = qn.Jc,
        j_r = qn.Jr,
        j_tot = qn.J,
        f_tot = qn.F,
        species = String(species),
    )
    return ket
end

function get_ket(qn::fjQuantumNumbers, species::Symbol)
    rydstate = pyimport("rydstate")
    ket = rydstate.angular.AngularKetFJ(;
        s_c = qn.sc,
        l_c = qn.lc,
        l_r = qn.lr,
        j_c = qn.Jc,
        f_c = qn.Fc,
        j_r = qn.Jr,
        f_tot = qn.F,
        species = String(species),
    )
    return ket
end