diff --git a/src/energy_degradation.jl b/src/energy_degradation.jl index cf458f5..64453f0 100644 --- a/src/energy_degradation.jl +++ b/src/energy_degradation.jl @@ -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) @@ -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 @@ -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 @@ -140,11 +83,9 @@ 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)) @@ -152,8 +93,8 @@ function add_ionization_collisions!(Q, Ie, h_atm, t, n, σ, E_levels, cascading, 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 @@ -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)) @@ -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 @@ -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] + @@ -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 @@ -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 @@ -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