add angular.jl

Author: johannes-moegerle

Project.toml

@@ -8,4 +8,5 @@ LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MQDT = "1bfeae05-7215-4237-8fdd-b424babe7f02"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Parquet2 = "98572fba-bba0-415d-956f-fa77e587d26d"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

angular.jl

@@ -0,0 +1,87 @@
+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 = Vector{Any}(undef, size(channels))
+    for i in eachindex(channels)
+        qn = channels[i]
+        kets[i] = get_ket(qn, basisstate.species)
+    end
+    state = rydstate.angular.AngularState(basisstate.coeff, 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

utils.jl

@@ -1,4 +1,6 @@
 import MQDT
+include("angular.jl")
+
 
 const START_ID = 0
 
@@ -12,7 +14,7 @@ macro timelog(expr)
 end
 
 
-function all_matrix_element(B::BasisArray, parameters::MQDT.Parameters)
+function all_matrix_element(B::MQDT.BasisArray, parameters::MQDT.Parameters)
     """Calculate all relevant matrix elements for a given basis array B.
 
     This means dipole, quadrupole, magnetic, and diamagnetic matrix elements.
@@ -81,26 +83,35 @@ function rcv_to_df(row_col_value::Vector{Tuple{Int64,Int64,Float64}})
     return df
 end
 
-function basis_to_df(T::BasisArray, P::Parameters)
+
+function get_n(T::MQDT.BasisArray)
+    # TODO for now just return round(nu)
+    # later calculate radial overlap with different sqdt states and take the corresponding n
+    nu = MQDT.get_nu(T)
+    return round.(Int, nu)
+end
+
+
+function basis_to_df(T::MQDT.BasisArray, P::MQDT.Parameters)
     df = DataFrame(
         id = collect(START_ID:(size(T)-1+START_ID)),
         energy = MQDT.get_e(T, P) / 219474.6313632,  # convert 1/cm to atomic units
         parity = MQDT.get_p(T),
-        n = MQDT.get_n(T, P),
+        n = get_n(T),
         nu = MQDT.get_nu(T),
         f = MQDT.get_f(T),
         exp_nui = MQDT.exp_nui(T),
-        exp_l = MQDT.calc_exp_qn(T, "l_tot"),
-        exp_j = MQDT.calc_exp_qn(T, "j_tot"),
-        exp_s = MQDT.calc_exp_qn(T, "s_tot"),
-        exp_l_ryd = MQDT.calc_exp_qn(T, "l_r"),
-        exp_j_ryd = MQDT.calc_exp_qn(T, "j_r"),
+        exp_l = calc_exp_qn(T, "l_tot"),
+        exp_j = calc_exp_qn(T, "j_tot"),
+        exp_s = calc_exp_qn(T, "s_tot"),
+        exp_l_ryd = calc_exp_qn(T, "l_r"),
+        exp_j_ryd = calc_exp_qn(T, "j_r"),
         std_nui = MQDT.std_nui(T),
-        std_l = MQDT.calc_std_qn(T, "l_tot"),
-        std_j = MQDT.calc_std_qn(T, "j_tot"),
-        std_s = MQDT.calc_std_qn(T, "s_tot"),
-        std_l_ryd = MQDT.calc_std_qn(T, "l_r"),
-        std_j_ryd = MQDT.calc_std_qn(T, "j_r"),
+        std_l = calc_std_qn(T, "l_tot"),
+        std_j = calc_std_qn(T, "j_tot"),
+        std_s = calc_std_qn(T, "s_tot"),
+        std_l_ryd = calc_std_qn(T, "l_r"),
+        std_j_ryd = calc_std_qn(T, "j_r"),
         is_j_total_momentum = MQDT.is_J(T, P),
         is_calculated_with_mqdt = MQDT.is_mqdt(T),
         underspecified_channel_contribution = MQDT.get_neg(T),