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 diff --git a/examples/compton_jones/compton_jones.py b/examples/compton_jones/compton_jones.py new file mode 100644 index 000000000..505b777bf --- /dev/null +++ b/examples/compton_jones/compton_jones.py @@ -0,0 +1,63 @@ +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" +plt.rcParams["mathtext.fontset"] = "stix" + +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..8230a04c5 --- /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 = 100.0 + + [[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 = 5.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 = 9.0 + + [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..b507964c5 --- /dev/null +++ b/examples/compton_kompaneets/compton_kompaneets.py @@ -0,0 +1,66 @@ +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" +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, t in enumerate([0, 0.5 * tC, tC, 3 * tC]): + ax1.plot( + 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"] +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=(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"] / 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 / 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") + +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..a92699968 --- /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 = 60.0 + + [[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.1 + + [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_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 diff --git a/input.example.toml b/input.example.toml index 741d5c2d2..bb031d193 100644 --- a/input.example.toml +++ b/input.example.toml @@ -272,6 +272,45 @@ # @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 = "" + # Interval in timesteps between checking for the interaction + # @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 # @type: ushort [>= 0] diff --git a/src/archetypes/particle_injector.h b/src/archetypes/particle_injector.h index 415da851e..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,150 @@ 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 + * @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 diff --git a/src/archetypes/qed/compton.h b/src/archetypes/qed/compton.h new file mode 100644 index 000000000..8776d4a12 --- /dev/null +++ b/src/archetypes/qed/compton.h @@ -0,0 +1,305 @@ +/** + * @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); + static constexpr real_t Thomson_limit = static_cast(1e-3); + + const real_t nominal_probability_density; + random_number_pool_t random_pool; + + ComptonScattering(const prm::Parameters& params, + random_number_pool_t& random_pool) + : nominal_probability_density { params.template get( + "compton_scattering.nominal_probability_density") } + , random_pool { random_pool } { + if (nominal_probability_density <= ZERO) { + raise::Error("nominal_probability must be > 0", 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 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(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 static_cast(1.0 - 2.0 * e_ + 5.2 * SQR(e_)); + } else { + 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 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_); + 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; + } + + /* + * 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_); + 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 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), + 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 f_KN = KNCrossSection(photon_energy_); + + auto gen = random_pool.get_state(); + const auto rnd = Random(gen); + random_pool.free_state(gen); + + 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]; + } + } + }; + +} // namespace arch::qed + +#endif // ARCHETYPES_QED_COMPTON_H diff --git a/src/archetypes/utils.h b/src/archetypes/utils.h index 1ae775309..827598ddc 100644 --- a/src/archetypes/utils.h +++ b/src/archetypes/utils.h @@ -146,7 +146,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]; diff --git a/src/engines/engine.hpp b/src/engines/engine.hpp index b20e163ca..bd77c83a4 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; } }; @@ -248,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..125cb45a5 100644 --- a/src/engines/reporter.cpp +++ b/src/engines/reporter.cpp @@ -6,8 +6,10 @@ #include "utils/formatting.h" #include "utils/reporter.h" +#include "framework/parameters/extra.h" #include "framework/parameters/parameters.h" +#include #include #include @@ -149,6 +151,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 new file mode 100644 index 000000000..c5d04ca66 --- /dev/null +++ b/src/engines/srpic/twobody.h @@ -0,0 +1,101 @@ +#ifndef ENGINES_SRPIC_TWOBODY_H +#define ENGINES_SRPIC_TWOBODY_H + +#include "enums.h" + +#include "traits/metric.h" +#include "utils/error.h" +#include "utils/formatting.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 { + + 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.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()); + + 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); + } + } + } + } + } // namespace srpic +} // namespace ntt + +#endif // ENGINES_SRPIC_TWOBODY_H diff --git a/src/framework/parameters/extra.cpp b/src/framework/parameters/extra.cpp index dbee39c7c..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 { @@ -114,6 +116,29 @@ 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", + "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), + .tile_size = toml::find_or(tbint, "tile_size", 4u), + .recoil1 = toml::find_or(tbint, "recoil1", true), + .recoil2 = toml::find_or(tbint, "recoil2", true) }); + } } void Extra::setParams(const std::map& extra, @@ -159,6 +184,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.interaction", twobody_interactions); } } // namespace params } // namespace ntt diff --git a/src/framework/parameters/extra.h b/src/framework/parameters/extra.h index 29de4527a..c3891429f 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,16 @@ namespace ntt { namespace params { + struct TwoBodyInteractionParams { + TwoBodyInteractionFlag type; + std::vector group1; + std::vector group2; + timestep_t interval; + ncells_t tile_size; + bool recoil1; + bool recoil2; + }; + struct Extra { // radiative drag parameters std::optional synchrotron_gamma_rad; @@ -45,6 +56,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 e1387677e..8c8c54f7b 100644 --- a/src/global/defaults.h +++ b/src/global/defaults.h @@ -117,6 +117,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 5700a3a28..b88f2da95 100644 --- a/src/global/enums.h +++ b/src/global/enums.h @@ -402,6 +402,43 @@ 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"; + } + } + + 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.c_str()), + HERE); + return NONE; + } + } + } // namespace TwoBodyInteraction + + using TwoBodyInteractionFlag = uint8_t; + namespace OutputSmoothingType { enum OutputSmoothingTypeFlag_ : uint8_t { NONE = 0, diff --git a/src/global/traits/policies.h b/src/global/traits/policies.h index 7f6d98985..fdca30ef3 100644 --- a/src/global/traits/policies.h +++ b/src/global/traits/policies.h @@ -138,4 +138,40 @@ concept CustomParticleUpdatePolicyClass = } or traits::custom_prtl_update::IsNoPolicy; +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; + #endif // TRAITS_POLICIES_H diff --git a/src/kernels/injectors.hpp b/src/kernels/injectors.hpp index 5d7904cd9..49ef50a21 100644 --- a/src/kernels/injectors.hpp +++ b/src/kernels/injectors.hpp @@ -5,6 +5,8 @@ * - kernel::UniformInjector_kernel<> * - kernel::GlobalInjector_kernel<> * - kernel::NonUniformInjector_kernel<> + * - kernel::SingleSpeciesNonUniformInjector_kernel<> + * - kernel::SingleSpeciesUniformInjector_kernel<> * @namespaces: * - kernel:: */ @@ -858,6 +860,363 @@ 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 { + + 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/particle_moments.hpp b/src/kernels/particle_moments.hpp index 1244ba21d..68461f768 100644 --- a/src/kernels/particle_moments.hpp +++ b/src/kernels/particle_moments.hpp @@ -252,7 +252,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 { diff --git a/src/kernels/twobody_interactions.hpp b/src/kernels/twobody_interactions.hpp new file mode 100644 index 000000000..901b43acf --- /dev/null +++ b/src/kernels/twobody_interactions.hpp @@ -0,0 +1,387 @@ +/** + * @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 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); + 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); + 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); + ncells_on_tile *= (i3_max - i3_min); + } + return ncells_on_tile; + } + + 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, + 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); + 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 cell 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; + + // 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].tag(p1) == + ParticleTag::alive) + ? interaction_policy.species[sp1 - 1].weight(p1) + : ZERO; + }, + 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].tag(p2) == + ParticleTag::alive) + ? interaction_policy.species[sp2 - 1].weight(p2) + : ZERO; + }, + 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)) { + 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))); + } + + array_t interaction_pairs { "interaction_pairs", + combined_idx1.extent(0) }; + array_t counter { "counter" }; + + Kokkos::parallel_for( + "PopulateInteractionPairs", + Kokkos::TeamPolicy<>(num_tiles, Kokkos::AUTO), + Lambda(const Kokkos::TeamPolicy<>::member_type& team) { + const ncells_t t = team.league_rank(); + 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); + 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)); + 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); + } + } + }); + }); + 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 + +#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..14350cbf4 --- /dev/null +++ b/tests/archetypes/qed_compton.cpp @@ -0,0 +1,280 @@ +#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 } + }; + const auto ppc0 = static_cast(npart) / (nx1 * nx2); + + prm::Parameters params; + params.set("compton_scattering.nominal_probability_density", + static_cast(1e-3)); + + 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, + ppc0, + 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; +} 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; +}