External force extension

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<real_t>;
   };
 
+  // force components which additionally depend on the particle index
+  // (e.g., to access per-particle properties captured by the setter itself)
+  template <class T, Dimension D>
+  concept HasFx1WithIndex = requires(const T& t, const coord_t<D>& x_Ph, prtlidx_t p) {
+    { t.fx1(x_Ph, p) } -> std::convertible_to<real_t>;
+  };
+
+  template <class T, Dimension D>
+  concept HasFx2WithIndex = requires(const T& t, const coord_t<D>& x_Ph, prtlidx_t p) {
+    { t.fx2(x_Ph, p) } -> std::convertible_to<real_t>;
+  };
+
+  template <class T, Dimension D>
+  concept HasFx3WithIndex = requires(const T& t, const coord_t<D>& x_Ph, prtlidx_t p) {
+    { t.fx3(x_Ph, p) } -> std::convertible_to<real_t>;
+  };
+
   template <class T, Dimension D>
   concept HasEx1 = requires(const T& t, const coord_t<D>& x_Ph) {
     { t.ex1(x_Ph) } -> std::convertible_to<real_t>;

src/global/traits/policies.h

@@ -109,10 +109,13 @@ namespace traits::extfields {
 template <class F, Dimension D>
 concept ExtFieldsPolicyClass =
   (::traits::fieldsetter::HasFx1<F, D> or ::traits::fieldsetter::HasFx2<F, D> or
-   ::traits::fieldsetter::HasFx3<F, D> or ::traits::fieldsetter::HasEx1<F, D> or
-   ::traits::fieldsetter::HasEx2<F, D> or ::traits::fieldsetter::HasEx3<F, D> or
-   ::traits::fieldsetter::HasBx1<F, D> or ::traits::fieldsetter::HasBx2<F, D> or
-   ::traits::fieldsetter::HasBx3<F, D>) or
+   ::traits::fieldsetter::HasFx3<F, D> or
+   ::traits::fieldsetter::HasFx1WithIndex<F, D> or
+   ::traits::fieldsetter::HasFx2WithIndex<F, D> or
+   ::traits::fieldsetter::HasFx3WithIndex<F, D> or
+   ::traits::fieldsetter::HasEx1<F, D> or ::traits::fieldsetter::HasEx2<F, D> or
+   ::traits::fieldsetter::HasEx3<F, D> or ::traits::fieldsetter::HasBx1<F, D> or
+   ::traits::fieldsetter::HasBx2<F, D> or ::traits::fieldsetter::HasBx3<F, D>) or
   ::traits::extfields::IsNoPolicy<F>;
 
 namespace traits::custom_prtl_update {

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<F, D>;
-    static constexpr auto HasExtFx2    = ::traits::fieldsetter::HasFx2<F, D>;
-    static constexpr auto HasExtFx3    = ::traits::fieldsetter::HasFx3<F, D>;
-    static constexpr auto HasExtForce  = HasExtFx1 or HasExtFx2 or HasExtFx3;
+    static constexpr auto D         = M::Dim;
+    static constexpr auto HasExtFx1 = ::traits::fieldsetter::HasFx1<F, D>;
+    static constexpr auto HasExtFx2 = ::traits::fieldsetter::HasFx2<F, D>;
+    static constexpr auto HasExtFx3 = ::traits::fieldsetter::HasFx3<F, D>;
+    static constexpr auto HasExtFx1WithIndex =
+      ::traits::fieldsetter::HasFx1WithIndex<F, D>;
+    static constexpr auto HasExtFx2WithIndex =
+      ::traits::fieldsetter::HasFx2WithIndex<F, D>;
+    static constexpr auto HasExtFx3WithIndex =
+      ::traits::fieldsetter::HasFx3WithIndex<F, D>;
+    static constexpr auto HasExtForce = HasExtFx1 or HasExtFx2 or HasExtFx3 or
+                                        HasExtFx1WithIndex or
+                                        HasExtFx2WithIndex or HasExtFx3WithIndex;
     static constexpr auto HasExtEx1    = ::traits::fieldsetter::HasEx1<F, D>;
     static constexpr auto HasExtEx2    = ::traits::fieldsetter::HasEx2<F, D>;
     static constexpr auto HasExtEx3    = ::traits::fieldsetter::HasEx3<F, D>;
@@ -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<M::PrtlDim>& xp_Cd,
                                  const coord_t<M::PrtlDim>& xp_Ph,
+                                 [[maybe_unused]] prtlidx_t p,
                                  vec_t<Dim::_3D>& 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) {

tests/global/traits_archetypes.cpp

@@ -102,6 +102,20 @@ struct WithFx1Fx2Fx3 {
   }
 };
 
+struct WithIndexedFx1Fx2Fx3 {
+  real_t fx1(const coord_t<Dimension::_2D>&, prtlidx_t) const {
+    return ZERO;
+  }
+
+  real_t fx2(const coord_t<Dimension::_2D>&, prtlidx_t) const {
+    return ZERO;
+  }
+
+  real_t fx3(const coord_t<Dimension::_2D>&, prtlidx_t) const {
+    return ZERO;
+  }
+};
+
 struct WithConditionalEx1Bx2 {
   Kokkos::pair<bool, real_t> ex1(const coord_t<Dimension::_2D>&,
                                  const vec_t<Dim::_3D>&,
@@ -135,6 +149,20 @@ static_assert(traits::fieldsetter::HasFx1<WithFx1Fx2Fx3, Dimension::_2D>);
 static_assert(traits::fieldsetter::HasFx2<WithFx1Fx2Fx3, Dimension::_2D>);
 static_assert(traits::fieldsetter::HasFx3<WithFx1Fx2Fx3, Dimension::_2D>);
 
+// the indexed variants are a separate, opt-in signature: neither form
+// satisfies the other's trait
+static_assert(
+  traits::fieldsetter::HasFx1WithIndex<WithIndexedFx1Fx2Fx3, Dimension::_2D>);
+static_assert(
+  traits::fieldsetter::HasFx2WithIndex<WithIndexedFx1Fx2Fx3, Dimension::_2D>);
+static_assert(
+  traits::fieldsetter::HasFx3WithIndex<WithIndexedFx1Fx2Fx3, Dimension::_2D>);
+static_assert(
+  not traits::fieldsetter::HasFx1WithIndex<WithFx1Fx2Fx3, Dimension::_2D>);
+static_assert(
+  not traits::fieldsetter::HasFx1<WithIndexedFx1Fx2Fx3, Dimension::_2D>);
+static_assert(not traits::fieldsetter::HasFx1WithIndex<Empty, Dimension::_2D>);
+
 // conditional variants require 3-arg signature returning Kokkos::pair<bool, real_t>
 static_assert(
   traits::fieldsetter::HasConditionalEx1<WithConditionalEx1Bx2, Dimension::_2D>);

tests/global/traits_policies.cpp

@@ -110,9 +110,16 @@ struct WithFx1 {
   }
 };
 
+struct WithIndexedFx1 {
+  real_t fx1(const coord_t<Dimension::_2D>&, prtlidx_t) const {
+    return ZERO;
+  }
+};
+
 struct Empty {};
 
 static_assert(ExtFieldsPolicyClass<WithFx1, Dimension::_2D>);
+static_assert(ExtFieldsPolicyClass<WithIndexedFx1, Dimension::_2D>);
 static_assert(not ExtFieldsPolicyClass<Empty, Dimension::_2D>);
 
 // --- CustomParticleUpdatePolicyClass with a real updater ---

tests/kernels/ext_force.cpp

@@ -24,6 +24,7 @@
 #include <cmath>
 #include <iostream>
 #include <string>
+#include <type_traits>
 #include <vector>
 
 using namespace ntt;
@@ -51,6 +52,13 @@ void put_value(array_t<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 <SimEngine::type S, typename M>
+// 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<Dim::_3D>&, prtlidx_t p) const -> real_t {
+    return prtls.weight(p) * force * math::sin(ONE) * math::sin(ONE);
+  }
+
+  Inline auto fx2(const coord_t<Dim::_3D>&, prtlidx_t p) const -> real_t {
+    return prtls.weight(p) * force * math::sin(ONE) * math::cos(ONE);
+  }
+
+  Inline auto fx3(const coord_t<Dim::_3D>&, prtlidx_t p) const -> real_t {
+    return prtls.weight(p) * force * math::cos(ONE);
+  }
+
+private:
+  const real_t         force;
+  const ParticleArrays prtls;
+};
+
+template <SimEngine::type S, typename M, bool Indexed>
 void testPusher(const std::vector<ncells_t>& 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<ncells_t>& res) {
   put_value<real_t>(ux1, ux1_0, 0);
   put_value<real_t>(ux2, ux2_0, 0);
   put_value<real_t>(ux3, ux3_0, 0);
+  put_value<real_t>(weight, weight_of(0), 0);
   put_value<short>(tag, ParticleTag::alive, 0);
 
   put_value<int>(i1, (int)(x1_0), 1);
@@ -152,12 +185,11 @@ void testPusher(const std::vector<ncells_t>& res) {
   put_value<real_t>(ux1, -ux1_0, 1);
   put_value<real_t>(ux2, -ux2_0, 1);
   put_value<real_t>(ux3, -ux3_0, 1);
+  put_value<real_t>(weight, weight_of(1), 1);
   put_value<short>(tag, ParticleTag::alive, 1);
 
   const real_t eps = std::is_same_v<real_t, float> ? 1e-4 : 1e-6;
 
-  const auto ext_force = Force { f_mag };
-
   static plog::RollingFileAppender<plog::NttInfoFormatter> file_appender(
     "pusher_log.csv");
   plog::init(plog::verbose, &file_appender);
@@ -186,8 +218,18 @@ void testPusher(const std::vector<ncells_t>& 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<Indexed, IndexedForce, Force>;
+  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<Minkowski<Dim::_3D>,
                                ::traits::emission::NoPolicy_t,
@@ -253,24 +295,28 @@ void testPusher(const std::vector<ncells_t>& 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");
       check_value(t, ux3_(0), ux3_expect, eps, "Particle #1 ux3");
     }
 
     {
-      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<SimEngine::SRPIC, Minkowski<Dim::_3D>>({ 10, 10, 10 });
+    testPusher<SimEngine::SRPIC, Minkowski<Dim::_3D>, false>({ 10, 10, 10 });
+    testPusher<SimEngine::SRPIC, Minkowski<Dim::_3D>, true>({ 10, 10, 10 });
 
   } catch (std::exception& e) {
     std::cerr << e.what() << '\n';