Skip to content

Commit

Permalink
fix: sum order
Browse files Browse the repository at this point in the history
  • Loading branch information
@lucifer1004
lucifer1004 committed Apr 3, 2023
1 parent cdc1b85 commit 08802e8
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 32 deletions.
49 changes: 34 additions & 15 deletions src/common/AbstractTransitionMatrix.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -435,6 +435,7 @@ function expansion_coefficients(𝐓::AbstractTransitionMatrix{CT, N}, λ;
T₃ = OffsetArray(zeros(ComplexF64, 2N + 1, N, 2N + 1, N), (-N):N, 1:N, (-N):N, 1:N)
T₄ = OffsetArray(zeros(ComplexF64, 2N + 1, N, 2N + 1, N), (-N):N, 1:N, (-N):N, 1:N)

@debug "Calculating T..."
Threads.@threads for (n, m) in collect(it)
for (n′, m′) in it
T₁[m, n, m′, n′] = 𝐓[m, n, m′, n′, 1, 1] + 𝐓[m, n, m′, n′, 1, 2] +
Expand All @@ -448,6 +449,12 @@ function expansion_coefficients(𝐓::AbstractTransitionMatrix{CT, N}, λ;
end
end

@debug "Calculating B..."
A₁ = OffsetArray(zeros(ComplexF64, N, 2N + 1), 1:N, 0:(2N))
A₂ = OffsetArray(zeros(ComplexF64, N, 2N + 1), 1:N, 0:(2N))
A₃ = OffsetArray(zeros(ComplexF64, N, 2N + 1), 1:N, 0:(2N))
A₄ = OffsetArray(zeros(ComplexF64, N, 2N + 1), 1:N, 0:(2N))

B₁ = OffsetArray(zeros(ComplexF64, 2N + 1, 2N + 1, N, 2N + 1), (-N):N, (-N):N, 1:N,
0:(2N))
B₂ = OffsetArray(zeros(ComplexF64, 2N + 1, 2N + 1, N, 2N + 1), (-N):N, (-N):N, 1:N,
Expand All @@ -458,12 +465,9 @@ function expansion_coefficients(𝐓::AbstractTransitionMatrix{CT, N}, λ;
0:(2N))

Threads.@threads for n₁ in 0:(2N)
for n in 1:N, m in (-N):N, k in (-N):N
b₁ = 0.0
b₂ = 0.0
b₃ = 0.0
b₄ = 0.0
@debug "n₁ = $n₁..."

for n in 1:N, k in (-N):N
for n′ in max(1, abs(n - n₁)):min(n + n₁, N)
a₁ = 0.0
a₂ = 0.0
Expand All @@ -478,20 +482,35 @@ function expansion_coefficients(𝐓::AbstractTransitionMatrix{CT, N}, λ;
a₄ += cg * T₄[m₁, n, m₁ + k, n′]
end

coeff = clebschgordan(n, m, n₁, 1 - m, n′) * ci[n′ - n] * a[n′]
b₁ += coeff * a₁
b₂ += coeff * a₂
b₃ += coeff * a₃
b₄ += coeff * a₄
A₁[n′, n₁] = a₁
A₂[n′, n₁] = a₂
A₃[n′, n₁] = a₃
A₄[n′, n₁] = a₄
end

B₁[k, m, n, n₁] = b₁
B₂[k, m, n, n₁] = b₂
B₃[k, m, n, n₁] = b₃
B₄[k, m, n, n₁] = b₄
for m in (-N):N
b₁ = 0.0
b₂ = 0.0
b₃ = 0.0
b₄ = 0.0

for n′ in max(1, abs(n - n₁)):min(n + n₁, N)
coeff = clebschgordan(n, m, n₁, 1 - m, n′) * ci[n′ - n] * a[n′]
b₁ += coeff * A₁[n′, n₁]
b₂ += coeff * A₂[n′, n₁]
b₃ += coeff * A₃[n′, n₁]
b₄ += coeff * A₄[n′, n₁]
end

B₁[k, m, n, n₁] = b₁
B₂[k, m, n, n₁] = b₂
B₃[k, m, n, n₁] = b₃
B₄[k, m, n, n₁] = b₄
end
end
end

@debug "Calculating D..."
D₀₀ = OffsetArray(zeros(ComplexF64, 2N + 1, N, N), (-N):N, 1:N, 1:N)
D₀₋₀ = OffsetArray(zeros(ComplexF64, 2N + 1, N, N), (-N):N, 1:N, 1:N)
D₋₀₋₀ = OffsetArray(zeros(ComplexF64, 2N + 1, N, N), (-N):N, 1:N, 1:N)
Expand Down Expand Up @@ -556,7 +575,7 @@ function expansion_coefficients(𝐓::AbstractTransitionMatrix{CT, N}, λ;
1:N,
1:N)

# Calculate g
@debug "Calculating g..."
g₀₀ = OffsetArray(zeros(ComplexF64, 2N + 1), 0:(2N))
g₀₋₀ = OffsetArray(zeros(ComplexF64, 2N + 1), 0:(2N))
g₋₀₋₀ = OffsetArray(zeros(ComplexF64, 2N + 1), 0:(2N))
Expand Down
42 changes: 25 additions & 17 deletions src/common/AxisymmetricTransitionMatrix.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,8 @@ function expansion_coefficients(𝐓::AxisymmetricTransitionMatrix{CT, N, V, T},

T1 = OffsetArray(zeros(ComplexF64, 2N + 1, N, N), (-N):N, 1:N, 1:N)
T2 = OffsetArray(zeros(ComplexF64, 2N + 1, N, N), (-N):N, 1:N, 1:N)
A1 = zeros(ComplexF64, N, N)
A2 = zeros(ComplexF64, N, N)
B1 = OffsetArray(zeros(ComplexF64, 2N + 1, 2N + 1, N), 0:(2N), (-N):N, 1:N)
B2 = OffsetArray(zeros(ComplexF64, 2N + 1, 2N + 1, N), 0:(2N), (-N):N, 1:N)

Expand All @@ -150,7 +152,10 @@ function expansion_coefficients(𝐓::AxisymmetricTransitionMatrix{CT, N, V, T},
@tspawnat i wig_thread_temp_init(4N)
end

@debug "Calculating B..."
Threads.@threads for n in 1:N
@debug "n = $n..."

# Calculate T1 and T2
for n′ in 1:N
for m in 0:min(n, n′)
Expand All @@ -169,34 +174,37 @@ function expansion_coefficients(𝐓::AxisymmetricTransitionMatrix{CT, N, V, T},
end

for n₁ in 0:(N + n)
# Calculate A1 and A2
for n′ in max(1, abs(n - n₁)):min(N, n₁ + n)
a₁ = 0.0im
a₂ = 0.0im
for m₁ in (-min(n, n′)):min(n, n′)
cg = clebschgordan(n, m₁, n₁, 0, n′)
a₁ += cg * T1[m₁, n′, n]
a₂ += cg * T2[m₁, n′, n]
end
a₁ *= ci[n′ - n] / ss[n′]
a₂ *= ci[n′ - n] / ss[n′]
A1[n′, n] = a₁
A2[n′, n] = a₂
end

# Calculate B1 and B2
for m in max(1 - n₁, -n):min(n₁ + 1, n)
# Calculate B1 and B2
b₁ = 0.0im
b₂ = 0.0im

for n′ in max(1, abs(n - n₁)):min(N, n₁ + n)
a₁ = 0.0im
a₂ = 0.0im
for m₁ in (-min(n, n′)):min(n, n′)
cg1 = clebschgordan(n, m₁, n₁, 0, n′)
a₁ += cg1 * T1[m₁, n′, n]
a₂ += cg1 * T2[m₁, n′, n]
end
a₁ *= ci[n′ - n] / ss[n′]
a₂ *= ci[n′ - n] / ss[n′]

cg = clebschgordan(n, m, n₁, 1 - m, n′)
b₁ += cg * a₁
b₂ += cg * a₂
b₁ += cg * A1[n′, n]
b₂ += cg * A2[n′, n]
end

B1[n₁, m, n] = b₁
B2[n₁, m, n] = b₂
end
end
end

# Calculate D
@debug "Calculating D..."
D₀₀ = OffsetArray(zeros(2N + 1, N, N), (-N):N, 1:N, 1:N)
D₀₋₀ = OffsetArray(zeros(2N + 1, N, N), (-N):N, 1:N, 1:N)
D₂₂ = OffsetArray(zeros(2N + 1, N, N), (-N):N, 1:N, 1:N)
Expand Down Expand Up @@ -229,7 +237,7 @@ function expansion_coefficients(𝐓::AxisymmetricTransitionMatrix{CT, N, V, T},
1:N,
1:N)

# Calculate g
@debug "Calculating g..."
g₀₀ = OffsetArray(zeros(2N + 1), 0:(2N))
g₀₋₀ = OffsetArray(zeros(2N + 1), 0:(2N))
g₂₂ = OffsetArray(zeros(2N + 1), 0:(2N))
Expand Down

0 comments on commit 08802e8

Please sign in to comment.