Add multithreading to electric field calculation

Author: fhagemann

src/ElectricField/ElectricField.jl

@@ -49,20 +49,20 @@ Base.convert(T::Type{ElectricField}, x::NamedTuple) = T(x)
 
 
 
-function ElectricField(epot::ElectricPotential{T, 3, S}, point_types::PointTypes{T}) where {T, S}
-    return ElectricField{T, 3, S, typeof(grid.axes)}(get_electric_field_from_potential( epot, point_types ), epot.grid)
+function ElectricField(epot::ElectricPotential{T, 3, S}, point_types::PointTypes{T}; use_nthreads::Int = Base.Threads.nthreads()) where {T, S}
+    return ElectricField{T, 3, S, typeof(grid.axes)}(get_electric_field_from_potential( epot, point_types; use_nthreads ), epot.grid)
 end
 
 
-function get_electric_field_from_potential(epot::ElectricPotential{T, 3, Cylindrical}, point_types::PointTypes{T}, fieldvector_coordinates=:xyz)::ElectricField{T, 3, Cylindrical} where {T <: SSDFloat}
+function get_electric_field_from_potential(epot::ElectricPotential{T, 3, Cylindrical}, point_types::PointTypes{T}, fieldvector_coordinates=:xyz; use_nthreads::Int = Base.Threads.threadid())::ElectricField{T, 3, Cylindrical} where {T <: SSDFloat}
     p = epot.data
     axr::Vector{T} = collect(epot.grid.axes[1])
     axφ::Vector{T} = collect(epot.grid.axes[2])
     axz::Vector{T} = collect(epot.grid.axes[3])
 
     cyclic::T = epot.grid.axes[2].interval.right
     ef = Array{SVector{3, T}}(undef, size(p)...)
-    for iz in 1:size(ef, 3)
+    @onthreads 1:use_nthreads for ix in workpart(1:size(ef, 3), 1:use_nthreads, Base.Threads.threadid())
         for iφ in 1:size(ef, 2)
             for ir in 1:size(ef, 1)
                 ### r ###
@@ -212,7 +212,7 @@ end
 
 
 
-function get_electric_field_from_potential(epot::ElectricPotential{T, 3, Cartesian}, point_types::PointTypes{T})::ElectricField{T, 3, Cartesian} where {T <: SSDFloat}
+function get_electric_field_from_potential(epot::ElectricPotential{T, 3, Cartesian}, point_types::PointTypes{T}; use_nthreads::Int = Base.Threads.nthreads())::ElectricField{T, 3, Cartesian} where {T <: SSDFloat}
     axx::Vector{T} = collect(epot.grid.axes[1])
     axy::Vector{T} = collect(epot.grid.axes[2])
     axz::Vector{T} = collect(epot.grid.axes[3])
@@ -222,7 +222,7 @@ function get_electric_field_from_potential(epot::ElectricPotential{T, 3, Cartesi
 
     ef::Array{SVector{3, T}} = Array{SVector{3, T}}(undef, size(epot.data))
 
-    for ix in eachindex(axx)
+    @onthreads 1:use_nthreads for ix in workpart(eachindex(axx), 1:use_nthreads, Base.Threads.threadid())
         for iy in eachindex(axy)
             for iz in eachindex(axz)
                 if ix - 1 < 1

src/Simulation/Simulation.jl

@@ -1167,7 +1167,7 @@ Calculates the [`ElectricField`](@ref) from the [`ElectricPotential`](@ref) stor
 !!! note 
     This method only works if `sim.electric_potential` has already been calculated and is not `missing`.
 """
-function calculate_electric_field!(sim::Simulation{T, CS}; n_points_in_φ::Union{Missing, Int} = missing)::Nothing where {T <: SSDFloat, CS}
+function calculate_electric_field!(sim::Simulation{T, CS}; n_points_in_φ::Union{Missing, Int} = missing, use_nthreads::Int = Base.Threads.nthreads())::Nothing where {T <: SSDFloat, CS}
     @assert !ismissing(sim.electric_potential) "Electric potential has not been calculated yet. Please run `calculate_electric_potential!(sim)` first."
     periodicity::T = width(sim.world.intervals[2])
     e_pot, point_types = if CS == Cylindrical && periodicity == T(0) # 2D, only one point in φ
@@ -1181,15 +1181,15 @@ function calculate_electric_field!(sim::Simulation{T, CS}; n_points_in_φ::Union
             end
         end
         get_2π_potential(sim.electric_potential, n_points_in_φ = n_points_in_φ),
-        get_2π_potential(sim.point_types,  n_points_in_φ = n_points_in_φ);
+        get_2π_potential(sim.point_types,  n_points_in_φ = n_points_in_φ)
     elseif CS == Cylindrical
         get_2π_potential(sim.electric_potential),
         get_2π_potential(sim.point_types)
     else
         sim.electric_potential,
         sim.point_types
     end
-    sim.electric_field = get_electric_field_from_potential(e_pot, point_types);
+    sim.electric_field = get_electric_field_from_potential(e_pot, point_types; use_nthreads)
     nothing
 end