From 5b7d06854bfca3667c4f96495fb325ce254daa40 Mon Sep 17 00:00:00 2001 From: hayk Date: Tue, 12 May 2026 01:47:27 -0400 Subject: [PATCH 01/25] twobody interactions as input --- input.example.toml | 23 +++ src/archetypes/qed.h | 160 --------------- src/archetypes/qed/compton.h | 285 +++++++++++++++++++++++++++ src/engines/srpic/twobody.h | 8 + src/global/enums.h | 23 +++ src/kernels/twobody_interactions.hpp | 9 + tests/archetypes/CMakeLists.txt | 3 +- tests/archetypes/qed_compton.cpp | 278 ++++++++++++++++++++++++++ 8 files changed, 628 insertions(+), 161 deletions(-) delete mode 100644 src/archetypes/qed.h create mode 100644 src/archetypes/qed/compton.h create mode 100644 src/engines/srpic/twobody.h create mode 100644 tests/archetypes/qed_compton.cpp diff --git a/input.example.toml b/input.example.toml index f890143a6..aa68d28ff 100644 --- a/input.example.toml +++ b/input.example.toml @@ -272,6 +272,29 @@ # @from: `.gamma_qed` # @value: `(1 / gamma_qed)^2` +[two_body] + # Nominal Thomson optical depth: `tau = n0 * sigma_T * 1` over a distance of 1 in physical units (n0 = nominal density) + # @type: float + # @default: 1.0 + thomson_optical_depth = "" + + [[two_body.interaction]] + # Type of the two-body interaction + # @required + # @type: string + # @enum: "Compton" + type = "" + # First group of species indices participating in the interaction + # @required + # @type: array + # @note: array indexing starting at 1 + group1 = "" + # Second group of species indices participating in the interaction + # @type: array + # @note: array indexing starting at 1 + # @note: For interactions between particles of the same group, leave `group2` empty + group2 = "" + [algorithms] # Number of current smoothing passes # @type: ushort [>= 0] diff --git a/src/archetypes/qed.h b/src/archetypes/qed.h deleted file mode 100644 index 1844a88fc..000000000 --- a/src/archetypes/qed.h +++ /dev/null @@ -1,160 +0,0 @@ -#ifndef ARCHETYPES_QED_H -#define ARCHETYPES_QED_H - -#include "enums.h" -#include "global.h" - -#include "arch/kokkos_aliases.h" -#include "utils/comparators.h" -#include "utils/error.h" -#include "utils/param_container.h" - -#include "framework/containers/particles.h" - -namespace arch { - using namespace ntt; - - constexpr spidx_t MAXSP = 16u; - - template - struct TwoBodyComptonScattering { - ParticleArrays species1[MAXSP]; - ParticleArrays species2[MAXSP]; - - const real_t nominal_probability_density; - const real_t Thomson_limit; - random_number_pool_t random_pool; - - TwoBodyComptonScattering(const prm::Parameters& params, - random_number_pool_t& random_pool) - : nominal_probability_density { params.template get( - "two_body.compton_scattering.nominal_probability_density") } - , Thomson_limit { params.template get( - "two_body.compton_scattering.Thomson_limit") } - , random_pool { random_pool } { - if (nominal_probability_density <= ZERO) { - raise::Error("nominal_probability must be in the range (0, 1]", HERE); - } - if (Thomson_limit <= ZERO or Thomson_limit > 2e-3) { - raise::Error("Thomson_limit must be in the range (0, 2e-3]", HERE); - } - } - - Inline void operator()(spidx_t sp1, - npart_t p1, - spidx_t sp2, - npart_t p2, - real_t tile_volume) const { - // values with "_" are in the lepton rest-frame - const auto lepton_ux1 = species1[sp1 - 1].ux1(p1); - const auto lepton_ux2 = species1[sp1 - 1].ux2(p1); - const auto lepton_ux3 = species1[sp1 - 1].ux3(p1); - const auto lepton_gamma = U2GAMMA(lepton_ux1, lepton_ux2, lepton_ux3); - const auto lepton_weight = species1[sp1 - 1].weight(p1); - - const auto photon_px1 = species2[sp2 - 1].ux1(p2); - const auto photon_px2 = species2[sp2 - 1].ux2(p2); - const auto photon_px3 = species2[sp2 - 1].ux3(p2); - const auto photon_energy = NORM(photon_px1, photon_px2, photon_px3); - const auto photon_weight = species2[sp2 - 1].weight(p2); - - // boost photon momentum to lepton rest frame - real_t photon_px1_ { ZERO }, photon_px2_ { ZERO }, photon_px3_ { ZERO }, - photon_energy_ { ZERO }; - { - const real_t p_dot_k = DOT(lepton_ux1, - lepton_ux2, - lepton_ux3, - photon_px1, - photon_px2, - photon_px3); - - photon_energy_ = lepton_gamma * photon_energy - p_dot_k; - - photon_px1_ = photon_px1 + - (p_dot_k / (ONE + lepton_gamma) - photon_energy) * lepton_ux1; - photon_px2_ = photon_px2 + - (p_dot_k / (ONE + lepton_gamma) - photon_energy) * lepton_ux2; - photon_px3_ = photon_px3 + - (p_dot_k / (ONE + lepton_gamma) - photon_energy) * lepton_ux3; - } - const bool KN_regime = photon_energy_ > Thomson_limit; - real_t f_KN { ONE }; - if (KN_regime) { - if (photon_energy_ < static_cast(2e-3)) { - // correctly handle the eph_RF << 1 limit using 2nd order expansion of f_KN - f_KN = ONE - TWO * photon_energy_ + - static_cast(5.2) * SQR(photon_energy_); - } else { - f_KN = static_cast(0.375) * - ((ONE - TWO / photon_energy_ - TWO / SQR(photon_energy_)) * - math::log(ONE + TWO * photon_energy_) + - HALF + FOUR / photon_energy_ - - HALF / SQR(ONE + TWO * photon_energy_)) / - photon_energy_; - } - } - const auto scattering_probability = nominal_probability_density * f_KN * - photon_energy_ * lepton_weight * - photon_weight / - (photon_energy * lepton_gamma * - tile_volume); - auto gen = random_pool.get_state(); - const auto rnd = Random(gen); - random_pool.free_state(gen); - - if (rnd < scattering_probability) { - // Define an orthonormal basis: {a, b, c} in the lepton frame - const auto ax1_ { photon_px1_ / photon_energy_ }; - const auto ax2_ { photon_px2_ / photon_energy_ }; - const auto ax3_ { photon_px3_ / photon_energy_ }; - real_t bx1_ { ONE }, bx2_ { ZERO }, bx3_ { ZERO }; - if (not cmp::AlmostZero(ax1_)) { - bx1_ = -ax2_ / ax1_; - bx2_ = ONE / math::sqrt(ONE + SQR(bx1_)); - bx1_ /= math::sqrt(ONE + SQR(bx1_)); - } - const auto cx1_ { CROSS_x1(ax1_, ax2_, ax3_, bx1_, bx2_, bx3_) }; - const auto cx2_ { CROSS_x2(ax1_, ax2_, ax3_, bx1_, bx2_, bx3_) }; - const auto cx3_ { CROSS_x3(ax1_, ax2_, ax3_, bx1_, bx2_, bx3_) }; - - auto gen_ = random_pool.get_state(); - const auto rnd_ = Random(gen_); - random_pool.free_state(gen_); - - real_t costheta_ { ZERO }; - if (not KN_regime) { - costheta_ = math::pow( - FOUR * rnd_ - TWO + - math::sqrt(FIVE + static_cast(16) * rnd_ * (rnd_ - ONE)), - THIRD); - costheta_ -= ONE / costheta_; - } else { - // iter = 0 - // converged = .false. - // c0 = 1.0d0 + 2.0d0 * eph_RF - // c1 = eph_RF / c0 - // c2 = eph_RF**2 - 2.0d0 * eph_RF - 2.0d0 - // c3 = eph_RF - 1.0d0 - 0.5d0 * c1**2 - // c4 = 1.0d0 / (4.0d0 * eph_RF + 2.0d0 * eph_RF * (1.0d0 + eph_RF) * c1**2 + c2 * log(c0)) - // u = 2.0d0 * rnd - 1.0d0 - // do while (iter .lt. max_iter) - // iter = iter + 1 - // du = du_KN_Newt(eph_RF, u, rnd, c0, c1, c2, c3, c4) - // u = u + du - // if (u .gt. 1.0d0) u = 1.0d0 - // if (u .lt. -1.0d0) u = -1.0d0 - // if (abs(du) .lt. thresh) then - // converged = .true. - // exit - // end if - // end do - } - // costheta_RF = u - } // not interacting - } - }; - -} // namespace arch - -#endif // ARCHETYPES_QED_H diff --git a/src/archetypes/qed/compton.h b/src/archetypes/qed/compton.h new file mode 100644 index 000000000..b54c268bc --- /dev/null +++ b/src/archetypes/qed/compton.h @@ -0,0 +1,285 @@ +/** + * @file archetypes/qed/compton.h + * @brief Two-body collision policy of Compton scattering between leptons and photons + * @implements + * - arch::qed::ComptonScattering<> + * @namespaces: + * - arch::qed:: + */ +#ifndef ARCHETYPES_QED_COMPTON_H +#define ARCHETYPES_QED_COMPTON_H + +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "utils/comparators.h" +#include "utils/error.h" +#include "utils/numeric.h" +#include "utils/param_container.h" + +#include "framework/containers/particles.h" + +#include + +namespace arch::qed { + using namespace ntt; + + template + struct ComptonScattering { + static constexpr spidx_t MAXSP = 16u; + static constexpr int MAX_ITER = 10; + + ParticleArrays species[MAXSP]; + static constexpr real_t low_energy_limit = static_cast(2e-3); + + const real_t nominal_probability_density; + const real_t Thomson_limit; + random_number_pool_t random_pool; + + ComptonScattering(const prm::Parameters& params, + random_number_pool_t& random_pool) + : nominal_probability_density { params.template get( + "qed.compton_scattering.nominal_probability_density") } + , Thomson_limit { params.template get( + "qed.compton_scattering.Thomson_limit") } + , random_pool { random_pool } { + if (nominal_probability_density <= ZERO) { + raise::Error("nominal_probability must be in the range (0, 1]", HERE); + } + if (Thomson_limit <= ZERO or Thomson_limit > low_energy_limit) { + raise::Error( + "Thomson_limit must be in the range (0, small_energy_limit]", + HERE); + } + } + + /* + * Lorentz boost a 4-momentum p of the photon to the frame moving with u + * @param u: 4-velocity of the boost frame + * @param gamma: Lorentz factor of the boost frame + * @param p: 4-momentum of the photon in the lab frame + * @param e: energy of the photon in the lab frame + * @return: 4-momentum of the photon in the boost frame + * @return: energy of the photon in the boost frame + */ + Inline void LorentzBoost(const vec_t& u, + real_t gamma, + const vec_t& p, + real_t e, + vec_t& p_, + real_t& e_) const { + const auto u_dot_p = DOT(u[0], u[1], u[2], p[0], p[1], p[2]); + + e_ = gamma * e - u_dot_p; + p_[0] = p[0] + (u_dot_p / (ONE + gamma) - e) * u[0]; + p_[1] = p[1] + (u_dot_p / (ONE + gamma) - e) * u[1]; + p_[2] = p[2] + (u_dot_p / (ONE + gamma) - e) * u[2]; + } + + /* + * Calculate the Klein-Nishina cross section for a photon with energy e_ in the lepton rest frame + * @param e_: photon energy in the lepton rest frame + * @return: pair of (is_KN_regime, f_KN) where + * - is_KN_regime: whether the photon energy is in the Klein-Nishina regime (e_ > Thomson_limit) + * - f_KN: the Klein-Nishina cross section normalized to the Thomson cross section + * @note for e_ > low_energy_limit, full Klein-Nishina formula + * @note for Thomson_limit < e_ <= low_energy_limit, 2nd order expansion of the Klein-Nishina formula + * @note for e_ <= Thomson_limit, return 1 (Thomson limit) + */ + Inline auto KNCrossSection(real_t e_) const -> Kokkos::pair { + if (e_ > Thomson_limit) { + if (e_ < low_energy_limit) { + // correctly handle the e_ << 1 limit using 2nd order expansion of f_KN + return { true, ONE - TWO * e_ + static_cast(5.2) * SQR(e_) }; + } else { + return { true, + static_cast(0.375) * + ((ONE - TWO / e_ - TWO / SQR(e_)) * math::log(ONE + TWO * e_) + + HALF + FOUR / e_ - HALF / SQR(ONE + TWO * e_)) / + e_ }; + } + } else { + return { false, ONE }; + } + } + + Inline auto RandomCosTheta_Th() const -> real_t { + auto gen_ = random_pool.get_state(); + const auto rnd_ = Random(gen_); + random_pool.free_state(gen_); + const auto u = math::pow( + FOUR * rnd_ - TWO + + math::sqrt(FIVE + static_cast(16) * rnd_ * (rnd_ - ONE)), + THIRD); + return u - ONE / u; + } + + Inline auto RandomCosTheta_KN(double e_) const -> real_t { + auto gen_ = random_pool.get_state(); + const auto rnd_ = Random(gen_); + random_pool.free_state(gen_); + + auto u = 2.0 * rnd_ - 1.0; + bool converged = false; + for (int iter = 0; iter < MAX_ITER; ++iter) { + const auto CDF = (-((2.0 + e_ * (4.0 + e_ - 4.0 * (-1.0 + u) * u * e_ + + 2.0 * CUBE(-1.0 + u) * SQR(e_))) / + SQR(1.0 + e_ - u * e_)) + + (2.0 + e_ * (4.0 - e_ * (7.0 + 16.0 * e_))) / + SQR(1.0 + 2.0 * e_) + + 2.0 * (-2.0 + (-2.0 + e_) * e_) * + math::log((1.0 + e_ - u * e_) / (1.0 + 2.0 * e_))) / + ((-4.0 * e_ * (2.0 + e_ * (1.0 + e_) * (8.0 + e_))) / + SQR(1.0 + 2.0 * e_) + + (4.0 - 2.0 * (-2.0 + e_) * e_) * + math::log(1.0 + 2.0 * e_)); + const auto dCDF_du = -((CUBE(e_) * SQR(1.0 + 2.0 * e_) * + (1.0 + SQR(u) - (-1.0 + u) * (1.0 + SQR(u)) * e_ + + SQR(-1.0 + u) * SQR(e_))) / + (CUBE(-1.0 + (-1.0 + u) * e_) * + (2.0 * e_ * (2.0 + e_ * (1.0 + e_) * (8.0 + e_)) + + SQR(1.0 + 2.0 * e_) * (-2.0 + (-2.0 + e_) * e_) * + math::log(1.0 + 2.0 * e_)))); + + const auto du = (rnd_ - CDF) / dCDF_du; + + u += du; + if (u > 1.0) { + u = 1.0; + } else if (u < -1.0) { + u = -1.0; + } + if (math::abs(du) < 1e-3) { + converged = true; + break; + } + } // iterative loop for u + return static_cast(u); + } + + /* + * Scatter a photon with initial momentum p_ and energy e_ in the lepton + * rest frame to a new momentum pnew_ and energy enew_ + * @param KN_regime: whether the photon energy is in the Klein-Nishina regime + * @param p_: initial photon momentum in the lepton rest frame + * @param e_: initial photon energy in the lepton rest frame + * @return pnew_: output photon momentum after scattering in the lepton rest frame + * @return enew_: output photon energy after scattering in the lepton rest frame + * @note the scattering angle is sampled from the Klein-Nishina differential + * cross section if KN_regime is true, otherwise it is sampled from the Thomson limit + */ + Inline void ScatterPhoton(bool KN_regime, + const vec_t& p_, + real_t e_, + vec_t& pnew_, + real_t& enew_) const { + auto rand_costheta_ { ZERO }; + if (not KN_regime) { + rand_costheta_ = RandomCosTheta_Th(); + } else { + rand_costheta_ = RandomCosTheta_KN(e_); + } + const auto rand_sintheta_ = math::sqrt(ONE - SQR(rand_costheta_)); + + auto gen_ = random_pool.get_state(); + const auto rand_phi_ = static_cast(constant::TWO_PI) * + Random(gen_); + random_pool.free_state(gen_); + const auto rand_cosphi_ = math::cos(rand_phi_); + const auto rand_sinphi_ = math::sin(rand_phi_); + + // Define an orthonormal basis: {a_, b_, c_} in the lepton frame + const vec_t a_ { p_[0] / e_, p_[1] / e_, p_[2] / e_ }; + vec_t b_ { ONE, ZERO, ZERO }; + if (not cmp::AlmostZero(a_[0])) { + b_[0] = -a_[1] / a_[0]; + b_[1] = ONE / math::sqrt(ONE + SQR(b_[0])); + b_[0] /= math::sqrt(ONE + SQR(b_[0])); + } + const vec_t c_ { + CROSS_x1(a_[0], a_[1], a_[2], b_[0], b_[1], b_[2]), + CROSS_x2(a_[0], a_[1], a_[2], b_[0], b_[1], b_[2]), + CROSS_x3(a_[0], a_[1], a_[2], b_[0], b_[1], b_[2]) + }; + + enew_ = e_ / (ONE + e_ * (ONE - rand_costheta_)); + + pnew_[0] = enew_ * (rand_costheta_ * a_[0] + + rand_sintheta_ * rand_cosphi_ * b_[0] + + rand_sintheta_ * rand_sinphi_ * c_[0]); + pnew_[1] = enew_ * (rand_costheta_ * a_[1] + + rand_sintheta_ * rand_cosphi_ * b_[1] + + rand_sintheta_ * rand_sinphi_ * c_[1]); + pnew_[2] = enew_ * (rand_costheta_ * a_[2] + + rand_sintheta_ * rand_cosphi_ * b_[2] + + rand_sintheta_ * rand_sinphi_ * c_[2]); + } + + Inline void operator()(spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2, + real_t tile_volume) const { + // @TODO coord/vec conversion + // values with "_" are in the lepton rest-frame + const vec_t lepton_u { species[sp1 - 1].ux1(p1), + species[sp1 - 1].ux2(p1), + species[sp1 - 1].ux3(p1) }; + const auto lepton_gamma = U2GAMMA(lepton_u[0], lepton_u[1], lepton_u[2]); + const auto lepton_weight = species[sp1 - 1].weight(p1); + + const vec_t photon_p { species[sp2 - 1].ux1(p2), + species[sp2 - 1].ux2(p2), + species[sp2 - 1].ux3(p2) }; + const auto photon_energy = NORM(photon_p[0], photon_p[1], photon_p[2]); + const auto photon_weight = species[sp2 - 1].weight(p2); + + // boost photon momentum to lepton rest frame + vec_t photon_p_ { ZERO, ZERO, ZERO }; + real_t photon_energy_ { ZERO }; + + LorentzBoost(lepton_u, lepton_gamma, photon_p, photon_energy, photon_p_, photon_energy_); + + const auto [KN_regime, f_KN] = KNCrossSection(photon_energy_); + const auto scattering_probability = nominal_probability_density * f_KN * + photon_energy_ * lepton_weight * + photon_weight / + (photon_energy * lepton_gamma * + tile_volume); + auto gen = random_pool.get_state(); + const auto rnd = Random(gen); + random_pool.free_state(gen); + + if (rnd < scattering_probability) { + vec_t photon_pnew_ { ZERO, ZERO, ZERO }, + photon_pnew { ZERO, ZERO, ZERO }; + real_t photon_energy_new_ { ZERO }, photon_energy_new { ZERO }; + + ScatterPhoton(KN_regime, + photon_p_, + photon_energy_, + photon_pnew_, + photon_energy_new_); + LorentzBoost({ -lepton_u[0], -lepton_u[1], -lepton_u[2] }, + lepton_gamma, + photon_pnew_, + photon_energy_new_, + photon_pnew, + photon_energy_new); + species[sp1 - 1].ux1(p1) += (photon_p[0] - photon_pnew[0]) * + photon_weight / lepton_weight; + species[sp1 - 1].ux2(p1) += (photon_p[1] - photon_pnew[1]) * + photon_weight / lepton_weight; + species[sp1 - 1].ux3(p1) += (photon_p[2] - photon_pnew[2]) * + photon_weight / lepton_weight; + + species[sp2 - 1].ux1(p2) = photon_pnew[0]; + species[sp2 - 1].ux2(p2) = photon_pnew[1]; + species[sp2 - 1].ux3(p2) = photon_pnew[2]; + } // not interacting + } + }; + +} // namespace arch::qed + +#endif // ARCHETYPES_QED_COMPTON_H diff --git a/src/engines/srpic/twobody.h b/src/engines/srpic/twobody.h new file mode 100644 index 000000000..a6e034c24 --- /dev/null +++ b/src/engines/srpic/twobody.h @@ -0,0 +1,8 @@ +#ifndef ENGINES_SRPIC_TWOBODY_H +#define ENGINES_SRPIC_TWOBODY_H + +namespace ntt { + namespace srpic {} +} // namespace ntt + +#endif // ENGINES_SRPIC_TWOBODY_H \ No newline at end of file diff --git a/src/global/enums.h b/src/global/enums.h index 68eefb75b..41fdd551b 100644 --- a/src/global/enums.h +++ b/src/global/enums.h @@ -401,6 +401,29 @@ namespace ntt { using EmissionTypeFlag = uint8_t; + namespace TwoBodyInteraction { + enum TwoBodyInteractionFlag_ : uint8_t { + NONE = 0, + COMPTON = 1, + CUSTOM = 2, + }; + + inline auto to_string(uint8_t flags) -> std::string { + switch (flags) { + case NONE: + return "none"; + case COMPTON: + return "compton"; + case CUSTOM: + return "custom"; + default: + return "unknown"; + } + } + } // namespace TwoBodyInteraction + + using TwoBodyInteractionFlag = uint8_t; + } // namespace ntt #endif // GLOBAL_ENUMS_H diff --git a/src/kernels/twobody_interactions.hpp b/src/kernels/twobody_interactions.hpp index 5af886959..44385e885 100644 --- a/src/kernels/twobody_interactions.hpp +++ b/src/kernels/twobody_interactions.hpp @@ -1,3 +1,12 @@ +/** + * @file kernels/twobody_interactions.hpp + * @brief Generic two-body interaction kernel that can be used to implement various + * types of collisions between species, e.g. Compton scattering, Breit-Wheeler pair production, etc. + * @implements + * - kernel::mink::TwoBodyInteraction<> + * @namespaces: + * - arch::mink:: + */ #ifndef KERNELS_TWOBODY_INTERACTIONS_HPP #define KERNELS_TWOBODY_INTERACTIONS_HPP diff --git a/tests/archetypes/CMakeLists.txt b/tests/archetypes/CMakeLists.txt index 4a5b501e1..4c1aad515 100644 --- a/tests/archetypes/CMakeLists.txt +++ b/tests/archetypes/CMakeLists.txt @@ -16,7 +16,7 @@ function(gen_test title) set(src ${title}.cpp) add_executable(${exec} ${src}) - set(libs ntt_archetypes ntt_global ntt_metrics) + set(libs ntt_archetypes ntt_framework ntt_global ntt_metrics) add_dependencies(${exec} ${libs}) target_link_libraries(${exec} PRIVATE ${libs}) @@ -28,3 +28,4 @@ gen_test(spatial_dist) gen_test(field_setter) gen_test(powerlaw) gen_test(pgen) +gen_test(qed_compton) diff --git a/tests/archetypes/qed_compton.cpp b/tests/archetypes/qed_compton.cpp new file mode 100644 index 000000000..265439a1b --- /dev/null +++ b/tests/archetypes/qed_compton.cpp @@ -0,0 +1,278 @@ +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" + +#include "archetypes/qed/compton.h" +#include "framework/containers/particles.h" +#include "kernels/twobody_interactions.hpp" + +#include + +#include +#include +#include + +using namespace ntt; + +void fill_random(array_t& i1, + array_t& i2, + array_t& ux1, + array_t& ux2, + array_t& ux3, + array_t& weight, + array_t& tag, + npart_t npart, + ncells_t nx1, + ncells_t nx2, + random_number_pool_t& rpool) { + Kokkos::parallel_for( + "FillRandom", + npart, + KOKKOS_LAMBDA(const npart_t p) { + auto gen = rpool.get_state(); + i1(p) = static_cast(gen.urand() % static_cast(nx1)); + i2(p) = static_cast(gen.urand() % static_cast(nx2)); + ux1(p) = Random(gen) * TWO - ONE; + ux2(p) = Random(gen) * TWO - ONE; + ux3(p) = Random(gen) * TWO - ONE; + weight(p) = ONE; + tag(p) = ParticleTag::alive; + rpool.free_state(gen); + }); + Kokkos::fence(); +} + +auto get_total_energy(bool is_massive, + array_t& ux1, + array_t& ux2, + array_t& ux3, + npart_t npart) -> real_t { + real_t total_energy = ZERO; + Kokkos::parallel_reduce( + "TotalEnergy", + npart, + Lambda(const npart_t p, real_t& local_sum) { + if (is_massive) { + local_sum += U2GAMMA(ux1(p), ux2(p), ux3(p)); + + } else { + local_sum += NORM(ux1(p), ux2(p), ux3(p)); + } + }, + total_energy); + return total_energy; +} + +auto get_total_momentum_in(in dir, + array_t& ux1, + array_t& ux2, + array_t& ux3, + npart_t npart) -> real_t { + real_t total_momentum_in = ZERO; + Kokkos::parallel_reduce( + "TotalMomentumIn", + npart, + Lambda(const npart_t p, real_t& local_sum) { + if (dir == in::x1) { + local_sum += ux1(p); + } else if (dir == in::x2) { + local_sum += ux2(p); + } else if (dir == in::x3) { + local_sum += ux3(p); + } + }, + total_momentum_in); + return total_momentum_in; +} + +auto main(int argc, char* argv[]) -> int { + ntt::GlobalInitialize(argc, argv); + + try { + const ncells_t nx1 = 32u; + const ncells_t nx2 = 64u; + const ncells_t tile_size = 3u; + const std::vector ncells = { nx1, nx2 }; + const ncells_t ntx1 = static_cast( + math::ceil(static_cast(nx1) / static_cast(tile_size))); + const ncells_t ntx2 = static_cast( + math::ceil(static_cast(nx2) / static_cast(tile_size))); + const npart_t npart = 1000u; + random_number_pool_t random_pool { 12345u }; + + Particles sp1 { 1u, + "sp1", + 1.0f, + 1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + Particles sp2 { 2u, + "sp2", + 1.0f, + -1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + Particles sp3 { 3u, + "sp3", + 0.0f, + 0.0f, + npart, + 0u, + 0u, + ParticlePusher::PHOTON, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + + for (auto* sp : { &sp1, &sp2, &sp3 }) { + sp->set_npart(npart); + fill_random(sp->i1, + sp->i2, + sp->ux1, + sp->ux2, + sp->ux3, + sp->weight, + sp->tag, + npart, + nx1, + nx2, + random_pool); + } + + boundaries_t extent { + { ZERO, ONE }, + { -ONE, ONE } + }; + + prm::Parameters params; + params.set("qed.compton_scattering.nominal_probability_density", + static_cast(1e-3)); + params.set("qed.compton_scattering.Thomson_limit", static_cast(1e-4)); + + auto policy = arch::qed::ComptonScattering(params, random_pool); + policy.species[0] = static_cast(sp1); + policy.species[1] = static_cast(sp2); + policy.species[2] = static_cast(sp3); + + std::array init_energies { + get_total_energy(true, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_energy(true, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_energy(false, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array init_moms_x1 { + get_total_momentum_in(in::x1, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x1, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x1, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array init_moms_x2 { + get_total_momentum_in(in::x2, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x2, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x2, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array init_moms_x3 { + get_total_momentum_in(in::x3, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x3, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x3, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + + for (int i = 0; i < 1000; ++i) { + kernel::mink::TwoBodyInteraction({ &sp1, &sp2 }, + { &sp3 }, + ncells, + extent, + tile_size, + random_pool, + policy); + } + + { + std::array fin_energies { + get_total_energy(true, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_energy(true, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_energy(false, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array fin_moms_x1 { + get_total_momentum_in(in::x1, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x1, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x1, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array fin_moms_x2 { + get_total_momentum_in(in::x2, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x2, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x2, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array fin_moms_x3 { + get_total_momentum_in(in::x3, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x3, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x3, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + + const auto fin_energy = fin_energies[0] + fin_energies[1] + fin_energies[2]; + const auto init_energy = init_energies[0] + init_energies[1] + + init_energies[2]; + const auto fin_mom_x1 = fin_moms_x1[0] + fin_moms_x1[1] + fin_moms_x1[2]; + const auto init_mom_x1 = init_moms_x1[0] + init_moms_x1[1] + init_moms_x1[2]; + const auto fin_mom_x2 = fin_moms_x2[0] + fin_moms_x2[1] + fin_moms_x2[2]; + const auto init_mom_x2 = init_moms_x2[0] + init_moms_x2[1] + init_moms_x2[2]; + const auto fin_mom_x3 = fin_moms_x3[0] + fin_moms_x3[1] + fin_moms_x3[2]; + const auto init_mom_x3 = init_moms_x3[0] + init_moms_x3[1] + init_moms_x3[2]; + + const auto err_energy = (fin_energy - init_energy) / init_energy; + const auto err_mom_x1 = (fin_mom_x1 - init_mom_x1) / + (std::abs(init_mom_x1) + 1e-10); + const auto err_mom_x2 = (fin_mom_x2 - init_mom_x2) / + (std::abs(init_mom_x2) + 1e-10); + const auto err_mom_x3 = (fin_mom_x3 - init_mom_x3) / + (std::abs(init_mom_x3) + 1e-10); + + raise::ErrorIf(err_energy > 1e-5, + fmt::format("energy is not conserved %e -> %e [%e]", + init_energy, + fin_energy, + err_energy), + HERE); + raise::ErrorIf(err_mom_x1 > 1e-5, + fmt::format("x1 momentum is not conserved %e -> %e [%e]", + init_mom_x1, + fin_mom_x1, + err_mom_x1), + HERE); + raise::ErrorIf(err_mom_x2 > 1e-5, + fmt::format("x2 momentum is not conserved %e -> %e [%e]", + init_mom_x2, + fin_mom_x2, + err_mom_x2), + HERE); + raise::ErrorIf(err_mom_x3 > 1e-5, + fmt::format("x3 momentum is not conserved %e -> %e [%e]", + init_mom_x3, + fin_mom_x3, + err_mom_x3), + HERE); + } + + } catch (std::exception& e) { + std::cerr << e.what() << '\n'; + ntt::GlobalFinalize(); + return 1; + } + ntt::GlobalFinalize(); + return 0; +} From 554e1720e0911f7794dd2aeb3e36701f9e8881a5 Mon Sep 17 00:00:00 2001 From: hayk Date: Tue, 12 May 2026 02:27:47 -0400 Subject: [PATCH 02/25] qed input params + caller in srpic.hpp --- input.example.toml | 4 ++++ src/engines/engine.hpp | 1 + src/engines/srpic/twobody.h | 37 +++++++++++++++++++++++++++++- src/framework/parameters/extra.cpp | 24 +++++++++++++++++++ src/framework/parameters/extra.h | 12 ++++++++++ src/global/defaults.h | 4 ++++ src/global/enums.h | 13 +++++++++++ 7 files changed, 94 insertions(+), 1 deletion(-) diff --git a/input.example.toml b/input.example.toml index aa68d28ff..20197b08e 100644 --- a/input.example.toml +++ b/input.example.toml @@ -294,6 +294,10 @@ # @note: array indexing starting at 1 # @note: For interactions between particles of the same group, leave `group2` empty group2 = "" + # Interval in timesteps between checking for the interaction + # @type: uint + # @default: 1 + interval = "" [algorithms] # Number of current smoothing passes diff --git a/src/engines/engine.hpp b/src/engines/engine.hpp index b20e163ca..026cd5c50 100644 --- a/src/engines/engine.hpp +++ b/src/engines/engine.hpp @@ -127,6 +127,7 @@ namespace ntt { auto parameters = prm::Parameters {}; parameters.set("dt", static_cast(dt)); parameters.set("time", static_cast(time)); + parameters.set("step", static_cast(step)); return parameters; } }; diff --git a/src/engines/srpic/twobody.h b/src/engines/srpic/twobody.h index a6e034c24..b969604ea 100644 --- a/src/engines/srpic/twobody.h +++ b/src/engines/srpic/twobody.h @@ -1,8 +1,43 @@ #ifndef ENGINES_SRPIC_TWOBODY_H #define ENGINES_SRPIC_TWOBODY_H +#include "enums.h" + +#include "traits/metric.h" +#include "utils/error.h" +#include "utils/log.h" +#include "utils/param_container.h" + +#include "archetypes/qed/compton.h" +#include "framework/domain/domain.h" +#include "framework/parameters/extra.h" +#include "framework/parameters/parameters.h" +#include "kernels/twobody_interactions.hpp" + namespace ntt { - namespace srpic {} + namespace srpic { + + template + void TwoBodyInteractions(Domain& domain, + const prm::Parameters& engine_params, + const SimulationParams& params) { + logger::Checkpoint("Launching TwoBodyInteractions routines", HERE); + const auto dt = engine_params.get("dt"); + const auto step = engine_params.get("step"); + for (const auto& interaction : + params.template get>( + "two_body.interactions")) { + if (step % interaction.interval == 0u) { + if (interaction.type == TwoBodyInteraction::COMPTON) { + + printf("CALLING COMPTON\n"); + } else if (interaction.type == TwoBodyInteraction::CUSTOM) { + raise::Error("Custom two-body interactions not implemented yet", HERE); + } + } + } + } + } // namespace srpic } // namespace ntt #endif // ENGINES_SRPIC_TWOBODY_H \ No newline at end of file diff --git a/src/framework/parameters/extra.cpp b/src/framework/parameters/extra.cpp index dbee39c7c..4d7169700 100644 --- a/src/framework/parameters/extra.cpp +++ b/src/framework/parameters/extra.cpp @@ -114,6 +114,26 @@ namespace ntt { compton_photon_weight.value(); compton_nominal_photon_energy = ONE / SQR(compton_gamma_qed.value()); } + + twobody_thomson_optical_depth = toml::find_or( + toml_data, + "two_body", + "thomson_optical_depth", + defaults::twobody::thomson_optical_depth); + + // find two-body interactions + const auto twobody_tab = toml::find_or(toml_data, + "two_body", + "interactions", + toml::array {}); + for (const auto& tbint : twobody_tab) { + twobody_interactions.push_back(TwoBodyInteractionParams { + .type = TwoBodyInteraction::from_string( + toml::find(tbint, "type")), + .group1 = toml::find>(tbint, "group1"), + .group2 = toml::find_or>(tbint, "group2", {}), + .interval = toml::find_or(tbint, "interval", 1) }); + } } void Extra::setParams(const std::map& extra, @@ -159,6 +179,10 @@ namespace ntt { params->set("radiation.emission.compton.nominal_photon_energy", compton_nominal_photon_energy.value()); } + + params->set("two_body.thomson_optical_depth", + twobody_thomson_optical_depth.value()); + params->set("two_body.interactions", twobody_interactions); } } // namespace params } // namespace ntt diff --git a/src/framework/parameters/extra.h b/src/framework/parameters/extra.h index 29de4527a..357b480aa 100644 --- a/src/framework/parameters/extra.h +++ b/src/framework/parameters/extra.h @@ -12,6 +12,7 @@ #ifndef FRAMEWORK_PARAMETERS_EXTRA_H #define FRAMEWORK_PARAMETERS_EXTRA_H +#include "enums.h" #include "global.h" #include "framework/parameters/parameters.h" @@ -25,6 +26,13 @@ namespace ntt { namespace params { + struct TwoBodyInteractionParams { + TwoBodyInteractionFlag type; + std::vector group1; + std::vector group2; + timestep_t interval; + }; + struct Extra { // radiative drag parameters std::optional synchrotron_gamma_rad; @@ -45,6 +53,10 @@ namespace ntt { std::optional compton_nominal_probability; std::optional compton_nominal_photon_energy; + // two-body interaction parameters + std::optional twobody_thomson_optical_depth; + std::vector twobody_interactions; + void read(const std::map&, const toml::value&, const SimulationParams* const); diff --git a/src/global/defaults.h b/src/global/defaults.h index dbe37ef0f..11792b2cf 100644 --- a/src/global/defaults.h +++ b/src/global/defaults.h @@ -107,6 +107,10 @@ namespace ntt::defaults { const real_t gamma_rad = 1.0; const real_t gamma_qed = 10.0; } // namespace compton + + namespace twobody { + const real_t thomson_optical_depth = 1.0; + } // namespace twobody } // namespace ntt::defaults #endif // GLOBAL_DEFAULTS_H diff --git a/src/global/enums.h b/src/global/enums.h index 41fdd551b..d8fc509ef 100644 --- a/src/global/enums.h +++ b/src/global/enums.h @@ -420,6 +420,19 @@ namespace ntt { return "unknown"; } } + + inline auto from_string(const std::string& s) -> uint8_t { + if (fmt::toLower(s) == "none") { + return NONE; + } else if (fmt::toLower(s) == "compton") { + return COMPTON; + } else if (fmt::toLower(s) == "custom") { + return CUSTOM; + } else { + raise::Error(fmt::format("Invalid TwoBodyInteraction type: %s", s), HERE); + return NONE; + } + } } // namespace TwoBodyInteraction using TwoBodyInteractionFlag = uint8_t; From d9930b8cb5222e339efe88fb06ddb2272271e677 Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 14 May 2026 17:16:40 -0400 Subject: [PATCH 03/25] compton tests --- examples/compton_jones/compton_jones.py | 62 ++++++++ examples/compton_jones/compton_jones.toml | 89 +++++++++++ examples/compton_jones/pgen.hpp | 148 ++++++++++++++++++ .../compton_kompaneets/compton_kompaneets.py | 60 +++++++ .../compton_kompaneets.toml | 87 ++++++++++ examples/compton_kompaneets/pgen.hpp | 74 +++++++++ 6 files changed, 520 insertions(+) create mode 100644 examples/compton_jones/compton_jones.py create mode 100644 examples/compton_jones/compton_jones.toml create mode 100644 examples/compton_jones/pgen.hpp create mode 100644 examples/compton_kompaneets/compton_kompaneets.py create mode 100644 examples/compton_kompaneets/compton_kompaneets.toml create mode 100644 examples/compton_kompaneets/pgen.hpp diff --git a/examples/compton_jones/compton_jones.py b/examples/compton_jones/compton_jones.py new file mode 100644 index 000000000..873a01f35 --- /dev/null +++ b/examples/compton_jones/compton_jones.py @@ -0,0 +1,62 @@ +import nt2 +import matplotlib.pyplot as plt +import numpy as np + +data = nt2.Data("compton_jones") + +photons = data.particles.sel(sp=3).isel(t=-1).load() +photons = photons[np.sqrt(photons.ux**2 + photons.uy**2 + photons.uz**2) > 0.01] + +plt.rcParams["figure.dpi"] = 300 +plt.rcParams["font.family"] = "serif" + +fig = plt.figure(figsize=(9, 4)) +gs = fig.add_gridspec(1, 2, wspace=0.35) +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1]) + +gamma = np.sqrt(1 + data.attrs["setup.electron_4vel"] ** 2) +e0 = data.attrs["setup.photon_energy"] +Gamma = 4 * e0 * gamma +emax = gamma * Gamma / (1 + Gamma) + +es = data.spectra.E[1:-1] / emax + +dnde = data.spectra.N_3.isel(t=-1)[1:-1] +dnde /= np.trapezoid(dnde, es) +ax1.plot(es, dnde) + +es = np.linspace(es.values.min(), es.values.max(), 250) +qs = es / (1 + Gamma * (1 - es)) +dnde_th = ( + 2 * qs * np.log(qs) + + (1 + 2 * qs) * (1 - qs) + + 0.5 * Gamma**2 * qs**2 / (1 + Gamma * qs) * (1 - qs) +) + +dnde_th /= np.trapezoid(dnde_th, es) + +ax1.plot(es, dnde_th, c="k", ls=":") +ax1.set( + xlim=(0, 1), + ylim=(0, 4), + xlabel=r"$\varepsilon_{\rm ph} / \varepsilon_{\rm max}$", + ylabel=r"$dn_{\rm ph}/d\varepsilon_{\rm ph}$", +) + +plt.scatter( + photons.ux / emax, + photons.uy / emax, + s=1, + linewidth=0, +) +xs = np.linspace(0, 1, 100) +ys = 2 / gamma * xs +ax2.plot(xs, ys, c="k", ls="--", lw=0.5) +ax2.plot(xs, -ys, c="k", ls="--", lw=0.5) +ax2.set( + xlabel=r"$p_{\rm ph}^x / \varepsilon_{\rm max}$", + ylabel=r"$p_{\rm ph}^y / \varepsilon_{\rm max}$", +) + +plt.savefig("compton_jones.png", bbox_inches="tight") diff --git a/examples/compton_jones/compton_jones.toml b/examples/compton_jones/compton_jones.toml new file mode 100644 index 000000000..796739978 --- /dev/null +++ b/examples/compton_jones/compton_jones.toml @@ -0,0 +1,89 @@ +[simulation] + name = "compton_jones" + engine = "srpic" + runtime = 10.0 + +[grid] + resolution = [32, 32] + extent = [[0.0, 1.0], [0.0, 1.0]] + + [grid.metric] + metric = "minkowski" + + [grid.boundaries] + fields = [["PERIODIC"], ["PERIODIC"]] + particles = [["PERIODIC"], ["PERIODIC"]] + +[scales] + larmor0 = 1.0 + skindepth0 = 1.0 + +[two_body] + thomson_optical_depth = 0.5 + + [[two_body.interaction]] + type = "compton" + group1 = [1] + group2 = [2] + interval = 1 + tile_size = 5 + recoil1 = false + recoil2 = true + +[algorithms] + current_filters = 0 + + [algorithms.deposit] + enable = false + + [algorithms.fieldsolver] + enable = false + +[particles] + ppc0 = 1.0 + clear_interval = 1 + + [[particles.species]] + label = "e-" + mass = 1.0 + charge = -1.0 + maxnpart = 1e6 + + [[particles.species]] + label = "ph" + mass = 0.0 + charge = 0.0 + maxnpart = 1e6 + + [[particles.species]] + label = "ph_out" + mass = 0.0 + charge = 0.0 + maxnpart = 1e7 + pusher = "none" + +[setup] + electron_4vel = 999.9995 + photon_energy = 1e-2 + +[output] + interval_time = 0.01 + + [output.fields] + quantities = ["N_1", "N_2", "N_3"] + + [output.particles] + species = [1, 2, 3] + stride = 1 + + [output.spectra] + log_bins = false + e_min = 0 + e_max = 1100 + n_bins = 100 + + [output.stats] + quantities = ["T00_1", "T00_2", "T00_3"] + +[checkpoint] + keep = 0 diff --git a/examples/compton_jones/pgen.hpp b/examples/compton_jones/pgen.hpp new file mode 100644 index 000000000..f2da19802 --- /dev/null +++ b/examples/compton_jones/pgen.hpp @@ -0,0 +1,148 @@ +#ifndef PROBLEM_GENERATOR_H +#define PROBLEM_GENERATOR_H + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "traits/pgen.h" + +#include "archetypes/particle_injector.h" +#include "framework/domain/metadomain.h" + +namespace user { + using namespace ntt; + + template + struct DeltaDistribution { + const real_t energy0; + bool monodirectional; + random_number_pool_t random_pool; + + DeltaDistribution(real_t energy0, + bool monodirectional, + random_number_pool_t& random_pool) + : energy0 { energy0 } + , monodirectional { monodirectional } + , random_pool { random_pool } {} + + Inline void operator()(const coord_t&, vec_t& v) const { + if (not monodirectional) { + auto gen = random_pool.get_state(); + auto rnd1 = Random(gen); + auto rnd2 = Random(gen); + random_pool.free_state(gen); + // random direction + const auto phi = static_cast(constant::TWO_PI) * rnd1; + const auto ct = 2.0 * rnd2 - 1.0; + const auto st = math::sqrt(1.0 - ct * ct); + v[0] = energy0 * st * math::cos(phi); + v[1] = energy0 * st * math::sin(phi); + v[2] = energy0 * ct; + } else { + v[0] = energy0; + v[1] = 0.0; + v[2] = 0.0; + } + } + }; + + template + struct PGen { + + static constexpr auto engines { + ::traits::pgen::compatible_with {} + }; + static constexpr auto metrics { + ::traits::pgen::compatible_with {} + }; + static constexpr auto dimensions { + ::traits::pgen::compatible_with {} + }; + + const SimulationParams& params; + const Metadomain& metadomain; + + PGen(const SimulationParams& p, const Metadomain& m) + : params { p } + , metadomain { m } {} + + void InitPrtls(Domain& domain) { + auto delta_electrons = DeltaDistribution { + params.template get("setup.electron_4vel"), + true, + domain.random_pool() + }; + arch::InjectUniform(params, + domain, + 1u, + delta_electrons, + ONE); + } + + void CustomPostStep(timestep_t /*step*/, simtime_t /*time*/, Domain& domain) { + // copy all photons from species #2 (idx 1) to #3 (idx 2) with an offset + const auto offset = domain.species[2].npart(); + const auto new_copies = domain.species[1].npart(); + const auto new_size = offset + new_copies; + const auto from_slice = prtl_slice_t { 0, new_copies }; + const auto to_slice = prtl_slice_t { offset, new_size }; + + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i1, to_slice), + Kokkos::subview(domain.species[1].i1, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i1_prev, to_slice), + Kokkos::subview(domain.species[1].i1_prev, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx1, to_slice), + Kokkos::subview(domain.species[1].dx1, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx1_prev, to_slice), + Kokkos::subview(domain.species[1].dx1_prev, from_slice)); + if constexpr (M::Dim == Dim::_2D or M::Dim == Dim::_3D) { + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i2, to_slice), + Kokkos::subview(domain.species[1].i2, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i2_prev, to_slice), + Kokkos::subview(domain.species[1].i2_prev, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx2, to_slice), + Kokkos::subview(domain.species[1].dx2, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx2_prev, to_slice), + Kokkos::subview(domain.species[1].dx2_prev, from_slice)); + } + if constexpr (M::Dim == Dim::_3D) { + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i3, to_slice), + Kokkos::subview(domain.species[1].i3, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i3_prev, to_slice), + Kokkos::subview(domain.species[1].i3_prev, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx3, to_slice), + Kokkos::subview(domain.species[1].dx3, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx3_prev, to_slice), + Kokkos::subview(domain.species[1].dx3_prev, from_slice)); + } + Kokkos::deep_copy(Kokkos::subview(domain.species[2].ux1, to_slice), + Kokkos::subview(domain.species[1].ux1, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].ux2, to_slice), + Kokkos::subview(domain.species[1].ux2, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].ux3, to_slice), + Kokkos::subview(domain.species[1].ux3, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].weight, to_slice), + Kokkos::subview(domain.species[1].weight, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].tag, to_slice), + Kokkos::subview(domain.species[1].tag, from_slice)); + + domain.species[1].set_npart(0); + domain.species[2].set_npart(new_size); + + auto delta_photons = DeltaDistribution { + params.template get("setup.photon_energy"), + false, + domain.random_pool() + }; + arch::InjectUniform(params, + domain, + 2u, + delta_photons, + ONE); + } + }; + +} // namespace user + +#endif diff --git a/examples/compton_kompaneets/compton_kompaneets.py b/examples/compton_kompaneets/compton_kompaneets.py new file mode 100644 index 000000000..4c713cbad --- /dev/null +++ b/examples/compton_kompaneets/compton_kompaneets.py @@ -0,0 +1,60 @@ +import nt2 +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +data = nt2.Data("compton_kompaneets") +stats = pd.read_csv("compton_kompaneets/compton_kompaneets_stats.csv") +stats.columns = stats.columns.str.strip() + +plt.rcParams["figure.dpi"] = 300 +plt.rcParams["font.family"] = "serif" + +fig = plt.figure(figsize=(9, 4)) +gs = fig.add_gridspec(1, 2, wspace=0.3) +ax1 = fig.add_subplot(gs[0, 0]) + +tvals = len(data.spectra.t.values) + +nphot = data.spectra.N_3.isel(t=-1).sum().values[()] +for ti in range(0, tvals, 10): + ax1.plot( + data.spectra.E.values / data.attrs["setup.temperature"], + data.spectra.N_3.isel(t=ti).values, + c=plt.get_cmap("plasma")(ti / tvals), + lw=0.5, + ) + +es = data.spectra.E.values / data.attrs["setup.temperature"] +dndes = es**2 * np.exp(-es) +dndes /= np.sum(dndes) +dndes *= nphot +ax1.plot( + es, + dndes, + c="k", + ls=":", + label=r"$\propto \varepsilon_{\rm ph}^2 e^{-\varepsilon_{\rm ph} / T_\pm}$", +) + +ax1.set( + yscale="log", + ylim=(1e-1, 1e5), + xlim=(0, 10), + xlabel=r"$\varepsilon / T_\pm$", + ylabel=r"$dn_{\rm ph}/d\varepsilon$", +) +ax1.legend() + +ax2 = fig.add_subplot(gs[0, 1]) +ax2.plot(stats["time"], stats["T00_3"], c="C0") +ax2.set(ylabel=r"total photon energy", xlabel=r"$t$") +ax2.yaxis.label.set_color("C0") +ax2.tick_params(axis="y", labelcolor="C0") +ax2twin = ax2.twinx() +ax2twin.plot(data.spectra.t.values, data.spectra.N_3.sum("E"), c="C2") +ax2twin.set(ylabel=r"photon number") +ax2twin.yaxis.label.set_color("C2") +ax2twin.tick_params(axis="y", labelcolor="C2") + +plt.savefig("compton_kompaneets_plot.png", bbox_inches="tight") diff --git a/examples/compton_kompaneets/compton_kompaneets.toml b/examples/compton_kompaneets/compton_kompaneets.toml new file mode 100644 index 000000000..1a006102f --- /dev/null +++ b/examples/compton_kompaneets/compton_kompaneets.toml @@ -0,0 +1,87 @@ +[simulation] + name = "compton_kompaneets" + engine = "srpic" + runtime = 5.0 + +[grid] + resolution = [128, 128] + extent = [[0.0, 1.0], [0.0, 1.0]] + + [grid.metric] + metric = "minkowski" + + [grid.boundaries] + fields = [["PERIODIC"], ["PERIODIC"]] + particles = [["PERIODIC"], ["PERIODIC"]] + +[scales] + larmor0 = 1.0 + skindepth0 = 1.0 + +[two_body] + thomson_optical_depth = 0.25 + + [[two_body.interaction]] + type = "compton" + group1 = [1, 2] + group2 = [3] + interval = 1 + tile_size = 5 + recoil1 = false + recoil2 = true + +[algorithms] + current_filters = 0 + + [algorithms.deposit] + enable = false + + [algorithms.fieldsolver] + enable = false + +[particles] + ppc0 = 2.0 + clear_interval = 1 + + [[particles.species]] + label = "e-" + mass = 1.0 + charge = -1.0 + maxnpart = 1e6 + + [[particles.species]] + label = "e+" + mass = 1.0 + charge = 1.0 + maxnpart = 1e6 + + [[particles.species]] + label = "ph" + mass = 0.0 + charge = 0.0 + maxnpart = 1e6 + +[setup] + temperature = 0.01 + photon_energy = 1e-3 + +[output] + interval_time = 0.01 + + [output.fields] + enable = false + + [output.particles] + enable = false + + [output.spectra] + log_bins = false + e_min = 0.0 + e_max = 1.0 + n_bins = 500 + + [output.stats] + quantities = ["T00_1", "T00_2", "T00_3"] + +[checkpoint] + keep = 0 diff --git a/examples/compton_kompaneets/pgen.hpp b/examples/compton_kompaneets/pgen.hpp new file mode 100644 index 000000000..0b0158ada --- /dev/null +++ b/examples/compton_kompaneets/pgen.hpp @@ -0,0 +1,74 @@ +#ifndef PROBLEM_GENERATOR_H +#define PROBLEM_GENERATOR_H + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "traits/pgen.h" + +#include "archetypes/particle_injector.h" +#include "archetypes/utils.h" +#include "framework/domain/metadomain.h" + +namespace user { + using namespace ntt; + + template + struct DeltaDistribution { + const real_t photon_energy0; + random_number_pool_t random_pool; + + DeltaDistribution(real_t photon_energy0, random_number_pool_t& random_pool) + : photon_energy0 { photon_energy0 } + , random_pool { random_pool } {} + + Inline void operator()(const coord_t&, vec_t& v) const { + auto gen = random_pool.get_state(); + auto rnd1 = Random(gen); + auto rnd2 = Random(gen); + random_pool.free_state(gen); + // random direction + const auto phi = static_cast(constant::TWO_PI) * rnd1; + const auto ct = 2.0 * rnd2 - 1.0; + const auto st = math::sqrt(1.0 - ct * ct); + v[0] = photon_energy0 * st * math::cos(phi); + v[1] = photon_energy0 * st * math::sin(phi); + v[2] = photon_energy0 * ct; + } + }; + + template + struct PGen { + + static constexpr auto engines { + ::traits::pgen::compatible_with {} + }; + static constexpr auto metrics { + ::traits::pgen::compatible_with {} + }; + static constexpr auto dimensions { + ::traits::pgen::compatible_with {} + }; + + const SimulationParams& params; + const Metadomain& metadomain; + + PGen(const SimulationParams& p, const Metadomain& m) + : params { p } + , metadomain { m } {} + + void InitPrtls(Domain& domain) { + const auto temperature = params.template get("setup.temperature"); + arch::InjectUniformMaxwellian(params, domain, ONE, temperature, { 1u, 2u }); + + auto delta = DeltaDistribution { params.template get( + "setup.photon_energy"), + domain.random_pool() }; + arch::InjectUniform(params, domain, 3u, delta, ONE); + } + }; + +} // namespace user + +#endif From 388d9af6d811ec26d6dbe20f2c72c7673cbb9892 Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 14 May 2026 17:16:59 -0400 Subject: [PATCH 04/25] extra injector for single species --- src/archetypes/particle_injector.h | 67 ++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/src/archetypes/particle_injector.h b/src/archetypes/particle_injector.h index 7adb8c5b5..047a40183 100644 --- a/src/archetypes/particle_injector.h +++ b/src/archetypes/particle_injector.h @@ -389,6 +389,73 @@ namespace arch { } } + /** + * @brief Injects uniform number density of a single species everywhere in the domain + * @param domain Domain object + * @param species Species index + * @param energy_dist Energy distribution objects + * @param number_density Number density (in units of n0) + * @param use_weights Use weights + * @param box Region to inject the particles in global coords + * @tparam S Simulation engine type + * @tparam M Metric type + * @tparam ED Energy distribution type + */ + template ED> + inline void InjectUniform(const SimulationParams& params, + Domain& domain, + spidx_t species, + const ED& energy_dist, + real_t number_density, + bool use_weights = false, + const boundaries_t& box = {}) { + raise::ErrorIf((M::CoordType != Coord::Cartesian) && (not use_weights), + "Weights must be used for non-Cartesian coordinates", + HERE); + raise::ErrorIf((M::CoordType == Coord::Cartesian) && use_weights, + "Weights should not be used for Cartesian coordinates", + HERE); + raise::ErrorIf(params.template get("particles.use_weights") != use_weights, + "Weights must be enabled from the input file to use them in " + "the injector", + HERE); + if (domain.species[species - 1].charge() != 0.0f) { + raise::Warning("Charge of the injected species is non-zero", HERE); + } + + { + boundaries_t nonempty_box; + for (auto d { 0u }; d < M::Dim; ++d) { + if (d < box.size()) { + nonempty_box.emplace_back(box[d].first, box[d].second); + } else { + nonempty_box.push_back(Range::All); + } + } + const auto result = ComputeNumInject(params, domain, number_density, nonempty_box); + if (not std::get<0>(result)) { + return; + } + const auto nparticles = std::get<1>(result); + const auto xi_min = std::get<2>(result); + const auto xi_max = std::get<3>(result); + + Kokkos::parallel_for("InjectUniform", + nparticles, + kernel::SingleSpeciesUniformInjector_kernel( + domain.species[species - 1], + domain.index(), + domain.mesh.metric, + xi_min, + xi_max, + energy_dist, + ONE / params.template get("scales.V0"), + domain.random_pool())); + domain.species[species - 1].set_npart( + domain.species[species - 1].npart() + nparticles); + } + } + } // namespace arch #endif // ARCHETYPES_PARTICLE_INJECTOR_H From a09db82605e2a50597596f763371e102979b7619 Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 14 May 2026 17:17:30 -0400 Subject: [PATCH 05/25] qed incorporated --- input.example.toml | 12 +++ src/archetypes/qed/compton.h | 120 ++++++++++++--------- src/engines/engine.hpp | 4 +- src/engines/reporter.cpp | 43 ++++++++ src/engines/srpic/srpic.hpp | 7 ++ src/engines/srpic/twobody.h | 64 ++++++++++- src/framework/parameters/extra.cpp | 13 ++- src/framework/parameters/extra.h | 3 + src/global/enums.h | 3 +- src/global/traits/policies.h | 35 ++++-- src/kernels/injectors.hpp | 126 ++++++++++++++++++++++ src/kernels/twobody_interactions.hpp | 131 ++++++++++++++++++----- tests/archetypes/qed_compton.cpp | 8 +- tests/kernels/twobody_interactions.cpp | 141 ++++++++----------------- 14 files changed, 510 insertions(+), 200 deletions(-) diff --git a/input.example.toml b/input.example.toml index 20197b08e..d721c9275 100644 --- a/input.example.toml +++ b/input.example.toml @@ -298,6 +298,18 @@ # @type: uint # @default: 1 interval = "" + # Size of interaction tile in number of cells + # @type: uint + # @default: 4 + tile_size = "" + # Whether to apply recoil on species of group1 + # @type: bool + # @default: true + recoil1 = "" + # Whether to apply recoil on species of group2 + # @type: bool + # @default: true + recoil2 = "" [algorithms] # Number of current smoothing passes diff --git a/src/archetypes/qed/compton.h b/src/archetypes/qed/compton.h index b54c268bc..4f466579e 100644 --- a/src/archetypes/qed/compton.h +++ b/src/archetypes/qed/compton.h @@ -24,33 +24,26 @@ namespace arch::qed { using namespace ntt; - template + template struct ComptonScattering { static constexpr spidx_t MAXSP = 16u; static constexpr int MAX_ITER = 10; ParticleArrays species[MAXSP]; static constexpr real_t low_energy_limit = static_cast(2e-3); + static constexpr real_t Thomson_limit = static_cast(1e-3); const real_t nominal_probability_density; - const real_t Thomson_limit; random_number_pool_t random_pool; ComptonScattering(const prm::Parameters& params, random_number_pool_t& random_pool) : nominal_probability_density { params.template get( - "qed.compton_scattering.nominal_probability_density") } - , Thomson_limit { params.template get( - "qed.compton_scattering.Thomson_limit") } + "compton_scattering.nominal_probability_density") } , random_pool { random_pool } { if (nominal_probability_density <= ZERO) { raise::Error("nominal_probability must be in the range (0, 1]", HERE); } - if (Thomson_limit <= ZERO or Thomson_limit > low_energy_limit) { - raise::Error( - "Thomson_limit must be in the range (0, small_energy_limit]", - HERE); - } } /* @@ -79,30 +72,31 @@ namespace arch::qed { /* * Calculate the Klein-Nishina cross section for a photon with energy e_ in the lepton rest frame * @param e_: photon energy in the lepton rest frame - * @return: pair of (is_KN_regime, f_KN) where - * - is_KN_regime: whether the photon energy is in the Klein-Nishina regime (e_ > Thomson_limit) - * - f_KN: the Klein-Nishina cross section normalized to the Thomson cross section + * @return f_KN: the Klein-Nishina cross section normalized to the Thomson cross section * @note for e_ > low_energy_limit, full Klein-Nishina formula * @note for Thomson_limit < e_ <= low_energy_limit, 2nd order expansion of the Klein-Nishina formula * @note for e_ <= Thomson_limit, return 1 (Thomson limit) */ - Inline auto KNCrossSection(real_t e_) const -> Kokkos::pair { + Inline auto KNCrossSection(double e_) const -> real_t { if (e_ > Thomson_limit) { if (e_ < low_energy_limit) { // correctly handle the e_ << 1 limit using 2nd order expansion of f_KN - return { true, ONE - TWO * e_ + static_cast(5.2) * SQR(e_) }; + return static_cast(1.0 - 2.0 * e_ + 5.2 * SQR(e_)); } else { - return { true, - static_cast(0.375) * - ((ONE - TWO / e_ - TWO / SQR(e_)) * math::log(ONE + TWO * e_) + - HALF + FOUR / e_ - HALF / SQR(ONE + TWO * e_)) / - e_ }; + return static_cast( + 0.375 * + ((1.0 - 2.0 / e_ - 2.0 / SQR(e_)) * math::log(1.0 + 2.0 * e_) + + 0.5 + 4.0 / e_ - 0.5 / SQR(1.0 + 2.0 * e_)) / + e_); } } else { - return { false, ONE }; + return ONE; } } + /* + * Sample a cosine theta value from the Thomson scattering cross section + */ Inline auto RandomCosTheta_Th() const -> real_t { auto gen_ = random_pool.get_state(); const auto rnd_ = Random(gen_); @@ -114,6 +108,9 @@ namespace arch::qed { return u - ONE / u; } + /* + * Sample a cosine theta from the Klein-Nishina scattering cross section + */ Inline auto RandomCosTheta_KN(double e_) const -> real_t { auto gen_ = random_pool.get_state(); const auto rnd_ = Random(gen_); @@ -215,12 +212,11 @@ namespace arch::qed { rand_sintheta_ * rand_sinphi_ * c_[2]); } - Inline void operator()(spidx_t sp1, - npart_t p1, - spidx_t sp2, - npart_t p2, - real_t tile_volume) const { - // @TODO coord/vec conversion + Inline auto should_interact(spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2, + real_t tile_weight) const -> bool { // values with "_" are in the lepton rest-frame const vec_t lepton_u { species[sp1 - 1].ux1(p1), species[sp1 - 1].ux2(p1), @@ -240,43 +236,67 @@ namespace arch::qed { LorentzBoost(lepton_u, lepton_gamma, photon_p, photon_energy, photon_p_, photon_energy_); - const auto [KN_regime, f_KN] = KNCrossSection(photon_energy_); - const auto scattering_probability = nominal_probability_density * f_KN * - photon_energy_ * lepton_weight * - photon_weight / - (photon_energy * lepton_gamma * - tile_volume); + const auto f_KN = KNCrossSection(photon_energy_); + auto gen = random_pool.get_state(); const auto rnd = Random(gen); random_pool.free_state(gen); - if (rnd < scattering_probability) { - vec_t photon_pnew_ { ZERO, ZERO, ZERO }, - photon_pnew { ZERO, ZERO, ZERO }; - real_t photon_energy_new_ { ZERO }, photon_energy_new { ZERO }; - - ScatterPhoton(KN_regime, - photon_p_, - photon_energy_, - photon_pnew_, - photon_energy_new_); - LorentzBoost({ -lepton_u[0], -lepton_u[1], -lepton_u[2] }, - lepton_gamma, - photon_pnew_, - photon_energy_new_, - photon_pnew, - photon_energy_new); + return rnd < + (tile_weight * nominal_probability_density * f_KN * photon_energy_ * + lepton_weight * photon_weight / (photon_energy * lepton_gamma)); + } + + Inline void operator()(spidx_t sp1, npart_t p1, spidx_t sp2, npart_t p2) const { + // @TODO coord/vec conversion + // values with "_" are in the lepton rest-frame + const vec_t lepton_u { species[sp1 - 1].ux1(p1), + species[sp1 - 1].ux2(p1), + species[sp1 - 1].ux3(p1) }; + const auto lepton_gamma = U2GAMMA(lepton_u[0], lepton_u[1], lepton_u[2]); + const auto lepton_weight = species[sp1 - 1].weight(p1); + + const vec_t photon_p { species[sp2 - 1].ux1(p2), + species[sp2 - 1].ux2(p2), + species[sp2 - 1].ux3(p2) }; + const auto photon_energy = NORM(photon_p[0], photon_p[1], photon_p[2]); + const auto photon_weight = species[sp2 - 1].weight(p2); + + // boost photon momentum to lepton rest frame + vec_t photon_p_ { ZERO, ZERO, ZERO }; + real_t photon_energy_ { ZERO }; + + LorentzBoost(lepton_u, lepton_gamma, photon_p, photon_energy, photon_p_, photon_energy_); + + vec_t photon_pnew_ { ZERO, ZERO, ZERO }, + photon_pnew { ZERO, ZERO, ZERO }; + real_t photon_energy_new_ { ZERO }, photon_energy_new { ZERO }; + + ScatterPhoton(photon_energy_ > Thomson_limit, + photon_p_, + photon_energy_, + photon_pnew_, + photon_energy_new_); + LorentzBoost({ -lepton_u[0], -lepton_u[1], -lepton_u[2] }, + lepton_gamma, + photon_pnew_, + photon_energy_new_, + photon_pnew, + photon_energy_new); + if constexpr (R1) { species[sp1 - 1].ux1(p1) += (photon_p[0] - photon_pnew[0]) * photon_weight / lepton_weight; species[sp1 - 1].ux2(p1) += (photon_p[1] - photon_pnew[1]) * photon_weight / lepton_weight; species[sp1 - 1].ux3(p1) += (photon_p[2] - photon_pnew[2]) * photon_weight / lepton_weight; + } + if constexpr (R2) { species[sp2 - 1].ux1(p2) = photon_pnew[0]; species[sp2 - 1].ux2(p2) = photon_pnew[1]; species[sp2 - 1].ux3(p2) = photon_pnew[2]; - } // not interacting + } } }; diff --git a/src/engines/engine.hpp b/src/engines/engine.hpp index 026cd5c50..bd77c83a4 100644 --- a/src/engines/engine.hpp +++ b/src/engines/engine.hpp @@ -249,8 +249,8 @@ namespace ntt { "ParticlePusher", "FieldBoundaries", "ParticleBoundaries", "Communications", "Injector", "Custom", - "ParticleSort", "Output", - "Checkpoint" }, + "TwoBodyInteractions", "ParticleSort", + "Output", "Checkpoint" }, []() { Kokkos::fence(); }, diff --git a/src/engines/reporter.cpp b/src/engines/reporter.cpp index f56874d23..11ae740ed 100644 --- a/src/engines/reporter.cpp +++ b/src/engines/reporter.cpp @@ -6,6 +6,7 @@ #include "utils/formatting.h" #include "utils/reporter.h" +#include "framework/parameters/extra.h" #include "framework/parameters/parameters.h" #include @@ -149,6 +150,48 @@ namespace ntt { params.template get( "radiation.emission.synchrotron.nominal_photon_energy")); } + + report += "\n"; + const auto two_body_interactions = + params.template get>( + "two_body.interaction"); + if (not two_body_interactions.empty()) { + reporter::AddCategory(report, 4, "Two-body interactions"); + reporter::AddParam( + report, + 6, + "Thomson optical depth", + "%.3e", + params.template get("two_body.thomson_optical_depth")); + for (const auto& interaction : two_body_interactions) { + reporter::AddSubcategory( + report, + 6, + TwoBodyInteraction::to_string(interaction.type).c_str()); + std::vector group1(interaction.group1.size()); + std::vector group2(interaction.group2.size()); + for (size_t g1 = 0; g1 < interaction.group1.size(); ++g1) { + group1[g1] = interaction.group1[g1]; + } + for (size_t g2 = 0; g2 < interaction.group2.size(); ++g2) { + group2[g2] = interaction.group2[g2]; + } + reporter::AddParam(report, + 8, + "group #1 species", + "%s (recoil: %s)", + fmt::formatVector(group1).c_str(), + interaction.recoil1 ? "ON" : "OFF"); + reporter::AddParam(report, + 8, + "group #2 species", + "%s (recoil: %s)", + fmt::formatVector(group2).c_str(), + interaction.recoil2 ? "ON" : "OFF"); + reporter::AddParam(report, 8, "tile size [cells]", "%u", interaction.tile_size); + reporter::AddParam(report, 8, "interval [steps]", "%u", interaction.interval); + } + } return report; } diff --git a/src/engines/srpic/srpic.hpp b/src/engines/srpic/srpic.hpp index 94a8949c1..1afb269a2 100644 --- a/src/engines/srpic/srpic.hpp +++ b/src/engines/srpic/srpic.hpp @@ -24,6 +24,7 @@ #include "engines/srpic/fieldsolvers.h" #include "engines/srpic/particle_pusher.h" #include "engines/srpic/particles_bcs.h" +#include "engines/srpic/twobody.h" #include "framework/domain/domain.h" #include "framework/parameters/parameters.h" @@ -182,6 +183,12 @@ namespace ntt { timers.stop("Injector"); } + if constexpr (CartesianMetricClass) { + timers.start("TwoBodyInteractions"); + srpic::TwoBodyInteractions(dom, this->engineParams(), m_params); + timers.stop("TwoBodyInteractions"); + } + timers.start("ParticleSort"); m_metadomain.SortParticles(time, step, m_params, dom); timers.stop("ParticleSort"); diff --git a/src/engines/srpic/twobody.h b/src/engines/srpic/twobody.h index b969604ea..c5d04ca66 100644 --- a/src/engines/srpic/twobody.h +++ b/src/engines/srpic/twobody.h @@ -5,6 +5,7 @@ #include "traits/metric.h" #include "utils/error.h" +#include "utils/formatting.h" #include "utils/log.h" #include "utils/param_container.h" @@ -26,11 +27,68 @@ namespace ntt { const auto step = engine_params.get("step"); for (const auto& interaction : params.template get>( - "two_body.interactions")) { + "two_body.interaction")) { if (step % interaction.interval == 0u) { + const auto thomson_optical_depth = params.template get( + "two_body.thomson_optical_depth"); + const auto nominal_thomson_probability_density = thomson_optical_depth * + dt * + static_cast( + interaction.interval); if (interaction.type == TwoBodyInteraction::COMPTON) { + prm::Parameters compton_params; + compton_params.set("compton_scattering.nominal_probability_density", + nominal_thomson_probability_density); + auto recoil1 = interaction.recoil1; + auto recoil2 = interaction.recoil2; + auto launch = [&]() { + auto policy = arch::qed::ComptonScattering( + compton_params, + domain.random_pool()); - printf("CALLING COMPTON\n"); + std::vector*> group1_species; + std::vector*> group2_species; + + for (const auto& sp_lepton : interaction.group1) { + raise::ErrorIf( + domain.species[sp_lepton - 1].mass() == ZERO, + fmt::format( + "Species %u is massless but is in the lepton group " + "of a Compton interaction", + sp_lepton), + HERE); + group1_species.push_back(&domain.species[sp_lepton - 1]); + } + for (const auto& sp_photon : interaction.group2) { + raise::ErrorIf( + domain.species[sp_photon - 1].mass() != ZERO, + fmt::format( + "Species %u is massive but is in the photon group " + "of a Compton interaction", + sp_photon), + HERE); + group2_species.push_back(&domain.species[sp_photon - 1]); + } + + kernel::mink::TwoBodyInteraction( + group1_species, + group2_species, + domain.mesh.n_active(), + domain.mesh.extent(), + interaction.tile_size, + params.template get("particles.ppc0"), + domain.random_pool(), + policy); + }; + if (interaction.recoil1 and interaction.recoil2) { + launch.template operator()(); + } else if (interaction.recoil1 and not interaction.recoil2) { + launch.template operator()(); + } else if (not interaction.recoil1 and interaction.recoil2) { + launch.template operator()(); + } else { + launch.template operator()(); + } } else if (interaction.type == TwoBodyInteraction::CUSTOM) { raise::Error("Custom two-body interactions not implemented yet", HERE); } @@ -40,4 +98,4 @@ namespace ntt { } // namespace srpic } // namespace ntt -#endif // ENGINES_SRPIC_TWOBODY_H \ No newline at end of file +#endif // ENGINES_SRPIC_TWOBODY_H diff --git a/src/framework/parameters/extra.cpp b/src/framework/parameters/extra.cpp index 4d7169700..53eca4117 100644 --- a/src/framework/parameters/extra.cpp +++ b/src/framework/parameters/extra.cpp @@ -124,15 +124,18 @@ namespace ntt { // find two-body interactions const auto twobody_tab = toml::find_or(toml_data, "two_body", - "interactions", + "interaction", toml::array {}); for (const auto& tbint : twobody_tab) { twobody_interactions.push_back(TwoBodyInteractionParams { .type = TwoBodyInteraction::from_string( toml::find(tbint, "type")), - .group1 = toml::find>(tbint, "group1"), - .group2 = toml::find_or>(tbint, "group2", {}), - .interval = toml::find_or(tbint, "interval", 1) }); + .group1 = toml::find>(tbint, "group1"), + .group2 = toml::find_or>(tbint, "group2", {}), + .interval = toml::find_or(tbint, "interval", 1), + .tile_size = toml::find_or(tbint, "tile_size", 4u), + .recoil1 = toml::find_or(tbint, "recoil1", true), + .recoil2 = toml::find_or(tbint, "recoil2", true) }); } } @@ -182,7 +185,7 @@ namespace ntt { params->set("two_body.thomson_optical_depth", twobody_thomson_optical_depth.value()); - params->set("two_body.interactions", twobody_interactions); + params->set("two_body.interaction", twobody_interactions); } } // namespace params } // namespace ntt diff --git a/src/framework/parameters/extra.h b/src/framework/parameters/extra.h index 357b480aa..c3891429f 100644 --- a/src/framework/parameters/extra.h +++ b/src/framework/parameters/extra.h @@ -31,6 +31,9 @@ namespace ntt { std::vector group1; std::vector group2; timestep_t interval; + ncells_t tile_size; + bool recoil1; + bool recoil2; }; struct Extra { diff --git a/src/global/enums.h b/src/global/enums.h index d8fc509ef..bb4fbc272 100644 --- a/src/global/enums.h +++ b/src/global/enums.h @@ -429,7 +429,8 @@ namespace ntt { } else if (fmt::toLower(s) == "custom") { return CUSTOM; } else { - raise::Error(fmt::format("Invalid TwoBodyInteraction type: %s", s), HERE); + raise::Error(fmt::format("Invalid TwoBodyInteraction type: %s", s.c_str()), + HERE); return NONE; } } diff --git a/src/global/traits/policies.h b/src/global/traits/policies.h index b2109ffb0..fdca30ef3 100644 --- a/src/global/traits/policies.h +++ b/src/global/traits/policies.h @@ -141,18 +141,37 @@ concept CustomParticleUpdatePolicyClass = namespace traits::twobodyinteractions { template - concept IsValid = requires(const I& interaction_policy, - spidx_t sp1, - npart_t p1, - spidx_t sp2, - npart_t p2, - real_t tile_vol) { - { interaction_policy(sp1, p1, sp2, p2, tile_vol) } -> std::same_as; + concept HasSpecies = requires(I& interaction_policy) { + { interaction_policy.species } -> std::convertible_to; + }; + + template + concept HasShouldInteract = requires(const I& interaction_policy, + spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2, + real_t tile_weight) { + { + interaction_policy.should_interact(sp1, p1, sp2, p2, tile_weight) + } -> std::same_as; + }; + + template + concept HasInteraction = requires(const I& interaction_policy, + spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2) { + { interaction_policy(sp1, p1, sp2, p2) } -> std::same_as; }; } // namespace traits::twobodyinteractions template -concept TwoBodyInteractionPolicyClass = traits::twobodyinteractions::IsValid; +concept TwoBodyInteractionPolicyClass = + traits::twobodyinteractions::HasSpecies and + traits::twobodyinteractions::HasShouldInteract and + traits::twobodyinteractions::HasInteraction; #endif // TRAITS_POLICIES_H diff --git a/src/kernels/injectors.hpp b/src/kernels/injectors.hpp index 05d83a0d4..cf99a698a 100644 --- a/src/kernels/injectors.hpp +++ b/src/kernels/injectors.hpp @@ -5,6 +5,7 @@ * - kernel::UniformInjector_kernel<> * - kernel::GlobalInjector_kernel<> * - kernel::NonUniformInjector_kernel<> + * - kernel::SingleSpeciesUniformInjector_kernel<> * @namespaces: * - kernel:: */ @@ -853,6 +854,131 @@ namespace kernel { } }; // struct NonUniformInjector_kernel + template ED> + struct SingleSpeciesUniformInjector_kernel { + + ParticleArrays particles; + + const npart_t offset; + const npart_t domain_idx, cntr; + const bool use_tracking; + const M metric; + const array_t xi_min, xi_max; + const ED energy_dist; + const real_t inv_V0; + random_number_pool_t random_pool; + + SingleSpeciesUniformInjector_kernel(Particles& particles, + npart_t domain_idx, + const M& metric, + const array_t& xi_min, + const array_t& xi_max, + const ED& energy_dist, + real_t inv_V0, + random_number_pool_t& random_pool) + : particles { particles } + , offset { particles.npart() } + , domain_idx { domain_idx } + , cntr { particles.counter() } + , use_tracking { particles.use_tracking() } + , metric { metric } + , xi_min { xi_min } + , xi_max { xi_max } + , energy_dist { energy_dist } + , inv_V0 { inv_V0 } + , random_pool { random_pool } { + if (use_tracking) { +#if !defined(MPI_ENABLED) + raise::ErrorIf(particles.pld_i.extent(1) < 1, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 1", + HERE); +#else + raise::ErrorIf(particles.pld_i.extent(1) < 2, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 2", + HERE); +#endif + } + } + + Inline void operator()(prtlidx_t p) const { + coord_t x_Cd { ZERO }; + tuple_t xi_Cd { 0 }; + tuple_t dxi_Cd { static_cast(0) }; + vec_t v { ZERO, ZERO, 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::Cartesian) { + energy_dist(x_Ph, v); + } 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(x_Ph, v_Ph); + metric.template transform_xyz(x_Cd_, v_Ph, v); + } else if constexpr (S == SimEngine::GRPIC) { + vec_t v_Ph { ZERO, ZERO, ZERO }; + energy_dist(x_Ph, v_Ph); + metric.template transform(x_Cd, v_Ph, v); + } else { + raise::KernelError(HERE, "Unknown simulation engine"); + } + } + real_t weight = ONE; + if constexpr (M::CoordType != Coord::Cartesian) { + const auto sqrt_det_h = metric.sqrt_det_h(x_Cd); + weight = sqrt_det_h * inv_V0; + } + // clang-format off + if (not use_tracking) { + InjectParticle( + p + offset, + particles.i1, particles.i2, particles.i3, + particles.dx1, particles.dx2, particles.dx3, + particles.ux1, particles.ux2, particles.ux3, + particles.phi, particles.weight, particles.tag, particles.pld_i, + xi_Cd, dxi_Cd, v, weight, ZERO); + } else { + InjectParticle( + p + offset, + particles.i1, particles.i2, particles.i3, + particles.dx1, particles.dx2, particles.dx3, + particles.ux1, particles.ux2, particles.ux3, + particles.phi, particles.weight, particles.tag, particles.pld_i, + xi_Cd, dxi_Cd, v, weight, ZERO, + domain_idx, cntr + p); + } + // clang-format on + } + }; // struct SingleSpeciesUniformInjector_kernel + } // namespace kernel #endif // KERNELS_INJECTORS_HPP diff --git a/src/kernels/twobody_interactions.hpp b/src/kernels/twobody_interactions.hpp index 44385e885..e0a41b1fd 100644 --- a/src/kernels/twobody_interactions.hpp +++ b/src/kernels/twobody_interactions.hpp @@ -73,32 +73,32 @@ namespace kernel::mink { } template - Inline auto TileIdxToVolume(ncells_t tile_idx, - ncells_t tile_size, - ncells_t ntx2, - ncells_t ntx3, - ncells_t nx1, - ncells_t nx2, - ncells_t nx3) -> real_t { - real_t tile_volume { ONE }; + Inline auto NCellsOnTile(ncells_t tile_idx, + ncells_t tile_size, + ncells_t ntx2, + ncells_t ntx3, + ncells_t nx1, + ncells_t nx2, + ncells_t nx3) -> ncells_t { + ncells_t ncells_on_tile { 1 }; ncells_t ti { 0u }, tj { 0u }, tk { 0u }; UnravelTileIdx(tile_idx, ntx2, ntx3, ti, tj, tk); if constexpr ((D == Dim::_1D) or (D == Dim::_2D) or (D == Dim::_3D)) { const auto i1_min = ti * tile_size; const auto i1_max = math::min(i1_min + tile_size, nx1); - tile_volume *= static_cast(i1_max - i1_min); + ncells_on_tile *= (i1_max - i1_min); } if constexpr ((D == Dim::_2D) or (D == Dim::_3D)) { const auto i2_min = tj * tile_size; const auto i2_max = math::min(i2_min + tile_size, nx2); - tile_volume *= static_cast(i2_max - i2_min); + ncells_on_tile *= (i2_max - i2_min); } if constexpr (D == Dim::_3D) { const auto i3_min = tk * tile_size; const auto i3_max = math::min(i3_min + tile_size, nx3); - tile_volume *= static_cast(i3_max - i3_min); + ncells_on_tile *= (i3_max - i3_min); } - return tile_volume; + return ncells_on_tile; } template @@ -112,11 +112,10 @@ namespace kernel::mink { ncells_t num_tiles { 0u }; - CollisionGroup( - const std::vector*>& particles, - const std::vector& ncells, - ncells_t tile_size, - random_number_pool_t& random_pool) { + CollisionGroup(const std::vector*>& particles, + const std::vector& ncells, + ncells_t tile_size, + random_number_pool_t& random_pool) { for (const auto* species : particles) { const auto npart_s = species->npart(); array_t tileidx { "tile_idx", npart_s }; @@ -215,13 +214,14 @@ namespace kernel::mink { template void TwoBodyInteraction( - const std::vector*>& species1, - const std::vector*>& species2, - const std::vector& ncells, - const boundaries_t& domain_extent, - ncells_t tile_size, - random_number_pool_t& random_pool, - const I& interaction_policy) { + const std::vector*>& species1, + const std::vector*>& species2, + const std::vector& ncells, + const boundaries_t& domain_extent, + ncells_t tile_size, + real_t ppc0, + random_number_pool_t& random_pool, + I& interaction_policy) { raise::ErrorIf(species1.empty() or species2.empty(), "species groups must be non-empty", HERE); @@ -231,7 +231,7 @@ namespace kernel::mink { raise::ErrorIf(domain_extent.size() != static_cast(D), "domain_extent size must match D", HERE); - // compute base tile volume in physical units + // compute base cell volume in physical units real_t cell_volume { ONE }; for (int d = 0; d < static_cast(D); ++d) { cell_volume *= static_cast( @@ -253,6 +253,56 @@ namespace kernel::mink { const auto& tile_offsets1 = group1.tile_offsets; const auto& tile_offsets2 = group2.tile_offsets; + // fill species in the interaction policy + for (auto& sp1 : species1) { + interaction_policy.species[sp1->sp - 1] = static_cast( + *sp1); + } + for (auto& sp2 : species2) { + interaction_policy.species[sp2->sp - 1] = static_cast( + *sp2); + } + + // total particle weight on each tile + auto weights_on_tile1 = array_t { "weights_on_tile1", num_tiles }; + auto weights_on_tile2 = array_t { "weights_on_tile2", num_tiles }; + Kokkos::parallel_for( + "ComputeWeightsOnTiles", + Kokkos::TeamPolicy<>(num_tiles, Kokkos::AUTO), + Lambda(const Kokkos::TeamPolicy<>::member_type& team) { + const ncells_t t = team.league_rank(); + + const auto o1 = tile_offsets1(t); + const auto num_ppt1 = combined_num_ppt1(t); + real_t weight_on_tile1 = ZERO; + Kokkos::parallel_reduce( + Kokkos::TeamThreadRange(team, num_ppt1), + [&](prtlidx_t i, real_t& lsum) { + const auto sp1 = static_cast(combined_idx1(o1 + i) >> 56); + const auto p1 = static_cast(combined_idx1(o1 + i) & + ((1ull << 56) - 1)); + + lsum += interaction_policy.species[sp1 - 1].weight(p1); + }, + weight_on_tile1); + weights_on_tile1(t) = weight_on_tile1; + + const auto o2 = tile_offsets2(t); + const auto num_ppt2 = combined_num_ppt2(t); + real_t weight_on_tile2 = ZERO; + Kokkos::parallel_reduce( + Kokkos::TeamThreadRange(team, num_ppt2), + [&](prtlidx_t i, real_t& lsum) { + const auto sp2 = static_cast(combined_idx2(o2 + i) >> 56); + const auto p2 = static_cast(combined_idx2(o2 + i) & + ((1ull << 56) - 1)); + + lsum += interaction_policy.species[sp2 - 1].weight(p2); + }, + weight_on_tile2); + weights_on_tile2(t) = weight_on_tile2; + }); + // number of cells in each direction ncells_t nx1 { 1u }, nx2 { 1u }, nx3 { 1u }; if constexpr ((D == Dim::_1D) or (D == Dim::_2D) or (D == Dim::_3D)) { @@ -274,13 +324,18 @@ namespace kernel::mink { math::ceil(static_cast(nx3) / static_cast(tile_size))); } + array_t interaction_pairs { "interaction_pairs", + combined_idx1.extent(0) }; + array_t counter { "counter" }; + Kokkos::parallel_for( - "EmitPairs", + "PopulateInteractionPairs", Kokkos::TeamPolicy<>(num_tiles, Kokkos::AUTO), Lambda(const Kokkos::TeamPolicy<>::member_type& team) { const ncells_t t = team.league_rank(); - const auto tile_volume = - TileIdxToVolume(t, tile_size, ntx2, ntx3, nx1, nx2, nx3) * cell_volume; + const auto tile_weight = + math::max(weights_on_tile1(t), weights_on_tile2(t)) / + (NCellsOnTile(t, tile_size, ntx2, ntx3, nx1, nx2, nx3) * ppc0); const auto k = math::min(combined_num_ppt1(t), combined_num_ppt2(t)); const auto o1 = tile_offsets1(t); @@ -295,9 +350,27 @@ namespace kernel::mink { ((1ull << 56) - 1)); const auto p2 = static_cast(combined_idx2(o2 + i) & ((1ull << 56) - 1)); - interaction_policy(sp1, p1, sp2, p2, tile_volume); + if (interaction_policy.should_interact(sp1, p1, sp2, p2, tile_weight)) { + const auto idx = Kokkos::atomic_fetch_add(&counter(), 1); + interaction_pairs(idx, 0) = combined_idx1(o1 + i); + interaction_pairs(idx, 1) = combined_idx2(o2 + i); + } }); }); + auto counter_h = Kokkos::create_mirror_view(counter); + Kokkos::deep_copy(counter_h, counter); + Kokkos::parallel_for( + "ProcessInteractions", + counter_h(), + Lambda(prtlidx_t idx) { + const auto sp1 = static_cast(interaction_pairs(idx, 0) >> 56); + const auto sp2 = static_cast(interaction_pairs(idx, 1) >> 56); + const auto p1 = static_cast(interaction_pairs(idx, 0) & + ((1ull << 56) - 1)); + const auto p2 = static_cast(interaction_pairs(idx, 1) & + ((1ull << 56) - 1)); + interaction_policy(sp1, p1, sp2, p2); + }); } } // namespace kernel::mink diff --git a/tests/archetypes/qed_compton.cpp b/tests/archetypes/qed_compton.cpp index 265439a1b..14350cbf4 100644 --- a/tests/archetypes/qed_compton.cpp +++ b/tests/archetypes/qed_compton.cpp @@ -160,13 +160,14 @@ auto main(int argc, char* argv[]) -> int { { ZERO, ONE }, { -ONE, ONE } }; + const auto ppc0 = static_cast(npart) / (nx1 * nx2); prm::Parameters params; - params.set("qed.compton_scattering.nominal_probability_density", + params.set("compton_scattering.nominal_probability_density", static_cast(1e-3)); - params.set("qed.compton_scattering.Thomson_limit", static_cast(1e-4)); - auto policy = arch::qed::ComptonScattering(params, random_pool); + auto policy = arch::qed::ComptonScattering(params, + random_pool); policy.species[0] = static_cast(sp1); policy.species[1] = static_cast(sp2); policy.species[2] = static_cast(sp3); @@ -198,6 +199,7 @@ auto main(int argc, char* argv[]) -> int { ncells, extent, tile_size, + ppc0, random_pool, policy); } diff --git a/tests/kernels/twobody_interactions.cpp b/tests/kernels/twobody_interactions.cpp index 0ce562de3..8051c5805 100644 --- a/tests/kernels/twobody_interactions.cpp +++ b/tests/kernels/twobody_interactions.cpp @@ -4,7 +4,6 @@ #include "global.h" #include "arch/kokkos_aliases.h" -#include "utils/comparators.h" #include "utils/error.h" #include "framework/containers/particles.h" @@ -19,52 +18,33 @@ using namespace ntt; // Verifies that each paired particle from group1 and group2 lies in the same tile struct SameTilePolicy { - const array_t i1_1, i2_1; - const array_t i1_2, i2_2; - const array_t i1_3, i2_3; - const array_t i1_4, i2_4; - const ncells_t tile_size; - const ncells_t ncx1, ncx2; // number of cells in each direction - const ncells_t ntx1, ntx2; // numbers of tiles - array_t diff_tile_errors { "diff_tile_errors" }; - array_t tile_vol_errors { "tile_vol_errors" }; - - SameTilePolicy(const array_t& i1_1, - const array_t& i2_1, - const array_t& i1_2, - const array_t& i2_2, - const array_t& i1_3, - const array_t& i2_3, - const array_t& i1_4, - const array_t& i2_4, - ncells_t tile_size, - ncells_t ncx1, - ncells_t ncx2, - ncells_t ntx1, - ncells_t ntx2) - : i1_1 { i1_1 } - , i2_1 { i2_1 } - , i1_2 { i1_2 } - , i2_2 { i2_2 } - , i1_3 { i1_3 } - , i2_3 { i2_3 } - , i1_4 { i1_4 } - , i2_4 { i2_4 } - , tile_size { tile_size } + ParticleArrays species[4]; + const ncells_t tile_size; + const ncells_t ncx1, ncx2; // number of cells in each direction + const ncells_t ntx1, ntx2; // numbers of tiles + array_t diff_tile_errors { "diff_tile_errors" }; + + SameTilePolicy(ncells_t tile_size, + ncells_t ncx1, + ncells_t ncx2, + ncells_t ntx1, + ncells_t ntx2) + : tile_size { tile_size } , ncx1 { ncx1 } , ncx2 { ncx2 } , ntx1 { ntx1 } , ntx2 { ntx2 } {} - Inline void operator()(spidx_t sp1, - npart_t p1, - spidx_t sp2, - npart_t p2, - real_t tile_volume) const { - const auto x1_1 = (sp1 == 1u) ? i1_1(p1) : i1_2(p1); - const auto x2_1 = (sp1 == 1u) ? i2_1(p1) : i2_2(p1); - const auto x1_2 = (sp2 == 3u) ? i1_3(p2) : i1_4(p2); - const auto x2_2 = (sp2 == 3u) ? i2_3(p2) : i2_4(p2); + Inline auto should_interact(spidx_t, npart_t, spidx_t, npart_t, real_t) const + -> bool { + return true; + } + + Inline void operator()(spidx_t sp1, npart_t p1, spidx_t sp2, npart_t p2) const { + const auto x1_1 = species[sp1 - 1].i1(p1); + const auto x2_1 = species[sp1 - 1].i2(p1); + const auto x1_2 = species[sp2 - 1].i1(p2); + const auto x2_2 = species[sp2 - 1].i2(p2); const auto t1 = static_cast(x1_1 / tile_size) * ntx2 + static_cast(x2_1 / tile_size); const auto t2 = static_cast(x1_2 / tile_size) * ntx2 + @@ -72,37 +52,6 @@ struct SameTilePolicy { if (t1 != t2) { Kokkos::atomic_add(&diff_tile_errors(), 1); } - - real_t vol1 { ONE }, vol2 { ONE }; - { - const auto ti1 = t1 / ntx2; - const auto tj1 = t1 % ntx2; - const auto i1_min_1 = ti1 * tile_size; - const auto i1_max_1 = math::min(i1_min_1 + tile_size, ncx1); - const auto i2_min_1 = tj1 * tile_size; - const auto i2_max_1 = math::min(i2_min_1 + tile_size, ncx2); - - vol1 *= static_cast(i1_max_1 - i1_min_1); - vol1 *= static_cast(i2_max_1 - i2_min_1); - } - { - const auto ti2 = t2 / ntx2; - const auto tj2 = t2 % ntx2; - const auto i1_min_2 = ti2 * tile_size; - const auto i1_max_2 = math::min(i1_min_2 + tile_size, ncx1); - const auto i2_min_2 = tj2 * tile_size; - const auto i2_max_2 = math::min(i2_min_2 + tile_size, ncx2); - - vol2 *= static_cast(i1_max_2 - i1_min_2); - vol2 *= static_cast(i2_max_2 - i2_min_2); - } - vol1 *= SQR(0.03125); - vol2 *= SQR(0.03125); - - if (not cmp::AlmostEqual(tile_volume, vol1) or - not cmp::AlmostEqual(tile_volume, vol2)) { - Kokkos::atomic_add(&tile_vol_errors(), 1); - } } }; @@ -134,7 +83,11 @@ auto main(int argc, char* argv[]) -> int { const ncells_t nx2 = 64u; const ncells_t tile_size = 3u; const std::vector ncells = { nx1, nx2 }; - const ncells_t ntx1 = static_cast( + const boundaries_t extent = { + { ZERO, ONE }, + { ZERO, TWO } + }; + const ncells_t ntx1 = static_cast( math::ceil(static_cast(nx1) / static_cast(tile_size))); const ncells_t ntx2 = static_cast( math::ceil(static_cast(nx2) / static_cast(tile_size))); @@ -199,25 +152,21 @@ auto main(int argc, char* argv[]) -> int { fill_random(sp->i1, sp->i2, sp->tag, npart, nx1, nx2, random_pool); } - const std::vector*> group1 = { &sp1, - &sp2 }; - const std::vector*> group2 = { &sp3, - &sp4 }; - - auto policy = SameTilePolicy { sp1.i1, sp1.i2, sp2.i1, sp2.i2, sp3.i1, - sp3.i2, sp4.i1, sp4.i2, tile_size, nx1, - nx2, ntx1, ntx2 }; - - kernel::mink::TwoBodyInteraction(group1, - group2, - ncells, - { - { ZERO, ONE }, - { ZERO, TWO } - }, - tile_size, - random_pool, - policy); + const std::vector*> group1 = { &sp1, + &sp2 }; + const std::vector*> group2 = { &sp3, + &sp4 }; + + auto policy = SameTilePolicy { tile_size, nx1, nx2, ntx1, ntx2 }; + + kernel::mink::TwoBodyInteraction(group1, + group2, + ncells, + extent, + tile_size, + ONE, + random_pool, + policy); Kokkos::fence(); { @@ -228,12 +177,6 @@ auto main(int argc, char* argv[]) -> int { HERE); } - { - auto errors_h = Kokkos::create_mirror_view(policy.tile_vol_errors); - Kokkos::deep_copy(errors_h, policy.tile_vol_errors); - raise::ErrorIf(errors_h() != 0, "tile volume errors detected", HERE); - } - } catch (std::exception& e) { std::cerr << e.what() << '\n'; ntt::GlobalFinalize(); From d70b3932ef4d428146a846e2fb313d59a3c9f849 Mon Sep 17 00:00:00 2001 From: hayk Date: Tue, 12 May 2026 01:47:27 -0400 Subject: [PATCH 06/25] twobody interactions as input --- input.example.toml | 23 ++ src/archetypes/qed/compton.h | 285 +++++++++++++++++++++++++ src/engines/srpic/twobody.h | 8 + src/global/enums.h | 23 ++ src/kernels/twobody_interactions.hpp | 305 +++++++++++++++++++++++++++ tests/archetypes/CMakeLists.txt | 3 +- tests/archetypes/qed_compton.cpp | 278 ++++++++++++++++++++++++ 7 files changed, 924 insertions(+), 1 deletion(-) create mode 100644 src/archetypes/qed/compton.h create mode 100644 src/engines/srpic/twobody.h create mode 100644 src/kernels/twobody_interactions.hpp create mode 100644 tests/archetypes/qed_compton.cpp diff --git a/input.example.toml b/input.example.toml index f890143a6..aa68d28ff 100644 --- a/input.example.toml +++ b/input.example.toml @@ -272,6 +272,29 @@ # @from: `.gamma_qed` # @value: `(1 / gamma_qed)^2` +[two_body] + # Nominal Thomson optical depth: `tau = n0 * sigma_T * 1` over a distance of 1 in physical units (n0 = nominal density) + # @type: float + # @default: 1.0 + thomson_optical_depth = "" + + [[two_body.interaction]] + # Type of the two-body interaction + # @required + # @type: string + # @enum: "Compton" + type = "" + # First group of species indices participating in the interaction + # @required + # @type: array + # @note: array indexing starting at 1 + group1 = "" + # Second group of species indices participating in the interaction + # @type: array + # @note: array indexing starting at 1 + # @note: For interactions between particles of the same group, leave `group2` empty + group2 = "" + [algorithms] # Number of current smoothing passes # @type: ushort [>= 0] diff --git a/src/archetypes/qed/compton.h b/src/archetypes/qed/compton.h new file mode 100644 index 000000000..b54c268bc --- /dev/null +++ b/src/archetypes/qed/compton.h @@ -0,0 +1,285 @@ +/** + * @file archetypes/qed/compton.h + * @brief Two-body collision policy of Compton scattering between leptons and photons + * @implements + * - arch::qed::ComptonScattering<> + * @namespaces: + * - arch::qed:: + */ +#ifndef ARCHETYPES_QED_COMPTON_H +#define ARCHETYPES_QED_COMPTON_H + +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "utils/comparators.h" +#include "utils/error.h" +#include "utils/numeric.h" +#include "utils/param_container.h" + +#include "framework/containers/particles.h" + +#include + +namespace arch::qed { + using namespace ntt; + + template + struct ComptonScattering { + static constexpr spidx_t MAXSP = 16u; + static constexpr int MAX_ITER = 10; + + ParticleArrays species[MAXSP]; + static constexpr real_t low_energy_limit = static_cast(2e-3); + + const real_t nominal_probability_density; + const real_t Thomson_limit; + random_number_pool_t random_pool; + + ComptonScattering(const prm::Parameters& params, + random_number_pool_t& random_pool) + : nominal_probability_density { params.template get( + "qed.compton_scattering.nominal_probability_density") } + , Thomson_limit { params.template get( + "qed.compton_scattering.Thomson_limit") } + , random_pool { random_pool } { + if (nominal_probability_density <= ZERO) { + raise::Error("nominal_probability must be in the range (0, 1]", HERE); + } + if (Thomson_limit <= ZERO or Thomson_limit > low_energy_limit) { + raise::Error( + "Thomson_limit must be in the range (0, small_energy_limit]", + HERE); + } + } + + /* + * Lorentz boost a 4-momentum p of the photon to the frame moving with u + * @param u: 4-velocity of the boost frame + * @param gamma: Lorentz factor of the boost frame + * @param p: 4-momentum of the photon in the lab frame + * @param e: energy of the photon in the lab frame + * @return: 4-momentum of the photon in the boost frame + * @return: energy of the photon in the boost frame + */ + Inline void LorentzBoost(const vec_t& u, + real_t gamma, + const vec_t& p, + real_t e, + vec_t& p_, + real_t& e_) const { + const auto u_dot_p = DOT(u[0], u[1], u[2], p[0], p[1], p[2]); + + e_ = gamma * e - u_dot_p; + p_[0] = p[0] + (u_dot_p / (ONE + gamma) - e) * u[0]; + p_[1] = p[1] + (u_dot_p / (ONE + gamma) - e) * u[1]; + p_[2] = p[2] + (u_dot_p / (ONE + gamma) - e) * u[2]; + } + + /* + * Calculate the Klein-Nishina cross section for a photon with energy e_ in the lepton rest frame + * @param e_: photon energy in the lepton rest frame + * @return: pair of (is_KN_regime, f_KN) where + * - is_KN_regime: whether the photon energy is in the Klein-Nishina regime (e_ > Thomson_limit) + * - f_KN: the Klein-Nishina cross section normalized to the Thomson cross section + * @note for e_ > low_energy_limit, full Klein-Nishina formula + * @note for Thomson_limit < e_ <= low_energy_limit, 2nd order expansion of the Klein-Nishina formula + * @note for e_ <= Thomson_limit, return 1 (Thomson limit) + */ + Inline auto KNCrossSection(real_t e_) const -> Kokkos::pair { + if (e_ > Thomson_limit) { + if (e_ < low_energy_limit) { + // correctly handle the e_ << 1 limit using 2nd order expansion of f_KN + return { true, ONE - TWO * e_ + static_cast(5.2) * SQR(e_) }; + } else { + return { true, + static_cast(0.375) * + ((ONE - TWO / e_ - TWO / SQR(e_)) * math::log(ONE + TWO * e_) + + HALF + FOUR / e_ - HALF / SQR(ONE + TWO * e_)) / + e_ }; + } + } else { + return { false, ONE }; + } + } + + Inline auto RandomCosTheta_Th() const -> real_t { + auto gen_ = random_pool.get_state(); + const auto rnd_ = Random(gen_); + random_pool.free_state(gen_); + const auto u = math::pow( + FOUR * rnd_ - TWO + + math::sqrt(FIVE + static_cast(16) * rnd_ * (rnd_ - ONE)), + THIRD); + return u - ONE / u; + } + + Inline auto RandomCosTheta_KN(double e_) const -> real_t { + auto gen_ = random_pool.get_state(); + const auto rnd_ = Random(gen_); + random_pool.free_state(gen_); + + auto u = 2.0 * rnd_ - 1.0; + bool converged = false; + for (int iter = 0; iter < MAX_ITER; ++iter) { + const auto CDF = (-((2.0 + e_ * (4.0 + e_ - 4.0 * (-1.0 + u) * u * e_ + + 2.0 * CUBE(-1.0 + u) * SQR(e_))) / + SQR(1.0 + e_ - u * e_)) + + (2.0 + e_ * (4.0 - e_ * (7.0 + 16.0 * e_))) / + SQR(1.0 + 2.0 * e_) + + 2.0 * (-2.0 + (-2.0 + e_) * e_) * + math::log((1.0 + e_ - u * e_) / (1.0 + 2.0 * e_))) / + ((-4.0 * e_ * (2.0 + e_ * (1.0 + e_) * (8.0 + e_))) / + SQR(1.0 + 2.0 * e_) + + (4.0 - 2.0 * (-2.0 + e_) * e_) * + math::log(1.0 + 2.0 * e_)); + const auto dCDF_du = -((CUBE(e_) * SQR(1.0 + 2.0 * e_) * + (1.0 + SQR(u) - (-1.0 + u) * (1.0 + SQR(u)) * e_ + + SQR(-1.0 + u) * SQR(e_))) / + (CUBE(-1.0 + (-1.0 + u) * e_) * + (2.0 * e_ * (2.0 + e_ * (1.0 + e_) * (8.0 + e_)) + + SQR(1.0 + 2.0 * e_) * (-2.0 + (-2.0 + e_) * e_) * + math::log(1.0 + 2.0 * e_)))); + + const auto du = (rnd_ - CDF) / dCDF_du; + + u += du; + if (u > 1.0) { + u = 1.0; + } else if (u < -1.0) { + u = -1.0; + } + if (math::abs(du) < 1e-3) { + converged = true; + break; + } + } // iterative loop for u + return static_cast(u); + } + + /* + * Scatter a photon with initial momentum p_ and energy e_ in the lepton + * rest frame to a new momentum pnew_ and energy enew_ + * @param KN_regime: whether the photon energy is in the Klein-Nishina regime + * @param p_: initial photon momentum in the lepton rest frame + * @param e_: initial photon energy in the lepton rest frame + * @return pnew_: output photon momentum after scattering in the lepton rest frame + * @return enew_: output photon energy after scattering in the lepton rest frame + * @note the scattering angle is sampled from the Klein-Nishina differential + * cross section if KN_regime is true, otherwise it is sampled from the Thomson limit + */ + Inline void ScatterPhoton(bool KN_regime, + const vec_t& p_, + real_t e_, + vec_t& pnew_, + real_t& enew_) const { + auto rand_costheta_ { ZERO }; + if (not KN_regime) { + rand_costheta_ = RandomCosTheta_Th(); + } else { + rand_costheta_ = RandomCosTheta_KN(e_); + } + const auto rand_sintheta_ = math::sqrt(ONE - SQR(rand_costheta_)); + + auto gen_ = random_pool.get_state(); + const auto rand_phi_ = static_cast(constant::TWO_PI) * + Random(gen_); + random_pool.free_state(gen_); + const auto rand_cosphi_ = math::cos(rand_phi_); + const auto rand_sinphi_ = math::sin(rand_phi_); + + // Define an orthonormal basis: {a_, b_, c_} in the lepton frame + const vec_t a_ { p_[0] / e_, p_[1] / e_, p_[2] / e_ }; + vec_t b_ { ONE, ZERO, ZERO }; + if (not cmp::AlmostZero(a_[0])) { + b_[0] = -a_[1] / a_[0]; + b_[1] = ONE / math::sqrt(ONE + SQR(b_[0])); + b_[0] /= math::sqrt(ONE + SQR(b_[0])); + } + const vec_t c_ { + CROSS_x1(a_[0], a_[1], a_[2], b_[0], b_[1], b_[2]), + CROSS_x2(a_[0], a_[1], a_[2], b_[0], b_[1], b_[2]), + CROSS_x3(a_[0], a_[1], a_[2], b_[0], b_[1], b_[2]) + }; + + enew_ = e_ / (ONE + e_ * (ONE - rand_costheta_)); + + pnew_[0] = enew_ * (rand_costheta_ * a_[0] + + rand_sintheta_ * rand_cosphi_ * b_[0] + + rand_sintheta_ * rand_sinphi_ * c_[0]); + pnew_[1] = enew_ * (rand_costheta_ * a_[1] + + rand_sintheta_ * rand_cosphi_ * b_[1] + + rand_sintheta_ * rand_sinphi_ * c_[1]); + pnew_[2] = enew_ * (rand_costheta_ * a_[2] + + rand_sintheta_ * rand_cosphi_ * b_[2] + + rand_sintheta_ * rand_sinphi_ * c_[2]); + } + + Inline void operator()(spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2, + real_t tile_volume) const { + // @TODO coord/vec conversion + // values with "_" are in the lepton rest-frame + const vec_t lepton_u { species[sp1 - 1].ux1(p1), + species[sp1 - 1].ux2(p1), + species[sp1 - 1].ux3(p1) }; + const auto lepton_gamma = U2GAMMA(lepton_u[0], lepton_u[1], lepton_u[2]); + const auto lepton_weight = species[sp1 - 1].weight(p1); + + const vec_t photon_p { species[sp2 - 1].ux1(p2), + species[sp2 - 1].ux2(p2), + species[sp2 - 1].ux3(p2) }; + const auto photon_energy = NORM(photon_p[0], photon_p[1], photon_p[2]); + const auto photon_weight = species[sp2 - 1].weight(p2); + + // boost photon momentum to lepton rest frame + vec_t photon_p_ { ZERO, ZERO, ZERO }; + real_t photon_energy_ { ZERO }; + + LorentzBoost(lepton_u, lepton_gamma, photon_p, photon_energy, photon_p_, photon_energy_); + + const auto [KN_regime, f_KN] = KNCrossSection(photon_energy_); + const auto scattering_probability = nominal_probability_density * f_KN * + photon_energy_ * lepton_weight * + photon_weight / + (photon_energy * lepton_gamma * + tile_volume); + auto gen = random_pool.get_state(); + const auto rnd = Random(gen); + random_pool.free_state(gen); + + if (rnd < scattering_probability) { + vec_t photon_pnew_ { ZERO, ZERO, ZERO }, + photon_pnew { ZERO, ZERO, ZERO }; + real_t photon_energy_new_ { ZERO }, photon_energy_new { ZERO }; + + ScatterPhoton(KN_regime, + photon_p_, + photon_energy_, + photon_pnew_, + photon_energy_new_); + LorentzBoost({ -lepton_u[0], -lepton_u[1], -lepton_u[2] }, + lepton_gamma, + photon_pnew_, + photon_energy_new_, + photon_pnew, + photon_energy_new); + species[sp1 - 1].ux1(p1) += (photon_p[0] - photon_pnew[0]) * + photon_weight / lepton_weight; + species[sp1 - 1].ux2(p1) += (photon_p[1] - photon_pnew[1]) * + photon_weight / lepton_weight; + species[sp1 - 1].ux3(p1) += (photon_p[2] - photon_pnew[2]) * + photon_weight / lepton_weight; + + species[sp2 - 1].ux1(p2) = photon_pnew[0]; + species[sp2 - 1].ux2(p2) = photon_pnew[1]; + species[sp2 - 1].ux3(p2) = photon_pnew[2]; + } // not interacting + } + }; + +} // namespace arch::qed + +#endif // ARCHETYPES_QED_COMPTON_H diff --git a/src/engines/srpic/twobody.h b/src/engines/srpic/twobody.h new file mode 100644 index 000000000..a6e034c24 --- /dev/null +++ b/src/engines/srpic/twobody.h @@ -0,0 +1,8 @@ +#ifndef ENGINES_SRPIC_TWOBODY_H +#define ENGINES_SRPIC_TWOBODY_H + +namespace ntt { + namespace srpic {} +} // namespace ntt + +#endif // ENGINES_SRPIC_TWOBODY_H \ No newline at end of file diff --git a/src/global/enums.h b/src/global/enums.h index 68eefb75b..41fdd551b 100644 --- a/src/global/enums.h +++ b/src/global/enums.h @@ -401,6 +401,29 @@ namespace ntt { using EmissionTypeFlag = uint8_t; + namespace TwoBodyInteraction { + enum TwoBodyInteractionFlag_ : uint8_t { + NONE = 0, + COMPTON = 1, + CUSTOM = 2, + }; + + inline auto to_string(uint8_t flags) -> std::string { + switch (flags) { + case NONE: + return "none"; + case COMPTON: + return "compton"; + case CUSTOM: + return "custom"; + default: + return "unknown"; + } + } + } // namespace TwoBodyInteraction + + using TwoBodyInteractionFlag = uint8_t; + } // namespace ntt #endif // GLOBAL_ENUMS_H diff --git a/src/kernels/twobody_interactions.hpp b/src/kernels/twobody_interactions.hpp new file mode 100644 index 000000000..44385e885 --- /dev/null +++ b/src/kernels/twobody_interactions.hpp @@ -0,0 +1,305 @@ +/** + * @file kernels/twobody_interactions.hpp + * @brief Generic two-body interaction kernel that can be used to implement various + * types of collisions between species, e.g. Compton scattering, Breit-Wheeler pair production, etc. + * @implements + * - kernel::mink::TwoBodyInteraction<> + * @namespaces: + * - arch::mink:: + */ +#ifndef KERNELS_TWOBODY_INTERACTIONS_HPP +#define KERNELS_TWOBODY_INTERACTIONS_HPP + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "traits/policies.h" +#include "utils/error.h" +#include "utils/sorting.h" + +#include "framework/containers/particles.h" + +#include +#include + +#include +#include + +namespace kernel::mink { + using namespace ntt; + + namespace { + struct CollisionSpecies { + const spidx_t sp; + const npart_t npart; + ncells_t num_tiles { 0u }; + + array_t tileidx; + array_t num_ppt; + + CollisionSpecies(spidx_t sp, + npart_t npart, + const array_t& tileidx, + const array_t& num_ppt, + ncells_t num_tiles) + : sp { sp } + , npart { npart } + , tileidx { tileidx } + , num_ppt { num_ppt } + , num_tiles { num_tiles } {} + }; + + template + Inline void UnravelTileIdx(ncells_t tile_idx, + ncells_t ntx2, + ncells_t ntx3, + ncells_t& ti, + ncells_t& tj, + ncells_t& tk) { + if constexpr (D == Dim::_1D) { + ti = tile_idx; + } else if constexpr (D == Dim::_2D) { + ti = tile_idx / ntx2; + tj = tile_idx % ntx2; + } else if constexpr (D == Dim::_3D) { + ti = tile_idx / (ntx2 * ntx3); + const auto rem = tile_idx % (ntx2 * ntx3); + tj = rem / ntx3; + tk = rem % ntx3; + } else { + raise::KernelError(HERE, "Wrong D in TileIdxUnravel"); + } + } + + template + Inline auto TileIdxToVolume(ncells_t tile_idx, + ncells_t tile_size, + ncells_t ntx2, + ncells_t ntx3, + ncells_t nx1, + ncells_t nx2, + ncells_t nx3) -> real_t { + real_t tile_volume { ONE }; + ncells_t ti { 0u }, tj { 0u }, tk { 0u }; + UnravelTileIdx(tile_idx, ntx2, ntx3, ti, tj, tk); + if constexpr ((D == Dim::_1D) or (D == Dim::_2D) or (D == Dim::_3D)) { + const auto i1_min = ti * tile_size; + const auto i1_max = math::min(i1_min + tile_size, nx1); + tile_volume *= static_cast(i1_max - i1_min); + } + if constexpr ((D == Dim::_2D) or (D == Dim::_3D)) { + const auto i2_min = tj * tile_size; + const auto i2_max = math::min(i2_min + tile_size, nx2); + tile_volume *= static_cast(i2_max - i2_min); + } + if constexpr (D == Dim::_3D) { + const auto i3_min = tk * tile_size; + const auto i3_max = math::min(i3_min + tile_size, nx3); + tile_volume *= static_cast(i3_max - i3_min); + } + return tile_volume; + } + + template + struct CollisionGroup { + std::vector group; + + array_t combined_idx; + array_t combined_tileidx; + array_t combined_num_ppt; + array_t tile_offsets; + + ncells_t num_tiles { 0u }; + + CollisionGroup( + const std::vector*>& particles, + const std::vector& ncells, + ncells_t tile_size, + random_number_pool_t& random_pool) { + for (const auto* species : particles) { + const auto npart_s = species->npart(); + array_t tileidx { "tile_idx", npart_s }; + auto tile_indexing_kernel = sort::PositionToTileIndex( + species->i1, + species->i2, + species->i3, + species->tag, + tileidx, + ncells, + tile_size); + Kokkos::parallel_for("TileIndexing", species->npart(), tile_indexing_kernel); + group.emplace_back(species->sp, + npart_s, + tileidx, + tile_indexing_kernel.num_ppt, + tile_indexing_kernel.total_tiles); + if (num_tiles == 0u) { + num_tiles = group.back().num_tiles; + } else if (num_tiles != group.back().num_tiles) { + raise::Error("unequal num_tiles across species in group", HERE); + } + raise::ErrorIf(group.back().tileidx.extent(0) != species->npart(), + "tileidx must have the same extent as npart for all " + "species in group", + HERE); + } + + npart_t tot_npart = 0u; + for (const auto& species : group) { + tot_npart += species.npart; + } + + combined_idx = array_t { "combined_idx", tot_npart }; + combined_tileidx = array_t { "combined_tileidx", tot_npart }; + combined_num_ppt = array_t { "combined_num_ppt", num_tiles }; + tile_offsets = array_t { "tile_offsets", num_tiles }; + + { + // combine particle indices in the group & compute total number in each tile + npart_t offset = 0u; + for (const auto& species : group) { + Kokkos::parallel_for( + "CombineInGroup", + species.npart, + ClassLambda(const npart_t p) { + // pack species idx into top 8 bits + prtl index into the remaining 56 bits + combined_idx(offset + p) = (static_cast(species.sp) + << 56) | + static_cast(p); + combined_tileidx(offset + p) = species.tileidx(p); + }); + offset += species.npart; + Kokkos::parallel_for( + "CombineNumPpt", + species.num_tiles, + ClassLambda(const ncells_t t) { + combined_num_ppt(t) += species.num_ppt(t); + }); + Kokkos::fence(); + } + } + { + // randomly shuffle particles within each tile and sort by tiles + array_t shuffle_key { "shuffle_key", tot_npart }; + Kokkos::parallel_for( + "PackRandom", + tot_npart, + ClassLambda(const npart_t p) { + auto gen = random_pool.get_state(); + const auto rnd = static_cast(gen.urand()); + random_pool.free_state(gen); + const auto tile_idx = static_cast(combined_tileidx(p)); + // packing top 32 bits with tile index, and the rest -- random + shuffle_key(p) = (tile_idx << 32) | rnd; + }); + Kokkos::Experimental::sort_by_key(Kokkos::DefaultExecutionSpace {}, + shuffle_key, + combined_idx); + } + { + // compute index offsets for each tile + Kokkos::parallel_scan( + "TileOffsets", + num_tiles, + ClassLambda(cellidx_t t, npart_t & acc, const bool final) { + if (final) { + tile_offsets(t) = acc; + } + acc += combined_num_ppt(t); + }); + } + } + }; + } // namespace + + template + void TwoBodyInteraction( + const std::vector*>& species1, + const std::vector*>& species2, + const std::vector& ncells, + const boundaries_t& domain_extent, + ncells_t tile_size, + random_number_pool_t& random_pool, + const I& interaction_policy) { + raise::ErrorIf(species1.empty() or species2.empty(), + "species groups must be non-empty", + HERE); + raise::ErrorIf(ncells.size() != static_cast(D), + "ncells size must match D", + HERE); + raise::ErrorIf(domain_extent.size() != static_cast(D), + "domain_extent size must match D", + HERE); + // compute base tile volume in physical units + real_t cell_volume { ONE }; + for (int d = 0; d < static_cast(D); ++d) { + cell_volume *= static_cast( + domain_extent[d].second - domain_extent[d].first) / + static_cast(ncells[d]); + } + + const auto group1 = CollisionGroup(species1, ncells, tile_size, random_pool); + const auto group2 = CollisionGroup(species2, ncells, tile_size, random_pool); + raise::ErrorIf(group1.num_tiles != group2.num_tiles, + "number of tiles differ in group1 vs group2", + HERE); + const auto num_tiles = group1.num_tiles; + + const auto& combined_idx1 = group1.combined_idx; + const auto& combined_idx2 = group2.combined_idx; + const auto& combined_num_ppt1 = group1.combined_num_ppt; + const auto& combined_num_ppt2 = group2.combined_num_ppt; + const auto& tile_offsets1 = group1.tile_offsets; + const auto& tile_offsets2 = group2.tile_offsets; + + // number of cells in each direction + ncells_t nx1 { 1u }, nx2 { 1u }, nx3 { 1u }; + if constexpr ((D == Dim::_1D) or (D == Dim::_2D) or (D == Dim::_3D)) { + nx1 = ncells[0]; + } + if constexpr ((D == Dim::_2D) or (D == Dim::_3D)) { + nx2 = ncells[1]; + } + if constexpr (D == Dim::_3D) { + nx3 = ncells[2]; + } + ncells_t ntx2 { 1u }, ntx3 { 1u }; + if constexpr ((D == Dim::_2D) or (D == Dim::_3D)) { + ntx2 = static_cast( + math::ceil(static_cast(nx2) / static_cast(tile_size))); + } + if constexpr (D == Dim::_3D) { + ntx3 = static_cast( + math::ceil(static_cast(nx3) / static_cast(tile_size))); + } + + Kokkos::parallel_for( + "EmitPairs", + Kokkos::TeamPolicy<>(num_tiles, Kokkos::AUTO), + Lambda(const Kokkos::TeamPolicy<>::member_type& team) { + const ncells_t t = team.league_rank(); + const auto tile_volume = + TileIdxToVolume(t, tile_size, ntx2, ntx3, nx1, nx2, nx3) * cell_volume; + + const auto k = math::min(combined_num_ppt1(t), combined_num_ppt2(t)); + const auto o1 = tile_offsets1(t); + const auto o2 = tile_offsets2(t); + Kokkos::parallel_for(Kokkos::TeamThreadRange(team, k), [&](prtlidx_t i) { + // unpack the higher 8 bits + const auto sp1 = static_cast(combined_idx1(o1 + i) >> 56); + const auto sp2 = static_cast(combined_idx2(o2 + i) >> 56); + + // unpack the lower 56 bits + const auto p1 = static_cast(combined_idx1(o1 + i) & + ((1ull << 56) - 1)); + const auto p2 = static_cast(combined_idx2(o2 + i) & + ((1ull << 56) - 1)); + interaction_policy(sp1, p1, sp2, p2, tile_volume); + }); + }); + } + +} // namespace kernel::mink + +#endif // KERNELS_TWOBODY_INTERACTIONS_HPP diff --git a/tests/archetypes/CMakeLists.txt b/tests/archetypes/CMakeLists.txt index 4a5b501e1..4c1aad515 100644 --- a/tests/archetypes/CMakeLists.txt +++ b/tests/archetypes/CMakeLists.txt @@ -16,7 +16,7 @@ function(gen_test title) set(src ${title}.cpp) add_executable(${exec} ${src}) - set(libs ntt_archetypes ntt_global ntt_metrics) + set(libs ntt_archetypes ntt_framework ntt_global ntt_metrics) add_dependencies(${exec} ${libs}) target_link_libraries(${exec} PRIVATE ${libs}) @@ -28,3 +28,4 @@ gen_test(spatial_dist) gen_test(field_setter) gen_test(powerlaw) gen_test(pgen) +gen_test(qed_compton) diff --git a/tests/archetypes/qed_compton.cpp b/tests/archetypes/qed_compton.cpp new file mode 100644 index 000000000..265439a1b --- /dev/null +++ b/tests/archetypes/qed_compton.cpp @@ -0,0 +1,278 @@ +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" + +#include "archetypes/qed/compton.h" +#include "framework/containers/particles.h" +#include "kernels/twobody_interactions.hpp" + +#include + +#include +#include +#include + +using namespace ntt; + +void fill_random(array_t& i1, + array_t& i2, + array_t& ux1, + array_t& ux2, + array_t& ux3, + array_t& weight, + array_t& tag, + npart_t npart, + ncells_t nx1, + ncells_t nx2, + random_number_pool_t& rpool) { + Kokkos::parallel_for( + "FillRandom", + npart, + KOKKOS_LAMBDA(const npart_t p) { + auto gen = rpool.get_state(); + i1(p) = static_cast(gen.urand() % static_cast(nx1)); + i2(p) = static_cast(gen.urand() % static_cast(nx2)); + ux1(p) = Random(gen) * TWO - ONE; + ux2(p) = Random(gen) * TWO - ONE; + ux3(p) = Random(gen) * TWO - ONE; + weight(p) = ONE; + tag(p) = ParticleTag::alive; + rpool.free_state(gen); + }); + Kokkos::fence(); +} + +auto get_total_energy(bool is_massive, + array_t& ux1, + array_t& ux2, + array_t& ux3, + npart_t npart) -> real_t { + real_t total_energy = ZERO; + Kokkos::parallel_reduce( + "TotalEnergy", + npart, + Lambda(const npart_t p, real_t& local_sum) { + if (is_massive) { + local_sum += U2GAMMA(ux1(p), ux2(p), ux3(p)); + + } else { + local_sum += NORM(ux1(p), ux2(p), ux3(p)); + } + }, + total_energy); + return total_energy; +} + +auto get_total_momentum_in(in dir, + array_t& ux1, + array_t& ux2, + array_t& ux3, + npart_t npart) -> real_t { + real_t total_momentum_in = ZERO; + Kokkos::parallel_reduce( + "TotalMomentumIn", + npart, + Lambda(const npart_t p, real_t& local_sum) { + if (dir == in::x1) { + local_sum += ux1(p); + } else if (dir == in::x2) { + local_sum += ux2(p); + } else if (dir == in::x3) { + local_sum += ux3(p); + } + }, + total_momentum_in); + return total_momentum_in; +} + +auto main(int argc, char* argv[]) -> int { + ntt::GlobalInitialize(argc, argv); + + try { + const ncells_t nx1 = 32u; + const ncells_t nx2 = 64u; + const ncells_t tile_size = 3u; + const std::vector ncells = { nx1, nx2 }; + const ncells_t ntx1 = static_cast( + math::ceil(static_cast(nx1) / static_cast(tile_size))); + const ncells_t ntx2 = static_cast( + math::ceil(static_cast(nx2) / static_cast(tile_size))); + const npart_t npart = 1000u; + random_number_pool_t random_pool { 12345u }; + + Particles sp1 { 1u, + "sp1", + 1.0f, + 1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + Particles sp2 { 2u, + "sp2", + 1.0f, + -1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + Particles sp3 { 3u, + "sp3", + 0.0f, + 0.0f, + npart, + 0u, + 0u, + ParticlePusher::PHOTON, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + + for (auto* sp : { &sp1, &sp2, &sp3 }) { + sp->set_npart(npart); + fill_random(sp->i1, + sp->i2, + sp->ux1, + sp->ux2, + sp->ux3, + sp->weight, + sp->tag, + npart, + nx1, + nx2, + random_pool); + } + + boundaries_t extent { + { ZERO, ONE }, + { -ONE, ONE } + }; + + prm::Parameters params; + params.set("qed.compton_scattering.nominal_probability_density", + static_cast(1e-3)); + params.set("qed.compton_scattering.Thomson_limit", static_cast(1e-4)); + + auto policy = arch::qed::ComptonScattering(params, random_pool); + policy.species[0] = static_cast(sp1); + policy.species[1] = static_cast(sp2); + policy.species[2] = static_cast(sp3); + + std::array init_energies { + get_total_energy(true, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_energy(true, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_energy(false, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array init_moms_x1 { + get_total_momentum_in(in::x1, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x1, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x1, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array init_moms_x2 { + get_total_momentum_in(in::x2, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x2, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x2, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array init_moms_x3 { + get_total_momentum_in(in::x3, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x3, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x3, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + + for (int i = 0; i < 1000; ++i) { + kernel::mink::TwoBodyInteraction({ &sp1, &sp2 }, + { &sp3 }, + ncells, + extent, + tile_size, + random_pool, + policy); + } + + { + std::array fin_energies { + get_total_energy(true, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_energy(true, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_energy(false, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array fin_moms_x1 { + get_total_momentum_in(in::x1, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x1, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x1, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array fin_moms_x2 { + get_total_momentum_in(in::x2, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x2, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x2, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + std::array fin_moms_x3 { + get_total_momentum_in(in::x3, sp1.ux1, sp1.ux2, sp1.ux3, sp1.npart()), + get_total_momentum_in(in::x3, sp2.ux1, sp2.ux2, sp2.ux3, sp2.npart()), + get_total_momentum_in(in::x3, sp3.ux1, sp3.ux2, sp3.ux3, sp3.npart()) + }; + + const auto fin_energy = fin_energies[0] + fin_energies[1] + fin_energies[2]; + const auto init_energy = init_energies[0] + init_energies[1] + + init_energies[2]; + const auto fin_mom_x1 = fin_moms_x1[0] + fin_moms_x1[1] + fin_moms_x1[2]; + const auto init_mom_x1 = init_moms_x1[0] + init_moms_x1[1] + init_moms_x1[2]; + const auto fin_mom_x2 = fin_moms_x2[0] + fin_moms_x2[1] + fin_moms_x2[2]; + const auto init_mom_x2 = init_moms_x2[0] + init_moms_x2[1] + init_moms_x2[2]; + const auto fin_mom_x3 = fin_moms_x3[0] + fin_moms_x3[1] + fin_moms_x3[2]; + const auto init_mom_x3 = init_moms_x3[0] + init_moms_x3[1] + init_moms_x3[2]; + + const auto err_energy = (fin_energy - init_energy) / init_energy; + const auto err_mom_x1 = (fin_mom_x1 - init_mom_x1) / + (std::abs(init_mom_x1) + 1e-10); + const auto err_mom_x2 = (fin_mom_x2 - init_mom_x2) / + (std::abs(init_mom_x2) + 1e-10); + const auto err_mom_x3 = (fin_mom_x3 - init_mom_x3) / + (std::abs(init_mom_x3) + 1e-10); + + raise::ErrorIf(err_energy > 1e-5, + fmt::format("energy is not conserved %e -> %e [%e]", + init_energy, + fin_energy, + err_energy), + HERE); + raise::ErrorIf(err_mom_x1 > 1e-5, + fmt::format("x1 momentum is not conserved %e -> %e [%e]", + init_mom_x1, + fin_mom_x1, + err_mom_x1), + HERE); + raise::ErrorIf(err_mom_x2 > 1e-5, + fmt::format("x2 momentum is not conserved %e -> %e [%e]", + init_mom_x2, + fin_mom_x2, + err_mom_x2), + HERE); + raise::ErrorIf(err_mom_x3 > 1e-5, + fmt::format("x3 momentum is not conserved %e -> %e [%e]", + init_mom_x3, + fin_mom_x3, + err_mom_x3), + HERE); + } + + } catch (std::exception& e) { + std::cerr << e.what() << '\n'; + ntt::GlobalFinalize(); + return 1; + } + ntt::GlobalFinalize(); + return 0; +} From f687ac6284f771621a6d5d67e7987e0e90faf5a0 Mon Sep 17 00:00:00 2001 From: hayk Date: Tue, 12 May 2026 02:27:47 -0400 Subject: [PATCH 07/25] qed input params + caller in srpic.hpp --- input.example.toml | 4 ++++ src/engines/engine.hpp | 1 + src/engines/srpic/twobody.h | 37 +++++++++++++++++++++++++++++- src/framework/parameters/extra.cpp | 24 +++++++++++++++++++ src/framework/parameters/extra.h | 12 ++++++++++ src/global/defaults.h | 4 ++++ src/global/enums.h | 13 +++++++++++ 7 files changed, 94 insertions(+), 1 deletion(-) diff --git a/input.example.toml b/input.example.toml index aa68d28ff..20197b08e 100644 --- a/input.example.toml +++ b/input.example.toml @@ -294,6 +294,10 @@ # @note: array indexing starting at 1 # @note: For interactions between particles of the same group, leave `group2` empty group2 = "" + # Interval in timesteps between checking for the interaction + # @type: uint + # @default: 1 + interval = "" [algorithms] # Number of current smoothing passes diff --git a/src/engines/engine.hpp b/src/engines/engine.hpp index b20e163ca..026cd5c50 100644 --- a/src/engines/engine.hpp +++ b/src/engines/engine.hpp @@ -127,6 +127,7 @@ namespace ntt { auto parameters = prm::Parameters {}; parameters.set("dt", static_cast(dt)); parameters.set("time", static_cast(time)); + parameters.set("step", static_cast(step)); return parameters; } }; diff --git a/src/engines/srpic/twobody.h b/src/engines/srpic/twobody.h index a6e034c24..b969604ea 100644 --- a/src/engines/srpic/twobody.h +++ b/src/engines/srpic/twobody.h @@ -1,8 +1,43 @@ #ifndef ENGINES_SRPIC_TWOBODY_H #define ENGINES_SRPIC_TWOBODY_H +#include "enums.h" + +#include "traits/metric.h" +#include "utils/error.h" +#include "utils/log.h" +#include "utils/param_container.h" + +#include "archetypes/qed/compton.h" +#include "framework/domain/domain.h" +#include "framework/parameters/extra.h" +#include "framework/parameters/parameters.h" +#include "kernels/twobody_interactions.hpp" + namespace ntt { - namespace srpic {} + namespace srpic { + + template + void TwoBodyInteractions(Domain& domain, + const prm::Parameters& engine_params, + const SimulationParams& params) { + logger::Checkpoint("Launching TwoBodyInteractions routines", HERE); + const auto dt = engine_params.get("dt"); + const auto step = engine_params.get("step"); + for (const auto& interaction : + params.template get>( + "two_body.interactions")) { + if (step % interaction.interval == 0u) { + if (interaction.type == TwoBodyInteraction::COMPTON) { + + printf("CALLING COMPTON\n"); + } else if (interaction.type == TwoBodyInteraction::CUSTOM) { + raise::Error("Custom two-body interactions not implemented yet", HERE); + } + } + } + } + } // namespace srpic } // namespace ntt #endif // ENGINES_SRPIC_TWOBODY_H \ No newline at end of file diff --git a/src/framework/parameters/extra.cpp b/src/framework/parameters/extra.cpp index dbee39c7c..4d7169700 100644 --- a/src/framework/parameters/extra.cpp +++ b/src/framework/parameters/extra.cpp @@ -114,6 +114,26 @@ namespace ntt { compton_photon_weight.value(); compton_nominal_photon_energy = ONE / SQR(compton_gamma_qed.value()); } + + twobody_thomson_optical_depth = toml::find_or( + toml_data, + "two_body", + "thomson_optical_depth", + defaults::twobody::thomson_optical_depth); + + // find two-body interactions + const auto twobody_tab = toml::find_or(toml_data, + "two_body", + "interactions", + toml::array {}); + for (const auto& tbint : twobody_tab) { + twobody_interactions.push_back(TwoBodyInteractionParams { + .type = TwoBodyInteraction::from_string( + toml::find(tbint, "type")), + .group1 = toml::find>(tbint, "group1"), + .group2 = toml::find_or>(tbint, "group2", {}), + .interval = toml::find_or(tbint, "interval", 1) }); + } } void Extra::setParams(const std::map& extra, @@ -159,6 +179,10 @@ namespace ntt { params->set("radiation.emission.compton.nominal_photon_energy", compton_nominal_photon_energy.value()); } + + params->set("two_body.thomson_optical_depth", + twobody_thomson_optical_depth.value()); + params->set("two_body.interactions", twobody_interactions); } } // namespace params } // namespace ntt diff --git a/src/framework/parameters/extra.h b/src/framework/parameters/extra.h index 29de4527a..357b480aa 100644 --- a/src/framework/parameters/extra.h +++ b/src/framework/parameters/extra.h @@ -12,6 +12,7 @@ #ifndef FRAMEWORK_PARAMETERS_EXTRA_H #define FRAMEWORK_PARAMETERS_EXTRA_H +#include "enums.h" #include "global.h" #include "framework/parameters/parameters.h" @@ -25,6 +26,13 @@ namespace ntt { namespace params { + struct TwoBodyInteractionParams { + TwoBodyInteractionFlag type; + std::vector group1; + std::vector group2; + timestep_t interval; + }; + struct Extra { // radiative drag parameters std::optional synchrotron_gamma_rad; @@ -45,6 +53,10 @@ namespace ntt { std::optional compton_nominal_probability; std::optional compton_nominal_photon_energy; + // two-body interaction parameters + std::optional twobody_thomson_optical_depth; + std::vector twobody_interactions; + void read(const std::map&, const toml::value&, const SimulationParams* const); diff --git a/src/global/defaults.h b/src/global/defaults.h index dbe37ef0f..11792b2cf 100644 --- a/src/global/defaults.h +++ b/src/global/defaults.h @@ -107,6 +107,10 @@ namespace ntt::defaults { const real_t gamma_rad = 1.0; const real_t gamma_qed = 10.0; } // namespace compton + + namespace twobody { + const real_t thomson_optical_depth = 1.0; + } // namespace twobody } // namespace ntt::defaults #endif // GLOBAL_DEFAULTS_H diff --git a/src/global/enums.h b/src/global/enums.h index 41fdd551b..d8fc509ef 100644 --- a/src/global/enums.h +++ b/src/global/enums.h @@ -420,6 +420,19 @@ namespace ntt { return "unknown"; } } + + inline auto from_string(const std::string& s) -> uint8_t { + if (fmt::toLower(s) == "none") { + return NONE; + } else if (fmt::toLower(s) == "compton") { + return COMPTON; + } else if (fmt::toLower(s) == "custom") { + return CUSTOM; + } else { + raise::Error(fmt::format("Invalid TwoBodyInteraction type: %s", s), HERE); + return NONE; + } + } } // namespace TwoBodyInteraction using TwoBodyInteractionFlag = uint8_t; From 5b4198284ee3e218708387f5551e121aac97fdbe Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 14 May 2026 17:16:40 -0400 Subject: [PATCH 08/25] compton tests --- examples/compton_jones/compton_jones.py | 62 ++++++++ examples/compton_jones/compton_jones.toml | 89 +++++++++++ examples/compton_jones/pgen.hpp | 148 ++++++++++++++++++ .../compton_kompaneets/compton_kompaneets.py | 60 +++++++ .../compton_kompaneets.toml | 87 ++++++++++ examples/compton_kompaneets/pgen.hpp | 74 +++++++++ 6 files changed, 520 insertions(+) create mode 100644 examples/compton_jones/compton_jones.py create mode 100644 examples/compton_jones/compton_jones.toml create mode 100644 examples/compton_jones/pgen.hpp create mode 100644 examples/compton_kompaneets/compton_kompaneets.py create mode 100644 examples/compton_kompaneets/compton_kompaneets.toml create mode 100644 examples/compton_kompaneets/pgen.hpp diff --git a/examples/compton_jones/compton_jones.py b/examples/compton_jones/compton_jones.py new file mode 100644 index 000000000..873a01f35 --- /dev/null +++ b/examples/compton_jones/compton_jones.py @@ -0,0 +1,62 @@ +import nt2 +import matplotlib.pyplot as plt +import numpy as np + +data = nt2.Data("compton_jones") + +photons = data.particles.sel(sp=3).isel(t=-1).load() +photons = photons[np.sqrt(photons.ux**2 + photons.uy**2 + photons.uz**2) > 0.01] + +plt.rcParams["figure.dpi"] = 300 +plt.rcParams["font.family"] = "serif" + +fig = plt.figure(figsize=(9, 4)) +gs = fig.add_gridspec(1, 2, wspace=0.35) +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1]) + +gamma = np.sqrt(1 + data.attrs["setup.electron_4vel"] ** 2) +e0 = data.attrs["setup.photon_energy"] +Gamma = 4 * e0 * gamma +emax = gamma * Gamma / (1 + Gamma) + +es = data.spectra.E[1:-1] / emax + +dnde = data.spectra.N_3.isel(t=-1)[1:-1] +dnde /= np.trapezoid(dnde, es) +ax1.plot(es, dnde) + +es = np.linspace(es.values.min(), es.values.max(), 250) +qs = es / (1 + Gamma * (1 - es)) +dnde_th = ( + 2 * qs * np.log(qs) + + (1 + 2 * qs) * (1 - qs) + + 0.5 * Gamma**2 * qs**2 / (1 + Gamma * qs) * (1 - qs) +) + +dnde_th /= np.trapezoid(dnde_th, es) + +ax1.plot(es, dnde_th, c="k", ls=":") +ax1.set( + xlim=(0, 1), + ylim=(0, 4), + xlabel=r"$\varepsilon_{\rm ph} / \varepsilon_{\rm max}$", + ylabel=r"$dn_{\rm ph}/d\varepsilon_{\rm ph}$", +) + +plt.scatter( + photons.ux / emax, + photons.uy / emax, + s=1, + linewidth=0, +) +xs = np.linspace(0, 1, 100) +ys = 2 / gamma * xs +ax2.plot(xs, ys, c="k", ls="--", lw=0.5) +ax2.plot(xs, -ys, c="k", ls="--", lw=0.5) +ax2.set( + xlabel=r"$p_{\rm ph}^x / \varepsilon_{\rm max}$", + ylabel=r"$p_{\rm ph}^y / \varepsilon_{\rm max}$", +) + +plt.savefig("compton_jones.png", bbox_inches="tight") diff --git a/examples/compton_jones/compton_jones.toml b/examples/compton_jones/compton_jones.toml new file mode 100644 index 000000000..796739978 --- /dev/null +++ b/examples/compton_jones/compton_jones.toml @@ -0,0 +1,89 @@ +[simulation] + name = "compton_jones" + engine = "srpic" + runtime = 10.0 + +[grid] + resolution = [32, 32] + extent = [[0.0, 1.0], [0.0, 1.0]] + + [grid.metric] + metric = "minkowski" + + [grid.boundaries] + fields = [["PERIODIC"], ["PERIODIC"]] + particles = [["PERIODIC"], ["PERIODIC"]] + +[scales] + larmor0 = 1.0 + skindepth0 = 1.0 + +[two_body] + thomson_optical_depth = 0.5 + + [[two_body.interaction]] + type = "compton" + group1 = [1] + group2 = [2] + interval = 1 + tile_size = 5 + recoil1 = false + recoil2 = true + +[algorithms] + current_filters = 0 + + [algorithms.deposit] + enable = false + + [algorithms.fieldsolver] + enable = false + +[particles] + ppc0 = 1.0 + clear_interval = 1 + + [[particles.species]] + label = "e-" + mass = 1.0 + charge = -1.0 + maxnpart = 1e6 + + [[particles.species]] + label = "ph" + mass = 0.0 + charge = 0.0 + maxnpart = 1e6 + + [[particles.species]] + label = "ph_out" + mass = 0.0 + charge = 0.0 + maxnpart = 1e7 + pusher = "none" + +[setup] + electron_4vel = 999.9995 + photon_energy = 1e-2 + +[output] + interval_time = 0.01 + + [output.fields] + quantities = ["N_1", "N_2", "N_3"] + + [output.particles] + species = [1, 2, 3] + stride = 1 + + [output.spectra] + log_bins = false + e_min = 0 + e_max = 1100 + n_bins = 100 + + [output.stats] + quantities = ["T00_1", "T00_2", "T00_3"] + +[checkpoint] + keep = 0 diff --git a/examples/compton_jones/pgen.hpp b/examples/compton_jones/pgen.hpp new file mode 100644 index 000000000..f2da19802 --- /dev/null +++ b/examples/compton_jones/pgen.hpp @@ -0,0 +1,148 @@ +#ifndef PROBLEM_GENERATOR_H +#define PROBLEM_GENERATOR_H + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "traits/pgen.h" + +#include "archetypes/particle_injector.h" +#include "framework/domain/metadomain.h" + +namespace user { + using namespace ntt; + + template + struct DeltaDistribution { + const real_t energy0; + bool monodirectional; + random_number_pool_t random_pool; + + DeltaDistribution(real_t energy0, + bool monodirectional, + random_number_pool_t& random_pool) + : energy0 { energy0 } + , monodirectional { monodirectional } + , random_pool { random_pool } {} + + Inline void operator()(const coord_t&, vec_t& v) const { + if (not monodirectional) { + auto gen = random_pool.get_state(); + auto rnd1 = Random(gen); + auto rnd2 = Random(gen); + random_pool.free_state(gen); + // random direction + const auto phi = static_cast(constant::TWO_PI) * rnd1; + const auto ct = 2.0 * rnd2 - 1.0; + const auto st = math::sqrt(1.0 - ct * ct); + v[0] = energy0 * st * math::cos(phi); + v[1] = energy0 * st * math::sin(phi); + v[2] = energy0 * ct; + } else { + v[0] = energy0; + v[1] = 0.0; + v[2] = 0.0; + } + } + }; + + template + struct PGen { + + static constexpr auto engines { + ::traits::pgen::compatible_with {} + }; + static constexpr auto metrics { + ::traits::pgen::compatible_with {} + }; + static constexpr auto dimensions { + ::traits::pgen::compatible_with {} + }; + + const SimulationParams& params; + const Metadomain& metadomain; + + PGen(const SimulationParams& p, const Metadomain& m) + : params { p } + , metadomain { m } {} + + void InitPrtls(Domain& domain) { + auto delta_electrons = DeltaDistribution { + params.template get("setup.electron_4vel"), + true, + domain.random_pool() + }; + arch::InjectUniform(params, + domain, + 1u, + delta_electrons, + ONE); + } + + void CustomPostStep(timestep_t /*step*/, simtime_t /*time*/, Domain& domain) { + // copy all photons from species #2 (idx 1) to #3 (idx 2) with an offset + const auto offset = domain.species[2].npart(); + const auto new_copies = domain.species[1].npart(); + const auto new_size = offset + new_copies; + const auto from_slice = prtl_slice_t { 0, new_copies }; + const auto to_slice = prtl_slice_t { offset, new_size }; + + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i1, to_slice), + Kokkos::subview(domain.species[1].i1, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i1_prev, to_slice), + Kokkos::subview(domain.species[1].i1_prev, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx1, to_slice), + Kokkos::subview(domain.species[1].dx1, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx1_prev, to_slice), + Kokkos::subview(domain.species[1].dx1_prev, from_slice)); + if constexpr (M::Dim == Dim::_2D or M::Dim == Dim::_3D) { + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i2, to_slice), + Kokkos::subview(domain.species[1].i2, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i2_prev, to_slice), + Kokkos::subview(domain.species[1].i2_prev, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx2, to_slice), + Kokkos::subview(domain.species[1].dx2, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx2_prev, to_slice), + Kokkos::subview(domain.species[1].dx2_prev, from_slice)); + } + if constexpr (M::Dim == Dim::_3D) { + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i3, to_slice), + Kokkos::subview(domain.species[1].i3, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].i3_prev, to_slice), + Kokkos::subview(domain.species[1].i3_prev, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx3, to_slice), + Kokkos::subview(domain.species[1].dx3, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].dx3_prev, to_slice), + Kokkos::subview(domain.species[1].dx3_prev, from_slice)); + } + Kokkos::deep_copy(Kokkos::subview(domain.species[2].ux1, to_slice), + Kokkos::subview(domain.species[1].ux1, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].ux2, to_slice), + Kokkos::subview(domain.species[1].ux2, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].ux3, to_slice), + Kokkos::subview(domain.species[1].ux3, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].weight, to_slice), + Kokkos::subview(domain.species[1].weight, from_slice)); + Kokkos::deep_copy(Kokkos::subview(domain.species[2].tag, to_slice), + Kokkos::subview(domain.species[1].tag, from_slice)); + + domain.species[1].set_npart(0); + domain.species[2].set_npart(new_size); + + auto delta_photons = DeltaDistribution { + params.template get("setup.photon_energy"), + false, + domain.random_pool() + }; + arch::InjectUniform(params, + domain, + 2u, + delta_photons, + ONE); + } + }; + +} // namespace user + +#endif diff --git a/examples/compton_kompaneets/compton_kompaneets.py b/examples/compton_kompaneets/compton_kompaneets.py new file mode 100644 index 000000000..4c713cbad --- /dev/null +++ b/examples/compton_kompaneets/compton_kompaneets.py @@ -0,0 +1,60 @@ +import nt2 +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +data = nt2.Data("compton_kompaneets") +stats = pd.read_csv("compton_kompaneets/compton_kompaneets_stats.csv") +stats.columns = stats.columns.str.strip() + +plt.rcParams["figure.dpi"] = 300 +plt.rcParams["font.family"] = "serif" + +fig = plt.figure(figsize=(9, 4)) +gs = fig.add_gridspec(1, 2, wspace=0.3) +ax1 = fig.add_subplot(gs[0, 0]) + +tvals = len(data.spectra.t.values) + +nphot = data.spectra.N_3.isel(t=-1).sum().values[()] +for ti in range(0, tvals, 10): + ax1.plot( + data.spectra.E.values / data.attrs["setup.temperature"], + data.spectra.N_3.isel(t=ti).values, + c=plt.get_cmap("plasma")(ti / tvals), + lw=0.5, + ) + +es = data.spectra.E.values / data.attrs["setup.temperature"] +dndes = es**2 * np.exp(-es) +dndes /= np.sum(dndes) +dndes *= nphot +ax1.plot( + es, + dndes, + c="k", + ls=":", + label=r"$\propto \varepsilon_{\rm ph}^2 e^{-\varepsilon_{\rm ph} / T_\pm}$", +) + +ax1.set( + yscale="log", + ylim=(1e-1, 1e5), + xlim=(0, 10), + xlabel=r"$\varepsilon / T_\pm$", + ylabel=r"$dn_{\rm ph}/d\varepsilon$", +) +ax1.legend() + +ax2 = fig.add_subplot(gs[0, 1]) +ax2.plot(stats["time"], stats["T00_3"], c="C0") +ax2.set(ylabel=r"total photon energy", xlabel=r"$t$") +ax2.yaxis.label.set_color("C0") +ax2.tick_params(axis="y", labelcolor="C0") +ax2twin = ax2.twinx() +ax2twin.plot(data.spectra.t.values, data.spectra.N_3.sum("E"), c="C2") +ax2twin.set(ylabel=r"photon number") +ax2twin.yaxis.label.set_color("C2") +ax2twin.tick_params(axis="y", labelcolor="C2") + +plt.savefig("compton_kompaneets_plot.png", bbox_inches="tight") diff --git a/examples/compton_kompaneets/compton_kompaneets.toml b/examples/compton_kompaneets/compton_kompaneets.toml new file mode 100644 index 000000000..1a006102f --- /dev/null +++ b/examples/compton_kompaneets/compton_kompaneets.toml @@ -0,0 +1,87 @@ +[simulation] + name = "compton_kompaneets" + engine = "srpic" + runtime = 5.0 + +[grid] + resolution = [128, 128] + extent = [[0.0, 1.0], [0.0, 1.0]] + + [grid.metric] + metric = "minkowski" + + [grid.boundaries] + fields = [["PERIODIC"], ["PERIODIC"]] + particles = [["PERIODIC"], ["PERIODIC"]] + +[scales] + larmor0 = 1.0 + skindepth0 = 1.0 + +[two_body] + thomson_optical_depth = 0.25 + + [[two_body.interaction]] + type = "compton" + group1 = [1, 2] + group2 = [3] + interval = 1 + tile_size = 5 + recoil1 = false + recoil2 = true + +[algorithms] + current_filters = 0 + + [algorithms.deposit] + enable = false + + [algorithms.fieldsolver] + enable = false + +[particles] + ppc0 = 2.0 + clear_interval = 1 + + [[particles.species]] + label = "e-" + mass = 1.0 + charge = -1.0 + maxnpart = 1e6 + + [[particles.species]] + label = "e+" + mass = 1.0 + charge = 1.0 + maxnpart = 1e6 + + [[particles.species]] + label = "ph" + mass = 0.0 + charge = 0.0 + maxnpart = 1e6 + +[setup] + temperature = 0.01 + photon_energy = 1e-3 + +[output] + interval_time = 0.01 + + [output.fields] + enable = false + + [output.particles] + enable = false + + [output.spectra] + log_bins = false + e_min = 0.0 + e_max = 1.0 + n_bins = 500 + + [output.stats] + quantities = ["T00_1", "T00_2", "T00_3"] + +[checkpoint] + keep = 0 diff --git a/examples/compton_kompaneets/pgen.hpp b/examples/compton_kompaneets/pgen.hpp new file mode 100644 index 000000000..0b0158ada --- /dev/null +++ b/examples/compton_kompaneets/pgen.hpp @@ -0,0 +1,74 @@ +#ifndef PROBLEM_GENERATOR_H +#define PROBLEM_GENERATOR_H + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "traits/pgen.h" + +#include "archetypes/particle_injector.h" +#include "archetypes/utils.h" +#include "framework/domain/metadomain.h" + +namespace user { + using namespace ntt; + + template + struct DeltaDistribution { + const real_t photon_energy0; + random_number_pool_t random_pool; + + DeltaDistribution(real_t photon_energy0, random_number_pool_t& random_pool) + : photon_energy0 { photon_energy0 } + , random_pool { random_pool } {} + + Inline void operator()(const coord_t&, vec_t& v) const { + auto gen = random_pool.get_state(); + auto rnd1 = Random(gen); + auto rnd2 = Random(gen); + random_pool.free_state(gen); + // random direction + const auto phi = static_cast(constant::TWO_PI) * rnd1; + const auto ct = 2.0 * rnd2 - 1.0; + const auto st = math::sqrt(1.0 - ct * ct); + v[0] = photon_energy0 * st * math::cos(phi); + v[1] = photon_energy0 * st * math::sin(phi); + v[2] = photon_energy0 * ct; + } + }; + + template + struct PGen { + + static constexpr auto engines { + ::traits::pgen::compatible_with {} + }; + static constexpr auto metrics { + ::traits::pgen::compatible_with {} + }; + static constexpr auto dimensions { + ::traits::pgen::compatible_with {} + }; + + const SimulationParams& params; + const Metadomain& metadomain; + + PGen(const SimulationParams& p, const Metadomain& m) + : params { p } + , metadomain { m } {} + + void InitPrtls(Domain& domain) { + const auto temperature = params.template get("setup.temperature"); + arch::InjectUniformMaxwellian(params, domain, ONE, temperature, { 1u, 2u }); + + auto delta = DeltaDistribution { params.template get( + "setup.photon_energy"), + domain.random_pool() }; + arch::InjectUniform(params, domain, 3u, delta, ONE); + } + }; + +} // namespace user + +#endif From 2ef59cb0da3d421cb0a3aacb4231fc3324c273e2 Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 14 May 2026 17:16:59 -0400 Subject: [PATCH 09/25] extra injector for single species --- src/archetypes/particle_injector.h | 67 ++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/src/archetypes/particle_injector.h b/src/archetypes/particle_injector.h index 7adb8c5b5..047a40183 100644 --- a/src/archetypes/particle_injector.h +++ b/src/archetypes/particle_injector.h @@ -389,6 +389,73 @@ namespace arch { } } + /** + * @brief Injects uniform number density of a single species everywhere in the domain + * @param domain Domain object + * @param species Species index + * @param energy_dist Energy distribution objects + * @param number_density Number density (in units of n0) + * @param use_weights Use weights + * @param box Region to inject the particles in global coords + * @tparam S Simulation engine type + * @tparam M Metric type + * @tparam ED Energy distribution type + */ + template ED> + inline void InjectUniform(const SimulationParams& params, + Domain& domain, + spidx_t species, + const ED& energy_dist, + real_t number_density, + bool use_weights = false, + const boundaries_t& box = {}) { + raise::ErrorIf((M::CoordType != Coord::Cartesian) && (not use_weights), + "Weights must be used for non-Cartesian coordinates", + HERE); + raise::ErrorIf((M::CoordType == Coord::Cartesian) && use_weights, + "Weights should not be used for Cartesian coordinates", + HERE); + raise::ErrorIf(params.template get("particles.use_weights") != use_weights, + "Weights must be enabled from the input file to use them in " + "the injector", + HERE); + if (domain.species[species - 1].charge() != 0.0f) { + raise::Warning("Charge of the injected species is non-zero", HERE); + } + + { + boundaries_t nonempty_box; + for (auto d { 0u }; d < M::Dim; ++d) { + if (d < box.size()) { + nonempty_box.emplace_back(box[d].first, box[d].second); + } else { + nonempty_box.push_back(Range::All); + } + } + const auto result = ComputeNumInject(params, domain, number_density, nonempty_box); + if (not std::get<0>(result)) { + return; + } + const auto nparticles = std::get<1>(result); + const auto xi_min = std::get<2>(result); + const auto xi_max = std::get<3>(result); + + Kokkos::parallel_for("InjectUniform", + nparticles, + kernel::SingleSpeciesUniformInjector_kernel( + domain.species[species - 1], + domain.index(), + domain.mesh.metric, + xi_min, + xi_max, + energy_dist, + ONE / params.template get("scales.V0"), + domain.random_pool())); + domain.species[species - 1].set_npart( + domain.species[species - 1].npart() + nparticles); + } + } + } // namespace arch #endif // ARCHETYPES_PARTICLE_INJECTOR_H From 5cfc138bcf589caff34419fbefb4b9946ad232b5 Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 14 May 2026 17:17:30 -0400 Subject: [PATCH 10/25] qed incorporated --- input.example.toml | 12 ++ src/archetypes/qed/compton.h | 120 +++++++++------- src/engines/engine.hpp | 4 +- src/engines/reporter.cpp | 43 ++++++ src/engines/srpic/srpic.hpp | 7 + src/engines/srpic/twobody.h | 64 ++++++++- src/framework/parameters/extra.cpp | 13 +- src/framework/parameters/extra.h | 3 + src/global/enums.h | 3 +- src/global/traits/policies.h | 39 ++++++ src/kernels/injectors.hpp | 126 +++++++++++++++++ src/kernels/twobody_interactions.hpp | 131 +++++++++++++---- tests/archetypes/qed_compton.cpp | 8 +- tests/kernels/twobody_interactions.cpp | 187 +++++++++++++++++++++++++ 14 files changed, 667 insertions(+), 93 deletions(-) create mode 100644 tests/kernels/twobody_interactions.cpp diff --git a/input.example.toml b/input.example.toml index 20197b08e..d721c9275 100644 --- a/input.example.toml +++ b/input.example.toml @@ -298,6 +298,18 @@ # @type: uint # @default: 1 interval = "" + # Size of interaction tile in number of cells + # @type: uint + # @default: 4 + tile_size = "" + # Whether to apply recoil on species of group1 + # @type: bool + # @default: true + recoil1 = "" + # Whether to apply recoil on species of group2 + # @type: bool + # @default: true + recoil2 = "" [algorithms] # Number of current smoothing passes diff --git a/src/archetypes/qed/compton.h b/src/archetypes/qed/compton.h index b54c268bc..4f466579e 100644 --- a/src/archetypes/qed/compton.h +++ b/src/archetypes/qed/compton.h @@ -24,33 +24,26 @@ namespace arch::qed { using namespace ntt; - template + template struct ComptonScattering { static constexpr spidx_t MAXSP = 16u; static constexpr int MAX_ITER = 10; ParticleArrays species[MAXSP]; static constexpr real_t low_energy_limit = static_cast(2e-3); + static constexpr real_t Thomson_limit = static_cast(1e-3); const real_t nominal_probability_density; - const real_t Thomson_limit; random_number_pool_t random_pool; ComptonScattering(const prm::Parameters& params, random_number_pool_t& random_pool) : nominal_probability_density { params.template get( - "qed.compton_scattering.nominal_probability_density") } - , Thomson_limit { params.template get( - "qed.compton_scattering.Thomson_limit") } + "compton_scattering.nominal_probability_density") } , random_pool { random_pool } { if (nominal_probability_density <= ZERO) { raise::Error("nominal_probability must be in the range (0, 1]", HERE); } - if (Thomson_limit <= ZERO or Thomson_limit > low_energy_limit) { - raise::Error( - "Thomson_limit must be in the range (0, small_energy_limit]", - HERE); - } } /* @@ -79,30 +72,31 @@ namespace arch::qed { /* * Calculate the Klein-Nishina cross section for a photon with energy e_ in the lepton rest frame * @param e_: photon energy in the lepton rest frame - * @return: pair of (is_KN_regime, f_KN) where - * - is_KN_regime: whether the photon energy is in the Klein-Nishina regime (e_ > Thomson_limit) - * - f_KN: the Klein-Nishina cross section normalized to the Thomson cross section + * @return f_KN: the Klein-Nishina cross section normalized to the Thomson cross section * @note for e_ > low_energy_limit, full Klein-Nishina formula * @note for Thomson_limit < e_ <= low_energy_limit, 2nd order expansion of the Klein-Nishina formula * @note for e_ <= Thomson_limit, return 1 (Thomson limit) */ - Inline auto KNCrossSection(real_t e_) const -> Kokkos::pair { + Inline auto KNCrossSection(double e_) const -> real_t { if (e_ > Thomson_limit) { if (e_ < low_energy_limit) { // correctly handle the e_ << 1 limit using 2nd order expansion of f_KN - return { true, ONE - TWO * e_ + static_cast(5.2) * SQR(e_) }; + return static_cast(1.0 - 2.0 * e_ + 5.2 * SQR(e_)); } else { - return { true, - static_cast(0.375) * - ((ONE - TWO / e_ - TWO / SQR(e_)) * math::log(ONE + TWO * e_) + - HALF + FOUR / e_ - HALF / SQR(ONE + TWO * e_)) / - e_ }; + return static_cast( + 0.375 * + ((1.0 - 2.0 / e_ - 2.0 / SQR(e_)) * math::log(1.0 + 2.0 * e_) + + 0.5 + 4.0 / e_ - 0.5 / SQR(1.0 + 2.0 * e_)) / + e_); } } else { - return { false, ONE }; + return ONE; } } + /* + * Sample a cosine theta value from the Thomson scattering cross section + */ Inline auto RandomCosTheta_Th() const -> real_t { auto gen_ = random_pool.get_state(); const auto rnd_ = Random(gen_); @@ -114,6 +108,9 @@ namespace arch::qed { return u - ONE / u; } + /* + * Sample a cosine theta from the Klein-Nishina scattering cross section + */ Inline auto RandomCosTheta_KN(double e_) const -> real_t { auto gen_ = random_pool.get_state(); const auto rnd_ = Random(gen_); @@ -215,12 +212,11 @@ namespace arch::qed { rand_sintheta_ * rand_sinphi_ * c_[2]); } - Inline void operator()(spidx_t sp1, - npart_t p1, - spidx_t sp2, - npart_t p2, - real_t tile_volume) const { - // @TODO coord/vec conversion + Inline auto should_interact(spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2, + real_t tile_weight) const -> bool { // values with "_" are in the lepton rest-frame const vec_t lepton_u { species[sp1 - 1].ux1(p1), species[sp1 - 1].ux2(p1), @@ -240,43 +236,67 @@ namespace arch::qed { LorentzBoost(lepton_u, lepton_gamma, photon_p, photon_energy, photon_p_, photon_energy_); - const auto [KN_regime, f_KN] = KNCrossSection(photon_energy_); - const auto scattering_probability = nominal_probability_density * f_KN * - photon_energy_ * lepton_weight * - photon_weight / - (photon_energy * lepton_gamma * - tile_volume); + const auto f_KN = KNCrossSection(photon_energy_); + auto gen = random_pool.get_state(); const auto rnd = Random(gen); random_pool.free_state(gen); - if (rnd < scattering_probability) { - vec_t photon_pnew_ { ZERO, ZERO, ZERO }, - photon_pnew { ZERO, ZERO, ZERO }; - real_t photon_energy_new_ { ZERO }, photon_energy_new { ZERO }; - - ScatterPhoton(KN_regime, - photon_p_, - photon_energy_, - photon_pnew_, - photon_energy_new_); - LorentzBoost({ -lepton_u[0], -lepton_u[1], -lepton_u[2] }, - lepton_gamma, - photon_pnew_, - photon_energy_new_, - photon_pnew, - photon_energy_new); + return rnd < + (tile_weight * nominal_probability_density * f_KN * photon_energy_ * + lepton_weight * photon_weight / (photon_energy * lepton_gamma)); + } + + Inline void operator()(spidx_t sp1, npart_t p1, spidx_t sp2, npart_t p2) const { + // @TODO coord/vec conversion + // values with "_" are in the lepton rest-frame + const vec_t lepton_u { species[sp1 - 1].ux1(p1), + species[sp1 - 1].ux2(p1), + species[sp1 - 1].ux3(p1) }; + const auto lepton_gamma = U2GAMMA(lepton_u[0], lepton_u[1], lepton_u[2]); + const auto lepton_weight = species[sp1 - 1].weight(p1); + + const vec_t photon_p { species[sp2 - 1].ux1(p2), + species[sp2 - 1].ux2(p2), + species[sp2 - 1].ux3(p2) }; + const auto photon_energy = NORM(photon_p[0], photon_p[1], photon_p[2]); + const auto photon_weight = species[sp2 - 1].weight(p2); + + // boost photon momentum to lepton rest frame + vec_t photon_p_ { ZERO, ZERO, ZERO }; + real_t photon_energy_ { ZERO }; + + LorentzBoost(lepton_u, lepton_gamma, photon_p, photon_energy, photon_p_, photon_energy_); + + vec_t photon_pnew_ { ZERO, ZERO, ZERO }, + photon_pnew { ZERO, ZERO, ZERO }; + real_t photon_energy_new_ { ZERO }, photon_energy_new { ZERO }; + + ScatterPhoton(photon_energy_ > Thomson_limit, + photon_p_, + photon_energy_, + photon_pnew_, + photon_energy_new_); + LorentzBoost({ -lepton_u[0], -lepton_u[1], -lepton_u[2] }, + lepton_gamma, + photon_pnew_, + photon_energy_new_, + photon_pnew, + photon_energy_new); + if constexpr (R1) { species[sp1 - 1].ux1(p1) += (photon_p[0] - photon_pnew[0]) * photon_weight / lepton_weight; species[sp1 - 1].ux2(p1) += (photon_p[1] - photon_pnew[1]) * photon_weight / lepton_weight; species[sp1 - 1].ux3(p1) += (photon_p[2] - photon_pnew[2]) * photon_weight / lepton_weight; + } + if constexpr (R2) { species[sp2 - 1].ux1(p2) = photon_pnew[0]; species[sp2 - 1].ux2(p2) = photon_pnew[1]; species[sp2 - 1].ux3(p2) = photon_pnew[2]; - } // not interacting + } } }; diff --git a/src/engines/engine.hpp b/src/engines/engine.hpp index 026cd5c50..bd77c83a4 100644 --- a/src/engines/engine.hpp +++ b/src/engines/engine.hpp @@ -249,8 +249,8 @@ namespace ntt { "ParticlePusher", "FieldBoundaries", "ParticleBoundaries", "Communications", "Injector", "Custom", - "ParticleSort", "Output", - "Checkpoint" }, + "TwoBodyInteractions", "ParticleSort", + "Output", "Checkpoint" }, []() { Kokkos::fence(); }, diff --git a/src/engines/reporter.cpp b/src/engines/reporter.cpp index f56874d23..11ae740ed 100644 --- a/src/engines/reporter.cpp +++ b/src/engines/reporter.cpp @@ -6,6 +6,7 @@ #include "utils/formatting.h" #include "utils/reporter.h" +#include "framework/parameters/extra.h" #include "framework/parameters/parameters.h" #include @@ -149,6 +150,48 @@ namespace ntt { params.template get( "radiation.emission.synchrotron.nominal_photon_energy")); } + + report += "\n"; + const auto two_body_interactions = + params.template get>( + "two_body.interaction"); + if (not two_body_interactions.empty()) { + reporter::AddCategory(report, 4, "Two-body interactions"); + reporter::AddParam( + report, + 6, + "Thomson optical depth", + "%.3e", + params.template get("two_body.thomson_optical_depth")); + for (const auto& interaction : two_body_interactions) { + reporter::AddSubcategory( + report, + 6, + TwoBodyInteraction::to_string(interaction.type).c_str()); + std::vector group1(interaction.group1.size()); + std::vector group2(interaction.group2.size()); + for (size_t g1 = 0; g1 < interaction.group1.size(); ++g1) { + group1[g1] = interaction.group1[g1]; + } + for (size_t g2 = 0; g2 < interaction.group2.size(); ++g2) { + group2[g2] = interaction.group2[g2]; + } + reporter::AddParam(report, + 8, + "group #1 species", + "%s (recoil: %s)", + fmt::formatVector(group1).c_str(), + interaction.recoil1 ? "ON" : "OFF"); + reporter::AddParam(report, + 8, + "group #2 species", + "%s (recoil: %s)", + fmt::formatVector(group2).c_str(), + interaction.recoil2 ? "ON" : "OFF"); + reporter::AddParam(report, 8, "tile size [cells]", "%u", interaction.tile_size); + reporter::AddParam(report, 8, "interval [steps]", "%u", interaction.interval); + } + } return report; } diff --git a/src/engines/srpic/srpic.hpp b/src/engines/srpic/srpic.hpp index 94a8949c1..1afb269a2 100644 --- a/src/engines/srpic/srpic.hpp +++ b/src/engines/srpic/srpic.hpp @@ -24,6 +24,7 @@ #include "engines/srpic/fieldsolvers.h" #include "engines/srpic/particle_pusher.h" #include "engines/srpic/particles_bcs.h" +#include "engines/srpic/twobody.h" #include "framework/domain/domain.h" #include "framework/parameters/parameters.h" @@ -182,6 +183,12 @@ namespace ntt { timers.stop("Injector"); } + if constexpr (CartesianMetricClass) { + timers.start("TwoBodyInteractions"); + srpic::TwoBodyInteractions(dom, this->engineParams(), m_params); + timers.stop("TwoBodyInteractions"); + } + timers.start("ParticleSort"); m_metadomain.SortParticles(time, step, m_params, dom); timers.stop("ParticleSort"); diff --git a/src/engines/srpic/twobody.h b/src/engines/srpic/twobody.h index b969604ea..c5d04ca66 100644 --- a/src/engines/srpic/twobody.h +++ b/src/engines/srpic/twobody.h @@ -5,6 +5,7 @@ #include "traits/metric.h" #include "utils/error.h" +#include "utils/formatting.h" #include "utils/log.h" #include "utils/param_container.h" @@ -26,11 +27,68 @@ namespace ntt { const auto step = engine_params.get("step"); for (const auto& interaction : params.template get>( - "two_body.interactions")) { + "two_body.interaction")) { if (step % interaction.interval == 0u) { + const auto thomson_optical_depth = params.template get( + "two_body.thomson_optical_depth"); + const auto nominal_thomson_probability_density = thomson_optical_depth * + dt * + static_cast( + interaction.interval); if (interaction.type == TwoBodyInteraction::COMPTON) { + prm::Parameters compton_params; + compton_params.set("compton_scattering.nominal_probability_density", + nominal_thomson_probability_density); + auto recoil1 = interaction.recoil1; + auto recoil2 = interaction.recoil2; + auto launch = [&]() { + auto policy = arch::qed::ComptonScattering( + compton_params, + domain.random_pool()); - printf("CALLING COMPTON\n"); + std::vector*> group1_species; + std::vector*> group2_species; + + for (const auto& sp_lepton : interaction.group1) { + raise::ErrorIf( + domain.species[sp_lepton - 1].mass() == ZERO, + fmt::format( + "Species %u is massless but is in the lepton group " + "of a Compton interaction", + sp_lepton), + HERE); + group1_species.push_back(&domain.species[sp_lepton - 1]); + } + for (const auto& sp_photon : interaction.group2) { + raise::ErrorIf( + domain.species[sp_photon - 1].mass() != ZERO, + fmt::format( + "Species %u is massive but is in the photon group " + "of a Compton interaction", + sp_photon), + HERE); + group2_species.push_back(&domain.species[sp_photon - 1]); + } + + kernel::mink::TwoBodyInteraction( + group1_species, + group2_species, + domain.mesh.n_active(), + domain.mesh.extent(), + interaction.tile_size, + params.template get("particles.ppc0"), + domain.random_pool(), + policy); + }; + if (interaction.recoil1 and interaction.recoil2) { + launch.template operator()(); + } else if (interaction.recoil1 and not interaction.recoil2) { + launch.template operator()(); + } else if (not interaction.recoil1 and interaction.recoil2) { + launch.template operator()(); + } else { + launch.template operator()(); + } } else if (interaction.type == TwoBodyInteraction::CUSTOM) { raise::Error("Custom two-body interactions not implemented yet", HERE); } @@ -40,4 +98,4 @@ namespace ntt { } // namespace srpic } // namespace ntt -#endif // ENGINES_SRPIC_TWOBODY_H \ No newline at end of file +#endif // ENGINES_SRPIC_TWOBODY_H diff --git a/src/framework/parameters/extra.cpp b/src/framework/parameters/extra.cpp index 4d7169700..53eca4117 100644 --- a/src/framework/parameters/extra.cpp +++ b/src/framework/parameters/extra.cpp @@ -124,15 +124,18 @@ namespace ntt { // find two-body interactions const auto twobody_tab = toml::find_or(toml_data, "two_body", - "interactions", + "interaction", toml::array {}); for (const auto& tbint : twobody_tab) { twobody_interactions.push_back(TwoBodyInteractionParams { .type = TwoBodyInteraction::from_string( toml::find(tbint, "type")), - .group1 = toml::find>(tbint, "group1"), - .group2 = toml::find_or>(tbint, "group2", {}), - .interval = toml::find_or(tbint, "interval", 1) }); + .group1 = toml::find>(tbint, "group1"), + .group2 = toml::find_or>(tbint, "group2", {}), + .interval = toml::find_or(tbint, "interval", 1), + .tile_size = toml::find_or(tbint, "tile_size", 4u), + .recoil1 = toml::find_or(tbint, "recoil1", true), + .recoil2 = toml::find_or(tbint, "recoil2", true) }); } } @@ -182,7 +185,7 @@ namespace ntt { params->set("two_body.thomson_optical_depth", twobody_thomson_optical_depth.value()); - params->set("two_body.interactions", twobody_interactions); + params->set("two_body.interaction", twobody_interactions); } } // namespace params } // namespace ntt diff --git a/src/framework/parameters/extra.h b/src/framework/parameters/extra.h index 357b480aa..c3891429f 100644 --- a/src/framework/parameters/extra.h +++ b/src/framework/parameters/extra.h @@ -31,6 +31,9 @@ namespace ntt { std::vector group1; std::vector group2; timestep_t interval; + ncells_t tile_size; + bool recoil1; + bool recoil2; }; struct Extra { diff --git a/src/global/enums.h b/src/global/enums.h index d8fc509ef..bb4fbc272 100644 --- a/src/global/enums.h +++ b/src/global/enums.h @@ -429,7 +429,8 @@ namespace ntt { } else if (fmt::toLower(s) == "custom") { return CUSTOM; } else { - raise::Error(fmt::format("Invalid TwoBodyInteraction type: %s", s), HERE); + raise::Error(fmt::format("Invalid TwoBodyInteraction type: %s", s.c_str()), + HERE); return NONE; } } diff --git a/src/global/traits/policies.h b/src/global/traits/policies.h index 7f6d98985..f82220a27 100644 --- a/src/global/traits/policies.h +++ b/src/global/traits/policies.h @@ -138,4 +138,43 @@ concept CustomParticleUpdatePolicyClass = } or traits::custom_prtl_update::IsNoPolicy; +<<<<<<< HEAD +======= +namespace traits::twobodyinteractions { + + template + concept HasSpecies = requires(I& interaction_policy) { + { interaction_policy.species } -> std::convertible_to; + }; + + template + concept HasShouldInteract = requires(const I& interaction_policy, + spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2, + real_t tile_weight) { + { + interaction_policy.should_interact(sp1, p1, sp2, p2, tile_weight) + } -> std::same_as; + }; + + template + concept HasInteraction = requires(const I& interaction_policy, + spidx_t sp1, + npart_t p1, + spidx_t sp2, + npart_t p2) { + { interaction_policy(sp1, p1, sp2, p2) } -> std::same_as; + }; + +} // namespace traits::twobodyinteractions + +template +concept TwoBodyInteractionPolicyClass = + traits::twobodyinteractions::HasSpecies and + traits::twobodyinteractions::HasShouldInteract and + traits::twobodyinteractions::HasInteraction; + +>>>>>>> a09db826 (qed incorporated) #endif // TRAITS_POLICIES_H diff --git a/src/kernels/injectors.hpp b/src/kernels/injectors.hpp index 05d83a0d4..cf99a698a 100644 --- a/src/kernels/injectors.hpp +++ b/src/kernels/injectors.hpp @@ -5,6 +5,7 @@ * - kernel::UniformInjector_kernel<> * - kernel::GlobalInjector_kernel<> * - kernel::NonUniformInjector_kernel<> + * - kernel::SingleSpeciesUniformInjector_kernel<> * @namespaces: * - kernel:: */ @@ -853,6 +854,131 @@ namespace kernel { } }; // struct NonUniformInjector_kernel + template ED> + struct SingleSpeciesUniformInjector_kernel { + + ParticleArrays particles; + + const npart_t offset; + const npart_t domain_idx, cntr; + const bool use_tracking; + const M metric; + const array_t xi_min, xi_max; + const ED energy_dist; + const real_t inv_V0; + random_number_pool_t random_pool; + + SingleSpeciesUniformInjector_kernel(Particles& particles, + npart_t domain_idx, + const M& metric, + const array_t& xi_min, + const array_t& xi_max, + const ED& energy_dist, + real_t inv_V0, + random_number_pool_t& random_pool) + : particles { particles } + , offset { particles.npart() } + , domain_idx { domain_idx } + , cntr { particles.counter() } + , use_tracking { particles.use_tracking() } + , metric { metric } + , xi_min { xi_min } + , xi_max { xi_max } + , energy_dist { energy_dist } + , inv_V0 { inv_V0 } + , random_pool { random_pool } { + if (use_tracking) { +#if !defined(MPI_ENABLED) + raise::ErrorIf(particles.pld_i.extent(1) < 1, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 1", + HERE); +#else + raise::ErrorIf(particles.pld_i.extent(1) < 2, + "Particle tracking is enabled but the " + "particle integer payload size is less " + "than 2", + HERE); +#endif + } + } + + Inline void operator()(prtlidx_t p) const { + coord_t x_Cd { ZERO }; + tuple_t xi_Cd { 0 }; + tuple_t dxi_Cd { static_cast(0) }; + vec_t v { ZERO, ZERO, 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::Cartesian) { + energy_dist(x_Ph, v); + } 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(x_Ph, v_Ph); + metric.template transform_xyz(x_Cd_, v_Ph, v); + } else if constexpr (S == SimEngine::GRPIC) { + vec_t v_Ph { ZERO, ZERO, ZERO }; + energy_dist(x_Ph, v_Ph); + metric.template transform(x_Cd, v_Ph, v); + } else { + raise::KernelError(HERE, "Unknown simulation engine"); + } + } + real_t weight = ONE; + if constexpr (M::CoordType != Coord::Cartesian) { + const auto sqrt_det_h = metric.sqrt_det_h(x_Cd); + weight = sqrt_det_h * inv_V0; + } + // clang-format off + if (not use_tracking) { + InjectParticle( + p + offset, + particles.i1, particles.i2, particles.i3, + particles.dx1, particles.dx2, particles.dx3, + particles.ux1, particles.ux2, particles.ux3, + particles.phi, particles.weight, particles.tag, particles.pld_i, + xi_Cd, dxi_Cd, v, weight, ZERO); + } else { + InjectParticle( + p + offset, + particles.i1, particles.i2, particles.i3, + particles.dx1, particles.dx2, particles.dx3, + particles.ux1, particles.ux2, particles.ux3, + particles.phi, particles.weight, particles.tag, particles.pld_i, + xi_Cd, dxi_Cd, v, weight, ZERO, + domain_idx, cntr + p); + } + // clang-format on + } + }; // struct SingleSpeciesUniformInjector_kernel + } // namespace kernel #endif // KERNELS_INJECTORS_HPP diff --git a/src/kernels/twobody_interactions.hpp b/src/kernels/twobody_interactions.hpp index 44385e885..e0a41b1fd 100644 --- a/src/kernels/twobody_interactions.hpp +++ b/src/kernels/twobody_interactions.hpp @@ -73,32 +73,32 @@ namespace kernel::mink { } template - Inline auto TileIdxToVolume(ncells_t tile_idx, - ncells_t tile_size, - ncells_t ntx2, - ncells_t ntx3, - ncells_t nx1, - ncells_t nx2, - ncells_t nx3) -> real_t { - real_t tile_volume { ONE }; + Inline auto NCellsOnTile(ncells_t tile_idx, + ncells_t tile_size, + ncells_t ntx2, + ncells_t ntx3, + ncells_t nx1, + ncells_t nx2, + ncells_t nx3) -> ncells_t { + ncells_t ncells_on_tile { 1 }; ncells_t ti { 0u }, tj { 0u }, tk { 0u }; UnravelTileIdx(tile_idx, ntx2, ntx3, ti, tj, tk); if constexpr ((D == Dim::_1D) or (D == Dim::_2D) or (D == Dim::_3D)) { const auto i1_min = ti * tile_size; const auto i1_max = math::min(i1_min + tile_size, nx1); - tile_volume *= static_cast(i1_max - i1_min); + ncells_on_tile *= (i1_max - i1_min); } if constexpr ((D == Dim::_2D) or (D == Dim::_3D)) { const auto i2_min = tj * tile_size; const auto i2_max = math::min(i2_min + tile_size, nx2); - tile_volume *= static_cast(i2_max - i2_min); + ncells_on_tile *= (i2_max - i2_min); } if constexpr (D == Dim::_3D) { const auto i3_min = tk * tile_size; const auto i3_max = math::min(i3_min + tile_size, nx3); - tile_volume *= static_cast(i3_max - i3_min); + ncells_on_tile *= (i3_max - i3_min); } - return tile_volume; + return ncells_on_tile; } template @@ -112,11 +112,10 @@ namespace kernel::mink { ncells_t num_tiles { 0u }; - CollisionGroup( - const std::vector*>& particles, - const std::vector& ncells, - ncells_t tile_size, - random_number_pool_t& random_pool) { + CollisionGroup(const std::vector*>& particles, + const std::vector& ncells, + ncells_t tile_size, + random_number_pool_t& random_pool) { for (const auto* species : particles) { const auto npart_s = species->npart(); array_t tileidx { "tile_idx", npart_s }; @@ -215,13 +214,14 @@ namespace kernel::mink { template void TwoBodyInteraction( - const std::vector*>& species1, - const std::vector*>& species2, - const std::vector& ncells, - const boundaries_t& domain_extent, - ncells_t tile_size, - random_number_pool_t& random_pool, - const I& interaction_policy) { + const std::vector*>& species1, + const std::vector*>& species2, + const std::vector& ncells, + const boundaries_t& domain_extent, + ncells_t tile_size, + real_t ppc0, + random_number_pool_t& random_pool, + I& interaction_policy) { raise::ErrorIf(species1.empty() or species2.empty(), "species groups must be non-empty", HERE); @@ -231,7 +231,7 @@ namespace kernel::mink { raise::ErrorIf(domain_extent.size() != static_cast(D), "domain_extent size must match D", HERE); - // compute base tile volume in physical units + // compute base cell volume in physical units real_t cell_volume { ONE }; for (int d = 0; d < static_cast(D); ++d) { cell_volume *= static_cast( @@ -253,6 +253,56 @@ namespace kernel::mink { const auto& tile_offsets1 = group1.tile_offsets; const auto& tile_offsets2 = group2.tile_offsets; + // fill species in the interaction policy + for (auto& sp1 : species1) { + interaction_policy.species[sp1->sp - 1] = static_cast( + *sp1); + } + for (auto& sp2 : species2) { + interaction_policy.species[sp2->sp - 1] = static_cast( + *sp2); + } + + // total particle weight on each tile + auto weights_on_tile1 = array_t { "weights_on_tile1", num_tiles }; + auto weights_on_tile2 = array_t { "weights_on_tile2", num_tiles }; + Kokkos::parallel_for( + "ComputeWeightsOnTiles", + Kokkos::TeamPolicy<>(num_tiles, Kokkos::AUTO), + Lambda(const Kokkos::TeamPolicy<>::member_type& team) { + const ncells_t t = team.league_rank(); + + const auto o1 = tile_offsets1(t); + const auto num_ppt1 = combined_num_ppt1(t); + real_t weight_on_tile1 = ZERO; + Kokkos::parallel_reduce( + Kokkos::TeamThreadRange(team, num_ppt1), + [&](prtlidx_t i, real_t& lsum) { + const auto sp1 = static_cast(combined_idx1(o1 + i) >> 56); + const auto p1 = static_cast(combined_idx1(o1 + i) & + ((1ull << 56) - 1)); + + lsum += interaction_policy.species[sp1 - 1].weight(p1); + }, + weight_on_tile1); + weights_on_tile1(t) = weight_on_tile1; + + const auto o2 = tile_offsets2(t); + const auto num_ppt2 = combined_num_ppt2(t); + real_t weight_on_tile2 = ZERO; + Kokkos::parallel_reduce( + Kokkos::TeamThreadRange(team, num_ppt2), + [&](prtlidx_t i, real_t& lsum) { + const auto sp2 = static_cast(combined_idx2(o2 + i) >> 56); + const auto p2 = static_cast(combined_idx2(o2 + i) & + ((1ull << 56) - 1)); + + lsum += interaction_policy.species[sp2 - 1].weight(p2); + }, + weight_on_tile2); + weights_on_tile2(t) = weight_on_tile2; + }); + // number of cells in each direction ncells_t nx1 { 1u }, nx2 { 1u }, nx3 { 1u }; if constexpr ((D == Dim::_1D) or (D == Dim::_2D) or (D == Dim::_3D)) { @@ -274,13 +324,18 @@ namespace kernel::mink { math::ceil(static_cast(nx3) / static_cast(tile_size))); } + array_t interaction_pairs { "interaction_pairs", + combined_idx1.extent(0) }; + array_t counter { "counter" }; + Kokkos::parallel_for( - "EmitPairs", + "PopulateInteractionPairs", Kokkos::TeamPolicy<>(num_tiles, Kokkos::AUTO), Lambda(const Kokkos::TeamPolicy<>::member_type& team) { const ncells_t t = team.league_rank(); - const auto tile_volume = - TileIdxToVolume(t, tile_size, ntx2, ntx3, nx1, nx2, nx3) * cell_volume; + const auto tile_weight = + math::max(weights_on_tile1(t), weights_on_tile2(t)) / + (NCellsOnTile(t, tile_size, ntx2, ntx3, nx1, nx2, nx3) * ppc0); const auto k = math::min(combined_num_ppt1(t), combined_num_ppt2(t)); const auto o1 = tile_offsets1(t); @@ -295,9 +350,27 @@ namespace kernel::mink { ((1ull << 56) - 1)); const auto p2 = static_cast(combined_idx2(o2 + i) & ((1ull << 56) - 1)); - interaction_policy(sp1, p1, sp2, p2, tile_volume); + if (interaction_policy.should_interact(sp1, p1, sp2, p2, tile_weight)) { + const auto idx = Kokkos::atomic_fetch_add(&counter(), 1); + interaction_pairs(idx, 0) = combined_idx1(o1 + i); + interaction_pairs(idx, 1) = combined_idx2(o2 + i); + } }); }); + auto counter_h = Kokkos::create_mirror_view(counter); + Kokkos::deep_copy(counter_h, counter); + Kokkos::parallel_for( + "ProcessInteractions", + counter_h(), + Lambda(prtlidx_t idx) { + const auto sp1 = static_cast(interaction_pairs(idx, 0) >> 56); + const auto sp2 = static_cast(interaction_pairs(idx, 1) >> 56); + const auto p1 = static_cast(interaction_pairs(idx, 0) & + ((1ull << 56) - 1)); + const auto p2 = static_cast(interaction_pairs(idx, 1) & + ((1ull << 56) - 1)); + interaction_policy(sp1, p1, sp2, p2); + }); } } // namespace kernel::mink diff --git a/tests/archetypes/qed_compton.cpp b/tests/archetypes/qed_compton.cpp index 265439a1b..14350cbf4 100644 --- a/tests/archetypes/qed_compton.cpp +++ b/tests/archetypes/qed_compton.cpp @@ -160,13 +160,14 @@ auto main(int argc, char* argv[]) -> int { { ZERO, ONE }, { -ONE, ONE } }; + const auto ppc0 = static_cast(npart) / (nx1 * nx2); prm::Parameters params; - params.set("qed.compton_scattering.nominal_probability_density", + params.set("compton_scattering.nominal_probability_density", static_cast(1e-3)); - params.set("qed.compton_scattering.Thomson_limit", static_cast(1e-4)); - auto policy = arch::qed::ComptonScattering(params, random_pool); + auto policy = arch::qed::ComptonScattering(params, + random_pool); policy.species[0] = static_cast(sp1); policy.species[1] = static_cast(sp2); policy.species[2] = static_cast(sp3); @@ -198,6 +199,7 @@ auto main(int argc, char* argv[]) -> int { ncells, extent, tile_size, + ppc0, random_pool, policy); } diff --git a/tests/kernels/twobody_interactions.cpp b/tests/kernels/twobody_interactions.cpp new file mode 100644 index 000000000..8051c5805 --- /dev/null +++ b/tests/kernels/twobody_interactions.cpp @@ -0,0 +1,187 @@ +#include "kernels/twobody_interactions.hpp" + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "utils/error.h" + +#include "framework/containers/particles.h" +#include "kernels/twobody_interactions.hpp" + +#include + +#include +#include + +using namespace ntt; + +// Verifies that each paired particle from group1 and group2 lies in the same tile +struct SameTilePolicy { + ParticleArrays species[4]; + const ncells_t tile_size; + const ncells_t ncx1, ncx2; // number of cells in each direction + const ncells_t ntx1, ntx2; // numbers of tiles + array_t diff_tile_errors { "diff_tile_errors" }; + + SameTilePolicy(ncells_t tile_size, + ncells_t ncx1, + ncells_t ncx2, + ncells_t ntx1, + ncells_t ntx2) + : tile_size { tile_size } + , ncx1 { ncx1 } + , ncx2 { ncx2 } + , ntx1 { ntx1 } + , ntx2 { ntx2 } {} + + Inline auto should_interact(spidx_t, npart_t, spidx_t, npart_t, real_t) const + -> bool { + return true; + } + + Inline void operator()(spidx_t sp1, npart_t p1, spidx_t sp2, npart_t p2) const { + const auto x1_1 = species[sp1 - 1].i1(p1); + const auto x2_1 = species[sp1 - 1].i2(p1); + const auto x1_2 = species[sp2 - 1].i1(p2); + const auto x2_2 = species[sp2 - 1].i2(p2); + const auto t1 = static_cast(x1_1 / tile_size) * ntx2 + + static_cast(x2_1 / tile_size); + const auto t2 = static_cast(x1_2 / tile_size) * ntx2 + + static_cast(x2_2 / tile_size); + if (t1 != t2) { + Kokkos::atomic_add(&diff_tile_errors(), 1); + } + } +}; + +void fill_random(array_t& i1, + array_t& i2, + array_t& tag, + npart_t npart, + ncells_t nx1, + ncells_t nx2, + random_number_pool_t& rpool) { + Kokkos::parallel_for( + "FillRandom", + npart, + KOKKOS_LAMBDA(const npart_t p) { + auto gen = rpool.get_state(); + i1(p) = static_cast(gen.urand() % static_cast(nx1)); + i2(p) = static_cast(gen.urand() % static_cast(nx2)); + tag(p) = ParticleTag::alive; + rpool.free_state(gen); + }); + Kokkos::fence(); +} + +auto main(int argc, char* argv[]) -> int { + ntt::GlobalInitialize(argc, argv); + + try { + const ncells_t nx1 = 32u; + const ncells_t nx2 = 64u; + const ncells_t tile_size = 3u; + const std::vector ncells = { nx1, nx2 }; + const boundaries_t extent = { + { ZERO, ONE }, + { ZERO, TWO } + }; + const ncells_t ntx1 = static_cast( + math::ceil(static_cast(nx1) / static_cast(tile_size))); + const ncells_t ntx2 = static_cast( + math::ceil(static_cast(nx2) / static_cast(tile_size))); + const npart_t npart = 1000u; + random_number_pool_t random_pool { 12345u }; + + Particles sp1 { 1u, + "sp1", + 1.0f, + 1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + Particles sp2 { 2u, + "sp2", + 1.0f, + 1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + Particles sp3 { 3u, + "sp3", + 1.0f, + 1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + Particles sp4 { 4u, + "sp4", + 1.0f, + 1.0f, + npart, + 0u, + 0u, + ParticlePusher::BORIS, + false, + RadiativeDrag::NONE, + EmissionType::NONE, + 0u, + 0u }; + + for (auto* sp : { &sp1, &sp2, &sp3, &sp4 }) { + sp->set_npart(npart); + fill_random(sp->i1, sp->i2, sp->tag, npart, nx1, nx2, random_pool); + } + + const std::vector*> group1 = { &sp1, + &sp2 }; + const std::vector*> group2 = { &sp3, + &sp4 }; + + auto policy = SameTilePolicy { tile_size, nx1, nx2, ntx1, ntx2 }; + + kernel::mink::TwoBodyInteraction(group1, + group2, + ncells, + extent, + tile_size, + ONE, + random_pool, + policy); + Kokkos::fence(); + + { + auto errors_h = Kokkos::create_mirror_view(policy.diff_tile_errors); + Kokkos::deep_copy(errors_h, policy.diff_tile_errors); + raise::ErrorIf(errors_h() != 0, + "paired particles from different tiles detected", + HERE); + } + + } catch (std::exception& e) { + std::cerr << e.what() << '\n'; + ntt::GlobalFinalize(); + return 1; + } + ntt::GlobalFinalize(); + return 0; +} From 4c11843578239fca035d4a1427c2be41d112ee2f Mon Sep 17 00:00:00 2001 From: haykh Date: Fri, 15 May 2026 08:56:41 -0400 Subject: [PATCH 11/25] rebased --- src/global/traits/policies.h | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/global/traits/policies.h b/src/global/traits/policies.h index f82220a27..fdca30ef3 100644 --- a/src/global/traits/policies.h +++ b/src/global/traits/policies.h @@ -138,8 +138,6 @@ concept CustomParticleUpdatePolicyClass = } or traits::custom_prtl_update::IsNoPolicy; -<<<<<<< HEAD -======= namespace traits::twobodyinteractions { template @@ -176,5 +174,4 @@ concept TwoBodyInteractionPolicyClass = traits::twobodyinteractions::HasShouldInteract and traits::twobodyinteractions::HasInteraction; ->>>>>>> a09db826 (qed incorporated) #endif // TRAITS_POLICIES_H From 4f6d914c0c12d58c35ae8376f2387e45c9eaa81f Mon Sep 17 00:00:00 2001 From: haykh Date: Fri, 15 May 2026 11:15:58 -0400 Subject: [PATCH 12/25] tests for compton fixed (proper cross section) --- examples/compton_jones/compton_jones.py | 1 + examples/compton_jones/compton_jones.toml | 6 ++--- .../compton_kompaneets/compton_kompaneets.py | 24 ++++++++++++------- .../compton_kompaneets.toml | 6 ++--- src/archetypes/qed/compton.h | 2 +- src/kernels/twobody_interactions.hpp | 21 +++++++++++----- 6 files changed, 38 insertions(+), 22 deletions(-) diff --git a/examples/compton_jones/compton_jones.py b/examples/compton_jones/compton_jones.py index 873a01f35..505b777bf 100644 --- a/examples/compton_jones/compton_jones.py +++ b/examples/compton_jones/compton_jones.py @@ -9,6 +9,7 @@ plt.rcParams["figure.dpi"] = 300 plt.rcParams["font.family"] = "serif" +plt.rcParams["mathtext.fontset"] = "stix" fig = plt.figure(figsize=(9, 4)) gs = fig.add_gridspec(1, 2, wspace=0.35) diff --git a/examples/compton_jones/compton_jones.toml b/examples/compton_jones/compton_jones.toml index 796739978..8230a04c5 100644 --- a/examples/compton_jones/compton_jones.toml +++ b/examples/compton_jones/compton_jones.toml @@ -19,7 +19,7 @@ skindepth0 = 1.0 [two_body] - thomson_optical_depth = 0.5 + thomson_optical_depth = 100.0 [[two_body.interaction]] type = "compton" @@ -40,7 +40,7 @@ enable = false [particles] - ppc0 = 1.0 + ppc0 = 5.0 clear_interval = 1 [[particles.species]] @@ -67,7 +67,7 @@ photon_energy = 1e-2 [output] - interval_time = 0.01 + interval_time = 9.0 [output.fields] quantities = ["N_1", "N_2", "N_3"] diff --git a/examples/compton_kompaneets/compton_kompaneets.py b/examples/compton_kompaneets/compton_kompaneets.py index 4c713cbad..b507964c5 100644 --- a/examples/compton_kompaneets/compton_kompaneets.py +++ b/examples/compton_kompaneets/compton_kompaneets.py @@ -9,20 +9,26 @@ plt.rcParams["figure.dpi"] = 300 plt.rcParams["font.family"] = "serif" +plt.rcParams["mathtext.fontset"] = "stix" fig = plt.figure(figsize=(9, 4)) gs = fig.add_gridspec(1, 2, wspace=0.3) ax1 = fig.add_subplot(gs[0, 0]) +tC = 1 / ( + data.attrs["setup.temperature"] * data.attrs["two_body.thomson_optical_depth"] +) + tvals = len(data.spectra.t.values) nphot = data.spectra.N_3.isel(t=-1).sum().values[()] -for ti in range(0, tvals, 10): +for ti, t in enumerate([0, 0.5 * tC, tC, 3 * tC]): ax1.plot( - data.spectra.E.values / data.attrs["setup.temperature"], - data.spectra.N_3.isel(t=ti).values, - c=plt.get_cmap("plasma")(ti / tvals), - lw=0.5, + data.spectra.E.coarsen(E=2).mean().values / data.attrs["setup.temperature"], + data.spectra.N_3.sel(t=t, method="nearest").coarsen(E=2).mean().values, + c=plt.get_cmap("viridis")(ti / 3), + label=f"$y={t / tC:.1f}$", + lw=1, ) es = data.spectra.E.values / data.attrs["setup.temperature"] @@ -39,7 +45,7 @@ ax1.set( yscale="log", - ylim=(1e-1, 1e5), + ylim=(1, 1e5), xlim=(0, 10), xlabel=r"$\varepsilon / T_\pm$", ylabel=r"$dn_{\rm ph}/d\varepsilon$", @@ -47,12 +53,12 @@ ax1.legend() ax2 = fig.add_subplot(gs[0, 1]) -ax2.plot(stats["time"], stats["T00_3"], c="C0") -ax2.set(ylabel=r"total photon energy", xlabel=r"$t$") +ax2.plot(stats["time"] / tC, stats["T00_3"], c="C0") +ax2.set(ylabel=r"total photon energy", xlabel=r"$y\equiv t/t_C$") ax2.yaxis.label.set_color("C0") ax2.tick_params(axis="y", labelcolor="C0") ax2twin = ax2.twinx() -ax2twin.plot(data.spectra.t.values, data.spectra.N_3.sum("E"), c="C2") +ax2twin.plot(data.spectra.t.values / tC, data.spectra.N_3.sum("E"), c="C2") ax2twin.set(ylabel=r"photon number") ax2twin.yaxis.label.set_color("C2") ax2twin.tick_params(axis="y", labelcolor="C2") diff --git a/examples/compton_kompaneets/compton_kompaneets.toml b/examples/compton_kompaneets/compton_kompaneets.toml index 1a006102f..a92699968 100644 --- a/examples/compton_kompaneets/compton_kompaneets.toml +++ b/examples/compton_kompaneets/compton_kompaneets.toml @@ -19,7 +19,7 @@ skindepth0 = 1.0 [two_body] - thomson_optical_depth = 0.25 + thomson_optical_depth = 60.0 [[two_body.interaction]] type = "compton" @@ -66,7 +66,7 @@ photon_energy = 1e-3 [output] - interval_time = 0.01 + interval_time = 0.1 [output.fields] enable = false @@ -81,7 +81,7 @@ n_bins = 500 [output.stats] - quantities = ["T00_1", "T00_2", "T00_3"] + quantities = ["T00_3"] [checkpoint] keep = 0 diff --git a/src/archetypes/qed/compton.h b/src/archetypes/qed/compton.h index 4f466579e..8776d4a12 100644 --- a/src/archetypes/qed/compton.h +++ b/src/archetypes/qed/compton.h @@ -42,7 +42,7 @@ namespace arch::qed { "compton_scattering.nominal_probability_density") } , random_pool { random_pool } { if (nominal_probability_density <= ZERO) { - raise::Error("nominal_probability must be in the range (0, 1]", HERE); + raise::Error("nominal_probability must be > 0", HERE); } } diff --git a/src/kernels/twobody_interactions.hpp b/src/kernels/twobody_interactions.hpp index e0a41b1fd..901b43acf 100644 --- a/src/kernels/twobody_interactions.hpp +++ b/src/kernels/twobody_interactions.hpp @@ -282,7 +282,10 @@ namespace kernel::mink { const auto p1 = static_cast(combined_idx1(o1 + i) & ((1ull << 56) - 1)); - lsum += interaction_policy.species[sp1 - 1].weight(p1); + lsum += (interaction_policy.species[sp1 - 1].tag(p1) == + ParticleTag::alive) + ? interaction_policy.species[sp1 - 1].weight(p1) + : ZERO; }, weight_on_tile1); weights_on_tile1(t) = weight_on_tile1; @@ -297,7 +300,10 @@ namespace kernel::mink { const auto p2 = static_cast(combined_idx2(o2 + i) & ((1ull << 56) - 1)); - lsum += interaction_policy.species[sp2 - 1].weight(p2); + lsum += (interaction_policy.species[sp2 - 1].tag(p2) == + ParticleTag::alive) + ? interaction_policy.species[sp2 - 1].weight(p2) + : ZERO; }, weight_on_tile2); weights_on_tile2(t) = weight_on_tile2; @@ -350,10 +356,13 @@ namespace kernel::mink { ((1ull << 56) - 1)); const auto p2 = static_cast(combined_idx2(o2 + i) & ((1ull << 56) - 1)); - if (interaction_policy.should_interact(sp1, p1, sp2, p2, tile_weight)) { - const auto idx = Kokkos::atomic_fetch_add(&counter(), 1); - interaction_pairs(idx, 0) = combined_idx1(o1 + i); - interaction_pairs(idx, 1) = combined_idx2(o2 + i); + if ((interaction_policy.species[sp1 - 1].tag(p1) == ParticleTag::alive) and + (interaction_policy.species[sp2 - 1].tag(p2) == ParticleTag::alive)) { + if (interaction_policy.should_interact(sp1, p1, sp2, p2, tile_weight)) { + const auto idx = Kokkos::atomic_fetch_add(&counter(), 1); + interaction_pairs(idx, 0) = combined_idx1(o1 + i); + interaction_pairs(idx, 1) = combined_idx2(o2 + i); + } } }); }); From e6f4c48179fc7506ef2433f22886f9a578523df9 Mon Sep 17 00:00:00 2001 From: haykh Date: Fri, 15 May 2026 11:24:53 -0400 Subject: [PATCH 13/25] fmt script updated --- dev/scripts/format.sh | 62 +++++++++++++++++++++---------------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/dev/scripts/format.sh b/dev/scripts/format.sh index e12bef2ec..ca4aab0f4 100755 --- a/dev/scripts/format.sh +++ b/dev/scripts/format.sh @@ -3,43 +3,43 @@ verify=false for arg in "$@"; do - case $arg in - --verify) verify=true ;; - esac + case $arg in + --verify) verify=true ;; + esac done if $verify; then - diff_output="" + diff_output="" - if command -v cmake-format &>/dev/null; then - while IFS= read -r -d '' f; do - if ! diff -q <(cmake-format "$f") "$f" &>/dev/null; then - diff_output+=" $f\n" - fi - done < <(find cmake/ src/ minimal/ tests/ -type f \( -name "*.cmake" -o -name "*.txt" \) -print0) - fi + if command -v cmake-format &>/dev/null; then + while IFS= read -r -d '' f; do + if ! diff -q <(cmake-format "$f") "$f" &>/dev/null; then + diff_output+=" $f\n" + fi + done < <(find cmake/ src/ minimal/ tests/ -type f \( -name "*.cmake" -o -name "*.txt" \) -print0) + fi - if command -v clang-format &>/dev/null; then - while IFS= read -r -d '' f; do - if ! clang-format --style=file --dry-run --Werror "$f" &>/dev/null; then - diff_output+=" $f\n" - fi - done < <(find pgens/ src/ minimal/ tests/ -type f \( -name "*.cpp" -o -name "*.hpp" -o -name "*.h" \) -print0) - fi + if command -v clang-format &>/dev/null; then + while IFS= read -r -d '' f; do + if ! clang-format --style=file --dry-run --Werror "$f" &>/dev/null; then + diff_output+=" $f\n" + fi + done < <(find pgens/ examples/ tutorials/ src/ minimal/ tests/ -type f \( -name "*.cpp" -o -name "*.hpp" -o -name "*.h" \) -print0) + fi - if [ -n "$diff_output" ]; then - echo "Formatting check failed. The following files need formatting:" - printf "$diff_output" - exit 1 - else - echo "All files are properly formatted." - fi + if [ -n "$diff_output" ]; then + echo "Formatting check failed. The following files need formatting:" + printf '%s' "$diff_output" + exit 1 + else + echo "All files are properly formatted." + fi else - if command -v cmake-format &>/dev/null; then - find cmake/ src/ minimal/ tests/ -type f -name "*.cmake" -o -name "*.txt" | xargs cmake-format -i - fi + if command -v cmake-format &>/dev/null; then + find cmake/ src/ minimal/ tests/ \( -type f -name "*.cmake" -o -name "*.txt" \) -exec cmake-format -i {} \; + fi - if command -v clang-format &>/dev/null; then - find pgens/ src/ minimal/ tests/ -type f -name "*.cpp" -o -name "*.hpp" -o -name "*.h" | xargs clang-format --style=file -i - fi + if command -v clang-format &>/dev/null; then + find pgens/ src/ minimal/ tests/ examples/ pgens/ tutorials/ \( -type f -name "*.cpp" -o -name "*.hpp" -o -name "*.h" \) -exec clang-format --style=file -i {} \; + fi fi From 7f57abb69f17b7c77979ea8993d413bc721f4a26 Mon Sep 17 00:00:00 2001 From: haykh Date: Fri, 15 May 2026 13:12:30 -0400 Subject: [PATCH 14/25] clang-tidy recommendations --- src/engines/reporter.cpp | 1 + src/framework/parameters/extra.cpp | 2 ++ 2 files changed, 3 insertions(+) diff --git a/src/engines/reporter.cpp b/src/engines/reporter.cpp index 11ae740ed..125cb45a5 100644 --- a/src/engines/reporter.cpp +++ b/src/engines/reporter.cpp @@ -9,6 +9,7 @@ #include "framework/parameters/extra.h" #include "framework/parameters/parameters.h" +#include #include #include diff --git a/src/framework/parameters/extra.cpp b/src/framework/parameters/extra.cpp index 53eca4117..f0d9d127f 100644 --- a/src/framework/parameters/extra.cpp +++ b/src/framework/parameters/extra.cpp @@ -5,12 +5,14 @@ #include "utils/numeric.h" +#include "framework/parameters/extra.h" #include "framework/parameters/parameters.h" #include #include #include +#include namespace ntt { namespace params { From 5ca3bd2279fa530dcbb133bb4e8831ca5b040978 Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 21 May 2026 10:58:07 -0400 Subject: [PATCH 15/25] confine zerofill only to necessary components --- src/archetypes/utils.h | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/archetypes/utils.h b/src/archetypes/utils.h index e371b4e8c..c15ffabe3 100644 --- a/src/archetypes/utils.h +++ b/src/archetypes/utils.h @@ -143,7 +143,17 @@ namespace arch { const auto inv_n0 = ONE / params.template get("scales.n0"); const auto use_weights = params.template get("particles.use_weights"); - Kokkos::deep_copy(buffer, ZERO); + if constexpr (M::Dim == Dim::_1D) { + Kokkos::deep_copy(Kokkos::subview(buffer, Kokkos::ALL(), buffer_idx), ZERO); + } else if constexpr (M::Dim == Dim::_2D) { + Kokkos::deep_copy( + Kokkos::subview(buffer, Kokkos::ALL, Kokkos::ALL, buffer_idx), + ZERO); + } else if constexpr (M::Dim == Dim::_3D) { + Kokkos::deep_copy( + Kokkos::subview(buffer, Kokkos::ALL, Kokkos::ALL, Kokkos::ALL, buffer_idx), + ZERO); + } auto scatter_buff = Kokkos::Experimental::create_scatter_view(buffer); for (const auto sp : species) { const auto& prtl_spec = domain.species[sp - 1]; From 5722ca4905f3811910e7ccc549ffe3d490d04483 Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 21 May 2026 10:58:22 -0400 Subject: [PATCH 16/25] support for average V scalar --- src/kernels/particle_moments.hpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/kernels/particle_moments.hpp b/src/kernels/particle_moments.hpp index f2799978c..91359d6d9 100644 --- a/src/kernels/particle_moments.hpp +++ b/src/kernels/particle_moments.hpp @@ -240,7 +240,11 @@ namespace kernel { } else { u0 = math::sqrt(ONE + NORM_SQR(u_Phys[0], u_Phys[1], u_Phys[2])); } - return (mass == ZERO ? ONE : mass) * u_Phys[c1 - 1] / u0; + if (c1 > 0u) { + return (mass == ZERO ? ONE : mass) * u_Phys[c1 - 1] / u0; + } else { + return (mass == ZERO ? ONE : (mass * math::sqrt(ONE - SQR(ONE / u0)))); + } } Inline auto computeEckartVelocityFluxComponent(prtlidx_t p) const -> real_t { From 8db0b2bc3d5f30c6a381c38b198984820236c86e Mon Sep 17 00:00:00 2001 From: haykh Date: Thu, 21 May 2026 10:58:32 -0400 Subject: [PATCH 17/25] single species nonuniform injector --- src/archetypes/particle_injector.h | 85 ++++++++++- src/kernels/injectors.hpp | 233 +++++++++++++++++++++++++++++ 2 files changed, 312 insertions(+), 6 deletions(-) diff --git a/src/archetypes/particle_injector.h b/src/archetypes/particle_injector.h index d53de9f14..b12d4106b 100644 --- a/src/archetypes/particle_injector.h +++ b/src/archetypes/particle_injector.h @@ -326,12 +326,8 @@ namespace arch { "Weights must be used for non-Cartesian coordinates", HERE); raise::ErrorIf( - params.template get("particles.use_weights") and not use_weights, - "Weights are enabled in the input but not enabled in the injector", - HERE); - raise::ErrorIf( - not params.template get("particles.use_weights") and use_weights, - "Weights are not enabled in the input but enabled in the injector", + params.template get("particles.use_weights") != use_weights, + "Mismatch between use_weights in the input file and the injector", HERE); if (domain.species[species.first - 1].charge() + domain.species[species.second - 1].charge() != @@ -386,6 +382,83 @@ namespace arch { } } + /** + * @brief Injects particles based on spatial distribution function + * @param params Simulation parameters + * @param domain Local domain object + * @param species Species index + * @param energy_dist Energy distribution class + * @param spatial_dist Spatial distribution class + * @param number_density Number density (in units of n0) + * @param use_weights Use weights + * @param box Region to inject the particles in + * @tparam S Simulation engine type + * @tparam M Metric type + * @tparam ED Energy distribution type + * @tparam SD Spatial distribution type + */ + template ED, SpatialDistClass SD> + inline void InjectNonUniform(const SimulationParams& params, + Domain& domain, + spidx_t species, + const ED& energy_dist, + const SD& spatial_dist, + real_t number_density, + bool use_weights = (M::CoordType != Coord::Cartesian), + const boundaries_t& box = {}) { + raise::ErrorIf((M::CoordType != Coord::Cartesian) && (not use_weights), + "Weights must be used for non-Cartesian coordinates", + HERE); + raise::ErrorIf( + params.template get("particles.use_weights") != use_weights, + "Mismatch between use_weights in the input file and the injector", + HERE); + { + range_t cell_range; + if (box.empty()) { + cell_range = domain.mesh.rangeActiveCells(); + } else { + boundaries_t reduced_box(box); + if (reduced_box.size() > M::Dim) { + reduced_box.resize(M::Dim); + } + raise::ErrorIf(reduced_box.size() != M::Dim, + "Box must have the same dimension as the mesh", + HERE); + boundaries_t incl_ghosts; + for (auto d = 0; d < M::Dim; ++d) { + incl_ghosts.emplace_back(false, false); + } + const auto extent = domain.mesh.ExtentToRange(reduced_box, incl_ghosts); + tuple_t x_min { 0 }, x_max { 0 }; + for (auto d = 0; d < M::Dim; ++d) { + x_min[d] = extent[d].first; + x_max[d] = extent[d].second; + } + cell_range = CreateRangePolicy(x_min, x_max); + } + const auto ppc = number_density * + params.template get("particles.ppc0") * HALF; + auto injector_kernel = kernel::SingleSpeciesNonUniformInjector_kernel( + ppc, + domain.species[species - 1], + domain.index(), + domain.mesh.metric, + energy_dist, + spatial_dist, + ONE / params.template get("scales.V0"), + domain.random_pool()); + Kokkos::parallel_for("InjectSingleSpeciesNonUniformNumberDensity", + cell_range, + injector_kernel); + const auto n_inj = injector_kernel.number_injected(); + domain.species[species - 1].set_npart( + domain.species[species - 1].npart() + n_inj); + domain.species[species - 1].set_counter( + domain.species[species - 1].counter() + n_inj); + } + } + /** * @brief Injects uniform number density of a single species everywhere in the domain * @param domain Domain object diff --git a/src/kernels/injectors.hpp b/src/kernels/injectors.hpp index 3462e3b2e..49ef50a21 100644 --- a/src/kernels/injectors.hpp +++ b/src/kernels/injectors.hpp @@ -5,6 +5,7 @@ * - kernel::UniformInjector_kernel<> * - kernel::GlobalInjector_kernel<> * - kernel::NonUniformInjector_kernel<> + * - kernel::SingleSpeciesNonUniformInjector_kernel<> * - kernel::SingleSpeciesUniformInjector_kernel<> * @namespaces: * - kernel:: @@ -859,6 +860,238 @@ namespace kernel { } }; // struct NonUniformInjector_kernel + template ED, SpatialDistClass SD> + struct SingleSpeciesNonUniformInjector_kernel { + + const real_t ppc0; + + ParticleArrays particles; + + array_t idx { "idx" }; + + const npart_t offset; + const npart_t domain_idx, cntr; + const bool use_tracking; + const M metric; + const ED energy_dist; + const SD spatial_dist; + const real_t inv_V0; + random_number_pool_t random_pool; + + SingleSpeciesNonUniformInjector_kernel(real_t ppc0, + Particles& particles, + npart_t domain_idx, + const M& metric, + const ED& energy_dist, + const SD& spatial_dist, + real_t inv_V0, + random_number_pool_t& random_pool) + : ppc0 { ppc0 } + , particles { static_cast(particles) } + , offset { particles.npart() } + , domain_idx { domain_idx } + , cntr { particles.counter() } + , use_tracking { particles.use_tracking() } + , metric { metric } + , energy_dist { energy_dist } + , 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 + -> Kokkos::pair { + real_t ppc_real = ppc0, weight = ONE; + if constexpr (SimpleSpatialDistClass) { + ppc_real *= spatial_dist(x_Ph); + } else { + const auto sp_dist = spatial_dist(x_Ph); + ppc_real *= sp_dist.first; + weight = sp_dist.second; + } + 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, weight }; + } + + Inline void inject(const prtlidx_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) { + InjectParticle(index + offset, + particles.i1, particles.i2, particles.i3, + particles.dx1, particles.dx2, particles.dx3, + particles.ux1, particles.ux2, particles.ux3, + particles.phi, particles.weight, particles.tag, particles.pld_i, + xi_Cd, dxi_Cd, v_Cd, weight, ZERO); + } else { + InjectParticle(index + offset, + particles.i1, particles.i2, particles.i3, + particles.dx1, particles.dx2, particles.dx3, + particles.ux1, particles.ux2, particles.ux3, + particles.phi, particles.weight, particles.tag, particles.pld_i, + xi_Cd, dxi_Cd, v_Cd, weight, ZERO, + domain_idx, index + cntr); + } + // clang-format on + } + + Inline void operator()(cellidx_t i1) const { + if constexpr (M::Dim == Dim::_1D) { + const auto i1_ = COORD(i1); + const coord_t x_Cd { i1_ + HALF }; + coord_t x_Ph { ZERO }; + metric.template convert(x_Cd, x_Ph); + + auto [ppc, weight] = injected_ppc(x_Ph); + if (ppc == 0) { + return; + } + + if constexpr (M::CoordType != Coord::Cartesian) { + 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(x_Ph, v_T); + metric.template transform_xyz(x_Cd, v_T, v_XYZ); + } + inject(index, { static_cast(i1_) }, { dx1 }, v_XYZ, weight); + } + } else { + raise::KernelError( + HERE, + "SingleSpeciesNonUniformInjector_kernel 1D called for 2D/3D"); + } + } + + Inline void operator()(cellidx_t i1, cellidx_t i2) const { + if constexpr (M::Dim == Dim::_2D) { + const auto i1_ = COORD(i1); + const auto i2_ = COORD(i2); + const 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::Cartesian) { + x_Cd_[2] = ZERO; + } + metric.template convert(x_Cd, x_Ph); + + auto [ppc, weight] = injected_ppc(x_Ph); + if (ppc == 0) { + return; + } + + if constexpr (M::CoordType != Coord::Cartesian) { + 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(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); + } + } + inject(index, + { static_cast(i1_), static_cast(i2_) }, + { dx1, dx2 }, + v_Cd, + weight); + } + } + + else { + raise::KernelError( + HERE, + "SingleSpeciesNonUniformInjector_kernel 2D called for 1D/3D"); + } + } + + Inline void operator()(cellidx_t i1, cellidx_t i2, cellidx_t i3) const { + if constexpr (M::Dim == Dim::_3D) { + const auto i1_ = COORD(i1); + const auto i2_ = COORD(i2); + const auto i3_ = COORD(i3); + const coord_t x_Cd { i1_ + HALF, i2_ + HALF, i3_ + HALF }; + coord_t x_Ph { ZERO }; + metric.template convert(x_Cd, x_Ph); + + auto [ppc, weight] = injected_ppc(x_Ph); + if (ppc == 0) { + return; + } + + if constexpr (M::CoordType != Coord::Cartesian) { + 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(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); + } + } + inject( + index, + { static_cast(i1_), static_cast(i2_), static_cast(i3_) }, + { dx1, dx2, dx3 }, + v_Cd, + weight); + } + } else { + raise::KernelError( + HERE, + "SingleSpeciesNonUniformInjector_kernel 3D called for 1D/2D"); + } + } + }; // struct SingleSpeciesNonUniformInjector_kernel + template ED> struct SingleSpeciesUniformInjector_kernel { From 3beef3901f9ec3bbfbd9bf5de8c023e2a88787ca Mon Sep 17 00:00:00 2001 From: haykh Date: Sat, 20 Jun 2026 19:00:24 +0400 Subject: [PATCH 18/25] nix modules upd --- dev/nix/adios2.nix | 2 +- dev/nix/kokkos.nix | 2 +- dev/nix/shell.nix | 39 +++++++++++++++++++++++---------------- 3 files changed, 25 insertions(+), 18 deletions(-) diff --git a/dev/nix/adios2.nix b/dev/nix/adios2.nix index 810c228dc..8880d84bf 100644 --- a/dev/nix/adios2.nix +++ b/dev/nix/adios2.nix @@ -39,7 +39,7 @@ stdenv.mkDerivation { ]; propagatedBuildInputs = [ - pkgs.gcc13 + pkgs.libgcc ] ++ ( if hdf5 then diff --git a/dev/nix/kokkos.nix b/dev/nix/kokkos.nix index f51d2c685..29712293e 100644 --- a/dev/nix/kokkos.nix +++ b/dev/nix/kokkos.nix @@ -27,7 +27,7 @@ let ]; "NONE" = [ pkgs.clang-tools - pkgs.gcc13 + pkgs.libgcc ]; }; getArch = diff --git a/dev/nix/shell.nix b/dev/nix/shell.nix index addd470fa..753a29ddd 100644 --- a/dev/nix/shell.nix +++ b/dev/nix/shell.nix @@ -7,6 +7,7 @@ arch ? "NATIVE", hdf5 ? false, mpi ? false, + extra ? "", }: let @@ -22,6 +23,9 @@ let gpu = gpuUpper; } ); + extraPkgs = map (name: pkgs.${name}) ( + pkgs.lib.filter (s: s != "") (pkgs.lib.splitString "," extra) + ); envVars = { compiler = { NONE = { @@ -41,28 +45,31 @@ pkgs.mkShell { "${name}" + (if gpu != "NONE" then "-${pkgs.lib.toLower gpu}" else "") + (if mpi then "-mpi" else ""); - nativeBuildInputs = with pkgs; [ - zlib - cmake + nativeBuildInputs = + with pkgs; + [ + zlib + cmake - adios2Pkg - kokkosPkg + adios2Pkg + kokkosPkg - python314 + python314 - cmake-format - cmake-lint - neocmakelsp - black - pyright - taplo - vscode-langservers-extracted - ]; + cmake-format + cmake-lint + neocmakelsp + black + pyright + taplo + vscode-langservers-extracted + ] + ++ extraPkgs; - LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath ([ + LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath [ pkgs.stdenv.cc.cc pkgs.zlib - ]); + ]; shellHook = '' BLUE='\033[0;34m' From 8084fa6e0b4fe32c3c760f9bf1e02957b4d6160b Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 29 Jun 2026 10:34:02 -0400 Subject: [PATCH 19/25] rms --- .gitignore | 3 +- src/archetypes/utils.h | 47 ++++++ src/global/traits/pgen.h | 2 +- src/kernels/particle_moments.hpp | 262 +++++++++++++++++++++++++++++++ 4 files changed, 312 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index e683ed8c5..a363191dd 100644 --- a/.gitignore +++ b/.gitignore @@ -64,4 +64,5 @@ action-token ignore-* tombi/ tidy/ -.claude \ No newline at end of file +.claude +.understand-anything diff --git a/src/archetypes/utils.h b/src/archetypes/utils.h index 827598ddc..405459091 100644 --- a/src/archetypes/utils.h +++ b/src/archetypes/utils.h @@ -178,6 +178,53 @@ namespace arch { Kokkos::Experimental::contribute(buffer, scatter_buff); } + template + inline void ComputeMomentWithSpeciesNew( + const SimulationParams& params, + const Domain& domain, + const std::vector& species, + ndfield_t& buffer, + const Kokkos::Array& buff_indices, + const MF& func, + uint8_t smoothing_order = 0u, + OutputSmoothingTypeFlag smoothing_method = OutputSmoothingType::SPLINE) { + const auto ni2 = domain.mesh.n_active(in::x2); + const auto inv_n0 = ONE / params.template get("scales.n0"); + const auto use_weights = params.template get("particles.use_weights"); + + for (const auto idx : buff_indices) { + if constexpr (M::Dim == Dim::_1D) { + Kokkos::deep_copy(Kokkos::subview(buffer, Kokkos::ALL(), idx), ZERO); + } else if constexpr (M::Dim == Dim::_2D) { + Kokkos::deep_copy(Kokkos::subview(buffer, Kokkos::ALL, Kokkos::ALL, idx), + ZERO); + } else if constexpr (M::Dim == Dim::_3D) { + Kokkos::deep_copy( + Kokkos::subview(buffer, Kokkos::ALL, Kokkos::ALL, Kokkos::ALL, idx), + ZERO); + } + } + auto scatter_buff = Kokkos::Experimental::create_scatter_view(buffer); + for (const auto sp : species) { + const auto& prtl_spec = domain.species[sp - 1]; + Kokkos::parallel_for( + "ComputeMoment", + prtl_spec.rangeActiveParticles(), + kernel::ParticleMomentsNew_kernel(func, + scatter_buff, + buff_indices, + prtl_spec, + use_weights, + domain.mesh.metric, + domain.mesh.flds_bc(), + ni2, + inv_n0, + smoothing_order, + smoothing_method)); + } + Kokkos::Experimental::contribute(buffer, scatter_buff); + } + template inline void UpdateEMFields(Domain& domain, const F& fieldsetter) { if constexpr (S == SimEngine::SRPIC) { diff --git a/src/global/traits/pgen.h b/src/global/traits/pgen.h index 179c2dbc9..0f84699dd 100644 --- a/src/global/traits/pgen.h +++ b/src/global/traits/pgen.h @@ -134,4 +134,4 @@ namespace traits::pgen { } // namespace traits::pgen -#endif // TRAITS_PGEN_H \ No newline at end of file +#endif // TRAITS_PGEN_H diff --git a/src/kernels/particle_moments.hpp b/src/kernels/particle_moments.hpp index 68461f768..58aac6023 100644 --- a/src/kernels/particle_moments.hpp +++ b/src/kernels/particle_moments.hpp @@ -3,6 +3,7 @@ * @brief Algorithm for computing different moments from particle distribution * @implements * - kernel::ParticleMoments_kernel<> + * - kernel::ParticleMomentsNew_kernel<> * - kernel::NormalizeVectorByRho_kernel<> * - kernel::Normalize4VelocityByNorm_kernel<> * - kernel::Transform4VelocitySpatialToPhysical_kernel<> @@ -31,6 +32,29 @@ namespace kernel { using namespace ntt; + namespace particle_moment_functor { + + template + concept IsValid = requires(const MF& fm, + const ParticleArrays& prtls, + float mass, + float charge, + prtlidx_t p, + list_t& buffer) { + { fm(prtls, mass, charge, p, buffer) } -> std::same_as; + }; + + template + concept HasN = requires { + { MF::N } -> std::convertible_to; + }; + + } // namespace particle_moment_functor + + template + concept ParticleMomentsFunctor = particle_moment_functor::IsValid && + particle_moment_functor::HasN; + template auto get_contrib(float mass, float charge) -> real_t { if constexpr (F == FldsID::Rho) { @@ -411,6 +435,244 @@ namespace kernel { } }; + /** + * @brief Generic moment-deposition kernel. + * + * Computes a per-particle scalar via the supplied @p Func and deposits it onto + * a scatter buffer, automatically applying volume normalization, weighting, + * smoothing (shape function) and axis reflection. @p Func is any device- + * callable object with signature + * `Inline auto operator()(const ParticleArrays&, prtlidx_t, const M&) const + * -> real_t` + * i.e. it receives the particle container, the particle index, and the metric, + * and returns the raw (un-normalized, un-weighted, un-smoothed) contribution of + * that particle to the moment. + * + * @tparam M Metric type + * @tparam N Last dimension of the buffer + * @tparam Func Per-particle contribution functor (deduced) + */ + template + class ParticleMomentsNew_kernel { + static_assert( + MF::N <= N, + "Buffer size N must be >= number of components N deposited by Func"); + static constexpr auto D = M::Dim; + + const MF func; + scatter_ndfield_t Buff; + const Kokkos::Array buff_indices; + const ParticleArrays particles; + const float mass, charge; + const bool use_weights; + const bool apply_norm; + const M metric; + const int ni2; + const real_t inv_n0; + + const uint8_t order; + const uint8_t window; + const OutputSmoothingTypeFlag smoothing; + + bool is_axis_i2min { false }, is_axis_i2max { false }; + + public: + ParticleMomentsNew_kernel( + const MF& func, + const scatter_ndfield_t& scatter_buff, + const Kokkos::Array& buff_indices, + const Particles& particles, + bool use_weights, + const M& metric, + const boundaries_t& boundaries, + ncells_t ni2, + real_t inv_n0, + uint8_t order = 0u, + OutputSmoothingTypeFlag smoothing = OutputSmoothingType::SPLINE, + bool apply_norm = true) + : func { func } + , Buff { scatter_buff } + , buff_indices { buff_indices } + , particles { static_cast(particles) } + , mass { particles.mass() } + , charge { particles.charge() } + , use_weights { use_weights } + , apply_norm { apply_norm } + , metric { metric } + , ni2 { static_cast(ni2) } + , inv_n0 { inv_n0 } + , order { order } + , smoothing { smoothing } + , window { static_cast( + math::ceil(static_cast(order) / 2.0f)) } { + raise::ErrorIf(window > N_GHOSTS, "Window size too large", HERE); + if constexpr ((M::CoordType != Coord::Cartesian) && + ((D == Dim::_2D) || (D == Dim::_3D))) { + raise::ErrorIf(boundaries.size() < 2, "boundaries defined incorrectly", HERE); + is_axis_i2min = (boundaries[1].first == FldsBC::AXIS); + is_axis_i2max = (boundaries[1].second == FldsBC::AXIS); + } + } + + Inline auto shapeFunction(real_t delta_x) const -> real_t { + if (smoothing == OutputSmoothingType::SPLINE) { + if (order == 0) { + return ONE; + } else if (order == 1) { + return prtl_shape::S1(delta_x); + } else if (order == 2) { + return prtl_shape::S2(delta_x); + } else if (order == 3) { + return prtl_shape::S3(delta_x); + } else if (order == 4) { + return prtl_shape::S4(delta_x); + } else if (order == 5) { + return prtl_shape::S5(delta_x); + } else if (order == 6) { + return prtl_shape::S6(delta_x); + } else if (order == 7) { + return prtl_shape::S7(delta_x); + } else if (order == 8) { + return prtl_shape::S8(delta_x); + } else if (order == 9) { + return prtl_shape::S9(delta_x); + } else if (order == 10) { + return prtl_shape::S10(delta_x); + } else if (order == 11) { + return prtl_shape::S11(delta_x); + } else { + raise::KernelError(HERE, "Unsupported shape function order"); + return ZERO; + } + } else if (smoothing == OutputSmoothingType::CONST) { + return ONE / (TWO * static_cast(window) + ONE); + } else { + raise::KernelError(HERE, "Unsupported smoothing method"); + return ZERO; + } + } + + Inline void operator()(prtlidx_t p) const { + if (particles.tag(p) == ParticleTag::dead) { + return; + } + list_t contributions { ZERO }; + func(particles, mass, charge, p, contributions); + for (uint8_t i = 0; i < MF::N; ++i) { + // apply volume normalization and (optionally) particle weights; + // skipped e.g. for nppc, which counts raw particles per cell + if constexpr (D == Dim::_1D) { + contributions[i] *= inv_n0 / + metric.sqrt_det_h( + { static_cast(particles.i1(p)) + HALF }); + } else if constexpr (D == Dim::_2D) { + contributions[i] *= inv_n0 / + metric.sqrt_det_h( + { static_cast(particles.i1(p)) + HALF, + static_cast(particles.i2(p)) + HALF }); + } else if constexpr (D == Dim::_3D) { + contributions[i] *= inv_n0 / + metric.sqrt_det_h( + { static_cast(particles.i1(p)) + HALF, + static_cast(particles.i2(p)) + HALF, + static_cast(particles.i3(p)) + HALF }); + } + if (use_weights) { + contributions[i] *= particles.weight(p); + } + } + auto buff_access = Buff.access(); + for (uint8_t i = 0; i < MF::N; ++i) { + const auto coeff = contributions[i]; + const auto buff_idx = buff_indices[i]; + if constexpr (D == Dim::_1D) { + for (auto di1 { -window }; di1 <= window; ++di1) { + const real_t delta_x1 = math::abs(static_cast(particles.dx1(p)) - + (static_cast(di1) + HALF)); + buff_access(particles.i1(p) + di1 + N_GHOSTS, + buff_idx) += coeff * shapeFunction(delta_x1); + } + } else if constexpr (D == Dim::_2D) { + for (auto di2 { -window }; di2 <= window; ++di2) { + for (auto di1 { -window }; di1 <= window; ++di1) { + const real_t delta_x1 = math::abs( + static_cast(particles.dx1(p)) - + (static_cast(di1) + HALF)); + const real_t delta_x2 = math::abs( + static_cast(particles.dx2(p)) - + (static_cast(di2) + HALF)); + const auto shape_coeff = shapeFunction(delta_x1) * + shapeFunction(delta_x2); + if constexpr (M::CoordType == Coord::Cartesian) { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + particles.i2(p) + di2 + N_GHOSTS, + buff_idx) += coeff * shape_coeff; + } else { + // reflect contribution at axes + if (is_axis_i2min && (particles.i2(p) + di2 < 0)) { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + N_GHOSTS - (particles.i2(p) + di2), + buff_idx) += coeff * shape_coeff; + } else if (is_axis_i2max && (particles.i2(p) + di2 >= ni2)) { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + 2 * ni2 - (particles.i2(p) + di2) + N_GHOSTS, + buff_idx) += coeff * shape_coeff; + } else { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + particles.i2(p) + di2 + N_GHOSTS, + buff_idx) += coeff * shape_coeff; + } + } + } + } + } else if constexpr (D == Dim::_3D) { + for (auto di3 { -window }; di3 <= window; ++di3) { + for (auto di2 { -window }; di2 <= window; ++di2) { + for (auto di1 { -window }; di1 <= window; ++di1) { + const auto delta_x1 = math::abs( + static_cast(particles.dx1(p)) - + (static_cast(di1) + HALF)); + const auto delta_x2 = math::abs( + static_cast(particles.dx2(p)) - + (static_cast(di2) + HALF)); + const auto delta_x3 = math::abs( + static_cast(particles.dx3(p)) - + (static_cast(di3) + HALF)); + const auto shape_coeff = shapeFunction(delta_x1) * + shapeFunction(delta_x2) * + shapeFunction(delta_x3); + if constexpr (M::CoordType == Coord::Cartesian) { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + particles.i2(p) + di2 + N_GHOSTS, + particles.i3(p) + di3 + N_GHOSTS, + buff_idx) += coeff * shape_coeff; + } else { + // reflect contribution at axes + if (is_axis_i2min && (particles.i2(p) + di2 < 0)) { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + N_GHOSTS - (particles.i2(p) + di2), + particles.i3(p) + di3 + N_GHOSTS, + buff_idx) += coeff * shape_coeff; + } else if (is_axis_i2max && (particles.i2(p) + di2 >= ni2)) { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + 2 * ni2 - (particles.i2(p) + di2) + N_GHOSTS, + particles.i3(p) + di3 + N_GHOSTS, + buff_idx) += coeff * shape_coeff; + } else { + buff_access(particles.i1(p) + di1 + N_GHOSTS, + particles.i2(p) + di2 + N_GHOSTS, + particles.i3(p) + di3 + N_GHOSTS, + buff_idx) += coeff * shape_coeff; + } + } + } + } + } + } + } + } + }; + template class NormalizeVectorByRho_kernel { const ndfield_t Rho; From 6ed5c5a466a20f8b31ec4794eed17d2f2702b45e Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 29 Jun 2026 10:39:32 -0400 Subject: [PATCH 20/25] rms pgen --- pgens/rms/cfg.rms.toml | 98 +++++ pgens/rms/pgen.hpp | 972 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 1070 insertions(+) create mode 100644 pgens/rms/cfg.rms.toml create mode 100644 pgens/rms/pgen.hpp diff --git a/pgens/rms/cfg.rms.toml b/pgens/rms/cfg.rms.toml new file mode 100644 index 000000000..81b319d03 --- /dev/null +++ b/pgens/rms/cfg.rms.toml @@ -0,0 +1,98 @@ +[simulation] + name = "rms" + engine = "srpic" + runtime = 20000.0 + +[grid] + resolution = [32768] + extent = [[0.0, 4000.0]] + + [grid.metric] + metric = "minkowski" + + [grid.boundaries] + fields = [["CONDUCTOR", "MATCH"]] + particles = [["REFLECT", "ABSORB"]] + +[scales] + larmor0 = 3.33333 + skindepth0 = 1.0 + +[algorithms] + current_filters = 8 + + [algorithms.timestep] + CFL = 0.7 + +[two_body] + thomson_optical_depth = 1e-3 + + [[two_body.interaction]] + type = "compton" + group1 = [1] + group2 = [3] + interval = 1 + tile_size = 8 + +[particles] + ppc0 = 8.0 + + [[particles.species]] + label = "e-" + mass = 1.0 + charge = -1.0 + maxnpart = 2e7 + + [[particles.species]] + label = "i" + mass = 25.0 + charge = 1.0 + maxnpart = 2e7 + + [[particles.species]] + label = "ph" + mass = 0.0 + charge = 0.0 + maxnpart = 2e7 + +[setup] + drift_ux = 0.1 # speed towards the wall [c] + temperature = 0.0001 # temperature of maxwell distribution [T / (m0 c^2)] + temperature_ratio = 1.0 # temperature ratio of electrons to protons + Bmag = 1.0 # magnetic field strength as fraction of magnetisation + Btheta = 60.0 # magnetic field angle in the plane + Bphi = 0.0 # magnetic field angle out of plane + filling_fraction = 1.0 # fraction of the shock piston filled with plasma + injector_velocity = 0.0 # speed of injector [c] + injection_start = 0.0 # start time of moving injector + injection_frequency = 100 + photon_injection_rate = 0.01 + +[output] + interval_time = 100.0 + + [output.fields] + quantities = ["N_1", "N_2", "N_3", "B", "E"] + custom = [ + "Temperature_e", + "Temperature_i", + "Temperature_ph", + "Vbx", + "Vby", + "Vbz", + ] + mom_smooth = 4 + + [output.particles] + enable = true + stride = 10 + + [output.spectra] + enable = true + +[checkpoint] + interval = 1000 + keep = 2 + +[diagnostics] + colored_stdout = true diff --git a/pgens/rms/pgen.hpp b/pgens/rms/pgen.hpp new file mode 100644 index 000000000..7dc3f73e5 --- /dev/null +++ b/pgens/rms/pgen.hpp @@ -0,0 +1,972 @@ +#ifndef PROBLEM_GENERATOR_H +#define PROBLEM_GENERATOR_H + +#include "enums.h" +#include "global.h" + +#include "traits/pgen.h" +#include "utils/error.h" +#include "utils/numeric.h" + +#include "archetypes/utils.h" +#include "framework/containers/particles.h" +#include "framework/domain/metadomain.h" + +#include +#include + +namespace user { + using namespace ntt; + + enum Component : idx_t { + comp_n = 0u, + comp_rho = 1u, + comp_rho_vx = 2u, + comp_rho_vy = 3u, + comp_rho_vz = 4u, + comp_n_t = 5u, + }; + + struct ComputeRhoV { + static constexpr uint8_t N = 5; + + Inline void operator()(const ParticleArrays& prtls, + float mass, + float /* charge */, + prtlidx_t p, + list_t& contribs) const { + contribs[0] = ONE; + contribs[1] = mass; + const auto gamma = U2GAMMA(prtls.ux1(p), prtls.ux2(p), prtls.ux3(p)); + contribs[2] = mass * prtls.ux1(p) / gamma; + contribs[3] = mass * prtls.ux2(p) / gamma; + contribs[4] = mass * prtls.ux3(p) / gamma; + } + }; + + struct ComputePhotonTmunu { + static constexpr uint8_t N = 5; + + Inline void operator()(const ParticleArrays& prtls, + float /* mass */, + float /* charge */, + prtlidx_t p, + list_t& contribs) const { + const auto energy = math::sqrt( + SQR(prtls.ux1(p)) + SQR(prtls.ux2(p)) + SQR(prtls.ux3(p))); + contribs[0] = ONE; + contribs[1] = energy; + contribs[2] = prtls.ux1(p); + contribs[3] = prtls.ux2(p); + contribs[4] = prtls.ux3(p); + } + }; + + template + struct ComputeN_Rho_Vi { + static constexpr uint8_t N = 1; + + Inline void operator()(const ParticleArrays& prtls, + float mass, + float /* charge */, + prtlidx_t p, + list_t& contribs) const { + const auto gamma = U2GAMMA(prtls.ux1(p), prtls.ux2(p), prtls.ux3(p)); + if constexpr (C == -1) { + contribs[0] = ONE; + } else if constexpr (C == 0) { + contribs[0] = mass; + } else if constexpr (C == 1) { + contribs[0] = mass * prtls.ux1(p) / gamma; + } else if constexpr (C == 2) { + contribs[0] = mass * prtls.ux2(p) / gamma; + } else if constexpr (C == 3) { + contribs[0] = mass * prtls.ux3(p) / gamma; + } + } + }; + + template + struct Normalize { + ndfield_t buffer; + const idx_t comp, comp_norm; + + Normalize(const ndfield_t& buff, idx_t comp, idx_t comp_norm) + : buffer { buff } + , comp { comp } + , comp_norm { comp_norm } {} + + Inline void operator()(cellidx_t i1) const { + if constexpr (D == Dim::_1D) { + if (buffer(i1, comp_norm) < 1e-6) { + buffer(i1, comp) = ZERO; + } else { + buffer(i1, comp) /= buffer(i1, comp_norm); + } + } else { + raise::KernelError(HERE, "Normalize is only implemented for 1D"); + } + } + + Inline void operator()(cellidx_t i1, cellidx_t i2) const { + if constexpr (D == Dim::_2D) { + if (buffer(i1, i2, comp_norm) < 1e-6) { + buffer(i1, i2, comp) = ZERO; + } else { + buffer(i1, i2, comp) /= buffer(i1, i2, comp_norm); + } + } else { + raise::KernelError(HERE, "Normalize is only implemented for 2D"); + } + } + + Inline void operator()(cellidx_t i1, cellidx_t i2, cellidx_t i3) const { + if constexpr (D == Dim::_3D) { + if (buffer(i1, i2, i3, comp_norm) < 1e-6) { + buffer(i1, i2, i3, comp) = ZERO; + } else { + buffer(i1, i2, i3, comp) /= buffer(i1, i2, i3, comp_norm); + } + } else { + raise::KernelError(HERE, "Normalize is only implemented for 3D"); + } + } + }; + + template + struct ComputePressure { + static constexpr uint8_t N = 1; + ndfield_t buffer; + + const idx_t rho, rho_vx, rho_vy, rho_vz; + + ComputePressure(const ndfield_t& buff, + idx_t rho = comp_rho, + idx_t rho_vx = comp_rho_vx, + idx_t rho_vy = comp_rho_vy, + idx_t rho_vz = comp_rho_vz) + : buffer { buff } + , rho { rho } + , rho_vx { rho_vx } + , rho_vy { rho_vy } + , rho_vz { rho_vz } {} + + Inline void operator()(const ParticleArrays& prtls, + float mass, + float, + prtlidx_t p, + list_t& contribs) const { + real_t Vx1 { ZERO }, Vx2 { ZERO }, Vx3 { ZERO }; + if constexpr (D == Dim::_1D) { + Vx1 = buffer(prtls.i1(p) + N_GHOSTS, rho_vx) / + buffer(prtls.i1(p) + N_GHOSTS, rho); + Vx2 = buffer(prtls.i1(p) + N_GHOSTS, rho_vy) / + buffer(prtls.i1(p) + N_GHOSTS, rho); + Vx3 = buffer(prtls.i1(p) + N_GHOSTS, rho_vz) / + buffer(prtls.i1(p) + N_GHOSTS, rho); + } else if constexpr (D == Dim::_2D) { + Vx1 = buffer(prtls.i1(p) + N_GHOSTS, prtls.i2(p) + N_GHOSTS, rho_vx) / + buffer(prtls.i1(p) + N_GHOSTS, prtls.i2(p) + N_GHOSTS, rho); + Vx2 = buffer(prtls.i1(p) + N_GHOSTS, prtls.i2(p) + N_GHOSTS, rho_vy) / + buffer(prtls.i1(p) + N_GHOSTS, prtls.i2(p) + N_GHOSTS, rho); + Vx3 = buffer(prtls.i1(p) + N_GHOSTS, prtls.i2(p) + N_GHOSTS, rho_vz) / + buffer(prtls.i1(p) + N_GHOSTS, prtls.i2(p) + N_GHOSTS, rho); + } else if constexpr (D == Dim::_3D) { + Vx1 = buffer(prtls.i1(p) + N_GHOSTS, + prtls.i2(p) + N_GHOSTS, + prtls.i3(p) + N_GHOSTS, + rho_vx) / + buffer(prtls.i1(p) + N_GHOSTS, + prtls.i2(p) + N_GHOSTS, + prtls.i3(p) + N_GHOSTS, + rho); + Vx2 = buffer(prtls.i1(p) + N_GHOSTS, + prtls.i2(p) + N_GHOSTS, + prtls.i3(p) + N_GHOSTS, + rho_vy) / + buffer(prtls.i1(p) + N_GHOSTS, + prtls.i2(p) + N_GHOSTS, + prtls.i3(p) + N_GHOSTS, + rho); + Vx3 = buffer(prtls.i1(p) + N_GHOSTS, + prtls.i2(p) + N_GHOSTS, + prtls.i3(p) + N_GHOSTS, + rho_vz) / + buffer(prtls.i1(p) + N_GHOSTS, + prtls.i2(p) + N_GHOSTS, + prtls.i3(p) + N_GHOSTS, + rho); + } + const auto Gamma = ONE / math::sqrt(ONE - SQR(Vx1) - SQR(Vx2) - SQR(Vx3)); + const auto gamma = U2GAMMA(prtls.ux1(p), prtls.ux2(p), prtls.ux3(p)); + contribs[0] = mass * + (SQR(Gamma) * SQR(gamma - Vx1 * prtls.ux1(p) - + Vx2 * prtls.ux2(p) - Vx3 * prtls.ux3(p)) - + ONE) / + (THREE * gamma); + } + }; + + template + struct ComputePhotonTemperature { + const ndfield_t t_array; + const idx_t comp_t, comp_n, comp_t00, comp_t01, comp_t02, comp_t03; + + ComputePhotonTemperature(ndfield_t& t_array, + idx_t comp_t, + idx_t comp_n, + idx_t comp_t00, + idx_t comp_t01, + idx_t comp_t02, + idx_t comp_t03) + : t_array { t_array } + , comp_t { comp_t } + , comp_n { comp_n } + , comp_t00 { comp_t00 } + , comp_t01 { comp_t01 } + , comp_t02 { comp_t02 } + , comp_t03 { comp_t03 } {} + + Inline auto Gamma(real_t T00_Sqr, real_t T0i_Sqr) const -> real_t { + return HALF * math::sqrt(T0i_Sqr) * + math::sqrt( + ONE / (T0i_Sqr + math::sqrt(T00_Sqr) * + (math::sqrt(FOUR * T00_Sqr - THREE * T0i_Sqr) - + TWO * math::sqrt(T00_Sqr)))); + } + + Inline void operator()(cellidx_t i1) const { + if constexpr (D == Dim::_1D) { + const auto T0i_Sqr = (SQR(t_array(i1, comp_t01)) + + SQR(t_array(i1, comp_t02)) + + SQR(t_array(i1, comp_t03))); + const auto T00_Sqr = SQR(t_array(i1, comp_t00)); + t_array(i1, comp_t) = (math::sqrt(FOUR * T00_Sqr - THREE * T0i_Sqr) - + math::sqrt(T00_Sqr)) / + (t_array(i1, comp_n) * Gamma(T00_Sqr, T0i_Sqr) * + static_cast(2.7)); + } else { + raise::KernelError(HERE, "ComputePhotonTemperature is only implemented for 1D"); + } + } + + Inline void operator()(cellidx_t i1, cellidx_t i2) const { + if constexpr (D == Dim::_2D) { + const auto T0i_Sqr = (SQR(t_array(i1, i2, comp_t01)) + + SQR(t_array(i1, i2, comp_t02)) + + SQR(t_array(i1, i2, comp_t03))); + const auto T00_Sqr = SQR(t_array(i1, i2, comp_t00)); + t_array(i1, i2, comp_t) = (math::sqrt(FOUR * T00_Sqr - THREE * T0i_Sqr) - + math::sqrt(T00_Sqr)) / + (t_array(i1, i2, comp_n) * + Gamma(T00_Sqr, T0i_Sqr) * + static_cast(2.7)); + } else { + raise::KernelError(HERE, "ComputePhotonTemperature is only implemented for 2D"); + } + } + + Inline void operator()(cellidx_t i1, cellidx_t i2, cellidx_t i3) const { + if constexpr (D == Dim::_3D) { + const auto T0i_Sqr = (SQR(t_array(i1, i2, i3, comp_t01)) + + SQR(t_array(i1, i2, i3, comp_t02)) + + SQR(t_array(i1, i2, i3, comp_t03))); + const auto T00_Sqr = SQR(t_array(i1, i2, i3, comp_t00)); + t_array(i1, i2, i3, comp_t) = (math::sqrt(FOUR * T00_Sqr - THREE * T0i_Sqr) - + math::sqrt(T00_Sqr)) / + (t_array(i1, i2, i3, comp_n) * + Gamma(T00_Sqr, T0i_Sqr) * + static_cast(2.7)); + } else { + raise::KernelError(HERE, "ComputePhotonTemperature is only implemented for 3D"); + } + } + }; + + template + struct PhotonSpatialDistribution { + const ndfield_t n_array; + + const M metric; + const real_t nmax { 4.0 }; + + PhotonSpatialDistribution(const ndfield_t& n_array, const M& metric) + : n_array { n_array } + , metric { metric } {} + + Inline auto operator()(const coord_t& x_Ph) const -> real_t { + coord_t x_Cd { ZERO }; + metric.template convert(x_Ph, x_Cd); + if constexpr (M::Dim == Dim::_1D) { + return n_array(static_cast(x_Cd[0]) + N_GHOSTS, comp_n) / nmax; + } else if constexpr (M::Dim == Dim::_2D) { + return n_array(static_cast(x_Cd[0]) + N_GHOSTS, + static_cast(x_Cd[1]) + N_GHOSTS, + comp_n) / + nmax; + } else if constexpr (M::Dim == Dim::_3D) { + return n_array(static_cast(x_Cd[0]) + N_GHOSTS, + static_cast(x_Cd[1]) + N_GHOSTS, + static_cast(x_Cd[2]) + N_GHOSTS, + comp_n) / + nmax; + } + } + }; + + template + struct PlanckDistribution { + const ndfield_t nt_array; + + const M metric; + random_number_pool_t random_pool; + + PlanckDistribution(const ndfield_t& nt_array, + const M& metric, + random_number_pool_t& pool) + : nt_array { nt_array } + , metric { metric } + , random_pool { pool } {} + + Inline void operator()(const coord_t& x_Ph, vec_t& v) const { + coord_t x_Cd { ZERO }; + metric.template convert(x_Ph, x_Cd); + real_t temperature { ZERO }; + if constexpr (M::Dim == Dim::_1D) { + const auto i1_ = static_cast(x_Cd[0]) + N_GHOSTS; + temperature = nt_array(i1_, comp_n_t) / nt_array(i1_, comp_n); + } else if constexpr (M::Dim == Dim::_2D) { + const auto i1_ = static_cast(x_Cd[0]) + N_GHOSTS; + const auto i2_ = static_cast(x_Cd[1]) + N_GHOSTS; + temperature = nt_array(i1_, i2_, comp_n_t) / nt_array(i1_, i2_, comp_n); + } else if constexpr (M::Dim == Dim::_3D) { + const auto i1_ = static_cast(x_Cd[0]) + N_GHOSTS; + const auto i2_ = static_cast(x_Cd[1]) + N_GHOSTS; + const auto i3_ = static_cast(x_Cd[2]) + N_GHOSTS; + temperature = nt_array(i1_, i2_, i3_, comp_n_t) / + nt_array(i1_, i2_, i3_, comp_n); + } + + real_t prob { ZERO }, n { ZERO }; + auto gen = random_pool.get_state(); + const auto rnd = Random(gen); + const auto rnd1 = Random(gen); + const auto rnd2 = Random(gen); + const auto rnd3 = Random(gen); + const auto rndth = Random(gen); + const auto rndph = Random(gen); + random_pool.free_state(gen); + + while ((prob < rnd) and (n < 40)) { + n += ONE; + prob += ONE / (static_cast(1.20206) * CUBE(n)); + } + const auto energy = -static_cast(2.7) * temperature * + math::log( + rnd1 * rnd2 * rnd3 + static_cast(1e-16)) / + n; + const auto costh = TWO * rndth - ONE; + const auto phi = static_cast(constant::TWO_PI) * rndph; + + v[0] = energy * math::sqrt(ONE - SQR(costh)) * math::cos(phi); + v[1] = energy * math::sqrt(ONE - SQR(costh)) * math::sin(phi); + v[2] = energy * costh; + } + }; + + template + struct InitFields { + /* + Sets up magnetic and electric field components for the simulation. + Must satisfy E = -v x B for Lorentz Force to be zero. + + @param bmag: magnetic field scaling + @param btheta: magnetic field polar angle + @param bphi: magnetic field azimuthal angle + @param drift_ux: drift velocity in the x direction + */ + InitFields(real_t bmag, real_t btheta, real_t bphi, real_t drift_ux) + : Bmag { bmag } + , Btheta { btheta * static_cast(convert::deg2rad) } + , Bphi { bphi * static_cast(convert::deg2rad) } + , Vx { drift_ux } {} + + // magnetic field components + Inline auto bx1(const coord_t&) const -> real_t { + return Bmag * math::cos(Btheta); + } + + Inline auto bx2(const coord_t&) const -> real_t { + return Bmag * math::sin(Btheta) * math::sin(Bphi); + } + + Inline auto bx3(const coord_t&) const -> real_t { + return Bmag * math::sin(Btheta) * math::cos(Bphi); + } + + // electric field components + Inline auto ex1(const coord_t&) const -> real_t { + return ZERO; + } + + Inline auto ex2(const coord_t&) const -> real_t { + return -Vx * Bmag * math::sin(Btheta) * math::cos(Bphi); + } + + Inline auto ex3(const coord_t&) const -> real_t { + return Vx * Bmag * math::sin(Btheta) * math::sin(Bphi); + } + + private: + const real_t Btheta, Bphi, Vx, Bmag; + }; + + template + struct PGen { + static constexpr auto D { M::Dim }; + // compatibility traits for the problem generator + static constexpr auto engines { + ::traits::pgen::compatible_with {} + }; + static constexpr auto metrics { + ::traits::pgen::compatible_with {} + }; + static constexpr auto dimensions { + ::traits::pgen::compatible_with {} + }; + const SimulationParams& params; + Metadomain& metadomain; + + // domain properties + const real_t global_xmin, global_xmax; + // gas properties + const real_t drift_ux, temperature, temperature_ratio, filling_fraction; + // magnetic field properties + real_t Btheta, Bphi, Bmag; + // injector properties + const real_t injector_velocity, injection_start; + const int injection_frequency; + InitFields init_flds; + + PGen(const SimulationParams& p, Metadomain& m) + : params { p } + , metadomain { m } + , global_xmin { metadomain.mesh().extent(in::x1).first } + , global_xmax { metadomain.mesh().extent(in::x1).second } + , drift_ux { params.template get("setup.drift_ux") } + , temperature { params.template get("setup.temperature") } + , temperature_ratio { params.template get( + "setup.temperature_ratio", + ONE) } + , Bmag { params.template get("setup.Bmag", ZERO) } + , Btheta { params.template get("setup.Btheta", ZERO) } + , Bphi { params.template get("setup.Bphi", ZERO) } + , init_flds { Bmag, Btheta, Bphi, drift_ux } + , filling_fraction { params.template get("setup.filling_fraction", + 1.0) } + , injector_velocity { params.template get( + "setup.injector_velocity", + 1.0) } + , injection_start { params.template get("setup.injection_start", 0.0) } + , injection_frequency { + params.template get("setup.injection_frequency", 100) + } {} + + auto MatchFields(simtime_t) const -> InitFields { + return init_flds; + } + + auto FixFieldsConst(const bc_in&, const em& comp) const + -> std::pair { + if (comp == em::ex1) { + return { init_flds.ex1({ ZERO }), true }; + } else if ((comp == em::ex2) or (comp == em::ex3)) { + return { ZERO, true }; + } else if (comp == em::bx1) { + return { init_flds.bx1({ ZERO }), true }; + } else if (comp == em::bx2) { + return { init_flds.bx2({ ZERO }), true }; + } else if (comp == em::bx3) { + return { init_flds.bx3({ ZERO }), true }; + } else { + raise::Error("Invalid component", HERE); + return { ZERO, false }; + } + } + + void InitPrtls(Domain& domain) { + /* + * Plasma setup as partially filled box + * + * Plasma setup: + * + * global_xmin global_xmax + * | | + * V V + * |:::::::::::|..........................| + * ^ + * | + * filling_fraction + */ + + // minimum and maximum position of particles + real_t xg_min = global_xmin; + // real_t xg_max = global_xmin + filling_fraction * (global_xmax - global_xmin); + real_t xg_max = global_xmax; + + // define box to inject into + boundaries_t box; + // loop over all dimensions + for (auto d { 0u }; d < (unsigned int)M::Dim; ++d) { + // compute the range for the x-direction + if (d == static_cast(in::x1)) { + box.emplace_back(xg_min, xg_max); + } else { + // inject into full range in other directions + box.push_back(Range::All); + } + } + + // define temperatures of species + const auto temperatures = std::make_pair(temperature, + temperature_ratio * temperature); + // define drift speed of species + const auto drifts = std::make_pair( + std::vector { -drift_ux, ZERO, ZERO }, + std::vector { -drift_ux, ZERO, ZERO }); + + // inject particles + arch::InjectUniformMaxwellians(params, + domain, + ONE, + temperatures, + { 1, 2 }, + drifts, + false, + box); + } + + void CustomPostStep(timestep_t /* step */, + simtime_t /* time */, + Domain& domain) { + const auto dt = params.template get("algorithms.timestep.dt"); + // if (step % injection_frequency == 0) { + // /* + // * Replenish plasma in a moving injector + // * + // * Injector setup: + // * + // * global_xmin purge/replenish global_xmax + // * | x_init | | + // * V v V V + // * |:::::::::::;::::::::::|\\\\\\\\|......| + // * xmin xmax + // * ^ + // * | + // * moving injector + // */ + // + // // initial position of injector + // const auto x_init = global_xmin + + // filling_fraction * (global_xmax - global_xmin); + // + // // compute the position of the injector after the current timestep + // auto xmax = x_init + + // injector_velocity * + // (std::max(time - injection_start, ZERO) + dt); + // if (xmax >= global_xmax) { + // xmax = global_xmax; + // } + // + // // compute the beginning of the injected region + // auto xmin = xmax - injection_frequency * dt; + // if (xmin <= global_xmin) { + // xmin = global_xmin; + // } + // + // // define indice range to reset fields + // boundaries_t incl_ghosts; + // for (auto d = 0; d < M::Dim; ++d) { + // incl_ghosts.emplace_back(false, false); + // } + // + // // define box to reset fields + // boundaries_t purge_box; + // // loop over all dimension + // for (auto d = 0u; d < M::Dim; ++d) { + // if (d == 0) { + // purge_box.emplace_back(xmin, global_xmax); + // } else { + // purge_box.push_back(Range::All); + // } + // } + // + // const auto extent = domain.mesh.ExtentToRange(purge_box, incl_ghosts); + // tuple_t x_min { 0 }, x_max { 0 }; + // for (auto d = 0; d < M::Dim; ++d) { + // x_min[d] = extent[d].first; + // x_max[d] = extent[d].second; + // } + // + // Kokkos::parallel_for("ResetFields", + // CreateRangePolicy(x_min, x_max), + // arch::SetEMFields_kernel { + // domain.fields.em, + // init_flds, + // domain.mesh.metric }); + // metadomain.CommunicateFields(domain, Comm::E | Comm::B); + // + // /* + // tag particles inside the injection zone as dead + // */ + // const auto& mesh = domain.mesh; + // + // // loop over particle species + // for (auto s { 0u }; s < 2; ++s) { + // // get particle properties + // auto& species = domain.species[s]; + // auto i1 = species.i1; + // auto dx1 = species.dx1; + // auto tag = species.tag; + // + // Kokkos::parallel_for( + // "RemoveParticles", + // species.rangeActiveParticles(), + // Lambda(prtlidx_t p) { + // // check if the particle is already dead + // if (tag(p) == ParticleTag::dead) { + // return; + // } + // const auto x_Cd = static_cast(i1(p)) + + // static_cast(dx1(p)); + // const auto x_Ph = mesh.metric.template convert<1, Crd::Cd, Crd::XYZ>( + // x_Cd); + // + // if (x_Ph > xmin) { + // tag(p) = ParticleTag::dead; + // } + // }); + // } + // + // // define box to inject into + // boundaries_t inj_box; + // // loop over all dimension + // for (auto d = 0u; d < M::Dim; ++d) { + // if (d == 0) { + // inj_box.emplace_back(xmin, xmax); + // } else { + // inj_box.push_back(Range::All); + // } + // } + // + // // same maxwell distribution as above + // const auto temperatures = std::make_pair(temperature, + // temperature_ratio * temperature); + // const auto drifts = std::make_pair( + // std::vector { -drift_ux, ZERO, ZERO }, + // std::vector { -drift_ux, ZERO, ZERO }); + // arch::InjectUniformMaxwellians(params, + // domain, + // ONE, + // temperatures, + // { 1, 2 }, + // drifts, + // false, + // inj_box); + // } + + { + /* + * Inject photons + */ + auto compute_rho_v = ComputeRhoV {}; + auto compute_pressure = ComputePressure { domain.fields.bckp }; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + domain.fields.bckp, + { comp_n, comp_rho, comp_rho_vx, comp_rho_vy, comp_rho_vz }, + compute_rho_v); + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + domain.fields.bckp, + { comp_n_t }, + compute_pressure); + + // inject photons with a Planck distribution in energy and spatial distribution following the plasma density + const auto energy_dist = PlanckDistribution(domain.fields.bckp, + domain.mesh.metric, + domain.random_pool()); + const auto spatial_dist = PhotonSpatialDistribution(domain.fields.bckp, + domain.mesh.metric); + const auto photon_inj_rate = params.template get( + "setup.photon_injection_rate", + ZERO); + arch::InjectNonUniform( + params, + domain, + 3, + energy_dist, + spatial_dist, + static_cast(photon_inj_rate * dt)); + } + } + + void CustomFieldOutput(const std::string& label, + ndfield_t& buff, + cellidx_t buff_idx, + timestep_t, + simtime_t, + const Domain& domain) { + const uint8_t smoothing_order = 2u * N_GHOSTS; + if (label == "Vbx") { + /** + * buff_idx + 1 -> rho_e + rho_i + */ + auto compute_rho = ComputeN_Rho_Vi<0> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)((buff_idx + 1) % 6) }, + compute_rho, + smoothing_order); + /** + * buff_idx -> rho_e * vxe + rho_i * vxi + */ + auto compute_rho_vx = ComputeN_Rho_Vi<1> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)buff_idx }, + compute_rho_vx, + smoothing_order); + /** + * buff_idx -> vx = (rho_e * vxe + rho_i * vxi) / (rho_e + rho_i) + */ + Kokkos::parallel_for( + "ComputeVx", + domain.mesh.rangeActiveCells(), + Normalize { buff, (idx_t)(buff_idx), (idx_t)((buff_idx + 1) % 6) }); + } else if (label == "Vby") { + /** + * buff_idx + 1 -> rho_e + rho_i + */ + auto compute_rho = ComputeN_Rho_Vi<0> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)((buff_idx + 1) % 6) }, + compute_rho, + smoothing_order); + /** + * buff_idx -> rho_e * vye + rho_i * vyi + */ + auto compute_rho_vy = ComputeN_Rho_Vi<2> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)(buff_idx) }, + compute_rho_vy, + smoothing_order); + /** + * buff_idx -> vy = (rho_e * vye + rho_i * vyi) / (rho_e + rho_i) + */ + Kokkos::parallel_for( + "ComputeVy", + domain.mesh.rangeActiveCells(), + Normalize { buff, (idx_t)(buff_idx), (idx_t)((buff_idx + 1) % 6) }); + } else if (label == "Vbz") { + /** + * buff_idx + 1 -> rho_e + rho_i + */ + auto compute_rho = ComputeN_Rho_Vi<0> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)((buff_idx + 1) % 6) }, + compute_rho, + smoothing_order); + /** + * buff_idx -> rho_e * vye + rho_i * vyi + */ + auto compute_rho_vz = ComputeN_Rho_Vi<3> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)buff_idx }, + compute_rho_vz, + smoothing_order); + /** + * buff_idx -> vy = (rho_e * vye + rho_i * vyi) / (rho_e + rho_i) + */ + Kokkos::parallel_for( + "ComputeVy", + domain.mesh.rangeActiveCells(), + Normalize { buff, (idx_t)(buff_idx), (idx_t)((buff_idx + 1) % 6) }); + } else if (label == "Temperature_e") { + /** + * buff_idx + 1 -> n_e + n_i + * buff_idx + 2 -> rho_e + rho_i + * buff_idx + 3 -> rho_e * vxe + rho_i * vxi + * buff_idx + 4 -> rho_e * vye + rho_i * vyi + * buff_idx + 5 -> rho_e * vze + rho_i * vzi + */ + auto compute_rho_v = ComputeRhoV {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)((buff_idx + 1) % 6), + (idx_t)((buff_idx + 2) % 6), + (idx_t)((buff_idx + 3) % 6), + (idx_t)((buff_idx + 4) % 6), + (idx_t)((buff_idx + 5) % 6) }, + compute_rho_v, + smoothing_order); + /** + * buff_idx -> n_e * T_e + */ + auto compute_pressure = ComputePressure { buff, + (idx_t)((buff_idx + 2) % 6), + (idx_t)((buff_idx + 3) % 6), + (idx_t)((buff_idx + 4) % 6), + (idx_t)((buff_idx + 5) % 6) }; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1 }, + buff, + { (idx_t)buff_idx }, + compute_pressure, + smoothing_order); + /** + * buff_idx + 1 -> n_e + */ + auto compute_n = ComputeN_Rho_Vi<-1> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1 }, + buff, + { (idx_t)((buff_idx + 1) % 6) }, + compute_n, + smoothing_order); + /** + * buff_idx -> T_e + */ + Kokkos::parallel_for( + "ComputeTe", + domain.mesh.rangeActiveCells(), + Normalize { buff, (idx_t)(buff_idx), (idx_t)((buff_idx + 1) % 6) }); + } else if (label == "Temperature_i") { + /** + * buff_idx + 1 -> n_e + n_i + * buff_idx + 2 -> rho_e + rho_i + * buff_idx + 3 -> rho_e * vxe + rho_i * vxi + * buff_idx + 4 -> rho_e * vye + rho_i * vyi + * buff_idx + 5 -> rho_e * vze + rho_i * vzi + */ + auto compute_rho_v = ComputeRhoV {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 1, 2 }, + buff, + { (idx_t)((buff_idx + 1) % 6), + (idx_t)((buff_idx + 2) % 6), + (idx_t)((buff_idx + 3) % 6), + (idx_t)((buff_idx + 4) % 6), + (idx_t)((buff_idx + 5) % 6) }, + compute_rho_v, + smoothing_order); + /** + * buff_idx -> n_i * T_i + */ + auto compute_pressure = ComputePressure { buff, + (idx_t)((buff_idx + 2) % 6), + (idx_t)((buff_idx + 3) % 6), + (idx_t)((buff_idx + 4) % 6), + (idx_t)((buff_idx + 5) % 6) }; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 2 }, + buff, + { (idx_t)buff_idx }, + compute_pressure, + smoothing_order); + /** + * buff_idx + 1 -> n_i + */ + auto compute_n = ComputeN_Rho_Vi<-1> {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 2 }, + buff, + { (idx_t)((buff_idx + 1) % 6) }, + compute_n, + smoothing_order); + /** + * buff_idx -> T_i + */ + Kokkos::parallel_for( + "ComputeTi", + domain.mesh.rangeActiveCells(), + Normalize { buff, (idx_t)buff_idx, (idx_t)((buff_idx + 1) % 6) }); + } else if (label == "Temperature_ph") { + /** + * buff_idx + 1 -> n_ph + * buff_idx + 2 -> T^00_ph + * buff_idx + 3 -> T^0x_ph + * buff_idx + 4 -> T^0y_ph + * buff_idx + 5 -> T^0z_ph + */ + auto compute_t_munu = ComputePhotonTmunu {}; + arch::ComputeMomentWithSpeciesNew( + params, + domain, + { 3 }, + buff, + { (idx_t)((buff_idx + 1) % 6), + (idx_t)((buff_idx + 2) % 6), + (idx_t)((buff_idx + 3) % 6), + (idx_t)((buff_idx + 4) % 6), + (idx_t)((buff_idx + 5) % 6) }, + compute_t_munu, + smoothing_order); + /** + * buff_idx -> T_ph + */ + Kokkos::parallel_for( + "ComputeT_ph", + domain.mesh.rangeActiveCells(), + ComputePhotonTemperature { buff, + (idx_t)buff_idx, + (idx_t)((buff_idx + 1) % 6), + (idx_t)((buff_idx + 2) % 6), + (idx_t)((buff_idx + 3) % 6), + (idx_t)((buff_idx + 4) % 6), + (idx_t)((buff_idx + 5) % 6) }); + } + } + }; +} // namespace user + +#endif // PROBLEM_GENERATOR_H From 6a622029219ca4c06648bdb20bc7f91bedbf13b8 Mon Sep 17 00:00:00 2001 From: hayk Date: Mon, 20 Jul 2026 01:52:53 -0400 Subject: [PATCH 21/25] minor fixes for nvcc --- src/engines/srpic/twobody.h | 8 +- src/kernels/twobody_interactions.hpp | 116 ++++++++++++++++++++------- 2 files changed, 89 insertions(+), 35 deletions(-) diff --git a/src/engines/srpic/twobody.h b/src/engines/srpic/twobody.h index c5d04ca66..8019cd20b 100644 --- a/src/engines/srpic/twobody.h +++ b/src/engines/srpic/twobody.h @@ -39,9 +39,7 @@ namespace ntt { prm::Parameters compton_params; compton_params.set("compton_scattering.nominal_probability_density", nominal_thomson_probability_density); - auto recoil1 = interaction.recoil1; - auto recoil2 = interaction.recoil2; - auto launch = [&]() { + auto launch = [&]() { auto policy = arch::qed::ComptonScattering( compton_params, domain.random_pool()); @@ -54,7 +52,7 @@ namespace ntt { domain.species[sp_lepton - 1].mass() == ZERO, fmt::format( "Species %u is massless but is in the lepton group " - "of a Compton interaction", + "of a Compton interaction", sp_lepton), HERE); group1_species.push_back(&domain.species[sp_lepton - 1]); @@ -64,7 +62,7 @@ namespace ntt { domain.species[sp_photon - 1].mass() != ZERO, fmt::format( "Species %u is massive but is in the photon group " - "of a Compton interaction", + "of a Compton interaction", sp_photon), HERE); group2_species.push_back(&domain.species[sp_photon - 1]); diff --git a/src/kernels/twobody_interactions.hpp b/src/kernels/twobody_interactions.hpp index 901b43acf..9a7df409a 100644 --- a/src/kernels/twobody_interactions.hpp +++ b/src/kernels/twobody_interactions.hpp @@ -101,6 +101,84 @@ namespace kernel::mink { return ncells_on_tile; } + namespace { + struct PackIndex { + const npart_t offset; + const CollisionSpecies species; + array_t combined_idx; + array_t combined_tileidx; + + PackIndex(npart_t offset, + const CollisionSpecies& species, + array_t& combined_idx, + array_t& combined_tileidx) + : offset { offset } + , species { species } + , combined_idx { combined_idx } + , combined_tileidx { combined_tileidx } {} + + Inline void operator()(prtlidx_t p) const { + // pack species idx into top 8 bits + prtl index into the remaining 56 bits + combined_idx(offset + p) = (static_cast(species.sp) << 56) | + static_cast(p); + combined_tileidx(offset + p) = species.tileidx(p); + } + }; + + struct CombineNumPpt { + const CollisionSpecies species; + array_t combined_num_ppt; + + CombineNumPpt(const CollisionSpecies& species, + array_t& combined_num_ppt) + : species { species } + , combined_num_ppt { combined_num_ppt } {} + + Inline void operator()(cellidx_t t) const { + combined_num_ppt(t) += species.num_ppt(t); + } + }; + + struct PackRandom { + const array_t combined_tileidx; + array_t shuffle_key; + random_number_pool_t random_pool; + + PackRandom(const array_t& combined_tileidx, + array_t& shuffle_key, + random_number_pool_t& random_pool) + : combined_tileidx { combined_tileidx } + , shuffle_key { shuffle_key } + , random_pool { random_pool } {} + + Inline void operator()(prtlidx_t p) const { + auto gen = random_pool.get_state(); + const auto rnd = static_cast(gen.urand()); + random_pool.free_state(gen); + const auto tile_idx = static_cast(combined_tileidx(p)); + // packing top 32 bits with tile index, and the rest -- random + shuffle_key(p) = (tile_idx << 32) | rnd; + } + }; + + struct TileOffsets { + array_t tile_offsets; + array_t combined_num_ppt; + + TileOffsets(array_t& tile_offsets, + const array_t& combined_num_ppt) + : tile_offsets { tile_offsets } + , combined_num_ppt { combined_num_ppt } {} + + Inline void operator()(cellidx_t t, npart_t& acc, const bool is_final) const { + if (is_final) { + tile_offsets(t) = acc; + } + acc += combined_num_ppt(t); + } + }; + } // namespace + template struct CollisionGroup { std::vector group; @@ -161,20 +239,11 @@ namespace kernel::mink { Kokkos::parallel_for( "CombineInGroup", species.npart, - ClassLambda(const npart_t p) { - // pack species idx into top 8 bits + prtl index into the remaining 56 bits - combined_idx(offset + p) = (static_cast(species.sp) - << 56) | - static_cast(p); - combined_tileidx(offset + p) = species.tileidx(p); - }); + PackIndex { offset, species, combined_idx, combined_tileidx }); offset += species.npart; - Kokkos::parallel_for( - "CombineNumPpt", - species.num_tiles, - ClassLambda(const ncells_t t) { - combined_num_ppt(t) += species.num_ppt(t); - }); + Kokkos::parallel_for("CombineNumPpt", + species.num_tiles, + CombineNumPpt { species, combined_num_ppt }); Kokkos::fence(); } } @@ -184,29 +253,16 @@ namespace kernel::mink { Kokkos::parallel_for( "PackRandom", tot_npart, - ClassLambda(const npart_t p) { - auto gen = random_pool.get_state(); - const auto rnd = static_cast(gen.urand()); - random_pool.free_state(gen); - const auto tile_idx = static_cast(combined_tileidx(p)); - // packing top 32 bits with tile index, and the rest -- random - shuffle_key(p) = (tile_idx << 32) | rnd; - }); + PackRandom { combined_tileidx, shuffle_key, random_pool }); Kokkos::Experimental::sort_by_key(Kokkos::DefaultExecutionSpace {}, shuffle_key, combined_idx); } { // compute index offsets for each tile - Kokkos::parallel_scan( - "TileOffsets", - num_tiles, - ClassLambda(cellidx_t t, npart_t & acc, const bool final) { - if (final) { - tile_offsets(t) = acc; - } - acc += combined_num_ppt(t); - }); + Kokkos::parallel_scan("TileOffsets", + num_tiles, + TileOffsets { tile_offsets, combined_num_ppt }); } } }; From 417beb4c76de54475c803e4183f6f146714074f0 Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 20 Jul 2026 09:25:50 -0400 Subject: [PATCH 22/25] nvcc 12.8 weirdness fix --- src/kernels/pushers/sr_policies.h | 40 +++++++++++++++---------------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/src/kernels/pushers/sr_policies.h b/src/kernels/pushers/sr_policies.h index b247b27b7..2050aac57 100644 --- a/src/kernels/pushers/sr_policies.h +++ b/src/kernels/pushers/sr_policies.h @@ -110,25 +110,23 @@ namespace kernel::sr { ntt::EmissionTypeFlag emission_type, bool atm, F&& callback) { - auto with_emission = [&](auto next) { + auto with_emission = [&](const PG& pg, D& dom, auto next) { switch (emission_type) { case ntt::EmissionType::SYNCHROTRON: - next(MakePusherPolicyEmission( - domain, + next(MakePusherPolicyEmission( + dom, params, pusher_ctx)); break; case ntt::EmissionType::COMPTON: - next(MakePusherPolicyEmission( - domain, + next(MakePusherPolicyEmission( + dom, params, pusher_ctx)); break; case ntt::EmissionType::CUSTOM: - if constexpr (::traits::pgen::HasEmissionPolicy) { - next(pgen.EmissionPolicy(pusher_ctx.time, - pusher_ctx.species_index, - domain)); + if constexpr (::traits::pgen::HasEmissionPolicy) { + next(pg.EmissionPolicy(pusher_ctx.time, pusher_ctx.species_index, dom)); } else { raise::Error("Custom emission policy flag is set but problem " "generator does not define an emission policy", @@ -142,22 +140,22 @@ namespace kernel::sr { } }; - auto with_custom_prtl_upd = [&](auto next) { - if constexpr (::traits::pgen::HasCustomPrtlUpdate) { - next(pgen.CustomParticleUpdate(pusher_ctx.time, - pusher_ctx.species_index, - domain)); + auto with_custom_prtl_upd = [&](const PG& pg, + D& dom, + auto next) { + if constexpr (::traits::pgen::HasCustomPrtlUpdate) { + next(pg.CustomParticleUpdate(pusher_ctx.time, pusher_ctx.species_index, dom)); } else { next(::traits::custom_prtl_update::NoPolicy_t {}); } }; - auto with_ext_fields = [&](auto next) { - if constexpr (::traits::pgen::HasExternalFields) { - const auto [apply_extfields, external_fields] = pgen.ExternalFields( + auto with_ext_fields = [&](const PG& pg, D& dom, auto next) { + if constexpr (::traits::pgen::HasExternalFields) { + const auto [apply_extfields, external_fields] = pg.ExternalFields( pusher_ctx.time, pusher_ctx.species_index, - domain); + dom); if (apply_extfields) { next(external_fields); } else { @@ -168,9 +166,9 @@ namespace kernel::sr { } }; - with_emission([&](auto ep) { - with_custom_prtl_upd([&](auto cpu) { - with_ext_fields([&](auto ef) { + with_emission(pgen, domain, [&](auto ep) { + with_custom_prtl_upd(pgen, domain, [&](auto cpu) { + with_ext_fields(pgen, domain, [&](auto ef) { using E = decltype(ep); using CPU = decltype(cpu); using EF = decltype(ef); From 4bcd2d04785f198dc3aab749d64697e629ea0f34 Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 7 Sep 2026 14:28:07 -0400 Subject: [PATCH 23/25] rms setup change --- pgens/rms/pgen.hpp | 415 +++++++++++++++++++++------------------------ 1 file changed, 189 insertions(+), 226 deletions(-) diff --git a/pgens/rms/pgen.hpp b/pgens/rms/pgen.hpp index 7dc3f73e5..2808d8fae 100644 --- a/pgens/rms/pgen.hpp +++ b/pgens/rms/pgen.hpp @@ -27,6 +27,18 @@ namespace user { comp_n_t = 5u, }; + struct ComputeN { + static constexpr uint8_t N = 1; + + Inline void operator()(const ParticleArrays& /* prtls */, + float /* mass */, + float /* charge */, + prtlidx_t /* p */, + list_t& contribs) const { + contribs[0] = ONE; + } + }; + struct ComputeRhoV { static constexpr uint8_t N = 5; @@ -314,39 +326,17 @@ namespace user { } }; - template + template struct PlanckDistribution { - const ndfield_t nt_array; + const real_t T_ph_inj; - const M metric; random_number_pool_t random_pool; - PlanckDistribution(const ndfield_t& nt_array, - const M& metric, - random_number_pool_t& pool) - : nt_array { nt_array } - , metric { metric } + PlanckDistribution(real_t T_ph_inj, random_number_pool_t& pool) + : T_ph_inj { T_ph_inj } , random_pool { pool } {} - Inline void operator()(const coord_t& x_Ph, vec_t& v) const { - coord_t x_Cd { ZERO }; - metric.template convert(x_Ph, x_Cd); - real_t temperature { ZERO }; - if constexpr (M::Dim == Dim::_1D) { - const auto i1_ = static_cast(x_Cd[0]) + N_GHOSTS; - temperature = nt_array(i1_, comp_n_t) / nt_array(i1_, comp_n); - } else if constexpr (M::Dim == Dim::_2D) { - const auto i1_ = static_cast(x_Cd[0]) + N_GHOSTS; - const auto i2_ = static_cast(x_Cd[1]) + N_GHOSTS; - temperature = nt_array(i1_, i2_, comp_n_t) / nt_array(i1_, i2_, comp_n); - } else if constexpr (M::Dim == Dim::_3D) { - const auto i1_ = static_cast(x_Cd[0]) + N_GHOSTS; - const auto i2_ = static_cast(x_Cd[1]) + N_GHOSTS; - const auto i3_ = static_cast(x_Cd[2]) + N_GHOSTS; - temperature = nt_array(i1_, i2_, i3_, comp_n_t) / - nt_array(i1_, i2_, i3_, comp_n); - } - + Inline void operator()(const coord_t&, vec_t& k) const { real_t prob { ZERO }, n { ZERO }; auto gen = random_pool.get_state(); const auto rnd = Random(gen); @@ -361,16 +351,16 @@ namespace user { n += ONE; prob += ONE / (static_cast(1.20206) * CUBE(n)); } - const auto energy = -static_cast(2.7) * temperature * + const auto energy = -static_cast(2.7) * T_ph_inj * math::log( rnd1 * rnd2 * rnd3 + static_cast(1e-16)) / n; const auto costh = TWO * rndth - ONE; const auto phi = static_cast(constant::TWO_PI) * rndph; - v[0] = energy * math::sqrt(ONE - SQR(costh)) * math::cos(phi); - v[1] = energy * math::sqrt(ONE - SQR(costh)) * math::sin(phi); - v[2] = energy * costh; + k[0] = energy * math::sqrt(ONE - SQR(costh)) * math::cos(phi); + k[1] = energy * math::sqrt(ONE - SQR(costh)) * math::sin(phi); + k[2] = energy * costh; } }; @@ -381,27 +371,25 @@ namespace user { Must satisfy E = -v x B for Lorentz Force to be zero. @param bmag: magnetic field scaling - @param btheta: magnetic field polar angle - @param bphi: magnetic field azimuthal angle - @param drift_ux: drift velocity in the x direction + @param thetaB: Bx = bmag * cos(thetaB) + @param beta_upstream: drift three-velocity in the x direction */ - InitFields(real_t bmag, real_t btheta, real_t bphi, real_t drift_ux) + InitFields(real_t bmag, real_t thetaB, real_t beta_upstream) : Bmag { bmag } - , Btheta { btheta * static_cast(convert::deg2rad) } - , Bphi { bphi * static_cast(convert::deg2rad) } - , Vx { drift_ux } {} + , thetaB { thetaB * static_cast(convert::deg2rad) } + , beta_upstream { beta_upstream } {} // magnetic field components Inline auto bx1(const coord_t&) const -> real_t { - return Bmag * math::cos(Btheta); + return Bmag * math::cos(thetaB); } Inline auto bx2(const coord_t&) const -> real_t { - return Bmag * math::sin(Btheta) * math::sin(Bphi); + return ZERO; } Inline auto bx3(const coord_t&) const -> real_t { - return Bmag * math::sin(Btheta) * math::cos(Bphi); + return Bmag * math::sin(thetaB); } // electric field components @@ -410,15 +398,15 @@ namespace user { } Inline auto ex2(const coord_t&) const -> real_t { - return -Vx * Bmag * math::sin(Btheta) * math::cos(Bphi); + return -beta_upstream * Bmag * math::sin(thetaB); } Inline auto ex3(const coord_t&) const -> real_t { - return Vx * Bmag * math::sin(Btheta) * math::sin(Bphi); + return ZERO; } private: - const real_t Btheta, Bphi, Vx, Bmag; + const real_t Bmag, thetaB, beta_upstream; }; template @@ -438,14 +426,18 @@ namespace user { Metadomain& metadomain; // domain properties - const real_t global_xmin, global_xmax; + const real_t global_xmin, global_xmax; // gas properties - const real_t drift_ux, temperature, temperature_ratio, filling_fraction; + const real_t beta_upstream, Te, Te_ovr_Ti; // magnetic field properties - real_t Btheta, Bphi, Bmag; - // injector properties - const real_t injector_velocity, injection_start; - const int injection_frequency; + const real_t Bmag, thetaB; + // photon properties + const real_t photon_inj_rate; // units of n0 / time + const real_t T_ph_inj; // photon injection temperature + // plasma injector properties + const real_t filling_fraction, beta_injector; + const int injection_interval; + InitFields init_flds; PGen(const SimulationParams& p, Metadomain& m) @@ -453,24 +445,18 @@ namespace user { , metadomain { m } , global_xmin { metadomain.mesh().extent(in::x1).first } , global_xmax { metadomain.mesh().extent(in::x1).second } - , drift_ux { params.template get("setup.drift_ux") } - , temperature { params.template get("setup.temperature") } - , temperature_ratio { params.template get( - "setup.temperature_ratio", - ONE) } + , beta_upstream { params.template get("setup.beta_upstream") } + , Te { params.template get("setup.Te") } + , Te_ovr_Ti { params.template get("setup.Te_ovr_Ti", ONE) } , Bmag { params.template get("setup.Bmag", ZERO) } - , Btheta { params.template get("setup.Btheta", ZERO) } - , Bphi { params.template get("setup.Bphi", ZERO) } - , init_flds { Bmag, Btheta, Bphi, drift_ux } + , thetaB { params.template get("setup.thetaB", ZERO) } + , photon_inj_rate { params.template get("setup.photon_inj_rate", ZERO) } + , T_ph_inj { params.template get("setup.T_ph_inj") } , filling_fraction { params.template get("setup.filling_fraction", 1.0) } - , injector_velocity { params.template get( - "setup.injector_velocity", - 1.0) } - , injection_start { params.template get("setup.injection_start", 0.0) } - , injection_frequency { - params.template get("setup.injection_frequency", 100) - } {} + , beta_injector { params.template get("setup.beta_injector", 1.0) } + , injection_interval { params.template get("setup.injection_interval", 100) } + , init_flds { Bmag, thetaB, beta_upstream } {} auto MatchFields(simtime_t) const -> InitFields { return init_flds; @@ -511,8 +497,7 @@ namespace user { // minimum and maximum position of particles real_t xg_min = global_xmin; - // real_t xg_max = global_xmin + filling_fraction * (global_xmax - global_xmin); - real_t xg_max = global_xmax; + real_t xg_max = global_xmin + filling_fraction * (global_xmax - global_xmin); // define box to inject into boundaries_t box; @@ -527,184 +512,162 @@ namespace user { } } - // define temperatures of species - const auto temperatures = std::make_pair(temperature, - temperature_ratio * temperature); - // define drift speed of species - const auto drifts = std::make_pair( - std::vector { -drift_ux, ZERO, ZERO }, - std::vector { -drift_ux, ZERO, ZERO }); + const auto gamma_upstream = ONE / math::sqrt(ONE - SQR(beta_upstream)); // inject particles - arch::InjectUniformMaxwellians(params, - domain, - ONE, - temperatures, - { 1, 2 }, - drifts, - false, - box); + arch::InjectUniformMaxwellians( + params, + domain, + TWO, + std::make_pair(Te, Te / Te_ovr_Ti), + { 1, 2 }, + std::make_pair( + std::vector { -gamma_upstream * beta_upstream, ZERO, ZERO }, + std::vector { -gamma_upstream * beta_upstream, ZERO, ZERO }), + false, + box); } - void CustomPostStep(timestep_t /* step */, - simtime_t /* time */, + void CustomPostStep(timestep_t step, + simtime_t time, Domain& domain) { const auto dt = params.template get("algorithms.timestep.dt"); - // if (step % injection_frequency == 0) { - // /* - // * Replenish plasma in a moving injector - // * - // * Injector setup: - // * - // * global_xmin purge/replenish global_xmax - // * | x_init | | - // * V v V V - // * |:::::::::::;::::::::::|\\\\\\\\|......| - // * xmin xmax - // * ^ - // * | - // * moving injector - // */ - // - // // initial position of injector - // const auto x_init = global_xmin + - // filling_fraction * (global_xmax - global_xmin); - // - // // compute the position of the injector after the current timestep - // auto xmax = x_init + - // injector_velocity * - // (std::max(time - injection_start, ZERO) + dt); - // if (xmax >= global_xmax) { - // xmax = global_xmax; - // } - // - // // compute the beginning of the injected region - // auto xmin = xmax - injection_frequency * dt; - // if (xmin <= global_xmin) { - // xmin = global_xmin; - // } - // - // // define indice range to reset fields - // boundaries_t incl_ghosts; - // for (auto d = 0; d < M::Dim; ++d) { - // incl_ghosts.emplace_back(false, false); - // } - // - // // define box to reset fields - // boundaries_t purge_box; - // // loop over all dimension - // for (auto d = 0u; d < M::Dim; ++d) { - // if (d == 0) { - // purge_box.emplace_back(xmin, global_xmax); - // } else { - // purge_box.push_back(Range::All); - // } - // } - // - // const auto extent = domain.mesh.ExtentToRange(purge_box, incl_ghosts); - // tuple_t x_min { 0 }, x_max { 0 }; - // for (auto d = 0; d < M::Dim; ++d) { - // x_min[d] = extent[d].first; - // x_max[d] = extent[d].second; - // } - // - // Kokkos::parallel_for("ResetFields", - // CreateRangePolicy(x_min, x_max), - // arch::SetEMFields_kernel { - // domain.fields.em, - // init_flds, - // domain.mesh.metric }); - // metadomain.CommunicateFields(domain, Comm::E | Comm::B); - // - // /* - // tag particles inside the injection zone as dead - // */ - // const auto& mesh = domain.mesh; - // - // // loop over particle species - // for (auto s { 0u }; s < 2; ++s) { - // // get particle properties - // auto& species = domain.species[s]; - // auto i1 = species.i1; - // auto dx1 = species.dx1; - // auto tag = species.tag; - // - // Kokkos::parallel_for( - // "RemoveParticles", - // species.rangeActiveParticles(), - // Lambda(prtlidx_t p) { - // // check if the particle is already dead - // if (tag(p) == ParticleTag::dead) { - // return; - // } - // const auto x_Cd = static_cast(i1(p)) + - // static_cast(dx1(p)); - // const auto x_Ph = mesh.metric.template convert<1, Crd::Cd, Crd::XYZ>( - // x_Cd); - // - // if (x_Ph > xmin) { - // tag(p) = ParticleTag::dead; - // } - // }); - // } - // - // // define box to inject into - // boundaries_t inj_box; - // // loop over all dimension - // for (auto d = 0u; d < M::Dim; ++d) { - // if (d == 0) { - // inj_box.emplace_back(xmin, xmax); - // } else { - // inj_box.push_back(Range::All); - // } - // } - // - // // same maxwell distribution as above - // const auto temperatures = std::make_pair(temperature, - // temperature_ratio * temperature); - // const auto drifts = std::make_pair( - // std::vector { -drift_ux, ZERO, ZERO }, - // std::vector { -drift_ux, ZERO, ZERO }); - // arch::InjectUniformMaxwellians(params, - // domain, - // ONE, - // temperatures, - // { 1, 2 }, - // drifts, - // false, - // inj_box); - // } - - { + + if (step % injection_interval == 0) { /* - * Inject photons + * Replenish plasma in a moving injector + * + * Injector setup: + * + * global_xmin purge/replenish global_xmax + * | x_init | | + * V v V V + * |:::::::::::;::::::::::|\\\\\\\\|......| + * xmin xmax + * ^ + * | + * moving injector */ - auto compute_rho_v = ComputeRhoV {}; - auto compute_pressure = ComputePressure { domain.fields.bckp }; - arch::ComputeMomentWithSpeciesNew( + + // initial position of injector + const auto x_init = global_xmin + filling_fraction * (global_xmax - global_xmin); + + // compute the position of the injector after the current timestep + const auto xmax = std::min(x_init + beta_injector * (step + 1) * dt, global_xmax); + + // compute the beginning of the injected region + const auto xmin = (step == 0) + ? std::max(x_init - beta_upstream * dt, global_xmin) + : xmax - injection_interval * dt * beta_injector - (injection_interval + 1) * dt * beta_upstream; + + // define indice range to reset fields + boundaries_t incl_ghosts; + for (auto d = 0; d < M::Dim; ++d) { + incl_ghosts.emplace_back(false, false); + } + + // define box to reset fields + boundaries_t purge_box; + // loop over all dimension + for (auto d = 0u; d < M::Dim; ++d) { + if (d == 0) { + purge_box.emplace_back(xmin, global_xmax); + } else { + purge_box.push_back(Range::All); + } + } + + const auto extent = domain.mesh.ExtentToRange(purge_box, incl_ghosts); + tuple_t x_min { 0 }, x_max { 0 }; + for (auto d = 0; d < M::Dim; ++d) { + x_min[d] = extent[d].first; + x_max[d] = extent[d].second; + } + + Kokkos::parallel_for("ResetFields", + CreateRangePolicy(x_min, x_max), + arch::SetEMFields_kernel { + domain.fields.em, + init_flds, + domain.mesh.metric }); + metadomain.CommunicateFields(domain, Comm::E | Comm::B); + + /* + tag particles inside the injection zone as dead + */ + // const auto& mesh = domain.mesh; + + // loop over particle species + // for (auto& species : domain.species) { + // // get particle properties + // auto i1 = species.i1; + // auto dx1 = species.dx1; + // auto tag = species.tag; + + // Kokkos::parallel_for( + // "RemoveParticles", + // species.rangeActiveParticles(), + // Lambda(prtlidx_t p) { + // // check if the particle is already dead + // if (tag(p) == ParticleTag::dead) { + // return; + // } + // const auto x_Cd = static_cast(i1(p)) + + // static_cast(dx1(p)); + // const auto x_Ph = mesh.metric.template convert<1, Crd::Cd, Crd::XYZ>( + // x_Cd); + + // if (x_Ph > xmin) { + // tag(p) = ParticleTag::dead; + // } + // }); + // } + + // define box to inject into + boundaries_t inj_box; + // loop over all dimension + for (auto d = 0u; d < M::Dim; ++d) { + if (d == 0) { + inj_box.emplace_back(xmin, xmax); + } else { + inj_box.push_back(Range::All); + } + } + + const auto gamma_upstream = ONE / math::sqrt(ONE - SQR(beta_upstream)); + + // same maxwell distribution as above + arch::InjectUniformMaxwellians( params, domain, + TWO, + std::make_pair(Te, Te / Te_ovr_Ti), { 1, 2 }, - domain.fields.bckp, - { comp_n, comp_rho, comp_rho_vx, comp_rho_vy, comp_rho_vz }, - compute_rho_v); - arch::ComputeMomentWithSpeciesNew( + std::make_pair( + std::vector { -gamma_upstream * beta_upstream, ZERO, ZERO }, + std::vector { -gamma_upstream * beta_upstream, ZERO, ZERO }), + false, + inj_box); + } + + { + /* + * Inject photons + */ + auto compute_n = ComputeN {}; + arch::ComputeMomentWithSpeciesNew( params, domain, { 1, 2 }, domain.fields.bckp, - { comp_n_t }, - compute_pressure); + { comp_n }, + compute_n); // inject photons with a Planck distribution in energy and spatial distribution following the plasma density - const auto energy_dist = PlanckDistribution(domain.fields.bckp, - domain.mesh.metric, - domain.random_pool()); + const auto energy_dist = PlanckDistribution(T_ph_inj, + domain.random_pool()); const auto spatial_dist = PhotonSpatialDistribution(domain.fields.bckp, domain.mesh.metric); - const auto photon_inj_rate = params.template get( - "setup.photon_injection_rate", - ZERO); arch::InjectNonUniform( params, domain, From 95a359022e54ec7d97bb2a2684ec9e35b4d2417b Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 14 Sep 2026 11:53:16 -0400 Subject: [PATCH 24/25] post merge bugs fixed --- src/kernels/pushers/sr_policies.h | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/kernels/pushers/sr_policies.h b/src/kernels/pushers/sr_policies.h index 0134fbbb9..cc34698a5 100644 --- a/src/kernels/pushers/sr_policies.h +++ b/src/kernels/pushers/sr_policies.h @@ -173,14 +173,14 @@ namespace kernel::sr { auto with_emission = [&](auto next) { switch (emission_type) { case ntt::EmissionType::SYNCHROTRON: - next(MakePusherPolicyEmission( - dom, + next(MakePusherPolicyEmission( + domain, params, pusher_ctx)); break; case ntt::EmissionType::COMPTON: - next(MakePusherPolicyEmission( - dom, + next(MakePusherPolicyEmission( + domain, params, pusher_ctx)); break; @@ -202,9 +202,9 @@ namespace kernel::sr { DispatchExternalFields(pgen, domain, pusher_ctx, next); }; - with_emission(pgen, domain, [&](auto ep) { - with_custom_prtl_upd(pgen, domain, [&](auto cpu) { - with_ext_fields(pgen, domain, [&](auto ef) { + with_emission([&](auto ep) { + with_custom_prtl_upd([&](auto cpu) { + with_ext_fields([&](auto ef) { using E = decltype(ep); using CPU = decltype(cpu); using EF = decltype(ef); From 277eb74f32e4b3e7b718206a1429defd970ccc61 Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 14 Sep 2026 13:34:44 -0400 Subject: [PATCH 25/25] devenv setup --- .gitignore | 4 + CODEGUIDE.md | 2 +- dev/nix/devenv.lock | 45 ++++++++++++ dev/nix/devenv.nix | 173 ++++++++++++++++++++++++++++++++++++++++++++ dev/nix/devenv.yaml | 4 + 5 files changed, 227 insertions(+), 1 deletion(-) create mode 100644 dev/nix/devenv.lock create mode 100644 dev/nix/devenv.nix create mode 100644 dev/nix/devenv.yaml diff --git a/.gitignore b/.gitignore index a363191dd..85331517d 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,7 @@ tombi/ tidy/ .claude .understand-anything + +# devenv +.devenv* +devenv.local.nix diff --git a/CODEGUIDE.md b/CODEGUIDE.md index 1e4dee68d..058edb5ef 100644 --- a/CODEGUIDE.md +++ b/CODEGUIDE.md @@ -16,7 +16,7 @@ entity │ ├── styling.cmake # styling functions │ └── tests.cmake # root cmake for tests ├── dev # developer-specific tools -│ ├── nix # nix-shells +│ ├── nix # nix-shell & devenv environments │ ├── runners # dockerfiles for github runners on different architectures │ ├── scripts # developer-specific scripts │ ├── Dockerfile.common # parent docker environment for development diff --git a/dev/nix/devenv.lock b/dev/nix/devenv.lock new file mode 100644 index 000000000..25a1cc90d --- /dev/null +++ b/dev/nix/devenv.lock @@ -0,0 +1,45 @@ +{ + "nodes": { + "devenv": { + "locked": { + "dir": "src/modules", + "lastModified": 1789340509, + "narHash": "sha256-It2AD16uFwiAe9GvuLbk+j+qTDrQxT3XPclL71ZnJ3I=", + "owner": "cachix", + "repo": "devenv", + "rev": "6e830506b517d6a373f2dec6ee8c7f4683908856", + "type": "github" + }, + "original": { + "dir": "src/modules", + "owner": "cachix", + "repo": "devenv", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1789286504, + "narHash": "sha256-eiEK7cKZORNEvX0GeF3RtNEF/JXhgf2RqSp3230q13E=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "ef34387ddd751e1ab8857adf4676492d32eb24ec", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "devenv": "devenv", + "nixpkgs": "nixpkgs" + } + } + }, + "root": "root", + "version": 7 +} \ No newline at end of file diff --git a/dev/nix/devenv.nix b/dev/nix/devenv.nix new file mode 100644 index 000000000..5da1b6a01 --- /dev/null +++ b/dev/nix/devenv.nix @@ -0,0 +1,173 @@ +# devenv counterpart of `shell.nix`; run from this directory: +# devenv shell # cpu-only +# devenv shell -P cuda -O entity.arch:string AMPERE80 # cuda +# devenv shell -P hip -O entity.arch:string AMD_GFX90A # hip +# devenv shell -P mpi -P hdf5 # adios2 with mpi + hdf5 +# persistent settings can be put into `devenv.local.nix` (gitignored). +{ + pkgs, + lib, + config, + inputs, + ... +}: + +let + cfg = config.entity; + + gpu = lib.toUpper cfg.gpu; + arch = lib.toUpper cfg.arch; + + # `shell.nix` imports nixpkgs with `allowUnfree`/`cudaSupport` decided by the + # requested backend. devenv instantiates its own `pkgs` before this module is + # evaluated, so it cannot be reconfigured from here -- import the same input + # ourselves and build everything from that instance. + nixpkgs = import inputs.nixpkgs { + inherit (pkgs.stdenv.hostPlatform) system; + config = { + allowUnfree = true; + cudaSupport = gpu == "CUDA"; + }; + }; + + adios2Pkg = nixpkgs.callPackage ./adios2.nix { + pkgs = nixpkgs; + inherit (cfg) hdf5 mpi; + }; + + kokkosPkg = nixpkgs.callPackage ./kokkos.nix { + pkgs = nixpkgs; + stdenv = nixpkgs.stdenv; + inherit arch gpu; + }; + + extraPkgs = map (name: nixpkgs.${name}) (lib.filter (s: s != "") (lib.splitString "," cfg.extra)); + + # compilers are picked by the backend; CUDA goes through kokkos' nvcc_wrapper + compilerEnv = + { + NONE = { + CXX = "g++"; + CC = "gcc"; + }; + HIP = { + CXX = "clang++"; + CC = "clang"; + }; + CUDA = { }; + } + .${gpu}; +in +{ + options.entity = { + gpu = lib.mkOption { + # case-insensitive, as in `shell.nix` + type = lib.types.enum [ + "NONE" + "none" + "CUDA" + "cuda" + "HIP" + "hip" + ]; + default = "NONE"; + description = "GPU backend to build Kokkos with."; + }; + + arch = lib.mkOption { + type = lib.types.str; + default = "NATIVE"; + example = "AMPERE80"; + description = '' + Kokkos architecture; mandatory when `gpu` is not `NONE`. See + https://kokkos.org/kokkos-core-wiki/get-started/configuration-guide.html#gpu-architectures + ''; + }; + + hdf5 = lib.mkOption { + type = lib.types.bool; + default = false; + description = "Build ADIOS2 with HDF5 support."; + }; + + mpi = lib.mkOption { + type = lib.types.bool; + default = false; + description = "Build ADIOS2 with MPI support."; + }; + + extra = lib.mkOption { + type = lib.types.str; + default = ""; + example = "gdb,valgrind"; + description = '' + Comma-separated nixpkgs attributes to add to the environment, kept for + parity with `shell.nix`. `-O packages:pkgs "gdb valgrind"` does the same + without going through this option. + ''; + }; + }; + + config = { + name = + "nt2" + (if gpu != "NONE" then "-${lib.toLower gpu}" else "") + (if cfg.mpi then "-mpi" else ""); + + profiles = { + cuda.module = { + entity.gpu = "CUDA"; + }; + hip.module = { + entity.gpu = "HIP"; + }; + mpi.module = { + entity.mpi = true; + }; + hdf5.module = { + entity.hdf5 = true; + }; + }; + + packages = + (with nixpkgs; [ + zlib + cmake + + adios2Pkg + kokkosPkg + + python314 + + cmake-format + cmake-lint + neocmakelsp + black + pyright + taplo + vscode-langservers-extracted + ]) + ++ extraPkgs; + + env = compilerEnv // { + LD_LIBRARY_PATH = lib.makeLibraryPath [ + nixpkgs.stdenv.cc.cc + nixpkgs.zlib + ]; + }; + + enterShell = '' + BLUE='\033[0;34m' + NC='\033[0m' + + echo "following environment variables are set:" + '' + + lib.concatStringsSep "" ( + lib.mapAttrsToList (name: value: '' + echo -e " ''${BLUE}${name}''${NC}=${value}" + '') compilerEnv + ) + + '' + echo "" + echo -e "${config.name} devenv activated" + ''; + }; +} diff --git a/dev/nix/devenv.yaml b/dev/nix/devenv.yaml new file mode 100644 index 000000000..b4c9128be --- /dev/null +++ b/dev/nix/devenv.yaml @@ -0,0 +1,4 @@ +# devenv inputs; update the lock with `devenv update --from path:./dev/nix` +inputs: + nixpkgs: + url: github:NixOS/nixpkgs/nixos-unstable