diff --git a/src/archetypes/spatial_dist.h b/src/archetypes/spatial_dist.h index 68477208..1a2d87f9 100644 --- a/src/archetypes/spatial_dist.h +++ b/src/archetypes/spatial_dist.h @@ -1,141 +1,152 @@ -/** - * @file archetypes/spatial_dist.hpp - * @brief Spatial distribution class passed to injectors - * @implements - * - arch::SpatialDistribution<> - * - arch::Uniform<> : arch::SpatialDistribution<> - * - arch::Replenish<> : arch::SpatialDistribution<> - * - arch::ReplenishUniform<> : arch::SpatialDistribution<> - * @namespace - * - arch:: - * @note - * Instances of these functors take coordinate position in code units - * and return a number between 0 and 1 that represents the spatial distribution - */ - -#ifndef ARCHETYPES_SPATIAL_DIST_HPP -#define ARCHETYPES_SPATIAL_DIST_HPP - -#include "enums.h" -#include "global.h" - -#include "arch/kokkos_aliases.h" -#include "utils/error.h" -#include "utils/numeric.h" - -namespace arch { - using namespace ntt; - - template - struct SpatialDistribution { - static constexpr bool is_spatial_dist { true }; - static_assert(M::is_metric, "M must be a metric class"); - - SpatialDistribution(const M& metric) : metric { metric } {} - - protected: - const M metric; - }; - - template - struct Uniform : public SpatialDistribution { - Uniform(const M& metric) : SpatialDistribution { metric } {} - - Inline auto operator()(const coord_t&) const -> real_t { - return ONE; - } - }; - - template - struct Replenish : public SpatialDistribution { - using SpatialDistribution::metric; - const ndfield_t density; - const idx_t idx; - - const T target_density; - const real_t target_max_density; - - Replenish(const M& metric, - const ndfield_t& density, - idx_t idx, - const T& target_density, - real_t target_max_density) - : SpatialDistribution { metric } - , density { density } - , idx { idx } - , target_density { target_density } - , target_max_density { target_max_density } {} - - Inline auto operator()(const coord_t& x_Ph) const -> real_t { - coord_t x_Cd { ZERO }; - metric.template convert(x_Ph, x_Cd); - real_t dens { ZERO }; - if constexpr (M::Dim == Dim::_1D) { - dens = density(static_cast(x_Cd[0]) + N_GHOSTS, idx); - } else if constexpr (M::Dim == Dim::_2D) { - dens = density(static_cast(x_Cd[0]) + N_GHOSTS, - static_cast(x_Cd[1]) + N_GHOSTS, - idx); - } else if constexpr (M::Dim == Dim::_3D) { - dens = density(static_cast(x_Cd[0]) + N_GHOSTS, - static_cast(x_Cd[1]) + N_GHOSTS, - static_cast(x_Cd[2]) + N_GHOSTS, - idx); - } else { - raise::KernelError(HERE, "Invalid dimension"); - } - const auto target = target_density(x_Ph); - if (0.9 * target > dens) { - return (target - dens) / target_max_density; - } else { - return ZERO; - } - } - }; - - template - struct ReplenishUniform : public SpatialDistribution { - using SpatialDistribution::metric; - const ndfield_t density; - const idx_t idx; - - const real_t target_density; - - ReplenishUniform(const M& metric, - const ndfield_t& density, - idx_t idx, - real_t target_density) - : SpatialDistribution { metric } - , density { density } - , idx { idx } - , target_density { target_density } {} - - Inline auto operator()(const coord_t& x_Ph) const -> real_t { - coord_t x_Cd { ZERO }; - metric.template convert(x_Ph, x_Cd); - real_t dens { ZERO }; - if constexpr (M::Dim == Dim::_1D) { - dens = density(static_cast(x_Cd[0]) + N_GHOSTS, idx); - } else if constexpr (M::Dim == Dim::_2D) { - dens = density(static_cast(x_Cd[0]) + N_GHOSTS, - static_cast(x_Cd[1]) + N_GHOSTS, - idx); - } else if constexpr (M::Dim == Dim::_3D) { - dens = density(static_cast(x_Cd[0]) + N_GHOSTS, - static_cast(x_Cd[1]) + N_GHOSTS, - static_cast(x_Cd[2]) + N_GHOSTS, - idx); - } else { - raise::KernelError(HERE, "Invalid dimension"); - } - if (0.9 * target_density > dens) { - return (target_density - dens) / target_density; - } else { - return ZERO; - } - } - }; - -} // namespace arch - -#endif // ARCHETYPES_SPATIAL_DIST_HPP +/** + * @file archetypes/spatial_dist.hpp + * @brief Spatial distribution class passed to injectors + * @implements + * - arch::SpatialDistribution<> + * - arch::Uniform<> : arch::SpatialDistribution<> + * - arch::Replenish<> : arch::SpatialDistribution<> + * - arch::ReplenishUniform<> : arch::SpatialDistribution<> + * @namespace + * - arch:: + * @note + * Instances of these functors take coordinate position in code units + * and return a number between 0 and 1 that represents the spatial distribution + */ + +#ifndef ARCHETYPES_SPATIAL_DIST_HPP +#define ARCHETYPES_SPATIAL_DIST_HPP + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "utils/error.h" +#include "utils/numeric.h" + +namespace arch { + using namespace ntt; + + template + struct SpatialDistribution { + static constexpr bool is_spatial_dist { true }; + static_assert(M::is_metric, "M must be a metric class"); + + SpatialDistribution(const M& metric) : metric { metric } {} + + protected: + const M metric; + }; + + template + struct Uniform : public SpatialDistribution { + Uniform(const M& metric) : SpatialDistribution { metric } {} + + Inline auto operator()(const coord_t&, + real_t& nppc_distribution, + real_t& weight_distribution) const { + nppc_distribution = ONE; + weight_distribution = ONE; + } + }; + + template + struct Replenish : public SpatialDistribution { + using SpatialDistribution::metric; + const ndfield_t density; + const idx_t idx; + + const T target_density; + const real_t target_max_density; + + Replenish(const M& metric, + const ndfield_t& density, + idx_t idx, + const T& target_density, + real_t target_max_density) + : SpatialDistribution { metric } + , density { density } + , idx { idx } + , target_density { target_density } + , target_max_density { target_max_density } {} + + Inline auto operator()(const coord_t& x_Ph, + real_t& nppc_distribution, + real_t& weight_distribution) const { + coord_t x_Cd { ZERO }; + metric.template convert(x_Ph, x_Cd); + real_t dens { ZERO }; + if constexpr (M::Dim == Dim::_1D) { + dens = density(static_cast(x_Cd[0]) + N_GHOSTS, idx); + } else if constexpr (M::Dim == Dim::_2D) { + dens = density(static_cast(x_Cd[0]) + N_GHOSTS, + static_cast(x_Cd[1]) + N_GHOSTS, + idx); + } else if constexpr (M::Dim == Dim::_3D) { + dens = density(static_cast(x_Cd[0]) + N_GHOSTS, + static_cast(x_Cd[1]) + N_GHOSTS, + static_cast(x_Cd[2]) + N_GHOSTS, + idx); + } else { + raise::KernelError(HERE, "Invalid dimension"); + } + const auto target = target_density(x_Ph); + if (0.9 * target > dens) { + nppc_distribution = (target - dens) / target_max_density; + weight_distribution = ONE; + } else { + nppc_distribution = ZERO; + weight_distribution = ONE; + } + } + }; + + template + struct ReplenishUniform : public SpatialDistribution { + using SpatialDistribution::metric; + const ndfield_t density; + const idx_t idx; + + const real_t target_density; + + ReplenishUniform(const M& metric, + const ndfield_t& density, + idx_t idx, + real_t target_density) + : SpatialDistribution { metric } + , density { density } + , idx { idx } + , target_density { target_density } {} + + Inline auto operator()(const coord_t& x_Ph, + real_t& nppc_distribution, + real_t& weight_distribution) const { + coord_t x_Cd { ZERO }; + metric.template convert(x_Ph, x_Cd); + real_t dens { ZERO }; + if constexpr (M::Dim == Dim::_1D) { + dens = density(static_cast(x_Cd[0]) + N_GHOSTS, idx); + } else if constexpr (M::Dim == Dim::_2D) { + dens = density(static_cast(x_Cd[0]) + N_GHOSTS, + static_cast(x_Cd[1]) + N_GHOSTS, + idx); + } else if constexpr (M::Dim == Dim::_3D) { + dens = density(static_cast(x_Cd[0]) + N_GHOSTS, + static_cast(x_Cd[1]) + N_GHOSTS, + static_cast(x_Cd[2]) + N_GHOSTS, + idx); + } else { + raise::KernelError(HERE, "Invalid dimension"); + } + if (0.9 * target_density > dens) { + nppc_distribution = (target_density - dens) / target_density; + weight_distribution = ONE; + } else { + nppc_distribution = ZERO; + weight_distribution = ONE; + } + } + }; + +} // namespace arch + +#endif // ARCHETYPES_SPATIAL_DIST_HPP diff --git a/src/engines/grpic.hpp b/src/engines/grpic.hpp index 35e544a9..a8bc8e67 100644 --- a/src/engines/grpic.hpp +++ b/src/engines/grpic.hpp @@ -1,1145 +1,1145 @@ -/** - * @file engines/grpic.hpp - * @brief Simulation engien class which specialized on GRPIC - * @implements - * - ntt::GRPICEngine<> : ntt::Engine<> - * @cpp: - * - grpic.cpp - * @namespaces: - * - ntt:: - */ - -#ifndef ENGINES_GRPIC_GRPIC_H -#define ENGINES_GRPIC_GRPIC_H - -#include "enums.h" -#include "global.h" - -#include "arch/kokkos_aliases.h" -#include "utils/log.h" -#include "utils/numeric.h" -#include "utils/timer.h" -#include "utils/toml.h" - -#include "framework/domain/domain.h" -#include "framework/parameters.h" - -#include "engines/engine.hpp" -#include "kernels/ampere_gr.hpp" -#include "kernels/aux_fields_gr.hpp" -#include "kernels/currents_deposit.hpp" -#include "kernels/digital_filter.hpp" -#include "kernels/faraday_gr.hpp" -#include "kernels/fields_bcs.hpp" -#include "kernels/particle_pusher_gr.hpp" -#include "pgen.hpp" - -#include -#include - -#include -#include - -namespace ntt { - - enum class gr_getE { - D0_B, - D_B0 - }; - enum class gr_getH { - D_B0, - D0_B0 - }; - enum class gr_faraday { - aux, - main - }; - enum class gr_ampere { - init, - aux, - main - }; - enum class gr_bc { - main, - aux, - curr - }; - - template - class GRPICEngine : public Engine { - using base_t = Engine; - using pgen_t = user::PGen; - using domain_t = Domain; - // constexprs - using base_t::pgen_is_ok; - // contents - using base_t::m_metadomain; - using base_t::m_params; - using base_t::m_pgen; - // methods - using base_t::init; - // variables - using base_t::dt; - using base_t::max_steps; - using base_t::runtime; - using base_t::step; - using base_t::time; - - public: - static constexpr auto S { SimEngine::GRPIC }; - - GRPICEngine(SimulationParams& params) : base_t { params } {} - - ~GRPICEngine() = default; - - void step_forward(timer::Timers& timers, domain_t& dom) override { - const auto fieldsolver_enabled = m_params.template get( - "algorithms.toggles.fieldsolver"); - const auto deposit_enabled = m_params.template get( - "algorithms.toggles.deposit"); - const auto clear_interval = m_params.template get( - "particles.clear_interval"); - - if (step == 0) { - if (fieldsolver_enabled) { - // communicate fields and apply BCs on the first timestep - /** - * Initially: em0::B -- - * em0::D -- - * em::B at -1/2 - * em::D at -1/2 - * - * cur0::J -- - * cur::J -- - * - * aux::E -- - * aux::H -- - * - * x_prtl at -1/2 - * u_prtl at -1/2 - */ - - /** - * em0::D, em::D, em0::B, em::B <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, - Comm::B | Comm::B0 | Comm::D | Comm::D0); - FieldBoundaries(dom, BC::B | BC::D, gr_bc::main); - - /** - * em0::B <- em::B - * em0::D <- em::D - * - * Now: em0::B & em0::D at -1/2 - */ - CopyFields(dom); - - /** - * aux::E <- alpha * em::D + beta x em0::B - * aux::H <- alpha * em::B0 - beta x em::D - * - * Now: aux::E & aux::H at -1/2 - */ - ComputeAuxE(dom, gr_getE::D_B0); - ComputeAuxH(dom, gr_getH::D_B0); - - /** - * aux::E, aux::H <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::H | Comm::E); - FieldBoundaries(dom, BC::H | BC::E, gr_bc::aux); - - /** - * em0::B <- (em0::B) <- -curl aux::E - * - * Now: em0::B at 0 - */ - Faraday(dom, gr_faraday::aux, HALF); - - /** - * em0::B, em::B <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); - FieldBoundaries(dom, BC::B, gr_bc::main); - - /** - * em::D <- (em0::D) <- curl aux::H - * - * Now: em::D at 0 - */ - Ampere(dom, gr_ampere::init, HALF); - - /** - * em0::D, em::D <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); - FieldBoundaries(dom, BC::D, gr_bc::main); - - /** - * aux::E <- alpha * em::D + beta x em0::B - * aux::H <- alpha * em0::B - beta x em::D - * - * Now: aux::E & aux::H at 0 - */ - ComputeAuxE(dom, gr_getE::D_B0); - ComputeAuxH(dom, gr_getH::D_B0); - - /** - * aux::E, aux::H <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::H | Comm::E); - FieldBoundaries(dom, BC::H | BC::E, gr_bc::aux); - - // !ADD: GR -- particles? - - /** - * em0::B <- (em::B) <- -curl aux::E - * - * Now: em0::B at 1/2 - */ - Faraday(dom, gr_faraday::main, ONE); - /** - * em0::B, em::B <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); - FieldBoundaries(dom, BC::B, gr_bc::main); - - /** - * em0::D <- (em0::D) <- curl aux::H - * - * Now: em0::D at 1/2 - */ - Ampere(dom, gr_ampere::aux, ONE); - /** - * em0::D, em::D <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); - FieldBoundaries(dom, BC::D, gr_bc::main); - - /** - * aux::H <- alpha * em0::B - beta x em0::D - * - * Now: aux::H at 1/2 - */ - ComputeAuxH(dom, gr_getH::D0_B0); - /** - * aux::H <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::H); - FieldBoundaries(dom, BC::H, gr_bc::aux); - - /** - * em0::D <- (em::D) <- curl aux::H - * - * Now: em0::D at 1 - * em::D at 0 - */ - Ampere(dom, gr_ampere::main, ONE); - /** - * em0::D, em::D <- boundary conditions - */ - m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); - FieldBoundaries(dom, BC::D, gr_bc::main); - - /** - * em::D <-> em0::D - * em::B <-> em0::B - * em::J <-> em0::J - */ - SwapFields(dom); - /** - * Finally: em0::B at -1/2 - * em0::D at 0 - * em::B at 1/2 - * em::D at 1 - * - * cur0::J -- - * cur::J -- - * - * aux::E -- - * aux::H -- - * - * x_prtl at 1 - * u_prtl at 1/2 - */ - } else { - /** - * em0::B <- em::B - * em0::D <- em::D - * - * Now: em0::B & em0::D at -1/2 - */ - CopyFields(dom); - } - } - - /** - * Initially: em0::B at n-3/2 - * em0::D at n-1 - * em::B at n-1/2 - * em::D at n - * - * cur0::J -- - * cur::J at n-1/2 - * - * aux::E -- - * aux::H -- - * - * x_prtl at n - * u_prtl at n-1/2 - */ - - if (fieldsolver_enabled) { - timers.start("FieldSolver"); - /** - * em0::D <- (em0::D + em::D) / 2 - * em0::B <- (em0::B + em::B) / 2 - * - * Now: em0::D at n-1/2 - * em0::B at n-1 - */ - TimeAverageDB(dom); - /** - * aux::E <- alpha * em0::D + beta x em::B - * - * Now: aux::E at n-1/2 - */ - ComputeAuxE(dom, gr_getE::D0_B); - timers.stop("FieldSolver"); - - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::E); - timers.stop("Communications"); - timers.start("FieldBoundaries"); - /** - * aux::E <- boundary conditions - */ - FieldBoundaries(dom, BC::E, gr_bc::aux); - timers.stop("FieldBoundaries"); - - timers.start("FieldSolver"); - /** - * em0::B <- (em0::B) <- -curl aux::E - * - * Now: em0::B at n - */ - Faraday(dom, gr_faraday::aux, ONE); - timers.stop("FieldSolver"); - - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); - timers.stop("Communications"); - /** - * em0::B, em::B <- boundary conditions - */ - timers.start("FieldBoundaries"); - FieldBoundaries(dom, BC::B, gr_bc::main); - timers.stop("FieldBoundaries"); - - timers.start("FieldSolver"); - /** - * aux::H <- alpha * em0::B - beta x em::D - * - * Now: aux::H at n - */ - ComputeAuxH(dom, gr_getH::D_B0); - timers.stop("FieldSolver"); - - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::H); - timers.stop("Communications"); - timers.start("FieldBoundaries"); - /** - * aux::H <- boundary conditions - */ - FieldBoundaries(dom, BC::H, gr_bc::aux); - timers.stop("FieldBoundaries"); - } - - { - /** - * x_prtl, u_prtl <- em::D, em0::B - * - * Now: x_prtl at n + 1, u_prtl at n + 1/2 - */ - timers.start("ParticlePusher"); - ParticlePush(dom); - timers.stop("ParticlePusher"); - - /** - * cur0::J <- current deposition - * - * Now: cur0::J at n+1/2 - */ - if (deposit_enabled) { - timers.start("CurrentDeposit"); - Kokkos::deep_copy(dom.fields.cur0, ZERO); - CurrentsDeposit(dom); - timers.stop("CurrentDeposit"); - - timers.start("Communications"); - m_metadomain.SynchronizeFields(dom, Comm::J); - m_metadomain.CommunicateFields(dom, Comm::J); - timers.stop("Communications"); - - timers.start("FieldBoundaries"); - FieldBoundaries(dom, BC::J, gr_bc::curr); - timers.stop("FieldBoundaries"); - - timers.start("CurrentFiltering"); - CurrentsFilter(dom); - timers.stop("CurrentFiltering"); - } - - timers.start("Communications"); - m_metadomain.CommunicateParticles(dom); - timers.stop("Communications"); - } - - if (fieldsolver_enabled) { - timers.start("FieldSolver"); - if (deposit_enabled) { - /** - * cur::J <- (cur0::J + cur::J) / 2 - * - * Now: cur::J at n - */ - TimeAverageJ(dom); - } - /** - * aux::Е <- alpha * em::D + beta x em0::B - * - * Now: aux::Е at n - */ - ComputeAuxE(dom, gr_getE::D_B0); - timers.stop("FieldSolver"); - - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::E); - timers.stop("Communications"); - timers.start("FieldBoundaries"); - /** - * aux::Е <- boundary conditions - */ - FieldBoundaries(dom, BC::E, gr_bc::aux); - timers.stop("FieldBoundaries"); - - timers.start("FieldSolver"); - /** - * em0::B <- (em::B) <- -curl aux::E - * - * Now: em0::B at n+1/2 - * em::B at n-1/2 - */ - Faraday(dom, gr_faraday::main, ONE); - timers.stop("FieldSolver"); - - /** - * em0::B, em::B <- boundary conditions - */ - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); - timers.stop("Communications"); - timers.start("FieldBoundaries"); - FieldBoundaries(dom, BC::B, gr_bc::main); - timers.stop("FieldBoundaries"); - - timers.start("FieldSolver"); - /** - * em0::D <- (em0::D) <- curl aux::H - * - * Now: em0::D at n+1/2 - */ - Ampere(dom, gr_ampere::aux, ONE); - timers.stop("FieldSolver"); - - if (deposit_enabled) { - timers.start("FieldSolver"); - /** - * em0::D <- (em0::D) <- cur::J - * - * Now: em0::D at n+1/2 - */ - AmpereCurrents(dom, gr_ampere::aux); - timers.stop("FieldSolver"); - } - - /** - * em0::D, em::D <- boundary conditions - */ - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); - timers.stop("Communications"); - timers.start("FieldBoundaries"); - FieldBoundaries(dom, BC::D, gr_bc::main); - timers.stop("FieldBoundaries"); - - timers.start("FieldSolver"); - /** - * aux::H <- alpha * em0::B - beta x em0::D - * - * Now: aux::H at n+1/2 - */ - ComputeAuxH(dom, gr_getH::D0_B0); - timers.stop("FieldSolver"); - - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::H); - timers.stop("Communications"); - timers.start("FieldBoundaries"); - /** - * aux::H <- boundary conditions - */ - FieldBoundaries(dom, BC::B, gr_bc::aux); - timers.stop("FieldBoundaries"); - - timers.start("FieldSolver"); - /** - * em0::D <- (em::D) <- curl aux::H - * - * Now: em0::D at n+1 - * em::D at n - */ - Ampere(dom, gr_ampere::main, ONE); - timers.stop("FieldSolver"); - - if (deposit_enabled) { - timers.start("FieldSolver"); - /** - * em0::D <- (em0::D) <- cur0::J - * - * Now: em0::D at n+1 - */ - AmpereCurrents(dom, gr_ampere::main); - timers.stop("FieldSolver"); - } - timers.start("FieldSolver"); - /** - * em::D <-> em0::D - * em::B <-> em0::B - * cur::J <-> cur0::J - */ - SwapFields(dom); - timers.stop("FieldSolver"); - - /** - * em0::D, em::D <- boundary conditions - */ - timers.start("Communications"); - m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); - timers.stop("Communications"); - timers.start("FieldBoundaries"); - FieldBoundaries(dom, BC::D, gr_bc::main); - timers.stop("FieldBoundaries"); - } - - if (clear_interval > 0 and step % clear_interval == 0 and step > 0) { - timers.start("PrtlClear"); - m_metadomain.RemoveDeadParticles(dom); - timers.stop("PrtlClear"); - } - - /** - * Finally: em0::B at n-1/2 - * em0::D at n - * em::B at n+1/2 - * em::D at n+1 - * - * cur0::J (at n) - * cur::J at n+1/2 - * - * aux::E (at n+1/2) - * aux::H (at n) - * - * x_prtl at n+1 - * u_prtl at n+1/2 - */ - } - - /* algorithm substeps --------------------------------------------------- */ - void FieldBoundaries(domain_t& domain, BCTags tags, const gr_bc& g) { - if (g == gr_bc::main) { - for (auto& direction : dir::Directions::orth) { - if (m_metadomain.mesh().flds_bc_in(direction) == FldsBC::MATCH) { - MatchFieldsIn(direction, domain, tags, g); - } else if (domain.mesh.flds_bc_in(direction) == FldsBC::AXIS) { - AxisFieldsIn(direction, domain, tags); - } else if (m_metadomain.mesh().flds_bc_in(direction) == FldsBC::CUSTOM) { - CustomFieldsIn(direction, domain, tags, g); - } else if (domain.mesh.flds_bc_in(direction) == FldsBC::HORIZON) { - HorizonFieldsIn(direction, domain, tags, g); - } - } // loop over directions - } else if (g == gr_bc::aux) { - for (auto& direction : dir::Directions::orth) { - if (domain.mesh.flds_bc_in(direction) == FldsBC::HORIZON) { - HorizonFieldsIn(direction, domain, tags, g); - } - } - } else if (g == gr_bc::curr) { - for (auto& direction : dir::Directions::orth) { - if (domain.mesh.prtl_bc_in(direction) == PrtlBC::ABSORB) { - MatchFieldsIn(direction, domain, tags, g); - } - } - } - } - - void MatchFieldsIn(dir::direction_t direction, - domain_t& domain, - BCTags tags, - const gr_bc& g) { - /** - * match boundaries - */ - const auto ds_array = m_params.template get>( - "grid.boundaries.match.ds"); - const auto dim = direction.get_dim(); - real_t xg_min, xg_max, xg_edge; - auto sign = direction.get_sign(); - - raise::ErrorIf(((dim != in::x1) or (sign < 0)) and (g == gr_bc::curr), - "Absorption of currents only possible in +x1 (+r)", - HERE); - - real_t ds; - if (sign > 0) { // + direction - ds = ds_array[(short)dim].second; - xg_max = m_metadomain.mesh().extent(dim).second; - xg_min = xg_max - ds; - xg_edge = xg_max; - } else { // - direction - ds = ds_array[(short)dim].first; - xg_min = m_metadomain.mesh().extent(dim).first; - xg_max = xg_min + ds; - xg_edge = xg_min; - } - boundaries_t box; - boundaries_t incl_ghosts; - for (unsigned short d { 0 }; d < M::Dim; ++d) { - if (d == static_cast(dim)) { - box.push_back({ xg_min, xg_max }); - incl_ghosts.push_back({ false, true }); - } else { - box.push_back(Range::All); - incl_ghosts.push_back({ true, true }); - } - } - if (not domain.mesh.Intersects(box)) { - return; - } - const auto intersect_range = domain.mesh.ExtentToRange(box, incl_ghosts); - tuple_t range_min { 0 }; - tuple_t range_max { 0 }; - - for (unsigned short d { 0 }; d < M::Dim; ++d) { - range_min[d] = intersect_range[d].first; - range_max[d] = intersect_range[d].second; - } - if (dim == in::x1) { - if (g != gr_bc::curr) { - Kokkos::parallel_for( - "MatchBoundaries", - CreateRangePolicy(range_min, range_max), - kernel::bc::MatchBoundaries_kernel( - domain.fields.em, - m_pgen.init_flds, - domain.mesh.metric, - xg_edge, - ds, - tags, - domain.mesh.flds_bc())); - Kokkos::parallel_for( - "MatchBoundaries", - CreateRangePolicy(range_min, range_max), - kernel::bc::MatchBoundaries_kernel( - domain.fields.em0, - m_pgen.init_flds, - domain.mesh.metric, - xg_edge, - ds, - tags, - domain.mesh.flds_bc())); - } else { - Kokkos::parallel_for( - "AbsorbCurrents", - CreateRangePolicy(range_min, range_max), - kernel::bc::gr::AbsorbCurrents_kernel(domain.fields.cur0, - domain.mesh.metric, - xg_edge, - ds)); - } - } else { - raise::Error("Invalid dimension", HERE); - } - } - - void HorizonFieldsIn(dir::direction_t direction, - domain_t& domain, - BCTags tags, - const gr_bc& g) { - /** - * open boundaries - */ - raise::ErrorIf(M::CoordType == Coord::Cart, - "Invalid coordinate type for horizon BCs", - HERE); - raise::ErrorIf(direction.get_dim() != in::x1, - "Invalid horizon direction, should be x1", - HERE); - const auto i1_min = domain.mesh.i_min(in::x1); - auto range = CreateRangePolicy({ domain.mesh.i_min(in::x2) }, - { domain.mesh.i_max(in::x2) + 1 }); - const auto nfilter = m_params.template get( - "algorithms.current_filters"); - if (g == gr_bc::main) { - Kokkos::parallel_for( - "OpenBCFields", - range, - kernel::bc::gr::HorizonBoundaries_kernel(domain.fields.em, - i1_min, - tags, - nfilter)); - Kokkos::parallel_for( - "OpenBCFields", - range, - kernel::bc::gr::HorizonBoundaries_kernel(domain.fields.em0, - i1_min, - tags, - nfilter)); - } - } - - void AxisFieldsIn(dir::direction_t direction, - domain_t& domain, - BCTags tags) { - /** - * axis boundaries - */ - raise::ErrorIf(M::CoordType == Coord::Cart, - "Invalid coordinate type for axis BCs", - HERE); - raise::ErrorIf(direction.get_dim() != in::x2, - "Invalid axis direction, should be x2", - HERE); - const auto i2_min = domain.mesh.i_min(in::x2); - const auto i2_max = domain.mesh.i_max(in::x2); - if (direction.get_sign() < 0) { - Kokkos::parallel_for( - "AxisBCFields", - domain.mesh.n_all(in::x1), - kernel::bc::AxisBoundaries_kernel(domain.fields.em, - i2_min, - tags)); - Kokkos::parallel_for( - "AxisBCFields", - domain.mesh.n_all(in::x1), - kernel::bc::AxisBoundaries_kernel(domain.fields.em0, - i2_min, - tags)); - } else { - Kokkos::parallel_for( - "AxisBCFields", - domain.mesh.n_all(in::x1), - kernel::bc::AxisBoundaries_kernel(domain.fields.em, - i2_max, - tags)); - Kokkos::parallel_for( - "AxisBCFields", - domain.mesh.n_all(in::x1), - kernel::bc::AxisBoundaries_kernel(domain.fields.em0, - i2_max, - tags)); - } - } - - void CustomFieldsIn(dir::direction_t direction, - domain_t& domain, - BCTags tags, - const gr_bc& g) { - (void)direction; - (void)domain; - (void)tags; - (void)g; - raise::Error("Custom boundaries not implemented", HERE); - // if constexpr ( - // traits::has_member::value) { - // const auto [box, custom_fields] = m_pgen.CustomFields(time); - // if (domain.mesh.Intersects(box)) { - // } - // - // } else { - // raise::Error("Custom boundaries not implemented", HERE); - // } - } - - /** - * @brief Swaps em and em0 fields, cur and cur0 currents. - */ - void SwapFields(domain_t& domain) { - std::swap(domain.fields.em, domain.fields.em0); - std::swap(domain.fields.cur, domain.fields.cur0); - } - - /** - * @brief Copies em fields into em0 - */ - void CopyFields(domain_t& domain) { - Kokkos::deep_copy(domain.fields.em0, domain.fields.em); - } - - void ComputeAuxE(domain_t& domain, const gr_getE& g) { - auto range = range_with_axis_BCs(domain); - if (g == gr_getE::D0_B) { - Kokkos::parallel_for( - "ComputeAuxE", - range, - kernel::gr::ComputeAuxE_kernel(domain.fields.em0, // D - domain.fields.em, // B - domain.fields.aux, // E - domain.mesh.metric)); - } else if (g == gr_getE::D_B0) { - Kokkos::parallel_for("ComputeAuxE", - range, - kernel::gr::ComputeAuxE_kernel(domain.fields.em, - domain.fields.em0, - domain.fields.aux, - domain.mesh.metric)); - } else { - raise::Error("Wrong option for `g`", HERE); - } - } - - void ComputeAuxH(domain_t& domain, const gr_getH& g) { - auto range = range_with_axis_BCs(domain); - if (g == gr_getH::D_B0) { - Kokkos::parallel_for( - "ComputeAuxH", - range, - kernel::gr::ComputeAuxH_kernel(domain.fields.em, // D - domain.fields.em0, // B - domain.fields.aux, // H - domain.mesh.metric)); - } else if (g == gr_getH::D0_B0) { - Kokkos::parallel_for("ComputeAuxH", - range, - kernel::gr::ComputeAuxH_kernel(domain.fields.em0, - domain.fields.em0, - domain.fields.aux, - domain.mesh.metric)); - } else { - raise::Error("Wrong option for `g`", HERE); - } - } - - auto range_with_axis_BCs(const domain_t& domain) -> range_t { - auto range = domain.mesh.rangeActiveCells(); - /** - * @brief taking one extra cell in the x1 and x2 directions if AXIS BCs - */ - if constexpr (M::Dim == Dim::_2D) { - if (domain.mesh.flds_bc_in({ 0, +1 }) == FldsBC::AXIS) { - range = CreateRangePolicy( - { domain.mesh.i_min(in::x1) - 1, domain.mesh.i_min(in::x2) }, - { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); - } else { - range = CreateRangePolicy( - { domain.mesh.i_min(in::x1) - 1, domain.mesh.i_min(in::x2) }, - { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) }); - } - } else if constexpr (M::Dim == Dim::_3D) { - raise::Error("Invalid dimension", HERE); - } - return range; - } - - void Faraday(domain_t& domain, const gr_faraday& g, real_t fraction = ONE) { - logger::Checkpoint("Launching Faraday kernel", HERE); - const auto dT = fraction * - m_params.template get( - "algorithms.timestep.correction") * - dt; - if (g == gr_faraday::aux) { - Kokkos::parallel_for( - "Faraday", - domain.mesh.rangeActiveCells(), - kernel::gr::Faraday_kernel(domain.fields.em0, // Bin - domain.fields.em0, // Bout - domain.fields.aux, // E - domain.mesh.metric, - dT, - domain.mesh.n_active(in::x2), - domain.mesh.flds_bc())); - } else if (g == gr_faraday::main) { - Kokkos::parallel_for( - "Faraday", - domain.mesh.rangeActiveCells(), - kernel::gr::Faraday_kernel(domain.fields.em, - domain.fields.em0, - domain.fields.aux, - domain.mesh.metric, - dT, - domain.mesh.n_active(in::x2), - domain.mesh.flds_bc())); - - } else { - raise::Error("Wrong option for `g`", HERE); - } - } - - void Ampere(domain_t& domain, const gr_ampere& g, real_t fraction = ONE) { - logger::Checkpoint("Launching Ampere kernel", HERE); - const auto dT = fraction * - m_params.template get( - "algorithms.timestep.correction") * - dt; - auto range = CreateRangePolicy( - { domain.mesh.i_min(in::x1), domain.mesh.i_min(in::x2) }, - { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); - const auto ni2 = domain.mesh.n_active(in::x2); - - if (g == gr_ampere::aux) { - // First push, updates D0 with J. - Kokkos::parallel_for("Ampere-1", - range, - kernel::gr::Ampere_kernel(domain.fields.em0, // Din - domain.fields.em0, // Dout - domain.fields.aux, - domain.mesh.metric, - dT, - ni2, - domain.mesh.flds_bc())); - } else if (g == gr_ampere::main) { - // Second push, updates D with J0 but assigns it to D0. - Kokkos::parallel_for("Ampere-2", - range, - kernel::gr::Ampere_kernel(domain.fields.em, - domain.fields.em0, - domain.fields.aux, - domain.mesh.metric, - dT, - ni2, - domain.mesh.flds_bc())); - } else if (g == gr_ampere::init) { - // Second push, updates D with J0 and assigns it to D. - Kokkos::parallel_for("Ampere-3", - range, - kernel::gr::Ampere_kernel(domain.fields.em, - domain.fields.em, - domain.fields.aux, - domain.mesh.metric, - dT, - ni2, - domain.mesh.flds_bc())); - } else { - raise::Error("Wrong option for `g`", HERE); - } - } - - void AmpereCurrents(domain_t& domain, const gr_ampere& g) { - logger::Checkpoint("Launching Ampere kernel for adding currents", HERE); - const auto q0 = m_params.template get("scales.q0"); - const auto B0 = m_params.template get("scales.B0"); - const auto coeff = -dt * q0 / B0; - auto range = CreateRangePolicy( - { domain.mesh.i_min(in::x1), domain.mesh.i_min(in::x2) }, - { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); - const auto ni2 = domain.mesh.n_active(in::x2); - - if (g == gr_ampere::aux) { - // Updates D0 with J: D0(n-1/2) -> (J(n)) -> D0(n+1/2) - Kokkos::parallel_for( - "AmpereCurrentsAux", - range, - kernel::gr::CurrentsAmpere_kernel(domain.fields.em0, - domain.fields.cur, - domain.mesh.metric, - coeff, - ni2, - domain.mesh.flds_bc())); - } else if (g == gr_ampere::main) { - // Updates D0 with J0: D0(n) -> (J0(n+1/2)) -> D0(n+1) - Kokkos::parallel_for( - "AmpereCurrentsMain", - range, - kernel::gr::CurrentsAmpere_kernel(domain.fields.em0, - domain.fields.cur0, - domain.mesh.metric, - coeff, - ni2, - domain.mesh.flds_bc())); - } else { - raise::Error("Wrong option for `g`", HERE); - } - } - - void TimeAverageDB(domain_t& domain) { - Kokkos::parallel_for("TimeAverageDB", - domain.mesh.rangeActiveCells(), - kernel::gr::TimeAverageDB_kernel(domain.fields.em, - domain.fields.em0, - domain.mesh.metric)); - } - - void TimeAverageJ(domain_t& domain) { - Kokkos::parallel_for("TimeAverageJ", - domain.mesh.rangeActiveCells(), - kernel::gr::TimeAverageJ_kernel(domain.fields.cur, - domain.fields.cur0, - domain.mesh.metric)); - } - - void CurrentsDeposit(domain_t& domain) { - auto scatter_cur0 = Kokkos::Experimental::create_scatter_view( - domain.fields.cur0); - for (auto& species : domain.species) { - logger::Checkpoint( - fmt::format("Launching currents deposit kernel for %d [%s] : %lu %f", - species.index(), - species.label().c_str(), - species.npart(), - (double)species.charge()), - HERE); - if (species.npart() == 0 || cmp::AlmostZero(species.charge())) { - continue; - } - Kokkos::parallel_for("CurrentsDeposit", - species.rangeActiveParticles(), - kernel::DepositCurrents_kernel( - scatter_cur0, - species.i1, - species.i2, - species.i3, - species.i1_prev, - species.i2_prev, - species.i3_prev, - species.dx1, - species.dx2, - species.dx3, - species.dx1_prev, - species.dx2_prev, - species.dx3_prev, - species.ux1, - species.ux2, - species.ux3, - species.phi, - species.weight, - species.tag, - domain.mesh.metric, - (real_t)(species.charge()), - dt)); - } - Kokkos::Experimental::contribute(domain.fields.cur0, scatter_cur0); - } - - void CurrentsFilter(domain_t& domain) { - logger::Checkpoint("Launching currents filtering kernels", HERE); - auto range = CreateRangePolicy( - { domain.mesh.i_min(in::x1), domain.mesh.i_min(in::x2) }, - { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); - const auto nfilter = m_params.template get( - "algorithms.current_filters"); - tuple_t size; - size[0] = domain.mesh.n_active(in::x1); - size[1] = domain.mesh.n_active(in::x2); - - // !TODO: this needs to be done more efficiently - for (unsigned short i = 0; i < nfilter; ++i) { - Kokkos::deep_copy(domain.fields.buff, domain.fields.cur0); - Kokkos::parallel_for("CurrentsFilter", - range, - kernel::DigitalFilter_kernel( - domain.fields.cur0, - domain.fields.buff, - size, - domain.mesh.flds_bc())); - m_metadomain.CommunicateFields(domain, Comm::J); // J0 - } - } - - void ParticlePush(domain_t& domain) { - for (auto& species : domain.species) { - species.set_unsorted(); - logger::Checkpoint( - fmt::format("Launching particle pusher kernel for %d [%s] : %lu", - species.index(), - species.label().c_str(), - species.npart()), - HERE); - if (species.npart() == 0) { - continue; - } - const auto q_ovr_m = species.mass() > ZERO - ? species.charge() / species.mass() - : ZERO; - // coeff = q / m (dt / 2) omegaB0 - const auto coeff = q_ovr_m * HALF * dt * - m_params.template get( - "algorithms.timestep.correction") * - m_params.template get("scales.omegaB0"); - const auto eps = m_params.template get( - "algorithms.gr.pusher_eps"); - const auto niter = m_params.template get( - "algorithms.gr.pusher_niter"); - // clang-format off - if (species.pusher() == PrtlPusher::PHOTON) { - auto range_policy = Kokkos::RangePolicy( - 0, - species.npart()); - - Kokkos::parallel_for( - "ParticlePusher", - range_policy, - kernel::gr::Pusher_kernel( - domain.fields.em, - domain.fields.em0, - species.i1, species.i2, species.i3, - species.i1_prev, species.i2_prev, species.i3_prev, - species.dx1, species.dx2, species.dx3, - species.dx1_prev, species.dx2_prev, species.dx3_prev, - species.ux1, species.ux2, species.ux3, - species.phi, species.tag, - domain.mesh.metric, - coeff, dt, - domain.mesh.n_active(in::x1), - domain.mesh.n_active(in::x2), - domain.mesh.n_active(in::x3), - eps, niter, - domain.mesh.prtl_bc() - )); - } else if (species.pusher() == PrtlPusher::BORIS) { - auto range_policy = Kokkos::RangePolicy( - 0, - species.npart()); - Kokkos::parallel_for( - "ParticlePusher", - range_policy, - kernel::gr::Pusher_kernel( - domain.fields.em, - domain.fields.em0, - species.i1, species.i2, species.i3, - species.i1_prev, species.i2_prev, species.i3_prev, - species.dx1, species.dx2, species.dx3, - species.dx1_prev, species.dx2_prev, species.dx3_prev, - species.ux1, species.ux2, species.ux3, - species.phi, species.tag, - domain.mesh.metric, - coeff, dt, - domain.mesh.n_active(in::x1), - domain.mesh.n_active(in::x2), - domain.mesh.n_active(in::x3), - eps, niter, - domain.mesh.prtl_bc() - )); - } else if (species.pusher() == PrtlPusher::NONE) { - // do nothing - } else { - raise::Error("not implemented", HERE); - } - // clang-format on - } - } - }; -} // namespace ntt - -#endif // ENGINES_GRPIC_GRPIC_H +/** + * @file engines/grpic.hpp + * @brief Simulation engien class which specialized on GRPIC + * @implements + * - ntt::GRPICEngine<> : ntt::Engine<> + * @cpp: + * - grpic.cpp + * @namespaces: + * - ntt:: + */ + +#ifndef ENGINES_GRPIC_GRPIC_H +#define ENGINES_GRPIC_GRPIC_H + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "utils/log.h" +#include "utils/numeric.h" +#include "utils/timer.h" +#include "utils/toml.h" + +#include "framework/domain/domain.h" +#include "framework/parameters.h" + +#include "engines/engine.hpp" +#include "kernels/ampere_gr.hpp" +#include "kernels/aux_fields_gr.hpp" +#include "kernels/currents_deposit.hpp" +#include "kernels/digital_filter.hpp" +#include "kernels/faraday_gr.hpp" +#include "kernels/fields_bcs.hpp" +#include "kernels/particle_pusher_gr.hpp" +#include "pgen.hpp" + +#include +#include + +#include +#include + +namespace ntt { + + enum class gr_getE { + D0_B, + D_B0 + }; + enum class gr_getH { + D_B0, + D0_B0 + }; + enum class gr_faraday { + aux, + main + }; + enum class gr_ampere { + init, + aux, + main + }; + enum class gr_bc { + main, + aux, + curr + }; + + template + class GRPICEngine : public Engine { + using base_t = Engine; + using pgen_t = user::PGen; + using domain_t = Domain; + // constexprs + using base_t::pgen_is_ok; + // contents + using base_t::m_metadomain; + using base_t::m_params; + using base_t::m_pgen; + // methods + using base_t::init; + // variables + using base_t::dt; + using base_t::max_steps; + using base_t::runtime; + using base_t::step; + using base_t::time; + + public: + static constexpr auto S { SimEngine::GRPIC }; + + GRPICEngine(SimulationParams& params) : base_t { params } {} + + ~GRPICEngine() = default; + + void step_forward(timer::Timers& timers, domain_t& dom) override { + const auto fieldsolver_enabled = m_params.template get( + "algorithms.fieldsolver.enable"); + const auto deposit_enabled = m_params.template get( + "algorithms.deposit.enable"); + const auto clear_interval = m_params.template get( + "particles.clear_interval"); + + if (step == 0) { + if (fieldsolver_enabled) { + // communicate fields and apply BCs on the first timestep + /** + * Initially: em0::B -- + * em0::D -- + * em::B at -1/2 + * em::D at -1/2 + * + * cur0::J -- + * cur::J -- + * + * aux::E -- + * aux::H -- + * + * x_prtl at -1/2 + * u_prtl at -1/2 + */ + + /** + * em0::D, em::D, em0::B, em::B <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, + Comm::B | Comm::B0 | Comm::D | Comm::D0); + FieldBoundaries(dom, BC::B | BC::D, gr_bc::main); + + /** + * em0::B <- em::B + * em0::D <- em::D + * + * Now: em0::B & em0::D at -1/2 + */ + CopyFields(dom); + + /** + * aux::E <- alpha * em::D + beta x em0::B + * aux::H <- alpha * em::B0 - beta x em::D + * + * Now: aux::E & aux::H at -1/2 + */ + ComputeAuxE(dom, gr_getE::D_B0); + ComputeAuxH(dom, gr_getH::D_B0); + + /** + * aux::E, aux::H <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::H | Comm::E); + FieldBoundaries(dom, BC::H | BC::E, gr_bc::aux); + + /** + * em0::B <- (em0::B) <- -curl aux::E + * + * Now: em0::B at 0 + */ + Faraday(dom, gr_faraday::aux, HALF); + + /** + * em0::B, em::B <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); + FieldBoundaries(dom, BC::B, gr_bc::main); + + /** + * em::D <- (em0::D) <- curl aux::H + * + * Now: em::D at 0 + */ + Ampere(dom, gr_ampere::init, HALF); + + /** + * em0::D, em::D <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); + FieldBoundaries(dom, BC::D, gr_bc::main); + + /** + * aux::E <- alpha * em::D + beta x em0::B + * aux::H <- alpha * em0::B - beta x em::D + * + * Now: aux::E & aux::H at 0 + */ + ComputeAuxE(dom, gr_getE::D_B0); + ComputeAuxH(dom, gr_getH::D_B0); + + /** + * aux::E, aux::H <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::H | Comm::E); + FieldBoundaries(dom, BC::H | BC::E, gr_bc::aux); + + // !ADD: GR -- particles? + + /** + * em0::B <- (em::B) <- -curl aux::E + * + * Now: em0::B at 1/2 + */ + Faraday(dom, gr_faraday::main, ONE); + /** + * em0::B, em::B <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); + FieldBoundaries(dom, BC::B, gr_bc::main); + + /** + * em0::D <- (em0::D) <- curl aux::H + * + * Now: em0::D at 1/2 + */ + Ampere(dom, gr_ampere::aux, ONE); + /** + * em0::D, em::D <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); + FieldBoundaries(dom, BC::D, gr_bc::main); + + /** + * aux::H <- alpha * em0::B - beta x em0::D + * + * Now: aux::H at 1/2 + */ + ComputeAuxH(dom, gr_getH::D0_B0); + /** + * aux::H <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::H); + FieldBoundaries(dom, BC::H, gr_bc::aux); + + /** + * em0::D <- (em::D) <- curl aux::H + * + * Now: em0::D at 1 + * em::D at 0 + */ + Ampere(dom, gr_ampere::main, ONE); + /** + * em0::D, em::D <- boundary conditions + */ + m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); + FieldBoundaries(dom, BC::D, gr_bc::main); + + /** + * em::D <-> em0::D + * em::B <-> em0::B + * em::J <-> em0::J + */ + SwapFields(dom); + /** + * Finally: em0::B at -1/2 + * em0::D at 0 + * em::B at 1/2 + * em::D at 1 + * + * cur0::J -- + * cur::J -- + * + * aux::E -- + * aux::H -- + * + * x_prtl at 1 + * u_prtl at 1/2 + */ + } else { + /** + * em0::B <- em::B + * em0::D <- em::D + * + * Now: em0::B & em0::D at -1/2 + */ + CopyFields(dom); + } + } + + /** + * Initially: em0::B at n-3/2 + * em0::D at n-1 + * em::B at n-1/2 + * em::D at n + * + * cur0::J -- + * cur::J at n-1/2 + * + * aux::E -- + * aux::H -- + * + * x_prtl at n + * u_prtl at n-1/2 + */ + + if (fieldsolver_enabled) { + timers.start("FieldSolver"); + /** + * em0::D <- (em0::D + em::D) / 2 + * em0::B <- (em0::B + em::B) / 2 + * + * Now: em0::D at n-1/2 + * em0::B at n-1 + */ + TimeAverageDB(dom); + /** + * aux::E <- alpha * em0::D + beta x em::B + * + * Now: aux::E at n-1/2 + */ + ComputeAuxE(dom, gr_getE::D0_B); + timers.stop("FieldSolver"); + + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::E); + timers.stop("Communications"); + timers.start("FieldBoundaries"); + /** + * aux::E <- boundary conditions + */ + FieldBoundaries(dom, BC::E, gr_bc::aux); + timers.stop("FieldBoundaries"); + + timers.start("FieldSolver"); + /** + * em0::B <- (em0::B) <- -curl aux::E + * + * Now: em0::B at n + */ + Faraday(dom, gr_faraday::aux, ONE); + timers.stop("FieldSolver"); + + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); + timers.stop("Communications"); + /** + * em0::B, em::B <- boundary conditions + */ + timers.start("FieldBoundaries"); + FieldBoundaries(dom, BC::B, gr_bc::main); + timers.stop("FieldBoundaries"); + + timers.start("FieldSolver"); + /** + * aux::H <- alpha * em0::B - beta x em::D + * + * Now: aux::H at n + */ + ComputeAuxH(dom, gr_getH::D_B0); + timers.stop("FieldSolver"); + + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::H); + timers.stop("Communications"); + timers.start("FieldBoundaries"); + /** + * aux::H <- boundary conditions + */ + FieldBoundaries(dom, BC::H, gr_bc::aux); + timers.stop("FieldBoundaries"); + } + + { + /** + * x_prtl, u_prtl <- em::D, em0::B + * + * Now: x_prtl at n + 1, u_prtl at n + 1/2 + */ + timers.start("ParticlePusher"); + ParticlePush(dom); + timers.stop("ParticlePusher"); + + /** + * cur0::J <- current deposition + * + * Now: cur0::J at n+1/2 + */ + if (deposit_enabled) { + timers.start("CurrentDeposit"); + Kokkos::deep_copy(dom.fields.cur0, ZERO); + CurrentsDeposit(dom); + timers.stop("CurrentDeposit"); + + timers.start("Communications"); + m_metadomain.SynchronizeFields(dom, Comm::J); + m_metadomain.CommunicateFields(dom, Comm::J); + timers.stop("Communications"); + + timers.start("FieldBoundaries"); + FieldBoundaries(dom, BC::J, gr_bc::curr); + timers.stop("FieldBoundaries"); + + timers.start("CurrentFiltering"); + CurrentsFilter(dom); + timers.stop("CurrentFiltering"); + } + + timers.start("Communications"); + m_metadomain.CommunicateParticles(dom); + timers.stop("Communications"); + } + + if (fieldsolver_enabled) { + timers.start("FieldSolver"); + if (deposit_enabled) { + /** + * cur::J <- (cur0::J + cur::J) / 2 + * + * Now: cur::J at n + */ + TimeAverageJ(dom); + } + /** + * aux::Е <- alpha * em::D + beta x em0::B + * + * Now: aux::Е at n + */ + ComputeAuxE(dom, gr_getE::D_B0); + timers.stop("FieldSolver"); + + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::E); + timers.stop("Communications"); + timers.start("FieldBoundaries"); + /** + * aux::Е <- boundary conditions + */ + FieldBoundaries(dom, BC::E, gr_bc::aux); + timers.stop("FieldBoundaries"); + + timers.start("FieldSolver"); + /** + * em0::B <- (em::B) <- -curl aux::E + * + * Now: em0::B at n+1/2 + * em::B at n-1/2 + */ + Faraday(dom, gr_faraday::main, ONE); + timers.stop("FieldSolver"); + + /** + * em0::B, em::B <- boundary conditions + */ + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::B | Comm::B0); + timers.stop("Communications"); + timers.start("FieldBoundaries"); + FieldBoundaries(dom, BC::B, gr_bc::main); + timers.stop("FieldBoundaries"); + + timers.start("FieldSolver"); + /** + * em0::D <- (em0::D) <- curl aux::H + * + * Now: em0::D at n+1/2 + */ + Ampere(dom, gr_ampere::aux, ONE); + timers.stop("FieldSolver"); + + if (deposit_enabled) { + timers.start("FieldSolver"); + /** + * em0::D <- (em0::D) <- cur::J + * + * Now: em0::D at n+1/2 + */ + AmpereCurrents(dom, gr_ampere::aux); + timers.stop("FieldSolver"); + } + + /** + * em0::D, em::D <- boundary conditions + */ + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); + timers.stop("Communications"); + timers.start("FieldBoundaries"); + FieldBoundaries(dom, BC::D, gr_bc::main); + timers.stop("FieldBoundaries"); + + timers.start("FieldSolver"); + /** + * aux::H <- alpha * em0::B - beta x em0::D + * + * Now: aux::H at n+1/2 + */ + ComputeAuxH(dom, gr_getH::D0_B0); + timers.stop("FieldSolver"); + + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::H); + timers.stop("Communications"); + timers.start("FieldBoundaries"); + /** + * aux::H <- boundary conditions + */ + FieldBoundaries(dom, BC::B, gr_bc::aux); + timers.stop("FieldBoundaries"); + + timers.start("FieldSolver"); + /** + * em0::D <- (em::D) <- curl aux::H + * + * Now: em0::D at n+1 + * em::D at n + */ + Ampere(dom, gr_ampere::main, ONE); + timers.stop("FieldSolver"); + + if (deposit_enabled) { + timers.start("FieldSolver"); + /** + * em0::D <- (em0::D) <- cur0::J + * + * Now: em0::D at n+1 + */ + AmpereCurrents(dom, gr_ampere::main); + timers.stop("FieldSolver"); + } + timers.start("FieldSolver"); + /** + * em::D <-> em0::D + * em::B <-> em0::B + * cur::J <-> cur0::J + */ + SwapFields(dom); + timers.stop("FieldSolver"); + + /** + * em0::D, em::D <- boundary conditions + */ + timers.start("Communications"); + m_metadomain.CommunicateFields(dom, Comm::D | Comm::D0); + timers.stop("Communications"); + timers.start("FieldBoundaries"); + FieldBoundaries(dom, BC::D, gr_bc::main); + timers.stop("FieldBoundaries"); + } + + if (clear_interval > 0 and step % clear_interval == 0 and step > 0) { + timers.start("PrtlClear"); + m_metadomain.RemoveDeadParticles(dom); + timers.stop("PrtlClear"); + } + + /** + * Finally: em0::B at n-1/2 + * em0::D at n + * em::B at n+1/2 + * em::D at n+1 + * + * cur0::J (at n) + * cur::J at n+1/2 + * + * aux::E (at n+1/2) + * aux::H (at n) + * + * x_prtl at n+1 + * u_prtl at n+1/2 + */ + } + + /* algorithm substeps --------------------------------------------------- */ + void FieldBoundaries(domain_t& domain, BCTags tags, const gr_bc& g) { + if (g == gr_bc::main) { + for (auto& direction : dir::Directions::orth) { + if (m_metadomain.mesh().flds_bc_in(direction) == FldsBC::MATCH) { + MatchFieldsIn(direction, domain, tags, g); + } else if (domain.mesh.flds_bc_in(direction) == FldsBC::AXIS) { + AxisFieldsIn(direction, domain, tags); + } else if (m_metadomain.mesh().flds_bc_in(direction) == FldsBC::CUSTOM) { + CustomFieldsIn(direction, domain, tags, g); + } else if (domain.mesh.flds_bc_in(direction) == FldsBC::HORIZON) { + HorizonFieldsIn(direction, domain, tags, g); + } + } // loop over directions + } else if (g == gr_bc::aux) { + for (auto& direction : dir::Directions::orth) { + if (domain.mesh.flds_bc_in(direction) == FldsBC::HORIZON) { + HorizonFieldsIn(direction, domain, tags, g); + } + } + } else if (g == gr_bc::curr) { + for (auto& direction : dir::Directions::orth) { + if (domain.mesh.prtl_bc_in(direction) == PrtlBC::ABSORB) { + MatchFieldsIn(direction, domain, tags, g); + } + } + } + } + + void MatchFieldsIn(dir::direction_t direction, + domain_t& domain, + BCTags tags, + const gr_bc& g) { + /** + * match boundaries + */ + const auto ds_array = m_params.template get>( + "grid.boundaries.match.ds"); + const auto dim = direction.get_dim(); + real_t xg_min, xg_max, xg_edge; + auto sign = direction.get_sign(); + + raise::ErrorIf(((dim != in::x1) or (sign < 0)) and (g == gr_bc::curr), + "Absorption of currents only possible in +x1 (+r)", + HERE); + + real_t ds; + if (sign > 0) { // + direction + ds = ds_array[(short)dim].second; + xg_max = m_metadomain.mesh().extent(dim).second; + xg_min = xg_max - ds; + xg_edge = xg_max; + } else { // - direction + ds = ds_array[(short)dim].first; + xg_min = m_metadomain.mesh().extent(dim).first; + xg_max = xg_min + ds; + xg_edge = xg_min; + } + boundaries_t box; + boundaries_t incl_ghosts; + for (unsigned short d { 0 }; d < M::Dim; ++d) { + if (d == static_cast(dim)) { + box.push_back({ xg_min, xg_max }); + incl_ghosts.push_back({ false, true }); + } else { + box.push_back(Range::All); + incl_ghosts.push_back({ true, true }); + } + } + if (not domain.mesh.Intersects(box)) { + return; + } + const auto intersect_range = domain.mesh.ExtentToRange(box, incl_ghosts); + tuple_t range_min { 0 }; + tuple_t range_max { 0 }; + + for (unsigned short d { 0 }; d < M::Dim; ++d) { + range_min[d] = intersect_range[d].first; + range_max[d] = intersect_range[d].second; + } + if (dim == in::x1) { + if (g != gr_bc::curr) { + Kokkos::parallel_for( + "MatchBoundaries", + CreateRangePolicy(range_min, range_max), + kernel::bc::MatchBoundaries_kernel( + domain.fields.em, + m_pgen.init_flds, + domain.mesh.metric, + xg_edge, + ds, + tags, + domain.mesh.flds_bc())); + Kokkos::parallel_for( + "MatchBoundaries", + CreateRangePolicy(range_min, range_max), + kernel::bc::MatchBoundaries_kernel( + domain.fields.em0, + m_pgen.init_flds, + domain.mesh.metric, + xg_edge, + ds, + tags, + domain.mesh.flds_bc())); + } else { + Kokkos::parallel_for( + "AbsorbCurrents", + CreateRangePolicy(range_min, range_max), + kernel::bc::gr::AbsorbCurrents_kernel(domain.fields.cur0, + domain.mesh.metric, + xg_edge, + ds)); + } + } else { + raise::Error("Invalid dimension", HERE); + } + } + + void HorizonFieldsIn(dir::direction_t direction, + domain_t& domain, + BCTags tags, + const gr_bc& g) { + /** + * open boundaries + */ + raise::ErrorIf(M::CoordType == Coord::Cart, + "Invalid coordinate type for horizon BCs", + HERE); + raise::ErrorIf(direction.get_dim() != in::x1, + "Invalid horizon direction, should be x1", + HERE); + const auto i1_min = domain.mesh.i_min(in::x1); + auto range = CreateRangePolicy({ domain.mesh.i_min(in::x2) }, + { domain.mesh.i_max(in::x2) + 1 }); + const auto nfilter = m_params.template get( + "algorithms.current_filters"); + if (g == gr_bc::main) { + Kokkos::parallel_for( + "OpenBCFields", + range, + kernel::bc::gr::HorizonBoundaries_kernel(domain.fields.em, + i1_min, + tags, + nfilter)); + Kokkos::parallel_for( + "OpenBCFields", + range, + kernel::bc::gr::HorizonBoundaries_kernel(domain.fields.em0, + i1_min, + tags, + nfilter)); + } + } + + void AxisFieldsIn(dir::direction_t direction, + domain_t& domain, + BCTags tags) { + /** + * axis boundaries + */ + raise::ErrorIf(M::CoordType == Coord::Cart, + "Invalid coordinate type for axis BCs", + HERE); + raise::ErrorIf(direction.get_dim() != in::x2, + "Invalid axis direction, should be x2", + HERE); + const auto i2_min = domain.mesh.i_min(in::x2); + const auto i2_max = domain.mesh.i_max(in::x2); + if (direction.get_sign() < 0) { + Kokkos::parallel_for( + "AxisBCFields", + domain.mesh.n_all(in::x1), + kernel::bc::AxisBoundaries_kernel(domain.fields.em, + i2_min, + tags)); + Kokkos::parallel_for( + "AxisBCFields", + domain.mesh.n_all(in::x1), + kernel::bc::AxisBoundaries_kernel(domain.fields.em0, + i2_min, + tags)); + } else { + Kokkos::parallel_for( + "AxisBCFields", + domain.mesh.n_all(in::x1), + kernel::bc::AxisBoundaries_kernel(domain.fields.em, + i2_max, + tags)); + Kokkos::parallel_for( + "AxisBCFields", + domain.mesh.n_all(in::x1), + kernel::bc::AxisBoundaries_kernel(domain.fields.em0, + i2_max, + tags)); + } + } + + void CustomFieldsIn(dir::direction_t direction, + domain_t& domain, + BCTags tags, + const gr_bc& g) { + (void)direction; + (void)domain; + (void)tags; + (void)g; + raise::Error("Custom boundaries not implemented", HERE); + // if constexpr ( + // traits::has_member::value) { + // const auto [box, custom_fields] = m_pgen.CustomFields(time); + // if (domain.mesh.Intersects(box)) { + // } + // + // } else { + // raise::Error("Custom boundaries not implemented", HERE); + // } + } + + /** + * @brief Swaps em and em0 fields, cur and cur0 currents. + */ + void SwapFields(domain_t& domain) { + std::swap(domain.fields.em, domain.fields.em0); + std::swap(domain.fields.cur, domain.fields.cur0); + } + + /** + * @brief Copies em fields into em0 + */ + void CopyFields(domain_t& domain) { + Kokkos::deep_copy(domain.fields.em0, domain.fields.em); + } + + void ComputeAuxE(domain_t& domain, const gr_getE& g) { + auto range = range_with_axis_BCs(domain); + if (g == gr_getE::D0_B) { + Kokkos::parallel_for( + "ComputeAuxE", + range, + kernel::gr::ComputeAuxE_kernel(domain.fields.em0, // D + domain.fields.em, // B + domain.fields.aux, // E + domain.mesh.metric)); + } else if (g == gr_getE::D_B0) { + Kokkos::parallel_for("ComputeAuxE", + range, + kernel::gr::ComputeAuxE_kernel(domain.fields.em, + domain.fields.em0, + domain.fields.aux, + domain.mesh.metric)); + } else { + raise::Error("Wrong option for `g`", HERE); + } + } + + void ComputeAuxH(domain_t& domain, const gr_getH& g) { + auto range = range_with_axis_BCs(domain); + if (g == gr_getH::D_B0) { + Kokkos::parallel_for( + "ComputeAuxH", + range, + kernel::gr::ComputeAuxH_kernel(domain.fields.em, // D + domain.fields.em0, // B + domain.fields.aux, // H + domain.mesh.metric)); + } else if (g == gr_getH::D0_B0) { + Kokkos::parallel_for("ComputeAuxH", + range, + kernel::gr::ComputeAuxH_kernel(domain.fields.em0, + domain.fields.em0, + domain.fields.aux, + domain.mesh.metric)); + } else { + raise::Error("Wrong option for `g`", HERE); + } + } + + auto range_with_axis_BCs(const domain_t& domain) -> range_t { + auto range = domain.mesh.rangeActiveCells(); + /** + * @brief taking one extra cell in the x1 and x2 directions if AXIS BCs + */ + if constexpr (M::Dim == Dim::_2D) { + if (domain.mesh.flds_bc_in({ 0, +1 }) == FldsBC::AXIS) { + range = CreateRangePolicy( + { domain.mesh.i_min(in::x1) - 1, domain.mesh.i_min(in::x2) }, + { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); + } else { + range = CreateRangePolicy( + { domain.mesh.i_min(in::x1) - 1, domain.mesh.i_min(in::x2) }, + { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) }); + } + } else if constexpr (M::Dim == Dim::_3D) { + raise::Error("Invalid dimension", HERE); + } + return range; + } + + void Faraday(domain_t& domain, const gr_faraday& g, real_t fraction = ONE) { + logger::Checkpoint("Launching Faraday kernel", HERE); + const auto dT = fraction * + m_params.template get( + "algorithms.timestep.correction") * + dt; + if (g == gr_faraday::aux) { + Kokkos::parallel_for( + "Faraday", + domain.mesh.rangeActiveCells(), + kernel::gr::Faraday_kernel(domain.fields.em0, // Bin + domain.fields.em0, // Bout + domain.fields.aux, // E + domain.mesh.metric, + dT, + domain.mesh.n_active(in::x2), + domain.mesh.flds_bc())); + } else if (g == gr_faraday::main) { + Kokkos::parallel_for( + "Faraday", + domain.mesh.rangeActiveCells(), + kernel::gr::Faraday_kernel(domain.fields.em, + domain.fields.em0, + domain.fields.aux, + domain.mesh.metric, + dT, + domain.mesh.n_active(in::x2), + domain.mesh.flds_bc())); + + } else { + raise::Error("Wrong option for `g`", HERE); + } + } + + void Ampere(domain_t& domain, const gr_ampere& g, real_t fraction = ONE) { + logger::Checkpoint("Launching Ampere kernel", HERE); + const auto dT = fraction * + m_params.template get( + "algorithms.timestep.correction") * + dt; + auto range = CreateRangePolicy( + { domain.mesh.i_min(in::x1), domain.mesh.i_min(in::x2) }, + { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); + const auto ni2 = domain.mesh.n_active(in::x2); + + if (g == gr_ampere::aux) { + // First push, updates D0 with J. + Kokkos::parallel_for("Ampere-1", + range, + kernel::gr::Ampere_kernel(domain.fields.em0, // Din + domain.fields.em0, // Dout + domain.fields.aux, + domain.mesh.metric, + dT, + ni2, + domain.mesh.flds_bc())); + } else if (g == gr_ampere::main) { + // Second push, updates D with J0 but assigns it to D0. + Kokkos::parallel_for("Ampere-2", + range, + kernel::gr::Ampere_kernel(domain.fields.em, + domain.fields.em0, + domain.fields.aux, + domain.mesh.metric, + dT, + ni2, + domain.mesh.flds_bc())); + } else if (g == gr_ampere::init) { + // Second push, updates D with J0 and assigns it to D. + Kokkos::parallel_for("Ampere-3", + range, + kernel::gr::Ampere_kernel(domain.fields.em, + domain.fields.em, + domain.fields.aux, + domain.mesh.metric, + dT, + ni2, + domain.mesh.flds_bc())); + } else { + raise::Error("Wrong option for `g`", HERE); + } + } + + void AmpereCurrents(domain_t& domain, const gr_ampere& g) { + logger::Checkpoint("Launching Ampere kernel for adding currents", HERE); + const auto q0 = m_params.template get("scales.q0"); + const auto B0 = m_params.template get("scales.B0"); + const auto coeff = -dt * q0 / B0; + auto range = CreateRangePolicy( + { domain.mesh.i_min(in::x1), domain.mesh.i_min(in::x2) }, + { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); + const auto ni2 = domain.mesh.n_active(in::x2); + + if (g == gr_ampere::aux) { + // Updates D0 with J: D0(n-1/2) -> (J(n)) -> D0(n+1/2) + Kokkos::parallel_for( + "AmpereCurrentsAux", + range, + kernel::gr::CurrentsAmpere_kernel(domain.fields.em0, + domain.fields.cur, + domain.mesh.metric, + coeff, + ni2, + domain.mesh.flds_bc())); + } else if (g == gr_ampere::main) { + // Updates D0 with J0: D0(n) -> (J0(n+1/2)) -> D0(n+1) + Kokkos::parallel_for( + "AmpereCurrentsMain", + range, + kernel::gr::CurrentsAmpere_kernel(domain.fields.em0, + domain.fields.cur0, + domain.mesh.metric, + coeff, + ni2, + domain.mesh.flds_bc())); + } else { + raise::Error("Wrong option for `g`", HERE); + } + } + + void TimeAverageDB(domain_t& domain) { + Kokkos::parallel_for("TimeAverageDB", + domain.mesh.rangeActiveCells(), + kernel::gr::TimeAverageDB_kernel(domain.fields.em, + domain.fields.em0, + domain.mesh.metric)); + } + + void TimeAverageJ(domain_t& domain) { + Kokkos::parallel_for("TimeAverageJ", + domain.mesh.rangeActiveCells(), + kernel::gr::TimeAverageJ_kernel(domain.fields.cur, + domain.fields.cur0, + domain.mesh.metric)); + } + + void CurrentsDeposit(domain_t& domain) { + auto scatter_cur0 = Kokkos::Experimental::create_scatter_view( + domain.fields.cur0); + for (auto& species : domain.species) { + logger::Checkpoint( + fmt::format("Launching currents deposit kernel for %d [%s] : %lu %f", + species.index(), + species.label().c_str(), + species.npart(), + (double)species.charge()), + HERE); + if (species.npart() == 0 || cmp::AlmostZero(species.charge())) { + continue; + } + Kokkos::parallel_for("CurrentsDeposit", + species.rangeActiveParticles(), + kernel::DepositCurrents_kernel( + scatter_cur0, + species.i1, + species.i2, + species.i3, + species.i1_prev, + species.i2_prev, + species.i3_prev, + species.dx1, + species.dx2, + species.dx3, + species.dx1_prev, + species.dx2_prev, + species.dx3_prev, + species.ux1, + species.ux2, + species.ux3, + species.phi, + species.weight, + species.tag, + domain.mesh.metric, + (real_t)(species.charge()), + dt)); + } + Kokkos::Experimental::contribute(domain.fields.cur0, scatter_cur0); + } + + void CurrentsFilter(domain_t& domain) { + logger::Checkpoint("Launching currents filtering kernels", HERE); + auto range = CreateRangePolicy( + { domain.mesh.i_min(in::x1), domain.mesh.i_min(in::x2) }, + { domain.mesh.i_max(in::x1), domain.mesh.i_max(in::x2) + 1 }); + const auto nfilter = m_params.template get( + "algorithms.current_filters"); + tuple_t size; + size[0] = domain.mesh.n_active(in::x1); + size[1] = domain.mesh.n_active(in::x2); + + // !TODO: this needs to be done more efficiently + for (unsigned short i = 0; i < nfilter; ++i) { + Kokkos::deep_copy(domain.fields.buff, domain.fields.cur0); + Kokkos::parallel_for("CurrentsFilter", + range, + kernel::DigitalFilter_kernel( + domain.fields.cur0, + domain.fields.buff, + size, + domain.mesh.flds_bc())); + m_metadomain.CommunicateFields(domain, Comm::J); // J0 + } + } + + void ParticlePush(domain_t& domain) { + for (auto& species : domain.species) { + species.set_unsorted(); + logger::Checkpoint( + fmt::format("Launching particle pusher kernel for %d [%s] : %lu", + species.index(), + species.label().c_str(), + species.npart()), + HERE); + if (species.npart() == 0) { + continue; + } + const auto q_ovr_m = species.mass() > ZERO + ? species.charge() / species.mass() + : ZERO; + // coeff = q / m (dt / 2) omegaB0 + const auto coeff = q_ovr_m * HALF * dt * + m_params.template get( + "algorithms.timestep.correction") * + m_params.template get("scales.omegaB0"); + const auto eps = m_params.template get( + "algorithms.gr.pusher_eps"); + const auto niter = m_params.template get( + "algorithms.gr.pusher_niter"); + // clang-format off + if (species.pusher() == PrtlPusher::PHOTON) { + auto range_policy = Kokkos::RangePolicy( + 0, + species.npart()); + + Kokkos::parallel_for( + "ParticlePusher", + range_policy, + kernel::gr::Pusher_kernel( + domain.fields.em, + domain.fields.em0, + species.i1, species.i2, species.i3, + species.i1_prev, species.i2_prev, species.i3_prev, + species.dx1, species.dx2, species.dx3, + species.dx1_prev, species.dx2_prev, species.dx3_prev, + species.ux1, species.ux2, species.ux3, + species.phi, species.tag, + domain.mesh.metric, + coeff, dt, + domain.mesh.n_active(in::x1), + domain.mesh.n_active(in::x2), + domain.mesh.n_active(in::x3), + eps, niter, + domain.mesh.prtl_bc() + )); + } else if (species.pusher() == PrtlPusher::BORIS) { + auto range_policy = Kokkos::RangePolicy( + 0, + species.npart()); + Kokkos::parallel_for( + "ParticlePusher", + range_policy, + kernel::gr::Pusher_kernel( + domain.fields.em, + domain.fields.em0, + species.i1, species.i2, species.i3, + species.i1_prev, species.i2_prev, species.i3_prev, + species.dx1, species.dx2, species.dx3, + species.dx1_prev, species.dx2_prev, species.dx3_prev, + species.ux1, species.ux2, species.ux3, + species.phi, species.tag, + domain.mesh.metric, + coeff, dt, + domain.mesh.n_active(in::x1), + domain.mesh.n_active(in::x2), + domain.mesh.n_active(in::x3), + eps, niter, + domain.mesh.prtl_bc() + )); + } else if (species.pusher() == PrtlPusher::NONE) { + // do nothing + } else { + raise::Error("not implemented", HERE); + } + // clang-format on + } + } + }; +} // namespace ntt + +#endif // ENGINES_GRPIC_GRPIC_H diff --git a/src/kernels/injectors.hpp b/src/kernels/injectors.hpp index cddb1883..e9926e36 100644 --- a/src/kernels/injectors.hpp +++ b/src/kernels/injectors.hpp @@ -1,860 +1,864 @@ -/** - * @file kernels/injectors.hpp - * @brief Kernels for injecting particles in different ways - * @implements - * - kernel::UniformInjector_kernel<> - * - kernel::GlobalInjector_kernel<> - * - kernel::NonUniformInjector_kernel<> - * @namespaces: - * - kernel:: - */ - -#ifndef KERNELS_INJECTORS_HPP -#define KERNELS_INJECTORS_HPP - -#include "enums.h" -#include "global.h" - -#include "arch/kokkos_aliases.h" -#include "utils/error.h" -#include "utils/numeric.h" - -#include "framework/containers/particles.h" -#include "framework/domain/domain.h" - -namespace kernel { - using namespace ntt; - - template - Inline void InjectParticle(npart_t p, - const array_t& i1_arr, - const array_t& i2_arr, - const array_t& i3_arr, - const array_t& dx1_arr, - const array_t& dx2_arr, - const array_t& dx3_arr, - const array_t& ux1_arr, - const array_t& ux2_arr, - const array_t& ux3_arr, - const array_t& phi_arr, - const array_t& weight_arr, - const array_t& tag_arr, - const array_t& pld_i_arr, - const tuple_t& xi_Cd, - const tuple_t& dxi_Cd, - const vec_t& v_Cd, - real_t weight = ONE, - real_t phi = ZERO, - npart_t domain_idx = 0u, - npart_t species_cntr = 0u) { - if constexpr (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { - i1_arr(p) = xi_Cd[0]; - dx1_arr(p) = dxi_Cd[0]; - } - if constexpr (D == Dim::_2D or D == Dim::_3D) { - i2_arr(p) = xi_Cd[1]; - dx2_arr(p) = dxi_Cd[1]; - } - if constexpr (D == Dim::_3D) { - i3_arr(p) = xi_Cd[2]; - dx3_arr(p) = dxi_Cd[2]; - } - if constexpr (D == Dim::_2D and C != Coord::Cart) { - phi_arr(p) = phi; - } - ux1_arr(p) = v_Cd[0]; - ux2_arr(p) = v_Cd[1]; - ux3_arr(p) = v_Cd[2]; - tag_arr(p) = ParticleTag::alive; - weight_arr(p) = weight; - if constexpr (T) { - pld_i_arr(p, pldi::spcCtr) = species_cntr; -#if defined(MPI_ENABLED) - pld_i_arr(p, pldi::domIdx) = domain_idx; -#endif - } - } - - template - struct UniformInjector_kernel { - static_assert(ED1::is_energy_dist, "ED1 must be an energy distribution class"); - static_assert(ED2::is_energy_dist, "ED2 must be an energy distribution class"); - static_assert(M::is_metric, "M must be a metric class"); - - array_t i1s_1, i2s_1, i3s_1; - array_t dx1s_1, dx2s_1, dx3s_1; - array_t ux1s_1, ux2s_1, ux3s_1; - array_t phis_1; - array_t weights_1; - array_t tags_1; - array_t pldis_1; - - array_t i1s_2, i2s_2, i3s_2; - array_t dx1s_2, dx2s_2, dx3s_2; - array_t ux1s_2, ux2s_2, ux3s_2; - array_t phis_2; - array_t weights_2; - array_t tags_2; - array_t pldis_2; - - const npart_t offset1, offset2; - const npart_t domain_idx, cntr1, cntr2; - const bool use_tracking_1, use_tracking_2; - const M metric; - const array_t xi_min, xi_max; - const ED1 energy_dist_1; - const ED2 energy_dist_2; - const real_t inv_V0; - random_number_pool_t random_pool; - - UniformInjector_kernel(Particles& species1, - Particles& species2, - npart_t inject_npart, - npart_t domain_idx, - const M& metric, - const array_t& xi_min, - const array_t& xi_max, - const ED1& energy_dist_1, - const ED2& energy_dist_2, - real_t inv_V0, - random_number_pool_t& random_pool) - : i1s_1 { species1.i1 } - , i2s_1 { species1.i2 } - , i3s_1 { species1.i3 } - , dx1s_1 { species1.dx1 } - , dx2s_1 { species1.dx2 } - , dx3s_1 { species1.dx3 } - , ux1s_1 { species1.ux1 } - , ux2s_1 { species1.ux2 } - , ux3s_1 { species1.ux3 } - , phis_1 { species1.phi } - , weights_1 { species1.weight } - , tags_1 { species1.tag } - , pldis_1 { species1.pld_i } - , i1s_2 { species2.i1 } - , i2s_2 { species2.i2 } - , i3s_2 { species2.i3 } - , dx1s_2 { species2.dx1 } - , dx2s_2 { species2.dx2 } - , dx3s_2 { species2.dx3 } - , ux1s_2 { species2.ux1 } - , ux2s_2 { species2.ux2 } - , ux3s_2 { species2.ux3 } - , phis_2 { species2.phi } - , weights_2 { species2.weight } - , tags_2 { species2.tag } - , pldis_2 { species2.pld_i } - , offset1 { species1.npart() } - , offset2 { species2.npart() } - , domain_idx { domain_idx } - , cntr1 { species1.counter() } - , cntr2 { species2.counter() } - , use_tracking_1 { species1.use_tracking() } - , use_tracking_2 { species2.use_tracking() } - , metric { metric } - , xi_min { xi_min } - , xi_max { xi_max } - , energy_dist_1 { energy_dist_1 } - , energy_dist_2 { energy_dist_2 } - , inv_V0 { inv_V0 } - , random_pool { random_pool } { - if (use_tracking_1) { - species1.set_counter(cntr1 + inject_npart); -#if !defined(MPI_ENABLED) - raise::ErrorIf(species1.pld_i.extent(1) < 1, - "Particle tracking is enabled but the " - "particle integer payload size is less " - "than 1", - HERE); -#else - raise::ErrorIf(species1.pld_i.extent(1) < 2, - "Particle tracking is enabled but the " - "particle integer payload size is less " - "than 2", - HERE); -#endif - } - if (use_tracking_2) { - species2.set_counter(cntr2 + inject_npart); -#if !defined(MPI_ENABLED) - raise::ErrorIf(species2.pld_i.extent(1) < 1, - "Particle tracking is enabled but the " - "particle integer payload size is less " - "than 1", - HERE); -#else - raise::ErrorIf(species2.pld_i.extent(1) < 2, - "Particle tracking is enabled but the " - "particle integer payload size is less " - "than 2", - HERE); -#endif - } - } - - Inline void operator()(index_t p) const { - coord_t x_Cd { ZERO }; - tuple_t xi_Cd { 0 }; - tuple_t dxi_Cd { static_cast(0) }; - vec_t v1 { ZERO }, v2 { ZERO }; - { // generate a random coordinate - auto rand_gen = random_pool.get_state(); - if constexpr (M::Dim == Dim::_1D or M::Dim == Dim::_2D or - M::Dim == Dim::_3D) { - x_Cd[0] = xi_min(0) + Random(rand_gen) * (xi_max(0) - xi_min(0)); - xi_Cd[0] = static_cast(x_Cd[0]); - dxi_Cd[0] = static_cast(x_Cd[0] - xi_Cd[0]); - } - if constexpr (M::Dim == Dim::_2D or M::Dim == Dim::_3D) { - x_Cd[1] = xi_min(1) + Random(rand_gen) * (xi_max(1) - xi_min(1)); - xi_Cd[1] = static_cast(x_Cd[1]); - xi_Cd[1] = static_cast(x_Cd[1]); - dxi_Cd[1] = static_cast(x_Cd[1] - xi_Cd[1]); - } - if constexpr (M::Dim == Dim::_3D) { - x_Cd[2] = xi_min(2) + Random(rand_gen) * (xi_max(2) - xi_min(2)); - xi_Cd[2] = static_cast(x_Cd[2]); - dxi_Cd[2] = static_cast(x_Cd[2] - xi_Cd[2]); - } - random_pool.free_state(rand_gen); - } - { // generate the velocity - coord_t x_Ph { ZERO }; - metric.template convert(x_Cd, x_Ph); - if constexpr (M::CoordType == Coord::Cart) { - energy_dist_1(x_Ph, v1); - energy_dist_2(x_Ph, v2); - } else if constexpr (S == SimEngine::SRPIC) { - coord_t x_Cd_ { ZERO }; - x_Cd_[0] = x_Cd[0]; - x_Cd_[1] = x_Cd[1]; - x_Cd_[2] = ZERO; // phi = 0 - vec_t v_Ph { ZERO }; - energy_dist_1(x_Ph, v_Ph); - metric.template transform_xyz(x_Cd_, v_Ph, v1); - energy_dist_2(x_Ph, v_Ph); - metric.template transform_xyz(x_Cd_, v_Ph, v2); - } else if constexpr (S == SimEngine::GRPIC) { - vec_t v_Ph { ZERO }; - energy_dist_1(x_Ph, v_Ph); - metric.template transform(x_Cd, v_Ph, v1); - energy_dist_2(x_Ph, v_Ph); - metric.template transform(x_Cd, v_Ph, v2); - } else { - raise::KernelError(HERE, "Unknown simulation engine"); - } - } - real_t weight = ONE; - if constexpr (M::CoordType != Coord::Cart) { - const auto sqrt_det_h = metric.sqrt_det_h(x_Cd); - weight = sqrt_det_h * inv_V0; - } - // clang-format off - if (not use_tracking_1) { - InjectParticle( - p + offset1, - i1s_1, i2s_1, i3s_1, - dx1s_1, dx2s_1, dx3s_1, - ux1s_1, ux2s_1, ux3s_1, - phis_1, weights_1, tags_1, pldis_1, - xi_Cd, dxi_Cd, v1, weight, ZERO); - } else { - InjectParticle( - p + offset1, - i1s_1, i2s_1, i3s_1, - dx1s_1, dx2s_1, dx3s_1, - ux1s_1, ux2s_1, ux3s_1, - phis_1, weights_1, tags_1, pldis_1, - xi_Cd, dxi_Cd, v1, weight, ZERO, - domain_idx, cntr1 + p); - } - if (not use_tracking_2) { - InjectParticle( - p + offset2, - i1s_2, i2s_2, i3s_2, - dx1s_2, dx2s_2, dx3s_2, - ux1s_2, ux2s_2, ux3s_2, - phis_2, weights_2, tags_2, pldis_2, - xi_Cd, dxi_Cd, v2, weight, ZERO); - } else { - InjectParticle( - p + offset2, - i1s_2, i2s_2, i3s_2, - dx1s_2, dx2s_2, dx3s_2, - ux1s_2, ux2s_2, ux3s_2, - phis_2, weights_2, tags_2, pldis_2, - xi_Cd, dxi_Cd, v2, weight, ZERO, - domain_idx, cntr2 + p); - } - // clang-format on - } - }; // struct UniformInjector_kernel - - template - struct GlobalInjector_kernel { - static_assert(M::is_metric, "M must be a metric class"); - static constexpr auto D = M::Dim; - - const bool use_weights; - - array_t in_ux1; - array_t in_ux2; - array_t in_ux3; - array_t in_x1; - array_t in_x2; - array_t in_x3; - array_t in_phi; - array_t in_wei; - - array_t idx { "idx" }; - array_t i1s, i2s, i3s; - array_t dx1s, dx2s, dx3s; - array_t ux1s, ux2s, ux3s; - array_t phis; - array_t weights; - array_t tags; - - const npart_t offset; - - M global_metric; - - real_t x1_min, x1_max, x2_min, x2_max, x3_min, x3_max; - ncells_t i1_offset, i2_offset, i3_offset; - - GlobalInjector_kernel(Particles& species, - const M& global_metric, - const Domain& local_domain, - const std::map>& data, - bool use_weights) - : use_weights { use_weights } - , i1s { species.i1 } - , i2s { species.i2 } - , i3s { species.i3 } - , dx1s { species.dx1 } - , dx2s { species.dx2 } - , dx3s { species.dx3 } - , ux1s { species.ux1 } - , ux2s { species.ux2 } - , ux3s { species.ux3 } - , phis { species.phi } - , weights { species.weight } - , tags { species.tag } - , offset { species.npart() } - , global_metric { global_metric } { - const auto n_inject = data.at("x1").size(); - - x1_min = local_domain.mesh.extent(in::x1).first; - x1_max = local_domain.mesh.extent(in::x1).second; - i1_offset = local_domain.offset_ncells()[0]; - - copy_from_vector("x1", in_x1, data, n_inject); - copy_from_vector("ux1", in_ux1, data, n_inject); - copy_from_vector("ux2", in_ux2, data, n_inject); - copy_from_vector("ux3", in_ux3, data, n_inject); - if (use_weights) { - copy_from_vector("weights", in_wei, data, n_inject); - } - if constexpr (D == Dim::_2D or D == Dim::_3D) { - x2_min = local_domain.mesh.extent(in::x2).first; - x2_max = local_domain.mesh.extent(in::x2).second; - i2_offset = local_domain.offset_ncells()[1]; - copy_from_vector("x2", in_x2, data, n_inject); - } - if constexpr (D == Dim::_2D and M::CoordType != Coord::Cart) { - copy_from_vector("phi", in_phi, data, n_inject); - } - if constexpr (D == Dim::_3D) { - x3_min = local_domain.mesh.extent(in::x3).first; - x3_max = local_domain.mesh.extent(in::x3).second; - i3_offset = local_domain.offset_ncells()[2]; - copy_from_vector("x3", in_x3, data, n_inject); - } - } - - void copy_from_vector(const std::string& name, - array_t& arr, - const std::map>& data, - npart_t n_inject) { - raise::ErrorIf(data.find(name) == data.end(), name + " not found in data", HERE); - raise::ErrorIf(data.at(name).size() != n_inject, "Inconsistent data size", HERE); - arr = array_t { name, n_inject }; - auto arr_h = Kokkos::create_mirror_view(arr); - for (auto i = 0u; i < data.at(name).size(); ++i) { - arr_h(i) = data.at(name)[i]; - } - Kokkos::deep_copy(arr, arr_h); - } - - auto number_injected() const -> npart_t { - auto idx_h = Kokkos::create_mirror_view(idx); - Kokkos::deep_copy(idx_h, idx); - return idx_h(); - } - - Inline void operator()(index_t p) const { - if constexpr (D == Dim::_1D) { - if (in_x1(p) >= x1_min and in_x1(p) < x1_max) { - coord_t x_Cd { ZERO }; - vec_t u_XYZ { ZERO }; - const vec_t u_Ph { in_ux1(p), in_ux2(p), in_ux3(p) }; - - auto index { offset + Kokkos::atomic_fetch_add(&idx(), 1) }; - global_metric.template convert({ in_x1(p) }, x_Cd); - global_metric.template transform_xyz(x_Cd, u_Ph, u_XYZ); - - const auto i1 = static_cast( - static_cast(x_Cd[0]) - i1_offset); - const auto dx1 = static_cast( - x_Cd[0] - static_cast(i1 + i1_offset)); - - i1s(index) = i1; - dx1s(index) = dx1; - ux1s(index) = u_XYZ[0]; - ux2s(index) = u_XYZ[1]; - ux3s(index) = u_XYZ[2]; - tags(index) = ParticleTag::alive; - if (use_weights) { - weights(index) = weights(p); - } else { - weights(index) = ONE; - } - } - } else if constexpr (D == Dim::_2D) { - if ((in_x1(p) >= x1_min and in_x1(p) < x1_max) and - (in_x2(p) >= x2_min and in_x2(p) < x2_max)) { - coord_t x_Cd { ZERO }; - vec_t u_Cd { ZERO }; - vec_t u_Ph { in_ux1(p), in_ux2(p), in_ux3(p) }; - coord_t x_Cd_ { ZERO }; - - auto index { offset + - Kokkos::atomic_fetch_add(&idx(), static_cast(1)) }; - global_metric.template convert({ in_x1(p), in_x2(p) }, - x_Cd); - x_Cd_[0] = x_Cd[0]; - x_Cd_[1] = x_Cd[1]; - if constexpr (S == SimEngine::SRPIC and M::CoordType != Coord::Cart) { - x_Cd_[2] = in_phi(p); - } - if constexpr (S == SimEngine::SRPIC) { - global_metric.template transform_xyz(x_Cd_, u_Ph, u_Cd); - } else if constexpr (S == SimEngine::GRPIC) { - global_metric.template transform(x_Cd, u_Ph, u_Cd); - } else { - raise::KernelError(HERE, "Unknown simulation engine"); - } - const auto i1 = static_cast( - static_cast(x_Cd[0]) - i1_offset); - const auto dx1 = static_cast( - x_Cd[0] - static_cast(i1 + i1_offset)); - const auto i2 = static_cast( - static_cast(x_Cd[1]) - i2_offset); - const auto dx2 = static_cast( - x_Cd[1] - static_cast(i2 + i2_offset)); - - i1s(index) = i1; - dx1s(index) = dx1; - i2s(index) = i2; - dx2s(index) = dx2; - ux1s(index) = u_Cd[0]; - ux2s(index) = u_Cd[1]; - ux3s(index) = u_Cd[2]; - if (M::CoordType != Coord::Cart) { - phis(index) = in_phi(p); - } - tags(index) = ParticleTag::alive; - if (use_weights) { - weights(index) = weights(p); - } else { - weights(index) = ONE; - } - } - } else { - if ((in_x1(p) >= x1_min and in_x1(p) < x1_max) and - (in_x2(p) >= x2_min and in_x2(p) < x2_max) and - (in_x3(p) >= x3_min and in_x3(p) < x3_max)) { - coord_t x_Cd { ZERO }; - vec_t u_Cd { ZERO }; - vec_t u_Ph { in_ux1(p), in_ux2(p), in_ux3(p) }; - - auto index { offset + Kokkos::atomic_fetch_add(&idx(), 1) }; - global_metric.template convert( - { in_x1(p), in_x2(p), in_x3(p) }, - x_Cd); - if constexpr (S == SimEngine::SRPIC) { - global_metric.template transform_xyz(x_Cd, u_Ph, u_Cd); - } else if constexpr (S == SimEngine::GRPIC) { - global_metric.template transform(x_Cd, u_Ph, u_Cd); - } else { - raise::KernelError(HERE, "Unknown simulation engine"); - } - const auto i1 = static_cast( - static_cast(x_Cd[0]) - i1_offset); - const auto dx1 = static_cast( - x_Cd[0] - static_cast(i1 + i1_offset)); - const auto i2 = static_cast( - static_cast(x_Cd[1]) - i2_offset); - const auto dx2 = static_cast( - x_Cd[1] - static_cast(i2 + i2_offset)); - const auto i3 = static_cast( - static_cast(x_Cd[2]) - i3_offset); - const auto dx3 = static_cast( - x_Cd[2] - static_cast(i3 + i3_offset)); - - i1s(index) = i1; - dx1s(index) = dx1; - i2s(index) = i2; - dx2s(index) = dx2; - i3s(index) = i3; - dx3s(index) = dx3; - ux1s(index) = u_Cd[0]; - ux2s(index) = u_Cd[1]; - ux3s(index) = u_Cd[2]; - tags(index) = ParticleTag::alive; - if (use_weights) { - weights(index) = weights(p); - } else { - weights(index) = ONE; - } - } - } - } - }; // struct GlobalInjector_kernel - - template - struct NonUniformInjector_kernel { - static_assert(M::is_metric, "M must be a metric class"); - static_assert(ED1::is_energy_dist, "ED1 must be an energy distribution class"); - static_assert(ED2::is_energy_dist, "ED2 must be an energy distribution class"); - static_assert(SD::is_spatial_dist, "SD must be a spatial distribution class"); - - const real_t ppc0; - - array_t i1s_1, i2s_1, i3s_1; - array_t dx1s_1, dx2s_1, dx3s_1; - array_t ux1s_1, ux2s_1, ux3s_1; - array_t phis_1; - array_t weights_1; - array_t tags_1; - array_t pldis_1; - - array_t i1s_2, i2s_2, i3s_2; - array_t dx1s_2, dx2s_2, dx3s_2; - array_t ux1s_2, ux2s_2, ux3s_2; - array_t phis_2; - array_t weights_2; - array_t tags_2; - array_t pldis_2; - - array_t idx { "idx" }; - - const npart_t offset1, offset2; - const npart_t domain_idx, cntr1, cntr2; - const bool use_tracking_1, use_tracking_2; - const M metric; - const ED1 energy_dist_1; - const ED2 energy_dist_2; - const SD spatial_dist; - const real_t inv_V0; - random_number_pool_t random_pool; - - NonUniformInjector_kernel(real_t ppc0, - Particles& species1, - Particles& species2, - npart_t domain_idx, - const M& metric, - const ED1& energy_dist_1, - const ED2& energy_dist_2, - const SD& spatial_dist, - real_t inv_V0, - random_number_pool_t& random_pool) - : ppc0 { ppc0 } - , i1s_1 { species1.i1 } - , i2s_1 { species1.i2 } - , i3s_1 { species1.i3 } - , dx1s_1 { species1.dx1 } - , dx2s_1 { species1.dx2 } - , dx3s_1 { species1.dx3 } - , ux1s_1 { species1.ux1 } - , ux2s_1 { species1.ux2 } - , ux3s_1 { species1.ux3 } - , phis_1 { species1.phi } - , weights_1 { species1.weight } - , tags_1 { species1.tag } - , pldis_1 { species1.pld_i } - , i1s_2 { species2.i1 } - , i2s_2 { species2.i2 } - , i3s_2 { species2.i3 } - , dx1s_2 { species2.dx1 } - , dx2s_2 { species2.dx2 } - , dx3s_2 { species2.dx3 } - , ux1s_2 { species2.ux1 } - , ux2s_2 { species2.ux2 } - , ux3s_2 { species2.ux3 } - , phis_2 { species2.phi } - , weights_2 { species2.weight } - , tags_2 { species2.tag } - , pldis_2 { species2.pld_i } - , offset1 { species1.npart() } - , offset2 { species2.npart() } - , domain_idx { domain_idx } - , cntr1 { species1.counter() } - , cntr2 { species2.counter() } - , use_tracking_1 { species1.use_tracking() } - , use_tracking_2 { species2.use_tracking() } - , metric { metric } - , energy_dist_1 { energy_dist_1 } - , energy_dist_2 { energy_dist_2 } - , spatial_dist { spatial_dist } - , inv_V0 { inv_V0 } - , random_pool { random_pool } {} - - auto number_injected() const -> npart_t { - auto idx_h = Kokkos::create_mirror_view(idx); - Kokkos::deep_copy(idx_h, idx); - return idx_h(); - } - - Inline auto injected_ppc(const coord_t& x_Ph) const -> npart_t { - const auto ppc_real = ppc0 * spatial_dist(x_Ph); - auto ppc = static_cast(ppc_real); - auto rand_gen = random_pool.get_state(); - if (Random(rand_gen) < (ppc_real - static_cast(ppc))) { - ppc += 1; - } - random_pool.free_state(rand_gen); - return ppc; - } - - Inline void inject1(const index_t index, - const tuple_t& xi_Cd, - const tuple_t& dxi_Cd, - const vec_t& v_Cd, - const real_t weight) const { - // clang-format off - if (not use_tracking_1) { - InjectParticle(index + offset1, - i1s_1, i2s_1, i3s_1, - dx1s_1, dx2s_1, dx3s_1, - ux1s_1, ux2s_1, ux3s_1, - phis_1, weights_1, tags_1, pldis_1, - xi_Cd, dxi_Cd, v_Cd, weight, ZERO); - } else { - InjectParticle(index + offset1, - i1s_1, i2s_1, i3s_1, - dx1s_1, dx2s_1, dx3s_1, - ux1s_1, ux2s_1, ux3s_1, - phis_1, weights_1, tags_1, pldis_1, - xi_Cd, dxi_Cd, v_Cd, weight, ZERO, - domain_idx, index + cntr1); - } - // clang-format on - } - - Inline void inject2(const index_t index, - const tuple_t& xi_Cd, - const tuple_t& dxi_Cd, - const vec_t& v_Cd, - const real_t weight) const { - // clang-format off - if (not use_tracking_2) { - InjectParticle(index + offset2, - i1s_2, i2s_2, i3s_2, - dx1s_2, dx2s_2, dx3s_2, - ux1s_2, ux2s_2, ux3s_2, - phis_2, weights_2, tags_2, pldis_2, - xi_Cd, dxi_Cd, v_Cd, weight, ZERO); - } else { - InjectParticle(index + offset2, - i1s_2, i2s_2, i3s_2, - dx1s_2, dx2s_2, dx3s_2, - ux1s_2, ux2s_2, ux3s_2, - phis_2, weights_2, tags_2, pldis_2, - xi_Cd, dxi_Cd, v_Cd, weight, ZERO, - domain_idx, index + cntr2); - } - // clang-format on - } - - Inline void operator()(index_t i1) const { - if constexpr (M::Dim == Dim::_1D) { - const auto i1_ = COORD(i1); - coord_t x_Cd { i1_ + HALF }; - coord_t x_Ph { ZERO }; - metric.template convert(x_Cd, x_Ph); - - const auto ppc = injected_ppc(x_Ph); - if (ppc == 0) { - return; - } - - auto weight = ONE; - if constexpr (M::CoordType != Coord::Cart) { - weight = metric.sqrt_det_h({ i1_ + HALF }) * inv_V0; - } - for (auto p { 0u }; p < ppc; ++p) { - const auto index = Kokkos::atomic_fetch_add(&idx(), 1); - - auto rand_gen = random_pool.get_state(); - const auto dx1 = Random(rand_gen); - random_pool.free_state(rand_gen); - - vec_t v_XYZ { ZERO }; - { - vec_t v_T { ZERO }; - energy_dist_1(x_Ph, v_T); - metric.template transform_xyz(x_Cd, v_T, v_XYZ); - } - inject1(index, { static_cast(i1_) }, { dx1 }, v_XYZ, weight); - - { - vec_t v_T { ZERO }; - energy_dist_2(x_Ph, v_T); - metric.template transform_xyz(x_Cd, v_T, v_XYZ); - } - inject2(index, { static_cast(i1_) }, { dx1 }, v_XYZ, weight); - } - } else { - raise::KernelError(HERE, "NonUniformInjector_kernel 1D called for 2D/3D"); - } - } - - Inline void operator()(index_t i1, index_t i2) const { - if constexpr (M::Dim == Dim::_2D) { - const auto i1_ = COORD(i1); - const auto i2_ = COORD(i2); - coord_t x_Cd { i1_ + HALF, i2_ + HALF }; - coord_t x_Ph { ZERO }; - coord_t x_Cd_ { ZERO }; - x_Cd_[0] = x_Cd[0]; - x_Cd_[1] = x_Cd[1]; - if constexpr (S == SimEngine::SRPIC and M::CoordType != Coord::Cart) { - x_Cd_[2] = ZERO; - } - metric.template convert(x_Cd, x_Ph); - - const auto ppc = injected_ppc(x_Ph); - if (ppc == 0) { - return; - } - - auto weight = ONE; - if constexpr (M::CoordType != Coord::Cart) { - weight = metric.sqrt_det_h({ i1_ + HALF, i2_ + HALF }) * inv_V0; - } - for (auto p { 0u }; p < ppc; ++p) { - const auto index = Kokkos::atomic_fetch_add(&idx(), 1); - - auto rand_gen = random_pool.get_state(); - const auto dx1 = Random(rand_gen); - const auto dx2 = Random(rand_gen); - random_pool.free_state(rand_gen); - - vec_t v_Cd { ZERO }; - { - vec_t v_T { ZERO }; - energy_dist_1(x_Ph, v_T); - if constexpr (S == SimEngine::SRPIC) { - metric.template transform_xyz(x_Cd_, v_T, v_Cd); - } else if constexpr (S == SimEngine::GRPIC) { - metric.template transform(x_Cd_, v_T, v_Cd); - } - } - inject1(index, - { static_cast(i1_), static_cast(i2_) }, - { dx1, dx2 }, - v_Cd, - weight); - - { - vec_t v_T { ZERO }; - energy_dist_2(x_Ph, v_T); - if constexpr (S == SimEngine::SRPIC) { - metric.template transform_xyz(x_Cd_, v_T, v_Cd); - } else if constexpr (S == SimEngine::GRPIC) { - metric.template transform(x_Cd_, v_T, v_Cd); - } - } - inject2(index, - { static_cast(i1_), static_cast(i2_) }, - { dx1, dx2 }, - v_Cd, - weight); - } - } - - else { - raise::KernelError(HERE, "NonUniformInjector_kernel 2D called for 1D/3D"); - } - } - - Inline void operator()(index_t i1, index_t i2, index_t i3) const { - if constexpr (M::Dim == Dim::_3D) { - const auto i1_ = COORD(i1); - const auto i2_ = COORD(i2); - const auto i3_ = COORD(i3); - coord_t x_Cd { i1_ + HALF, i2_ + HALF, i3_ + HALF }; - coord_t x_Ph { ZERO }; - metric.template convert(x_Cd, x_Ph); - - const auto ppc = injected_ppc(x_Ph); - if (ppc == 0) { - return; - } - - auto weight = ONE; - if constexpr (M::CoordType != Coord::Cart) { - weight = metric.sqrt_det_h({ i1_ + HALF, i2_ + HALF, i3_ + HALF }) * - inv_V0; - } - for (auto p { 0u }; p < ppc; ++p) { - const auto index = Kokkos::atomic_fetch_add(&idx(), 1); - - auto rand_gen = random_pool.get_state(); - const auto dx1 = Random(rand_gen); - const auto dx2 = Random(rand_gen); - const auto dx3 = Random(rand_gen); - random_pool.free_state(rand_gen); - - vec_t v_Cd { ZERO }; - { - vec_t v_T { ZERO }; - energy_dist_1(x_Ph, v_T); - if constexpr (S == SimEngine::SRPIC) { - metric.template transform_xyz(x_Cd, v_T, v_Cd); - } else if constexpr (S == SimEngine::GRPIC) { - metric.template transform(x_Cd, v_T, v_Cd); - } - } - inject1( - index, - { static_cast(i1_), static_cast(i2_), static_cast(i3_) }, - { dx1, dx2, dx3 }, - v_Cd, - weight); - - { - vec_t v_T { ZERO }; - energy_dist_2(x_Ph, v_T); - if constexpr (S == SimEngine::SRPIC) { - metric.template transform_xyz(x_Cd, v_T, v_Cd); - } else if constexpr (S == SimEngine::GRPIC) { - metric.template transform(x_Cd, v_T, v_Cd); - } - } - inject2( - index, - { static_cast(i1_), static_cast(i2_), static_cast(i3_) }, - { dx1, dx2, dx3 }, - v_Cd, - weight); - } - } else { - raise::KernelError(HERE, "NonUniformInjector_kernel 3D called for 1D/2D"); - } - } - }; // struct NonUniformInjector_kernel - -} // namespace kernel - -#endif // KERNELS_INJECTORS_HPP +/** + * @file kernels/injectors.hpp + * @brief Kernels for injecting particles in different ways + * @implements + * - kernel::UniformInjector_kernel<> + * - kernel::GlobalInjector_kernel<> + * - kernel::NonUniformInjector_kernel<> + * @namespaces: + * - kernel:: + */ + +#ifndef KERNELS_INJECTORS_HPP +#define KERNELS_INJECTORS_HPP + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "utils/error.h" +#include "utils/numeric.h" + +#include "framework/containers/particles.h" +#include "framework/domain/domain.h" + +namespace kernel { + using namespace ntt; + + template + Inline void InjectParticle(npart_t p, + const array_t& i1_arr, + const array_t& i2_arr, + const array_t& i3_arr, + const array_t& dx1_arr, + const array_t& dx2_arr, + const array_t& dx3_arr, + const array_t& ux1_arr, + const array_t& ux2_arr, + const array_t& ux3_arr, + const array_t& phi_arr, + const array_t& weight_arr, + const array_t& tag_arr, + const array_t& pld_i_arr, + const tuple_t& xi_Cd, + const tuple_t& dxi_Cd, + const vec_t& v_Cd, + real_t weight = ONE, + real_t phi = ZERO, + npart_t domain_idx = 0u, + npart_t species_cntr = 0u) { + if constexpr (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { + i1_arr(p) = xi_Cd[0]; + dx1_arr(p) = dxi_Cd[0]; + } + if constexpr (D == Dim::_2D or D == Dim::_3D) { + i2_arr(p) = xi_Cd[1]; + dx2_arr(p) = dxi_Cd[1]; + } + if constexpr (D == Dim::_3D) { + i3_arr(p) = xi_Cd[2]; + dx3_arr(p) = dxi_Cd[2]; + } + if constexpr (D == Dim::_2D and C != Coord::Cart) { + phi_arr(p) = phi; + } + ux1_arr(p) = v_Cd[0]; + ux2_arr(p) = v_Cd[1]; + ux3_arr(p) = v_Cd[2]; + tag_arr(p) = ParticleTag::alive; + weight_arr(p) = weight; + if constexpr (T) { + pld_i_arr(p, pldi::spcCtr) = species_cntr; +#if defined(MPI_ENABLED) + pld_i_arr(p, pldi::domIdx) = domain_idx; +#endif + } + } + + template + struct UniformInjector_kernel { + static_assert(ED1::is_energy_dist, "ED1 must be an energy distribution class"); + static_assert(ED2::is_energy_dist, "ED2 must be an energy distribution class"); + static_assert(M::is_metric, "M must be a metric class"); + + array_t i1s_1, i2s_1, i3s_1; + array_t dx1s_1, dx2s_1, dx3s_1; + array_t ux1s_1, ux2s_1, ux3s_1; + array_t phis_1; + array_t weights_1; + array_t tags_1; + array_t pldis_1; + + array_t i1s_2, i2s_2, i3s_2; + array_t dx1s_2, dx2s_2, dx3s_2; + array_t ux1s_2, ux2s_2, ux3s_2; + array_t phis_2; + array_t weights_2; + array_t tags_2; + array_t pldis_2; + + const npart_t offset1, offset2; + const npart_t domain_idx, cntr1, cntr2; + const bool use_tracking_1, use_tracking_2; + const M metric; + const array_t xi_min, xi_max; + const ED1 energy_dist_1; + const ED2 energy_dist_2; + const real_t inv_V0; + random_number_pool_t random_pool; + + UniformInjector_kernel(Particles& species1, + Particles& species2, + npart_t inject_npart, + npart_t domain_idx, + const M& metric, + const array_t& xi_min, + const array_t& xi_max, + const ED1& energy_dist_1, + const ED2& energy_dist_2, + real_t inv_V0, + random_number_pool_t& random_pool) + : i1s_1 { species1.i1 } + , i2s_1 { species1.i2 } + , i3s_1 { species1.i3 } + , dx1s_1 { species1.dx1 } + , dx2s_1 { species1.dx2 } + , dx3s_1 { species1.dx3 } + , ux1s_1 { species1.ux1 } + , ux2s_1 { species1.ux2 } + , ux3s_1 { species1.ux3 } + , phis_1 { species1.phi } + , weights_1 { species1.weight } + , tags_1 { species1.tag } + , pldis_1 { species1.pld_i } + , i1s_2 { species2.i1 } + , i2s_2 { species2.i2 } + , i3s_2 { species2.i3 } + , dx1s_2 { species2.dx1 } + , dx2s_2 { species2.dx2 } + , dx3s_2 { species2.dx3 } + , ux1s_2 { species2.ux1 } + , ux2s_2 { species2.ux2 } + , ux3s_2 { species2.ux3 } + , phis_2 { species2.phi } + , weights_2 { species2.weight } + , tags_2 { species2.tag } + , pldis_2 { species2.pld_i } + , offset1 { species1.npart() } + , offset2 { species2.npart() } + , domain_idx { domain_idx } + , cntr1 { species1.counter() } + , cntr2 { species2.counter() } + , use_tracking_1 { species1.use_tracking() } + , use_tracking_2 { species2.use_tracking() } + , metric { metric } + , xi_min { xi_min } + , xi_max { xi_max } + , energy_dist_1 { energy_dist_1 } + , energy_dist_2 { energy_dist_2 } + , inv_V0 { inv_V0 } + , random_pool { random_pool } { + if (use_tracking_1) { + species1.set_counter(cntr1 + inject_npart); +#if !defined(MPI_ENABLED) + raise::ErrorIf(species1.pld_i.extent(1) < 1, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 1", + HERE); +#else + raise::ErrorIf(species1.pld_i.extent(1) < 2, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 2", + HERE); +#endif + } + if (use_tracking_2) { + species2.set_counter(cntr2 + inject_npart); +#if !defined(MPI_ENABLED) + raise::ErrorIf(species2.pld_i.extent(1) < 1, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 1", + HERE); +#else + raise::ErrorIf(species2.pld_i.extent(1) < 2, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 2", + HERE); +#endif + } + } + + Inline void operator()(index_t p) const { + coord_t x_Cd { ZERO }; + tuple_t xi_Cd { 0 }; + tuple_t dxi_Cd { static_cast(0) }; + vec_t v1 { ZERO }, v2 { ZERO }; + { // generate a random coordinate + auto rand_gen = random_pool.get_state(); + if constexpr (M::Dim == Dim::_1D or M::Dim == Dim::_2D or + M::Dim == Dim::_3D) { + x_Cd[0] = xi_min(0) + Random(rand_gen) * (xi_max(0) - xi_min(0)); + xi_Cd[0] = static_cast(x_Cd[0]); + dxi_Cd[0] = static_cast(x_Cd[0] - xi_Cd[0]); + } + if constexpr (M::Dim == Dim::_2D or M::Dim == Dim::_3D) { + x_Cd[1] = xi_min(1) + Random(rand_gen) * (xi_max(1) - xi_min(1)); + xi_Cd[1] = static_cast(x_Cd[1]); + xi_Cd[1] = static_cast(x_Cd[1]); + dxi_Cd[1] = static_cast(x_Cd[1] - xi_Cd[1]); + } + if constexpr (M::Dim == Dim::_3D) { + x_Cd[2] = xi_min(2) + Random(rand_gen) * (xi_max(2) - xi_min(2)); + xi_Cd[2] = static_cast(x_Cd[2]); + dxi_Cd[2] = static_cast(x_Cd[2] - xi_Cd[2]); + } + random_pool.free_state(rand_gen); + } + { // generate the velocity + coord_t x_Ph { ZERO }; + metric.template convert(x_Cd, x_Ph); + if constexpr (M::CoordType == Coord::Cart) { + energy_dist_1(x_Ph, v1); + energy_dist_2(x_Ph, v2); + } else if constexpr (S == SimEngine::SRPIC) { + coord_t x_Cd_ { ZERO }; + x_Cd_[0] = x_Cd[0]; + x_Cd_[1] = x_Cd[1]; + x_Cd_[2] = ZERO; // phi = 0 + vec_t v_Ph { ZERO }; + energy_dist_1(x_Ph, v_Ph); + metric.template transform_xyz(x_Cd_, v_Ph, v1); + energy_dist_2(x_Ph, v_Ph); + metric.template transform_xyz(x_Cd_, v_Ph, v2); + } else if constexpr (S == SimEngine::GRPIC) { + vec_t v_Ph { ZERO }; + energy_dist_1(x_Ph, v_Ph); + metric.template transform(x_Cd, v_Ph, v1); + energy_dist_2(x_Ph, v_Ph); + metric.template transform(x_Cd, v_Ph, v2); + } else { + raise::KernelError(HERE, "Unknown simulation engine"); + } + } + real_t weight = ONE; + if constexpr (M::CoordType != Coord::Cart) { + const auto sqrt_det_h = metric.sqrt_det_h(x_Cd); + weight = sqrt_det_h * inv_V0; + } + // clang-format off + if (not use_tracking_1) { + InjectParticle( + p + offset1, + i1s_1, i2s_1, i3s_1, + dx1s_1, dx2s_1, dx3s_1, + ux1s_1, ux2s_1, ux3s_1, + phis_1, weights_1, tags_1, pldis_1, + xi_Cd, dxi_Cd, v1, weight, ZERO); + } else { + InjectParticle( + p + offset1, + i1s_1, i2s_1, i3s_1, + dx1s_1, dx2s_1, dx3s_1, + ux1s_1, ux2s_1, ux3s_1, + phis_1, weights_1, tags_1, pldis_1, + xi_Cd, dxi_Cd, v1, weight, ZERO, + domain_idx, cntr1 + p); + } + if (not use_tracking_2) { + InjectParticle( + p + offset2, + i1s_2, i2s_2, i3s_2, + dx1s_2, dx2s_2, dx3s_2, + ux1s_2, ux2s_2, ux3s_2, + phis_2, weights_2, tags_2, pldis_2, + xi_Cd, dxi_Cd, v2, weight, ZERO); + } else { + InjectParticle( + p + offset2, + i1s_2, i2s_2, i3s_2, + dx1s_2, dx2s_2, dx3s_2, + ux1s_2, ux2s_2, ux3s_2, + phis_2, weights_2, tags_2, pldis_2, + xi_Cd, dxi_Cd, v2, weight, ZERO, + domain_idx, cntr2 + p); + } + // clang-format on + } + }; // struct UniformInjector_kernel + + template + struct GlobalInjector_kernel { + static_assert(M::is_metric, "M must be a metric class"); + static constexpr auto D = M::Dim; + + const bool use_weights; + + array_t in_ux1; + array_t in_ux2; + array_t in_ux3; + array_t in_x1; + array_t in_x2; + array_t in_x3; + array_t in_phi; + array_t in_wei; + + array_t idx { "idx" }; + array_t i1s, i2s, i3s; + array_t dx1s, dx2s, dx3s; + array_t ux1s, ux2s, ux3s; + array_t phis; + array_t weights; + array_t tags; + + const npart_t offset; + + M global_metric; + + real_t x1_min, x1_max, x2_min, x2_max, x3_min, x3_max; + ncells_t i1_offset, i2_offset, i3_offset; + + GlobalInjector_kernel(Particles& species, + const M& global_metric, + const Domain& local_domain, + const std::map>& data, + bool use_weights) + : use_weights { use_weights } + , i1s { species.i1 } + , i2s { species.i2 } + , i3s { species.i3 } + , dx1s { species.dx1 } + , dx2s { species.dx2 } + , dx3s { species.dx3 } + , ux1s { species.ux1 } + , ux2s { species.ux2 } + , ux3s { species.ux3 } + , phis { species.phi } + , weights { species.weight } + , tags { species.tag } + , offset { species.npart() } + , global_metric { global_metric } { + const auto n_inject = data.at("x1").size(); + + x1_min = local_domain.mesh.extent(in::x1).first; + x1_max = local_domain.mesh.extent(in::x1).second; + i1_offset = local_domain.offset_ncells()[0]; + + copy_from_vector("x1", in_x1, data, n_inject); + copy_from_vector("ux1", in_ux1, data, n_inject); + copy_from_vector("ux2", in_ux2, data, n_inject); + copy_from_vector("ux3", in_ux3, data, n_inject); + if (use_weights) { + copy_from_vector("weights", in_wei, data, n_inject); + } + if constexpr (D == Dim::_2D or D == Dim::_3D) { + x2_min = local_domain.mesh.extent(in::x2).first; + x2_max = local_domain.mesh.extent(in::x2).second; + i2_offset = local_domain.offset_ncells()[1]; + copy_from_vector("x2", in_x2, data, n_inject); + } + if constexpr (D == Dim::_2D and M::CoordType != Coord::Cart) { + copy_from_vector("phi", in_phi, data, n_inject); + } + if constexpr (D == Dim::_3D) { + x3_min = local_domain.mesh.extent(in::x3).first; + x3_max = local_domain.mesh.extent(in::x3).second; + i3_offset = local_domain.offset_ncells()[2]; + copy_from_vector("x3", in_x3, data, n_inject); + } + } + + void copy_from_vector(const std::string& name, + array_t& arr, + const std::map>& data, + npart_t n_inject) { + raise::ErrorIf(data.find(name) == data.end(), name + " not found in data", HERE); + raise::ErrorIf(data.at(name).size() != n_inject, "Inconsistent data size", HERE); + arr = array_t { name, n_inject }; + auto arr_h = Kokkos::create_mirror_view(arr); + for (auto i = 0u; i < data.at(name).size(); ++i) { + arr_h(i) = data.at(name)[i]; + } + Kokkos::deep_copy(arr, arr_h); + } + + auto number_injected() const -> npart_t { + auto idx_h = Kokkos::create_mirror_view(idx); + Kokkos::deep_copy(idx_h, idx); + return idx_h(); + } + + Inline void operator()(index_t p) const { + if constexpr (D == Dim::_1D) { + if (in_x1(p) >= x1_min and in_x1(p) < x1_max) { + coord_t x_Cd { ZERO }; + vec_t u_XYZ { ZERO }; + const vec_t u_Ph { in_ux1(p), in_ux2(p), in_ux3(p) }; + + auto index { offset + Kokkos::atomic_fetch_add(&idx(), 1) }; + global_metric.template convert({ in_x1(p) }, x_Cd); + global_metric.template transform_xyz(x_Cd, u_Ph, u_XYZ); + + const auto i1 = static_cast( + static_cast(x_Cd[0]) - i1_offset); + const auto dx1 = static_cast( + x_Cd[0] - static_cast(i1 + i1_offset)); + + i1s(index) = i1; + dx1s(index) = dx1; + ux1s(index) = u_XYZ[0]; + ux2s(index) = u_XYZ[1]; + ux3s(index) = u_XYZ[2]; + tags(index) = ParticleTag::alive; + if (use_weights) { + weights(index) = weights(p); + } else { + weights(index) = ONE; + } + } + } else if constexpr (D == Dim::_2D) { + if ((in_x1(p) >= x1_min and in_x1(p) < x1_max) and + (in_x2(p) >= x2_min and in_x2(p) < x2_max)) { + coord_t x_Cd { ZERO }; + vec_t u_Cd { ZERO }; + vec_t u_Ph { in_ux1(p), in_ux2(p), in_ux3(p) }; + coord_t x_Cd_ { ZERO }; + + auto index { offset + + Kokkos::atomic_fetch_add(&idx(), static_cast(1)) }; + global_metric.template convert({ in_x1(p), in_x2(p) }, + x_Cd); + x_Cd_[0] = x_Cd[0]; + x_Cd_[1] = x_Cd[1]; + if constexpr (S == SimEngine::SRPIC and M::CoordType != Coord::Cart) { + x_Cd_[2] = in_phi(p); + } + if constexpr (S == SimEngine::SRPIC) { + global_metric.template transform_xyz(x_Cd_, u_Ph, u_Cd); + } else if constexpr (S == SimEngine::GRPIC) { + global_metric.template transform(x_Cd, u_Ph, u_Cd); + } else { + raise::KernelError(HERE, "Unknown simulation engine"); + } + const auto i1 = static_cast( + static_cast(x_Cd[0]) - i1_offset); + const auto dx1 = static_cast( + x_Cd[0] - static_cast(i1 + i1_offset)); + const auto i2 = static_cast( + static_cast(x_Cd[1]) - i2_offset); + const auto dx2 = static_cast( + x_Cd[1] - static_cast(i2 + i2_offset)); + + i1s(index) = i1; + dx1s(index) = dx1; + i2s(index) = i2; + dx2s(index) = dx2; + ux1s(index) = u_Cd[0]; + ux2s(index) = u_Cd[1]; + ux3s(index) = u_Cd[2]; + if (M::CoordType != Coord::Cart) { + phis(index) = in_phi(p); + } + tags(index) = ParticleTag::alive; + if (use_weights) { + weights(index) = weights(p); + } else { + weights(index) = ONE; + } + } + } else { + if ((in_x1(p) >= x1_min and in_x1(p) < x1_max) and + (in_x2(p) >= x2_min and in_x2(p) < x2_max) and + (in_x3(p) >= x3_min and in_x3(p) < x3_max)) { + coord_t x_Cd { ZERO }; + vec_t u_Cd { ZERO }; + vec_t u_Ph { in_ux1(p), in_ux2(p), in_ux3(p) }; + + auto index { offset + Kokkos::atomic_fetch_add(&idx(), 1) }; + global_metric.template convert( + { in_x1(p), in_x2(p), in_x3(p) }, + x_Cd); + if constexpr (S == SimEngine::SRPIC) { + global_metric.template transform_xyz(x_Cd, u_Ph, u_Cd); + } else if constexpr (S == SimEngine::GRPIC) { + global_metric.template transform(x_Cd, u_Ph, u_Cd); + } else { + raise::KernelError(HERE, "Unknown simulation engine"); + } + const auto i1 = static_cast( + static_cast(x_Cd[0]) - i1_offset); + const auto dx1 = static_cast( + x_Cd[0] - static_cast(i1 + i1_offset)); + const auto i2 = static_cast( + static_cast(x_Cd[1]) - i2_offset); + const auto dx2 = static_cast( + x_Cd[1] - static_cast(i2 + i2_offset)); + const auto i3 = static_cast( + static_cast(x_Cd[2]) - i3_offset); + const auto dx3 = static_cast( + x_Cd[2] - static_cast(i3 + i3_offset)); + + i1s(index) = i1; + dx1s(index) = dx1; + i2s(index) = i2; + dx2s(index) = dx2; + i3s(index) = i3; + dx3s(index) = dx3; + ux1s(index) = u_Cd[0]; + ux2s(index) = u_Cd[1]; + ux3s(index) = u_Cd[2]; + tags(index) = ParticleTag::alive; + if (use_weights) { + weights(index) = weights(p); + } else { + weights(index) = ONE; + } + } + } + } + }; // struct GlobalInjector_kernel + + template + struct NonUniformInjector_kernel { + static_assert(M::is_metric, "M must be a metric class"); + static_assert(ED1::is_energy_dist, "ED1 must be an energy distribution class"); + static_assert(ED2::is_energy_dist, "ED2 must be an energy distribution class"); + static_assert(SD::is_spatial_dist, "SD must be a spatial distribution class"); + + const real_t ppc0; + + array_t i1s_1, i2s_1, i3s_1; + array_t dx1s_1, dx2s_1, dx3s_1; + array_t ux1s_1, ux2s_1, ux3s_1; + array_t phis_1; + array_t weights_1; + array_t tags_1; + array_t pldis_1; + + array_t i1s_2, i2s_2, i3s_2; + array_t dx1s_2, dx2s_2, dx3s_2; + array_t ux1s_2, ux2s_2, ux3s_2; + array_t phis_2; + array_t weights_2; + array_t tags_2; + array_t pldis_2; + + array_t idx { "idx" }; + + const npart_t offset1, offset2; + const npart_t domain_idx, cntr1, cntr2; + const bool use_tracking_1, use_tracking_2; + const M metric; + const ED1 energy_dist_1; + const ED2 energy_dist_2; + const SD spatial_dist; + const real_t inv_V0; + random_number_pool_t random_pool; + + NonUniformInjector_kernel(real_t ppc0, + Particles& species1, + Particles& species2, + npart_t domain_idx, + const M& metric, + const ED1& energy_dist_1, + const ED2& energy_dist_2, + const SD& spatial_dist, + real_t inv_V0, + random_number_pool_t& random_pool) + : ppc0 { ppc0 } + , i1s_1 { species1.i1 } + , i2s_1 { species1.i2 } + , i3s_1 { species1.i3 } + , dx1s_1 { species1.dx1 } + , dx2s_1 { species1.dx2 } + , dx3s_1 { species1.dx3 } + , ux1s_1 { species1.ux1 } + , ux2s_1 { species1.ux2 } + , ux3s_1 { species1.ux3 } + , phis_1 { species1.phi } + , weights_1 { species1.weight } + , tags_1 { species1.tag } + , pldis_1 { species1.pld_i } + , i1s_2 { species2.i1 } + , i2s_2 { species2.i2 } + , i3s_2 { species2.i3 } + , dx1s_2 { species2.dx1 } + , dx2s_2 { species2.dx2 } + , dx3s_2 { species2.dx3 } + , ux1s_2 { species2.ux1 } + , ux2s_2 { species2.ux2 } + , ux3s_2 { species2.ux3 } + , phis_2 { species2.phi } + , weights_2 { species2.weight } + , tags_2 { species2.tag } + , pldis_2 { species2.pld_i } + , offset1 { species1.npart() } + , offset2 { species2.npart() } + , domain_idx { domain_idx } + , cntr1 { species1.counter() } + , cntr2 { species2.counter() } + , use_tracking_1 { species1.use_tracking() } + , use_tracking_2 { species2.use_tracking() } + , metric { metric } + , energy_dist_1 { energy_dist_1 } + , energy_dist_2 { energy_dist_2 } + , spatial_dist { spatial_dist } + , inv_V0 { inv_V0 } + , random_pool { random_pool } {} + + auto number_injected() const -> npart_t { + auto idx_h = Kokkos::create_mirror_view(idx); + Kokkos::deep_copy(idx_h, idx); + return idx_h(); + } + + Inline auto injected_ppc(const coord_t& x_Ph, + real_t& ppc_dist, + real_t& weight_dist) const -> npart_t { + spatial_dist(x_Ph, ppc_dist, weight_dist); + const auto ppc_real = ppc0 * ppc_dist; + auto ppc = static_cast(ppc_real); + auto rand_gen = random_pool.get_state(); + if (Random(rand_gen) < (ppc_real - static_cast(ppc))) { + ppc += 1; + } + random_pool.free_state(rand_gen); + return ppc; + } + + Inline void inject1(const index_t index, + const tuple_t& xi_Cd, + const tuple_t& dxi_Cd, + const vec_t& v_Cd, + const real_t weight) const { + // clang-format off + if (not use_tracking_1) { + InjectParticle(index + offset1, + i1s_1, i2s_1, i3s_1, + dx1s_1, dx2s_1, dx3s_1, + ux1s_1, ux2s_1, ux3s_1, + phis_1, weights_1, tags_1, pldis_1, + xi_Cd, dxi_Cd, v_Cd, weight, ZERO); + } else { + InjectParticle(index + offset1, + i1s_1, i2s_1, i3s_1, + dx1s_1, dx2s_1, dx3s_1, + ux1s_1, ux2s_1, ux3s_1, + phis_1, weights_1, tags_1, pldis_1, + xi_Cd, dxi_Cd, v_Cd, weight, ZERO, + domain_idx, index + cntr1); + } + // clang-format on + } + + Inline void inject2(const index_t index, + const tuple_t& xi_Cd, + const tuple_t& dxi_Cd, + const vec_t& v_Cd, + const real_t weight) const { + // clang-format off + if (not use_tracking_2) { + InjectParticle(index + offset2, + i1s_2, i2s_2, i3s_2, + dx1s_2, dx2s_2, dx3s_2, + ux1s_2, ux2s_2, ux3s_2, + phis_2, weights_2, tags_2, pldis_2, + xi_Cd, dxi_Cd, v_Cd, weight, ZERO); + } else { + InjectParticle(index + offset2, + i1s_2, i2s_2, i3s_2, + dx1s_2, dx2s_2, dx3s_2, + ux1s_2, ux2s_2, ux3s_2, + phis_2, weights_2, tags_2, pldis_2, + xi_Cd, dxi_Cd, v_Cd, weight, ZERO, + domain_idx, index + cntr2); + } + // clang-format on + } + + Inline void operator()(index_t i1) const { + if constexpr (M::Dim == Dim::_1D) { + const auto i1_ = COORD(i1); + coord_t x_Cd { i1_ + HALF }; + coord_t x_Ph { ZERO }; + metric.template convert(x_Cd, x_Ph); + auto ppc_dist { ZERO }, weight_dist { ONE }; + const auto ppc = injected_ppc(x_Ph, ppc_dist, weight_dist); + if (ppc == 0) { + return; + } + + auto weight = ONE * weight_dist; + if constexpr (M::CoordType != Coord::Cart) { + weight = metric.sqrt_det_h({ i1_ + HALF }) * inv_V0 * weight_dist; + } + for (auto p { 0u }; p < ppc; ++p) { + const auto index = Kokkos::atomic_fetch_add(&idx(), 1); + + auto rand_gen = random_pool.get_state(); + const auto dx1 = Random(rand_gen); + random_pool.free_state(rand_gen); + + vec_t v_XYZ { ZERO }; + { + vec_t v_T { ZERO }; + energy_dist_1(x_Ph, v_T); + metric.template transform_xyz(x_Cd, v_T, v_XYZ); + } + inject1(index, { static_cast(i1_) }, { dx1 }, v_XYZ, weight); + + { + vec_t v_T { ZERO }; + energy_dist_2(x_Ph, v_T); + metric.template transform_xyz(x_Cd, v_T, v_XYZ); + } + inject2(index, { static_cast(i1_) }, { dx1 }, v_XYZ, weight); + } + } else { + raise::KernelError(HERE, "NonUniformInjector_kernel 1D called for 2D/3D"); + } + } + + Inline void operator()(index_t i1, index_t i2) const { + if constexpr (M::Dim == Dim::_2D) { + const auto i1_ = COORD(i1); + const auto i2_ = COORD(i2); + coord_t x_Cd { i1_ + HALF, i2_ + HALF }; + coord_t x_Ph { ZERO }; + coord_t x_Cd_ { ZERO }; + x_Cd_[0] = x_Cd[0]; + x_Cd_[1] = x_Cd[1]; + if constexpr (S == SimEngine::SRPIC and M::CoordType != Coord::Cart) { + x_Cd_[2] = ZERO; + } + metric.template convert(x_Cd, x_Ph); + + auto ppc_dist { ZERO }, weight_dist { ONE }; + const auto ppc = injected_ppc(x_Ph, ppc_dist, weight_dist); + if (ppc == 0) { + return; + } + + auto weight = ONE * weight_dist; + if constexpr (M::CoordType != Coord::Cart) { + weight = metric.sqrt_det_h({ i1_ + HALF, i2_ + HALF }) * inv_V0 * weight_dist; + } + for (auto p { 0u }; p < ppc; ++p) { + const auto index = Kokkos::atomic_fetch_add(&idx(), 1); + + auto rand_gen = random_pool.get_state(); + const auto dx1 = Random(rand_gen); + const auto dx2 = Random(rand_gen); + random_pool.free_state(rand_gen); + + vec_t v_Cd { ZERO }; + { + vec_t v_T { ZERO }; + energy_dist_1(x_Ph, v_T); + if constexpr (S == SimEngine::SRPIC) { + metric.template transform_xyz(x_Cd_, v_T, v_Cd); + } else if constexpr (S == SimEngine::GRPIC) { + metric.template transform(x_Cd_, v_T, v_Cd); + } + } + inject1(index, + { static_cast(i1_), static_cast(i2_) }, + { dx1, dx2 }, + v_Cd, + weight); + + { + vec_t v_T { ZERO }; + energy_dist_2(x_Ph, v_T); + if constexpr (S == SimEngine::SRPIC) { + metric.template transform_xyz(x_Cd_, v_T, v_Cd); + } else if constexpr (S == SimEngine::GRPIC) { + metric.template transform(x_Cd_, v_T, v_Cd); + } + } + inject2(index, + { static_cast(i1_), static_cast(i2_) }, + { dx1, dx2 }, + v_Cd, + weight); + } + } + + else { + raise::KernelError(HERE, "NonUniformInjector_kernel 2D called for 1D/3D"); + } + } + + Inline void operator()(index_t i1, index_t i2, index_t i3) const { + if constexpr (M::Dim == Dim::_3D) { + const auto i1_ = COORD(i1); + const auto i2_ = COORD(i2); + const auto i3_ = COORD(i3); + coord_t x_Cd { i1_ + HALF, i2_ + HALF, i3_ + HALF }; + coord_t x_Ph { ZERO }; + metric.template convert(x_Cd, x_Ph); + auto ppc_dist { ZERO }, weight_dist { ONE }; + const auto ppc = injected_ppc(x_Ph, ppc_dist, weight_dist); + if (ppc == 0) { + return; + } + + auto weight = ONE * weight_dist; + if constexpr (M::CoordType != Coord::Cart) { + weight = metric.sqrt_det_h({ i1_ + HALF, i2_ + HALF, i3_ + HALF }) * + inv_V0 * weight_dist; + } + for (auto p { 0u }; p < ppc; ++p) { + const auto index = Kokkos::atomic_fetch_add(&idx(), 1); + + auto rand_gen = random_pool.get_state(); + const auto dx1 = Random(rand_gen); + const auto dx2 = Random(rand_gen); + const auto dx3 = Random(rand_gen); + random_pool.free_state(rand_gen); + + vec_t v_Cd { ZERO }; + { + vec_t v_T { ZERO }; + energy_dist_1(x_Ph, v_T); + if constexpr (S == SimEngine::SRPIC) { + metric.template transform_xyz(x_Cd, v_T, v_Cd); + } else if constexpr (S == SimEngine::GRPIC) { + metric.template transform(x_Cd, v_T, v_Cd); + } + } + inject1( + index, + { static_cast(i1_), static_cast(i2_), static_cast(i3_) }, + { dx1, dx2, dx3 }, + v_Cd, + weight); + + { + vec_t v_T { ZERO }; + energy_dist_2(x_Ph, v_T); + if constexpr (S == SimEngine::SRPIC) { + metric.template transform_xyz(x_Cd, v_T, v_Cd); + } else if constexpr (S == SimEngine::GRPIC) { + metric.template transform(x_Cd, v_T, v_Cd); + } + } + inject2( + index, + { static_cast(i1_), static_cast(i2_), static_cast(i3_) }, + { dx1, dx2, dx3 }, + v_Cd, + weight); + } + } else { + raise::KernelError(HERE, "NonUniformInjector_kernel 3D called for 1D/2D"); + } + } + }; // struct NonUniformInjector_kernel + +} // namespace kernel + +#endif // KERNELS_INJECTORS_HPP