update BasisState

Author: johannes-moegerle

angular.jl

@@ -37,7 +37,8 @@ function get_state(basisstate::BasisState)
         qn = channels[i]
         kets[i] = get_ket(qn, basisstate.species)
     end
-    state = rydstate.angular.AngularState(basisstate.coeff, kets; warn_if_not_normalized = false)
+    coeff = basisstate.coefficients[findall(basisstate.model.core)]
+    state = rydstate.angular.AngularState(coeff, kets; warn_if_not_normalized = false)
     return state
 end
 

utils.jl

@@ -1,3 +1,4 @@
+import LinearAlgebra
 import MQDT
 include("angular.jl")
 
@@ -14,6 +15,16 @@ macro timelog(expr)
 end
 
 
+function get_relevant_lr(state::MQDT.BasisState)
+    inds = findall(state.model.core)
+    return state.lr_list[inds]
+end
+
+function get_relevant_nu(state::MQDT.BasisState)
+    inds = findall(state.model.core)
+    return state.nu_list[inds]
+end
+
 function all_matrix_element(B::MQDT.BasisArray, parameters::MQDT.Parameters)
     """Calculate all relevant matrix elements for a given basis array B.
 
@@ -29,25 +40,32 @@ function all_matrix_element(B::MQDT.BasisArray, parameters::MQDT.Parameters)
     )
 
     states_indexed = [(ids - 1 + START_ID, state) for (ids, state) in enumerate(B.states)]
-    states_sorted =
-        sort(states_indexed, by = x -> (minimum(x[2].lr), minimum(x[2].nu), x[1]))
+    states_sorted = sort(
+        states_indexed,
+        by = x ->
+            (minimum(get_relevant_lr(x[2])), minimum(get_relevant_nu(x[2])), x[1]),
+    )
 
-    for (i1, (id1, b1)) in enumerate(states_sorted)
-        for (id2, b2) in states_sorted[i1:end]
+    for (i1, (id1, s1)) in enumerate(states_sorted)
+        lr1 = get_relevant_lr(s1)
+        nus1 = get_relevant_nu(s1)
+        for (id2, s2) in states_sorted[i1:end]
+            lr2 = get_relevant_lr(s2)
+            nus2 = get_relevant_nu(s2)
 
             # Skip if all contributions of the two states are far apart in angular momentum
-            if minimum(b2.lr) - maximum(b1.lr) > k_angular_max
+            if minimum(lr2) - maximum(lr1) > k_angular_max
                 continue
             end
 
             # Skip if all contributions of the two states are far apart in n and None of them is low-n
-            if all(abs(nu1-nu2) >= 11 for nu1 in b1.nu for nu2 in b2.nu) &&
-               all(nu1 > 25 for nu1 in b1.nu) &&
-               all(nu2 > 25 for nu2 in b2.nu)
+            if all(abs(nu1-nu2) >= 11 for nu1 in nus1 for nu2 in nus2) &&
+               all(nu1 > 25 for nu1 in nus1) &&
+               all(nu2 > 25 for nu2 in nus2)
                 continue
             end
 
-            m = MQDT.multipole_moments(b1, b2, parameters)
+            m = MQDT.multipole_moments(s1, s2, parameters)
             # multipole_moments returns the matrix elements in the following order
             # electric dipole, electric quadrupole, diamagnetic, magnetic
             table_keys = [
@@ -56,7 +74,7 @@ function all_matrix_element(B::MQDT.BasisArray, parameters::MQDT.Parameters)
                 "matrix_elements_q0",
                 "matrix_elements_mu",
             ]
-            prefactor_transposed = (-1)^(b2.f - b1.f)
+            prefactor_transposed = (-1)^(s2.f - s1.f)
 
             for (i, key) in enumerate(table_keys)
                 if m[i] != 0
@@ -84,37 +102,122 @@ function rcv_to_df(row_col_value::Vector{Tuple{Int64,Int64,Float64}})
 end
 
 
-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 = get_n(T),
+        n = get_n(T, P),
         nu = MQDT.get_nu(T),
         f = MQDT.get_f(T),
-        exp_nui = MQDT.exp_nui(T),
+        exp_nui = exp_nui(T),
         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_nui = std_nui(T),
         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),
+        is_j_total_momentum = repeat([iszero(P.spin)], size(T)),
+        is_calculated_with_mqdt = is_mqdt(T),
+        underspecified_channel_contribution = get_neg(T),
     )
     return df
 end
+
+
+function exp_nui(T::MQDT.BasisArray)
+    t = Vector{Float64}(undef, size(T))
+    for (i, state) in enumerate(T.states)
+        t[i] = exp_q(state.nu_list, state.coefficients)
+    end
+    return t
+end
+
+function std_nui(T::MQDT.BasisArray)
+    t = Vector{Float64}(undef, size(T))
+    for (i, state) in enumerate(T.states)
+        t[i] = std_q(state.nu_list, state.coefficients)
+    end
+    return t
+end
+
+function is_mqdt(T::MQDT.BasisArray)
+    t = Vector{Bool}(undef, size(T))
+    for (i, state) in enumerate(T.states)
+        t[i] = !isone(length(state.coefficients))
+    end
+    return t
+end
+
+
+function get_neg(T::MQDT.BasisArray)
+    t = Vector{Float64}(undef, size(T))
+    for (i, state) in enumerate(T.states)
+        irrel = findall(iszero, state.model.core)
+        t[i] = sum(state.coefficients[irrel] .^ 2)
+    end
+    return t
+end
+
+function exp_q(q::Vector, n::Vector)
+    if allequal(q)
+        return Float64(q[1])
+    else
+        m = n .^ 2
+        M = sum(m)
+        if M > 1
+            m /= M
+        end
+        return LinearAlgebra.dot(q, m)
+    end
+end
+
+function std_q(q::Vector, n::Vector)
+    if allequal(q)
+        return 0.0
+    else
+        m = n .^ 2
+        M = sum(m)
+        if M > 1
+            m /= M
+        end
+        e1 = LinearAlgebra.dot(q, m)^2
+        e2 = LinearAlgebra.dot(q .^ 2, m)
+        if abs(e1 - e2) < 1e-11
+            return 0.0
+        else
+            return sqrt(e2 - e1)
+        end
+    end
+end
+
+function get_n(T::MQDT.BasisArray, P::MQDT.Parameters)
+    nu = MQDT.get_nu(T)
+    l = round.(Int, calc_exp_qn(T, "l_r"))
+    return get_n(nu, l, P.species)
+end
+
+function get_n(nu::Vector{Float64}, l::Vector{Int}, species::Symbol)
+    i0 = findall(iszero, l)
+    i1 = findall(iszero, l .- 1)
+    i2 = findall(iszero, l .- 2)
+    i3 = findall(iszero, l .- 3)
+    j0 = findall(x->x<2, nu)
+    nu[j0] .+= 1
+    if occursin("Yb", String(species))
+        nu[i0] .+= 4
+        nu[i1] .+= 3
+        nu[i2] .+= 2
+        nu[i3] .+= 1
+    else
+        nu[i0] .+= 3
+        nu[i1] .+= 2
+        nu[i2] .+= 2
+    end
+    return ceil.(Int, nu)
+end