diff --git a/src/global/traits/archetypes.h b/src/global/traits/archetypes.h index e170170fb..0dfced45c 100644 --- a/src/global/traits/archetypes.h +++ b/src/global/traits/archetypes.h @@ -5,6 +5,8 @@ * - EnrgDistClass<> - checks if a class can be used as an energy distribution * - SpatialDistClass<> - checks if a class can be used as a spatial distribution * - traits::fieldsetter::HasFx1, ::HasFx2, ::HasFx3 - checks for F functions in field setter class + * - traits::fieldsetter::HasFx1WithIndex, ::HasFx2WithIndex, ::HasFx3WithIndex + * - checks for F functions taking an additional particle index * - traits::fieldsetter::HasEx1, ::HasEx2, ::HasEx3 - checks for E functions in field setter class * - traits::fieldsetter::HasBx1, ::HasBx2, ::HasBx3 - checks for B functions in field setter class * - traits::fieldsetter::HasDx1, ::HasDx2, ::HasDx3 - checks for D functions in field setter class @@ -64,6 +66,23 @@ namespace traits::fieldsetter { { t.fx3(x_Ph) } -> std::convertible_to; }; + // force components which additionally depend on the particle index + // (e.g., to access per-particle properties captured by the setter itself) + template + concept HasFx1WithIndex = requires(const T& t, const coord_t& x_Ph, prtlidx_t p) { + { t.fx1(x_Ph, p) } -> std::convertible_to; + }; + + template + concept HasFx2WithIndex = requires(const T& t, const coord_t& x_Ph, prtlidx_t p) { + { t.fx2(x_Ph, p) } -> std::convertible_to; + }; + + template + concept HasFx3WithIndex = requires(const T& t, const coord_t& x_Ph, prtlidx_t p) { + { t.fx3(x_Ph, p) } -> std::convertible_to; + }; + template concept HasEx1 = requires(const T& t, const coord_t& x_Ph) { { t.ex1(x_Ph) } -> std::convertible_to; diff --git a/src/global/traits/policies.h b/src/global/traits/policies.h index 7f6d98985..a317fd687 100644 --- a/src/global/traits/policies.h +++ b/src/global/traits/policies.h @@ -109,10 +109,13 @@ namespace traits::extfields { template concept ExtFieldsPolicyClass = (::traits::fieldsetter::HasFx1 or ::traits::fieldsetter::HasFx2 or - ::traits::fieldsetter::HasFx3 or ::traits::fieldsetter::HasEx1 or - ::traits::fieldsetter::HasEx2 or ::traits::fieldsetter::HasEx3 or - ::traits::fieldsetter::HasBx1 or ::traits::fieldsetter::HasBx2 or - ::traits::fieldsetter::HasBx3) or + ::traits::fieldsetter::HasFx3 or + ::traits::fieldsetter::HasFx1WithIndex or + ::traits::fieldsetter::HasFx2WithIndex or + ::traits::fieldsetter::HasFx3WithIndex or + ::traits::fieldsetter::HasEx1 or ::traits::fieldsetter::HasEx2 or + ::traits::fieldsetter::HasEx3 or ::traits::fieldsetter::HasBx1 or + ::traits::fieldsetter::HasBx2 or ::traits::fieldsetter::HasBx3) or ::traits::extfields::IsNoPolicy; namespace traits::custom_prtl_update { diff --git a/src/kernels/pushers/sr.hpp b/src/kernels/pushers/sr.hpp index 78b5c4916..6d691d3df 100644 --- a/src/kernels/pushers/sr.hpp +++ b/src/kernels/pushers/sr.hpp @@ -67,11 +67,19 @@ namespace kernel::sr { using F = typename P::ExternalFieldsPolicy; static constexpr auto Atm = P::ApplyAtmosphere; - static constexpr auto D = M::Dim; - static constexpr auto HasExtFx1 = ::traits::fieldsetter::HasFx1; - static constexpr auto HasExtFx2 = ::traits::fieldsetter::HasFx2; - static constexpr auto HasExtFx3 = ::traits::fieldsetter::HasFx3; - static constexpr auto HasExtForce = HasExtFx1 or HasExtFx2 or HasExtFx3; + static constexpr auto D = M::Dim; + static constexpr auto HasExtFx1 = ::traits::fieldsetter::HasFx1; + static constexpr auto HasExtFx2 = ::traits::fieldsetter::HasFx2; + static constexpr auto HasExtFx3 = ::traits::fieldsetter::HasFx3; + static constexpr auto HasExtFx1WithIndex = + ::traits::fieldsetter::HasFx1WithIndex; + static constexpr auto HasExtFx2WithIndex = + ::traits::fieldsetter::HasFx2WithIndex; + static constexpr auto HasExtFx3WithIndex = + ::traits::fieldsetter::HasFx3WithIndex; + static constexpr auto HasExtForce = HasExtFx1 or HasExtFx2 or HasExtFx3 or + HasExtFx1WithIndex or + HasExtFx2WithIndex or HasExtFx3WithIndex; static constexpr auto HasExtEx1 = ::traits::fieldsetter::HasEx1; static constexpr auto HasExtEx2 = ::traits::fieldsetter::HasEx2; static constexpr auto HasExtEx3 = ::traits::fieldsetter::HasEx3; @@ -245,7 +253,7 @@ namespace kernel::sr { // compute the external force either user-provided or from the atmosphere model if constexpr (HasExtForce or Atm) { - getExternalForce(xp_Cd, xp_Ph, external_force_Cart); + getExternalForce(xp_Cd, xp_Ph, p, external_force_Cart); } if (ctx.pusher_flags & ParticlePusher::GCA) { @@ -1425,17 +1433,24 @@ namespace kernel::sr { Inline void getExternalForce(const coord_t& xp_Cd, const coord_t& xp_Ph, + [[maybe_unused]] prtlidx_t p, vec_t& external_force_Cart) const requires(Atm or HasExtForce) { real_t f_x1 = ZERO, f_x2 = ZERO, f_x3 = ZERO; - if constexpr (HasExtFx1) { + if constexpr (HasExtFx1WithIndex) { + f_x1 = policies.external_fields_policy.fx1(xp_Ph, p); + } else if constexpr (HasExtFx1) { f_x1 = policies.external_fields_policy.fx1(xp_Ph); } - if constexpr (HasExtFx2) { + if constexpr (HasExtFx2WithIndex) { + f_x2 = policies.external_fields_policy.fx2(xp_Ph, p); + } else if constexpr (HasExtFx2) { f_x2 = policies.external_fields_policy.fx2(xp_Ph); } - if constexpr (HasExtFx3) { + if constexpr (HasExtFx3WithIndex) { + f_x3 = policies.external_fields_policy.fx3(xp_Ph, p); + } else if constexpr (HasExtFx3) { f_x3 = policies.external_fields_policy.fx3(xp_Ph); } if constexpr (Atm) { diff --git a/tests/global/traits_archetypes.cpp b/tests/global/traits_archetypes.cpp index d72545102..70ce3e9e2 100644 --- a/tests/global/traits_archetypes.cpp +++ b/tests/global/traits_archetypes.cpp @@ -102,6 +102,20 @@ struct WithFx1Fx2Fx3 { } }; +struct WithIndexedFx1Fx2Fx3 { + real_t fx1(const coord_t&, prtlidx_t) const { + return ZERO; + } + + real_t fx2(const coord_t&, prtlidx_t) const { + return ZERO; + } + + real_t fx3(const coord_t&, prtlidx_t) const { + return ZERO; + } +}; + struct WithConditionalEx1Bx2 { Kokkos::pair ex1(const coord_t&, const vec_t&, @@ -135,6 +149,20 @@ static_assert(traits::fieldsetter::HasFx1); static_assert(traits::fieldsetter::HasFx2); static_assert(traits::fieldsetter::HasFx3); +// the indexed variants are a separate, opt-in signature: neither form +// satisfies the other's trait +static_assert( + traits::fieldsetter::HasFx1WithIndex); +static_assert( + traits::fieldsetter::HasFx2WithIndex); +static_assert( + traits::fieldsetter::HasFx3WithIndex); +static_assert( + not traits::fieldsetter::HasFx1WithIndex); +static_assert( + not traits::fieldsetter::HasFx1); +static_assert(not traits::fieldsetter::HasFx1WithIndex); + // conditional variants require 3-arg signature returning Kokkos::pair static_assert( traits::fieldsetter::HasConditionalEx1); diff --git a/tests/global/traits_policies.cpp b/tests/global/traits_policies.cpp index 5f5a83f9f..7f212dab3 100644 --- a/tests/global/traits_policies.cpp +++ b/tests/global/traits_policies.cpp @@ -110,9 +110,16 @@ struct WithFx1 { } }; +struct WithIndexedFx1 { + real_t fx1(const coord_t&, prtlidx_t) const { + return ZERO; + } +}; + struct Empty {}; static_assert(ExtFieldsPolicyClass); +static_assert(ExtFieldsPolicyClass); static_assert(not ExtFieldsPolicyClass); // --- CustomParticleUpdatePolicyClass with a real updater --- diff --git a/tests/kernels/ext_force.cpp b/tests/kernels/ext_force.cpp index 69d773d10..0a705f19b 100644 --- a/tests/kernels/ext_force.cpp +++ b/tests/kernels/ext_force.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include using namespace ntt; @@ -51,6 +52,13 @@ void put_value(array_t& arr, T v, prtlidx_t p) { Kokkos::deep_copy(arr, h); } +// per-particle weight given to the two test particles; the indexed force +// setter reads it back from the particle arrays and scales the force by it +auto weight_of(prtlidx_t p) -> real_t { + return (p == 0) ? ONE : TWO; +} + +// force setter depending only on the position struct Force { Force(real_t force) : force { force } {} @@ -70,7 +78,31 @@ struct Force { const real_t force; }; -template +// force setter carrying its own copy of the particle arrays and reading +// per-particle properties through the particle index +struct IndexedForce { + IndexedForce(real_t force, const ParticleArrays& prtls) + : force { force } + , prtls { prtls } {} + + Inline auto fx1(const coord_t&, prtlidx_t p) const -> real_t { + return prtls.weight(p) * force * math::sin(ONE) * math::sin(ONE); + } + + Inline auto fx2(const coord_t&, prtlidx_t p) const -> real_t { + return prtls.weight(p) * force * math::sin(ONE) * math::cos(ONE); + } + + Inline auto fx3(const coord_t&, prtlidx_t p) const -> real_t { + return prtls.weight(p) * force * math::cos(ONE); + } + +private: + const real_t force; + const ParticleArrays prtls; +}; + +template void testPusher(const std::vector& res) { static_assert(M::Dim == 3); raise::ErrorIf(res.size() != M::Dim, "res.size() != M::Dim", HERE); @@ -141,6 +173,7 @@ void testPusher(const std::vector& res) { put_value(ux1, ux1_0, 0); put_value(ux2, ux2_0, 0); put_value(ux3, ux3_0, 0); + put_value(weight, weight_of(0), 0); put_value(tag, ParticleTag::alive, 0); put_value(i1, (int)(x1_0), 1); @@ -152,12 +185,11 @@ void testPusher(const std::vector& res) { put_value(ux1, -ux1_0, 1); put_value(ux2, -ux2_0, 1); put_value(ux3, -ux3_0, 1); + put_value(weight, weight_of(1), 1); put_value(tag, ParticleTag::alive, 1); const real_t eps = std::is_same_v ? 1e-4 : 1e-6; - const auto ext_force = Force { f_mag }; - static plog::RollingFileAppender file_appender( "pusher_log.csv"); plog::init(plog::verbose, &file_appender); @@ -186,8 +218,18 @@ void testPusher(const std::vector& res) { pusher_arrays.ux2 = ux2; pusher_arrays.ux3 = ux3; pusher_arrays.phi = phi; + pusher_arrays.weight = weight; pusher_arrays.tag = tag; + using ext_force_t = std::conditional_t; + const auto ext_force = [&]() -> ext_force_t { + if constexpr (Indexed) { + return { f_mag, pusher_arrays }; + } else { + return { f_mag }; + } + }(); + const auto pusher_policy = ::kernel::sr::PusherPolicy, ::traits::emission::NoPolicy_t, @@ -253,12 +295,16 @@ void testPusher(const std::vector& res) { ux2_(1), ux3_(1)); + // the indexed setter scales the force by the particle weight + const real_t scale_0 = Indexed ? weight_of(0) : ONE; + const real_t scale_1 = Indexed ? weight_of(1) : ONE; + { - const real_t ux1_expect = ux1_0 + (time + dt) * f_mag * std::sin(ONE) * - std::sin(ONE); - const real_t ux2_expect = ux2_0 + (time + dt) * f_mag * std::sin(ONE) * - std::cos(ONE); - const real_t ux3_expect = ux3_0 + (time + dt) * f_mag * std::cos(ONE); + const real_t f = (time + dt) * f_mag * scale_0; + + const real_t ux1_expect = ux1_0 + f * std::sin(ONE) * std::sin(ONE); + const real_t ux2_expect = ux2_0 + f * std::sin(ONE) * std::cos(ONE); + const real_t ux3_expect = ux3_0 + f * std::cos(ONE); check_value(t, ux1_(0), ux1_expect, eps, "Particle #1 ux1"); check_value(t, ux2_(0), ux2_expect, eps, "Particle #1 ux2"); @@ -266,11 +312,11 @@ void testPusher(const std::vector& res) { } { - const real_t ux1_expect = -ux1_0 + (time + dt) * f_mag * std::sin(ONE) * - std::sin(ONE); - const real_t ux2_expect = -ux2_0 + (time + dt) * f_mag * std::sin(ONE) * - std::cos(ONE); - const real_t ux3_expect = -ux3_0 + (time + dt) * f_mag * std::cos(ONE); + const real_t f = (time + dt) * f_mag * scale_1; + + const real_t ux1_expect = -ux1_0 + f * std::sin(ONE) * std::sin(ONE); + const real_t ux2_expect = -ux2_0 + f * std::sin(ONE) * std::cos(ONE); + const real_t ux3_expect = -ux3_0 + f * std::cos(ONE); check_value(t, ux1_(1), ux1_expect, eps, "Particle #2 ux1"); check_value(t, ux2_(1), ux2_expect, eps, "Particle #2 ux2"); @@ -285,7 +331,8 @@ auto main(int argc, char* argv[]) -> int { try { using namespace ntt; - testPusher>({ 10, 10, 10 }); + testPusher, false>({ 10, 10, 10 }); + testPusher, true>({ 10, 10, 10 }); } catch (std::exception& e) { std::cerr << e.what() << '\n';