Skip to content

Commit

Permalink
clean code and rename functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@egavazzi
egavazzi committed Oct 17, 2024
1 parent c064b22 commit fdf36a4
Showing 1 changed file with 45 additions and 102 deletions.
147 changes: 45 additions & 102 deletions src/energy_degradation.jl
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,7 @@
using SparseArrays



#################################################################################
# B2B matrix #
#################################################################################

using SparseArrays
function make_big_B2B_matrix(B2B_inelastic, n, h_atm)
idx1 = Vector{Float64}(undef, 0)
idx2 = Vector{Float64}(undef, 0)
Expand All @@ -26,82 +22,31 @@ end
#################################################################################
# Inelastic collisions #
#################################################################################

function add_inelastic_collisions!(Q, Ie, h_atm, n, σ, E_levels, B2B_inelastic, E, dE, iE)
Ie_degraded = Matrix{Float64}(undef, size(Ie, 1), size(Ie, 2))

# Multiply each element of B2B with n (density vector) and resize to get a matrix that can
# be multiplied with Ie
AB2B = make_big_B2B_matrix(B2B_inelastic, n, h_atm)
Ie_scatter = AB2B * @view(Ie[:, :, iE])

# Loop over the energy levels of the collisions with the i-th neutral species
for i_level in 2:size(E_levels, 1)
if E_levels[i_level, 2] <= 0 # these collisions should not produce secondary e-
# The flux of e- degraded from energy bin [E[iE], E[iE] + dE[iE]] to any lower energy
# bin by excitation of the E_levels[i_level] state of the current species.
# The second factor corrects for the case where the energy loss is maller than the width
# in energy of the energy bin. That is, when dE[iE] > E_levels[i_level,1], only the
# fraction E_levels[i_level,1] / dE[iE] is lost from the energy bin [E[iE], E[iE] + dE[iE]].

Ie_degraded .= (σ[i_level, iE] * min(1, E_levels[i_level, 1] / dE[iE])) .* Ie_scatter

# Find the energy bins where the e- in the current energy bin will degrade when losing
# E_levels[i_level, 1] eV
i_degrade = intersect( findall(x -> x > E[iE] - E_levels[i_level, 1], E + dE), # find lowest bin
findall(x -> x < E[iE] + dE[iE] - E_levels[i_level,1], E)) # find highest bin

partition_fraction = zeros(size(i_degrade)) # initialise
if !isempty(i_degrade) && i_degrade[1] < iE
# Distribute the degrading e- between those bins
partition_fraction[1] = min(1, (E[i_degrade[1]] .+ dE[i_degrade[1]] .-
E[iE] .+ E_levels[i_level, 1]) / dE[iE])
if length(i_degrade) > 2
partition_fraction[2:end-1] = min.(1, dE[i_degrade[2:end-1]] / dE[iE])
end
partition_fraction[end] = min(1, (E[iE] .+ dE[iE] .- E[i_degrade[end]] .-
E_levels[i_level, 1]) / dE[iE])
if i_degrade[end] == iE
partition_fraction[end] = 0
end

# normalise
partition_fraction = partition_fraction / sum(partition_fraction)

# and finally calculate the flux of degrading e-
Threads.@threads for i_u in findall(x -> x != 0, partition_fraction)
@view(Q[:, :, i_degrade[i_u]]) .+= max.(0, Ie_degraded) .* partition_fraction[i_u]
end
end
end
end
end

using LoopVectorization
using Polyester
function add_inelastic_collisions_turbo!(Q, Ie, h_atm, n, σ, E_levels, B2B_inelastic, E, dE, iE)
function add_inelastic_collisions!(Q, Ie, h_atm, n, σ, E_levels, B2B_inelastic, E, dE, iE)
Ie_degraded = Matrix{Float64}(undef, size(Ie, 1), size(Ie, 2))

# Multiply each element of B2B with n (density vector) and resize to get a matrix that can
# be multiplied with Ie
# Multiply each element of B2B with n (density vector) and resize to get a matrix that
# can be multiplied with Ie
AB2B = make_big_B2B_matrix(B2B_inelastic, n, h_atm)
Ie_scatter = AB2B * @view(Ie[:, :, iE])

# Loop over the energy levels of the collisions with the i-th neutral species
for i_level in 2:size(E_levels, 1)
if E_levels[i_level, 2] <= 0 # these collisions should not produce secondary e-
# The flux of e- degraded from energy bin [E[iE], E[iE] + dE[iE]] to any lower energy
# bin by excitation of the E_levels[i_level] state of the current species.
# The second factor corrects for the case where the energy loss is maller than the width
# in energy of the energy bin. That is, when dE[iE] > E_levels[i_level,1], only the
# fraction E_levels[i_level,1] / dE[iE] is lost from the energy bin [E[iE], E[iE] + dE[iE]].
# The flux of e- degraded from energy bin [E[iE], E[iE] + dE[iE]] to any lower
# energy bin by excitation of the E_levels[i_level] state of the current
# species. The second factor corrects for the case where the energy loss is
# smaller than the width in energy of the energy bin. That is, when dE[iE] >
# E_levels[i_level,1], only the fraction E_levels[i_level,1] / dE[iE] is lost
# from the energy bin [E[iE], E[iE] + dE[iE]].

Ie_degraded .= (σ[i_level, iE] * min(1, E_levels[i_level, 1] / dE[iE])) .* Ie_scatter

# Find the energy bins where the e- in the current energy bin will degrade when losing
# E_levels[i_level, 1] eV
i_degrade = intersect( findall(x -> x > E[iE] - E_levels[i_level, 1], E + dE), # find lowest bin
findall(x -> x < E[iE] + dE[iE] - E_levels[i_level,1], E)) # find highest bin
# Find the energy bins where the e- in the current energy bin will degrade when
# losing E_levels[i_level, 1] eV
i_degrade = intersect(findall(x -> x > E[iE] - E_levels[i_level, 1], E + dE), # find lowest bin
findall(x -> x < E[iE] + dE[iE] - E_levels[i_level,1], E)) # find highest bin

partition_fraction = zeros(size(i_degrade)) # initialise
if !isempty(i_degrade) && i_degrade[1] < iE
Expand All @@ -123,8 +68,6 @@ function add_inelastic_collisions_turbo!(Q, Ie, h_atm, n, σ, E_levels, B2B_inel
# and finally calculate the flux of degrading e-
@tturbo for i_u in eachindex(findall(x -> x != 0, partition_fraction))
for j in axes(Q, 2)
# @batch per=thread for i_u in eachindex(findall(x -> x != 0, partition_fraction))
# @turbo for j in axes(Q, 2)
for k in axes(Q, 1)
Q[k, j, i_degrade[i_u]] += max(0, Ie_degraded[k, j]) * partition_fraction[i_u]
end
Expand All @@ -140,20 +83,18 @@ end
#################################################################################
# Ionization collisions #
#################################################################################
function add_ionization_collisions!(Q, Ie, h_atm, t, n, σ, E_levels, cascading, E, dE, iE,
BeamWeight, μ_center)

using Polyester
function add_ionization_collisions!(Q, Ie, h_atm, t, n, σ, E_levels, cascading, E, dE, iE, BeamWeight, μ_center, Nthreads = 6)
# Nthreads is set to 6 by default as it seems to be optimal on my machine
# But on Revontuli, something like 20 is more optimal
Ionization = zeros(size(Ie, 1), size(Ie, 2))
Ionizing = Matrix{Float64}(undef, size(Ie, 1), size(Ie, 2))
n_repeated_over_μt = repeat(n, length(μ_center), length(t))
n_repeated_over_t = repeat(n, 1, length(t))

for i_level = 2:size(E_levels, 1)
if E_levels[i_level, 2] > 0 # these collisions should produce secondary e-
# Find the energy bins where the e- in the current energy bin will degrade when losing
# E_levels[i_level, 1] eV
# Find the energy bins where the e- in the current energy bin will degrade when
# losing E_levels[i_level, 1] eV
i_degrade = intersect(findall(x -> x > E[iE] - E_levels[i_level, 1], E + dE), # find lowest bin
findall(x -> x < E[iE] + dE[iE] - E_levels[i_level,1], E)) # find highest bin

Expand Down Expand Up @@ -192,26 +133,28 @@ function add_ionization_collisions!(Q, Ie, h_atm, t, n, σ, E_levels, cascading,
primary_e_spectra = primary_e_spectra ./ sum(primary_e_spectra)
end

# This is some kind of check that it is not empty?
e_ionized_distribution = max.(secondary_e_spectra, primary_e_spectra)
e_ionized_distribution[isnan.(e_ionized_distribution)] .= 0

# and finally add this to the flux of degrading e-
nbatch = Int(floor(iE / Nthreads)) - 1
nbatch < 1 ? nbatch = 1 : nothing
@batch minbatch=nbatch for iI in 1:(iE - 1)
@view(Q[:, :, iI]) .+= Ionization .* secondary_e_spectra[iI] .+
Ionizing .* primary_e_spectra[iI]
@tturbo for iI in 1:(iE - 1)
for j in axes(Q, 2)
for k in axes(Q, 1)
Q[k, j, iI] += Ionization[k, j] * secondary_e_spectra[iI] +
Ionizing[k, j] * primary_e_spectra[iI]
end
end
end
end
end
end
end



using LoopVectorization
function prepare_ionization_collisions_turbo!(Ie, h_atm, t, n, σ, E_levels, cascading, E, dE, iE,
BeamWeight, μ_center, Ionization_matrix, Ionizing_matrix,
secondary_vector, primary_vector, counter_ionization)
function prepare_ionization_collisions!(Ie, h_atm, t, n, σ, E_levels, cascading, E,
dE, iE, BeamWeight, μ_center,
Ionization_matrix, Ionizing_matrix,
secondary_vector, primary_vector,
counter_ionization)

Ionization = zeros(size(Ie, 1), size(Ie, 2))
Ionizing = Matrix{Float64}(undef, size(Ie, 1), size(Ie, 2))
Expand All @@ -220,8 +163,8 @@ function prepare_ionization_collisions_turbo!(Ie, h_atm, t, n, σ, E_levels, cas

for i_level = 2:size(E_levels, 1)
if E_levels[i_level, 2] > 0 # these collisions should produce secondary e-
# Find the energy bins where the e- in the current energy bin will degrade when losing
# E_levels[i_level, 1] eV
# Find the energy bins where the e- in the current energy bin will degrade when
# losing E_levels[i_level, 1] eV
i_degrade = intersect(findall(x -> x > E[iE] - E_levels[i_level, 1], E + dE), # find lowest bin
findall(x -> x < E[iE] + dE[iE] - E_levels[i_level,1], E)) # find highest bin

Expand Down Expand Up @@ -274,16 +217,16 @@ function prepare_ionization_collisions_turbo!(Ie, h_atm, t, n, σ, E_levels, cas
end
end

function add_ionization_collisions_turbo!(Q, iE, Ionization_matrix, Ionizing_matrix,


function add_ionization_collisions_batch!(Q, iE, Ionization_matrix, Ionizing_matrix,
secondary_vector, primary_vector)
# We need to add all the ionization collisions in three steps (15 different collisions
# split over groups of 5)
# split over groups of 5. This seems to be optimal on Revontuli)
for i_loop in 1:3
idx = (i_loop - 1) * 5
@tturbo for iI in 1:(iE - 1)
for j in axes(Q, 2)
# @batch per=thread for iI in 1:(iE - 1)
# @turbo for j in axes(Q, 2)
for k in axes(Q, 1)
Q[k, j, iI] += Ionization_matrix[idx + 1][k, j] * secondary_vector[idx + 1][iI] +
Ionizing_matrix[idx + 1][k, j] * primary_vector[idx + 1][iI] +
Expand All @@ -306,9 +249,9 @@ end
#################################################################################
# Update Q #
#################################################################################

function update_Q!(Q, Ie, h_atm, t, ne, Te, n_neutrals, σ_neutrals, E_levels_neutrals, B2B_inelastic_neutrals,
cascading_neutrals, E, dE, iE, BeamWeight, μ_center, Nthreads = 6)
function update_Q!(Q, Ie, h_atm, t, ne, Te, n_neutrals, σ_neutrals, E_levels_neutrals,
B2B_inelastic_neutrals, cascading_neutrals, E, dE, iE, BeamWeight,
μ_center)

# e-e collisions
@views if iE > 1
Expand All @@ -325,7 +268,7 @@ function update_Q!(Q, Ie, h_atm, t, ne, Te, n_neutrals, σ_neutrals, E_levels_ne
cascading = cascading_neutrals[i]; # Cascading function for the current i-th species

add_inelastic_collisions!(Q, Ie, h_atm, n, σ, E_levels, B2B_inelastic, E, dE, iE)
add_ionization_collisions!(Q, Ie, h_atm, t, n, σ, E_levels, cascading, E, dE, iE, BeamWeight, μ_center, Nthreads)
add_ionization_collisions!(Q, Ie, h_atm, t, n, σ, E_levels, cascading, E, dE, iE, BeamWeight, μ_center)
end
end

Expand All @@ -348,12 +291,12 @@ function update_Q_turbo!(Q, Ie, h_atm, t, ne, Te, n_neutrals, σ_neutrals, E_lev
B2B_inelastic = B2B_inelastic_neutrals[i]; # Array with the probablities of scattering from beam to beam
cascading = cascading_neutrals[i]; # Cascading function for the current i-th species

add_inelastic_collisions_turbo!(Q, Ie, h_atm, n, σ, E_levels, B2B_inelastic, E, dE, iE)
prepare_ionization_collisions_turbo!(Ie, h_atm, t, n, σ, E_levels, cascading, E, dE, iE,
add_inelastic_collisions!(Q, Ie, h_atm, n, σ, E_levels, B2B_inelastic, E, dE, iE)
prepare_ionization_collisions!(Ie, h_atm, t, n, σ, E_levels, cascading, E, dE, iE,
BeamWeight, μ_center, Ionization_matrix, Ionizing_matrix,
secondary_vector, primary_vector, counter_ionization)
end

add_ionization_collisions_turbo!(Q, iE, Ionization_matrix, Ionizing_matrix,
add_ionization_collisions_batch!(Q, iE, Ionization_matrix, Ionizing_matrix,
secondary_vector, primary_vector)
end

0 comments on commit fdf36a4

Please sign in to comment.