diff --git a/cmake/report.cmake b/cmake/report.cmake index 822d82650..e6ad4aaf8 100644 --- a/cmake/report.cmake +++ b/cmake/report.cmake @@ -146,7 +146,7 @@ printchoices( ON "${Green}" VENDOR_SORT_REPORT - 44) + 46) printchoices( "Debug mode" "DEBUG" @@ -234,6 +234,9 @@ string( " - DEVICES [${Magenta}Kokkos_ENABLE_***${ColorReset}]: " "${Kokkos_DEVICES}" "\n" + " " + ${VENDOR_SORT_REPORT} + "\n" " > Multi-node specs" " ${Dim}[requires mpi=ON]${ColorReset}" "\n" @@ -255,9 +258,6 @@ string( " " "- Deposit drift [${Magenta}team_policy_drift${ColorReset}]: " ${team_policy_drift} - "\n" - " " - ${VENDOR_SORT_REPORT} "\n") string( diff --git a/dependencies.py b/dependencies.py index 8f39bde85..54adf55b2 100755 --- a/dependencies.py +++ b/dependencies.py @@ -8,7 +8,6 @@ from dataclasses import dataclass, field from typing import Callable, List, Optional, Tuple - # ============================ # colors: edit these # ============================ @@ -54,7 +53,7 @@ class Settings: # versions kokkos_version: str = "5.2.1" - adios2_version: str = "2.11.0" + adios2_version: str = "2.12.1" # options kokkos_backend: str = "cpu" @@ -135,7 +134,9 @@ def InstallKokkosScriptModfile(settings: Settings) -> tuple[str, str]: [f"module load {module}" for module in settings.module_loads] ) src_path = os.path.join(prefix, "src", "kokkos") - install_path = os.path.join(prefix, "kokkos", version, backend, arch.lower() if arch else "") + install_path = os.path.join( + prefix, "kokkos", version, backend, arch.lower() if arch else "" + ) if os.path.exists(install_path) and not settings.overwrite: raise FileExistsError( f"Kokkos install path {install_path} already exists and overwrite is disabled" @@ -300,7 +301,7 @@ def InstallNt2pyScript(settings: Settings) -> str: }, "stellar": {"module_loads": []}, "perlmutter": { - "module_loads": ["gpu/1.0"], + "module_loads": ["gpu/1.0", "python/3.14-26.8.1"], "kokkos_backend": "cuda", "kokkos_arch": "AMPERE80", "extra_kokkos_flags": [ @@ -313,17 +314,19 @@ def InstallNt2pyScript(settings: Settings) -> str: ], }, "lumi": { - "module_loads": ["PrgEnv-cray", "cray-mpich", "craype-accel-amd-gfx90a", "rocm"], + "module_loads": [ + "PrgEnv-cray", + "cray-mpich", + "craype-accel-amd-gfx90a", + "rocm", + ], "kokkos_backend": "hip", "kokkos_arch": "AMD_GFX90A", "extra_kokkos_flags": [ "CMAKE_CXX_COMPILER=hipcc", "AMDGPU_TARGETS=gfx90a", ], - "extra_adios2_flags": [ - "CMAKE_CXX_COMPILER=CC", - "CMAKE_C_COMPILER=cc" - ] + "extra_adios2_flags": ["CMAKE_CXX_COMPILER=CC", "CMAKE_C_COMPILER=cc"], }, "frontier": {"module_loads": []}, "aurora": {"module_loads": []}, @@ -366,7 +369,7 @@ def on_install_confirmed(settings: Settings) -> None: "modules", "kokkos", settings.kokkos_version, - settings.kokkos_backend, + settings.kokkos_backend, settings.kokkos_arch.strip().lower(), ) os.makedirs(os.path.dirname(kokkos_modfile_file), exist_ok=True) @@ -1089,14 +1092,18 @@ def toggle(k: str): def menu_cluster(self) -> Tuple[str, str, List[MenuItem]]: def choose(name: str): - print ("CALLING:", name) + print("CALLING:", name) apply_preset(self.s, name) self.push("custom") return ( "cluster-specific", "pick a preset:", - [MenuItem(cluster, "apply preset", on_enter=lambda c=cluster: choose(c)) for cluster in list(PRESETS.keys())] + [MenuItem("back", "", on_enter=self.pop)], + [ + MenuItem(cluster, "apply preset", on_enter=lambda c=cluster: choose(c)) + for cluster in list(PRESETS.keys()) + ] + + [MenuItem("back", "", on_enter=self.pop)], ) def get_menu(self) -> Tuple[str, str, List[MenuItem]]: diff --git a/input.example.toml b/input.example.toml index 8d6f25117..8870b29b5 100644 --- a/input.example.toml +++ b/input.example.toml @@ -627,10 +627,14 @@ # @type: bool # @default: true log_bins = "" - # Number of bins for the spectra output + # Number of energy bins for the spectra output # @type: uint [> 0] # @default: 200 - n_bins = "" + num_energy_bins = "" + # Number of spatial bins for the spectra output + # @type: array [size 1 :->: 3] + # @default: [1, 1, 1] + num_spatial_bins = "" # Number of timesteps between spectra outputs # @type: uint # @default: 0 diff --git a/pgens/shock/shock.toml b/pgens/shock/shock.toml index 8cde86990..8d3270a74 100644 --- a/pgens/shock/shock.toml +++ b/pgens/shock/shock.toml @@ -40,35 +40,38 @@ maxnpart = 2e7 [setup] - drift_ux = 0.15 # speed towards the wall [c] - temperature = 0.168 # temperature of maxwell distribution [kB T / (m_i c^2)] - temperature_ratio = 1.0 # temperature ratio of electrons to protons - Bmag = 1.0 # magnetic field strength as fraction of magnetisation - Btheta = 63.0 # magnetic field angle in the plane - Bphi = 0.0 # magnetic field angle out of plane - filling_fraction = 0.99 # 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 + drift_ux = 0.15 # speed towards the wall [c] + temperature = 0.168 # temperature of maxwell distribution [kB T / (m_i c^2)] + temperature_ratio = 1.0 # temperature ratio of electrons to protons + Bmag = 1.0 # magnetic field strength as fraction of magnetisation + Btheta = 63.0 # magnetic field angle in the plane + Bphi = 0.0 # magnetic field angle out of plane + filling_fraction = 0.99 # 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 [output] interval_time = 1000.0 format = "BPFile" - + [output.fields] quantities = ["N", "B", "E"] [output.particles] enable = true - stride = 10 + stride = 10 [output.spectra] - enable = false + enable = true + e_min = 1e-2 + e_max = 1e1 + log_bins = true + num_spatial_bins = [40, 3, 1] [diagnostics] log_level = "WARNING" - blocking_timers = true [checkpoint] interval = 10000 - keep = 1 \ No newline at end of file + keep = 2 diff --git a/src/framework/domain/io/init.cpp b/src/framework/domain/io/init.cpp index c054dc9a5..472ba9618 100644 --- a/src/framework/domain/io/init.cpp +++ b/src/framework/domain/io/init.cpp @@ -93,7 +93,9 @@ namespace ntt { for (const auto& sp : species_params()) { spectra_species.push_back(sp.index()); } - g_writer.defineSpectraOutputs(spectra_species); + const auto num_spatial_bins = params.template get>( + "output.spectra.num_spatial_bins"); + g_writer.defineSpectraOutputs(spectra_species, num_spatial_bins); for (const auto& type : { "fields", "particles", "spectra" }) { g_writer.addTracker(type, params.template get( diff --git a/src/framework/domain/io/spectra.cpp b/src/framework/domain/io/spectra.cpp index d063b9d3d..f628ff05b 100644 --- a/src/framework/domain/io/spectra.cpp +++ b/src/framework/domain/io/spectra.cpp @@ -10,7 +10,7 @@ #include "framework/domain/metadomain.h" #include "framework/parameters/parameters.h" #include "framework/specialization_registry.h" -#include "kernels/particle_moments.hpp" +#include "kernels/particle_energy_distribution.hpp" #include #include @@ -20,6 +20,7 @@ #include #endif // MPI_ENABLED +#include #include #include #include @@ -27,53 +28,183 @@ namespace ntt { + template + void GenerateEnergyDistribution(const Particles& species, + array_t& dn, + const kernel::EnergyBinning& energy_binning, + const M& metric) { + auto dn_scatter = Kokkos::Experimental::create_scatter_view(dn); + Kokkos::parallel_for("ComputeSpectraSpatial", + species.rangeActiveParticles(), + kernel::ParticleDistribution_kernel { species, + dn_scatter, + energy_binning, + metric }); + Kokkos::Experimental::contribute(dn, dn_scatter); + } + + template + void GenerateSpatialEnergyDistribution( + const Particles& species, + nddata_t(M::Dim) + 1, real_t>& dn, + const array_t& nmin_i, + const array_t& dncells_i, + const array_t& nbins_i, + const kernel::EnergyBinning& energy_binning, + const M& metric) { + auto dn_scatter = Kokkos::Experimental::create_scatter_view(dn); + Kokkos::parallel_for( + "ComputeSpectra", + species.rangeActiveParticles(), + kernel::ParticleDistributionSpatial_kernel { species, + dn_scatter, + nmin_i, + dncells_i, + nbins_i, + energy_binning, + metric }); + Kokkos::Experimental::contribute(dn, dn_scatter); + } + + template + void GenerateSpatialBins(array_t& bins, + const M& metric, + real_t dncells, + size_t nbins) { + Kokkos::parallel_for( + "GenerateSpatialBins", + nbins + 1, + Lambda(uint32_t b) { + bins(b) = metric.template convert( + static_cast(b) * dncells); + }); + } + template void Metadomain::WriteSpectra(const SimulationParams& params, Domain* local_domain, timestep_t current_step, simtime_t current_time) { - + constexpr auto dim = static_cast(M::Dim); g_writer.beginWriting(WriteMode::Spectra, current_step, current_time); const auto log_bins = params.template get("output.spectra.log_bins"); - const auto n_bins = params.template get("output.spectra.n_bins"); - auto e_min = params.template get("output.spectra.e_min"); - auto e_max = params.template get("output.spectra.e_max"); + + const auto num_energy_bins = params.template get( + "output.spectra.num_energy_bins"); + auto e_min = params.template get("output.spectra.e_min"); + auto e_max = params.template get("output.spectra.e_max"); if (log_bins) { e_min = math::log10(e_min); e_max = math::log10(e_max); } - const array_t energy { "energy", n_bins + 1 }; - Kokkos::parallel_for( - "GenerateEnergyBins", - n_bins + 1, - Lambda(uint32_t e) { - if (log_bins) { - energy(e) = math::pow(static_cast(10), - e_min + (e_max - e_min) * static_cast(e) / - static_cast(n_bins)); + + const auto num_spatial_bins = params.template get>( + "output.spectra.num_spatial_bins"); + const auto spatial_binning_enabled = std::any_of(num_spatial_bins.begin(), + num_spatial_bins.end(), + [](const auto& n) { + return n != 1u; + }); + + // fractional number of cells per each direction in each bin + const array_t dncells_i { "dncells_i" }; + // left edge of the local domain in each direction in number of cells + const array_t nmin_i { "nmin_i" }; + const array_t nbins_i { "nbins_i" }; + if (spatial_binning_enabled) { + auto dncells_i_h = Kokkos::create_mirror_view(dncells_i); + auto nmin_i_h = Kokkos::create_mirror_view(nmin_i); + auto nbins_i_h = Kokkos::create_mirror_view(nbins_i); + for (auto d = 0u; d < dim; ++d) { + dncells_i_h(d) = static_cast(mesh().n_active(static_cast(d))) / + static_cast(num_spatial_bins[d]); + nmin_i_h(d) = static_cast(local_domain->offset_ncells()[d]); + nbins_i_h(d) = num_spatial_bins[d]; + } + Kokkos::deep_copy(dncells_i, dncells_i_h); + Kokkos::deep_copy(nmin_i, nmin_i_h); + Kokkos::deep_copy(nbins_i, nbins_i_h); + } + for (const auto& spec : g_writer.spectraWriters()) { + auto& species = local_domain->species[spec.species() - 1]; + if (not spatial_binning_enabled) { + array_t dn { "dn", num_energy_bins }; + GenerateEnergyDistribution(species, + dn, + { e_min, e_max, log_bins, num_energy_bins }, + local_domain->mesh.metric); + g_writer.writeSpectrum(dn, spec.name()); + } else { + nddata_t dn; + if constexpr (M::Dim == Dim::_1D) { + dn = { "dn", num_spatial_bins[0], num_energy_bins }; + } else if constexpr (M::Dim == Dim::_2D) { + dn = { "dn", num_spatial_bins[0], num_spatial_bins[1], num_energy_bins }; + } else if constexpr (M::Dim == Dim::_3D) { + dn = { "dn", + num_spatial_bins[0], + num_spatial_bins[1], + num_spatial_bins[2], + num_energy_bins }; } else { - energy(e) = e_min + (e_max - e_min) * static_cast(e) / - static_cast(n_bins); + raise::Error("invalid dimension", HERE); } - }); - for (const auto& spec : g_writer.spectraWriters()) { - auto& species = local_domain->species[spec.species() - 1]; - array_t dn { "dn", n_bins }; - auto dn_scatter = Kokkos::Experimental::create_scatter_view(dn); + GenerateSpatialEnergyDistribution( + species, + dn, + nmin_i, + dncells_i, + nbins_i, + { e_min, e_max, log_bins, num_energy_bins }, + local_domain->mesh.metric); + g_writer.writeSpectrumSpatial(dim + 1u)>(dn, + spec.name()); + } + } + + { + // energy bins + const array_t energy { "energy", num_energy_bins + 1 }; Kokkos::parallel_for( - "ComputeSpectra", - species.rangeActiveParticles(), - kernel::ParticleDistribution_kernel { species, - dn_scatter, - e_min, - e_max, - log_bins, - n_bins, - local_domain->mesh.metric }); - Kokkos::Experimental::contribute(dn, dn_scatter); - g_writer.writeSpectrum(dn, spec.name()); + "GenerateEnergyBins", + num_energy_bins + 1, + Lambda(uint32_t e) { + if (log_bins) { + energy(e) = math::pow(static_cast(10), + e_min + (e_max - e_min) * static_cast(e) / + static_cast(num_energy_bins)); + } else { + energy(e) = e_min + (e_max - e_min) * static_cast(e) / + static_cast(num_energy_bins); + } + }); + + g_writer.writeSpectrumBins(energy, "sEbn"); + } + if (spatial_binning_enabled) { + const auto metric = mesh().metric; + const auto n_active = mesh().n_active(); + for (auto d = 0u; d < dim; ++d) { + array_t xi { "xi", num_spatial_bins[d] + 1 }; + // # of cells per spatial bin in direction `d` (fractional) + const auto dncells = static_cast(n_active[d]) / + static_cast(num_spatial_bins[d]); + if (d == 0) { + GenerateSpatialBins<1>(xi, metric, dncells, num_spatial_bins[d]); + } else if (d == 1) { + if constexpr (dim > 1) { + GenerateSpatialBins<2>(xi, metric, dncells, num_spatial_bins[d]); + } + } else if (d == 2) { + if constexpr (dim > 2) { + GenerateSpatialBins<3>(xi, metric, dncells, num_spatial_bins[d]); + } + } else { + raise::Error("invalid dimension", HERE); + } + g_writer.writeSpectrumBins(xi, "sX" + std::to_string(d + 1) + "bn"); + } } - g_writer.writeSpectrumBins(energy, "sEbn"); g_writer.endWriting(WriteMode::Spectra); } diff --git a/src/framework/parameters/output.cpp b/src/framework/parameters/output.cpp index 6c1858e9d..d2a6a9c44 100644 --- a/src/framework/parameters/output.cpp +++ b/src/framework/parameters/output.cpp @@ -153,11 +153,38 @@ namespace ntt { "spectra", "log_bins", defaults::output::spec_log); - spectra_n_bins = toml::find_or(toml_data, - "output", - "spectra", - "n_bins", - defaults::output::spec_nbins); + if (toml_data.contains("output") and + toml_data.at("output").contains("spectra") and + toml_data.at("output").at("spectra").contains("n_bins")) { + spectra_num_energy_bins = toml::find(toml_data, + "output", + "spectra", + "n_bins"); + raise::Warning( + "`output.spectra.n_bins` is deprecated and will be removed in 1.6+ " + "versions, use `output.spectra.num_energy_bins` instead", + HERE); + } else { + spectra_num_energy_bins = toml::find_or(toml_data, + "output", + "spectra", + "num_energy_bins", + defaults::output::spec_num_e_bins); + } + spectra_num_spatial_bins = toml::find_or>( + toml_data, + "output", + "spectra", + "num_spatial_bins", + std::vector { 1, 1, 1 }); + if (spectra_num_spatial_bins->size() < static_cast(dim)) { + raise::Error("`output.spectra.num_spatial_bins` must have at least " + + std::to_string(static_cast(dim)) + " entries", + HERE); + } + spectra_num_spatial_bins->erase( + spectra_num_spatial_bins->begin() + static_cast(dim), + spectra_num_spatial_bins->end()); /* Stats ---------------------------------------------------------------- */ stats_quantities = toml::find_or(toml_data, @@ -207,7 +234,10 @@ namespace ntt { params->set("output.spectra.e_min", spectra_e_min.value()); params->set("output.spectra.e_max", spectra_e_max.value()); params->set("output.spectra.log_bins", spectra_log_bins.value()); - params->set("output.spectra.n_bins", spectra_n_bins.value()); + params->set("output.spectra.num_energy_bins", + spectra_num_energy_bins.value()); + params->set("output.spectra.num_spatial_bins", + spectra_num_spatial_bins.value()); params->set("output.stats.quantities", stats_quantities.value()); params->set("output.stats.custom", stats_custom_quantities.value()); diff --git a/src/framework/parameters/output.h b/src/framework/parameters/output.h index 27b077bc5..a51bf9290 100644 --- a/src/framework/parameters/output.h +++ b/src/framework/parameters/output.h @@ -50,10 +50,11 @@ namespace ntt { std::optional> particles_species; std::optional particles_stride; - std::optional spectra_e_min; - std::optional spectra_e_max; - std::optional spectra_log_bins; - std::optional spectra_n_bins; + std::optional spectra_e_min; + std::optional spectra_e_max; + std::optional spectra_log_bins; + std::optional spectra_num_energy_bins; + std::optional> spectra_num_spatial_bins; std::optional> stats_quantities; std::optional> stats_custom_quantities; diff --git a/src/global/arch/kokkos_aliases.h b/src/global/arch/kokkos_aliases.h index 7212b554e..b0ed48773 100644 --- a/src/global/arch/kokkos_aliases.h +++ b/src/global/arch/kokkos_aliases.h @@ -4,7 +4,7 @@ * @implements * - ClassLambda, Lambda, Function, Inline macros * - array_t, array_mirror_t, scatter_array_t - * - ndarray_t, ndfield_t + * - nddata_t, scatter_nddata_t, ndarray_t, ndfield_t * - ndfield_mirror_t, scatter_ndfield_t * - range_t, range_h_t * - CreateRangePolicy, CreateRangePolicyOnHost @@ -80,6 +80,38 @@ namespace kokkos_aliases_hidden { template using nddata_t = typename kokkos_aliases_hidden::nddata_impl::type; +// Scatter view for nddata +namespace kokkos_aliases_hidden { + // c++ magic + template + struct scatter_nddata_impl { + using type = void; + }; + + template + struct scatter_nddata_impl<1, T> { + using type = scatter_array_t; + }; + + template + struct scatter_nddata_impl<2, T> { + using type = scatter_array_t; + }; + + template + struct scatter_nddata_impl<3, T> { + using type = scatter_array_t; + }; + + template + struct scatter_nddata_impl<4, T> { + using type = scatter_array_t; + }; +} // namespace kokkos_aliases_hidden + +template +using scatter_nddata_t = typename kokkos_aliases_hidden::scatter_nddata_impl::type; + template using ndarray_t = typename kokkos_aliases_hidden::nddata_impl::type; diff --git a/src/global/defaults.h b/src/global/defaults.h index c4f9b67f4..82b4364cc 100644 --- a/src/global/defaults.h +++ b/src/global/defaults.h @@ -75,7 +75,7 @@ namespace ntt::defaults { const real_t spec_emin = 1e-3; const real_t spec_emax = 1e3; const bool spec_log = true; - const std::size_t spec_nbins = 200; + const std::size_t spec_num_e_bins = 200; const std::vector stats_quantities = { "B^2", "E^2", "ExB", diff --git a/src/kernels/particle_energy_distribution.hpp b/src/kernels/particle_energy_distribution.hpp new file mode 100644 index 000000000..eba88bac6 --- /dev/null +++ b/src/kernels/particle_energy_distribution.hpp @@ -0,0 +1,220 @@ +/** + * @file kernels/particle_moments.hpp + * @brief Kernels for computing particle distribution functions + * @implements + * - kernel::EnergyBinning + * - kernel::ParticleDistributionBase_kernel<> + * - kernel::ParticleDistribution_kernel<> + * - kernel::ParticleDistributionSpatial_kernel<> + * @namespaces: + * - kernel:: + */ + +#ifndef KERNELS_PARTICLE_ENERGY_DISTRIBUTION_HPP +#define KERNELS_PARTICLE_ENERGY_DISTRIBUTION_HPP + +#include "enums.h" +#include "global.h" + +#include "arch/kokkos_aliases.h" +#include "traits/metric.h" +#include "utils/error.h" +#include "utils/numeric.h" + +#include "framework/containers/particles.h" + +namespace kernel { + using namespace ntt; + + struct EnergyBinning { + const real_t e_min, e_max; + const bool log_bins; + const size_t n_bins; + + EnergyBinning(real_t e_min, real_t e_max, bool log_bins, size_t n_bins) + : e_min { e_min } + , e_max { e_max } + , log_bins { log_bins } + , n_bins { n_bins } {} + }; + + template + struct ParticleDistributionBase_kernel { + const ParticleArrays particles; + const bool is_massive; + const M metric; + + const EnergyBinning energy_binning; + + ParticleDistributionBase_kernel(const Particles& particles, + const EnergyBinning& energy_binning, + const M& metric) + : particles { static_cast(particles) } + , is_massive { (particles.mass() != 0.0f) } + , energy_binning { energy_binning } + , metric { metric } {} + + Inline auto EnergyBinIndex(prtlidx_t p) const -> size_t { + real_t en; + if constexpr (S == SimEngine::SRPIC) { + if (is_massive) { + en = U2GAMMA(particles.ux1(p), particles.ux2(p), particles.ux3(p)) - ONE; + } else { + en = NORM(particles.ux1(p), particles.ux2(p), particles.ux3(p)); + } + } else if constexpr (S == SimEngine::GRPIC) { + coord_t x_Code { ZERO }; + x_Code[0] = static_cast(particles.i1(p)) + + static_cast(particles.dx1(p)); + x_Code[1] = static_cast(particles.i2(p)) + + static_cast(particles.dx2(p)); + + // raise full covariant 4-vector to get correct contravariant u^0 + // u^i != h^{ij} u_j + const real_t u_0_cov { metric.u_0( + x_Code, + { particles.ux1(p), particles.ux2(p), particles.ux3(p) }, + (is_massive) ? ONE : ZERO) }; + vec_t u_cntrv_4d { ZERO }; + metric.template transform_4d( + x_Code, + { u_0_cov, particles.ux1(p), particles.ux2(p), particles.ux3(p) }, + u_cntrv_4d); + // in GR: u^0 = Gamma/alpha + const real_t Gamma { metric.alpha(x_Code) * u_cntrv_4d[0] }; + en = is_massive ? (Gamma - ONE) : Gamma; + } + if (energy_binning.log_bins) { + en = math::log10(en); + } + if (en <= energy_binning.e_min) { + return 0u; + } else if (en >= energy_binning.e_max) { + return energy_binning.n_bins - 1u; + } else { + return static_cast(static_cast(energy_binning.n_bins) * + (en - energy_binning.e_min) / + (energy_binning.e_max - energy_binning.e_min)); + } + } + }; + + template + class ParticleDistribution_kernel + : public ParticleDistributionBase_kernel { + scatter_array_t dn_scatter; + using ParticleDistributionBase_kernel::particles; + using ParticleDistributionBase_kernel::EnergyBinIndex; + + public: + ParticleDistribution_kernel(const Particles& particles, + const scatter_array_t& dn_scatter, + const EnergyBinning& energy_binning, + const M& metric) + : ParticleDistributionBase_kernel(particles, energy_binning, metric) + , dn_scatter { dn_scatter } {} + + Inline void operator()(prtlidx_t p) const { + if (particles.tag(p) != ParticleTag::alive) { + return; + } + const auto e_ind = EnergyBinIndex(p); + auto dn_acc = dn_scatter.access(); + dn_acc(e_ind) += particles.weight(p); + } + }; + + template + class ParticleDistributionSpatial_kernel + : public ParticleDistributionBase_kernel { + static constexpr auto D = M::Dim; + static constexpr auto Dim = static_cast(D); + static constexpr auto DimP1 = Dim + 1u; + using ParticleDistributionBase_kernel::particles; + using ParticleDistributionBase_kernel::EnergyBinIndex; + + scatter_nddata_t dn_scatter; + const array_t nmin_i; + const array_t dncells_i; + const array_t nbins_i; + + public: + ParticleDistributionSpatial_kernel( + const Particles& particles, + const scatter_nddata_t& dn_scatter, + const array_t nmin_i, + const array_t dncells_i, + const array_t nbins_i, + const EnergyBinning& energy_binning, + const M& metric) + : ParticleDistributionBase_kernel(particles, energy_binning, metric) + , dn_scatter { dn_scatter } + , nmin_i { nmin_i } + , dncells_i { dncells_i } + , nbins_i { nbins_i } {} + + Inline auto SpatialBinIndex(real_t i, real_t nmin, real_t dncells, size_t nbins) const + -> size_t { + const auto ni = i + nmin; + if (ni <= ZERO) { + return 0u; + } else if (ni >= dncells * static_cast(nbins)) { + return nbins - 1u; + } else { + return static_cast(ni / dncells); + } + } + + Inline void operator()(prtlidx_t p) const { + if (particles.tag(p) != ParticleTag::alive) { + return; + } + const auto e_ind = EnergyBinIndex(p); + if constexpr (D == Dim::_1D) { + const auto i1_ind = SpatialBinIndex(static_cast(particles.i1(p)) + + static_cast(particles.dx1(p)), + nmin_i(0), + dncells_i(0), + nbins_i(0)); + auto dn_acc = dn_scatter.access(); + dn_acc(i1_ind, e_ind) += particles.weight(p); + } else if constexpr (D == Dim::_2D) { + const auto i1_ind = SpatialBinIndex(static_cast(particles.i1(p)) + + static_cast(particles.dx1(p)), + nmin_i(0), + dncells_i(0), + nbins_i(0)); + const auto i2_ind = SpatialBinIndex(static_cast(particles.i2(p)) + + static_cast(particles.dx2(p)), + nmin_i(1), + dncells_i(1), + nbins_i(1)); + auto dn_acc = dn_scatter.access(); + dn_acc(i1_ind, i2_ind, e_ind) += particles.weight(p); + } else if constexpr (D == Dim::_3D) { + const auto i1_ind = SpatialBinIndex(static_cast(particles.i1(p)) + + static_cast(particles.dx1(p)), + nmin_i(0), + dncells_i(0), + nbins_i(0)); + const auto i2_ind = SpatialBinIndex(static_cast(particles.i2(p)) + + static_cast(particles.dx2(p)), + nmin_i(1), + dncells_i(1), + nbins_i(1)); + const auto i3_ind = SpatialBinIndex(static_cast(particles.i3(p)) + + static_cast(particles.dx3(p)), + nmin_i(2), + dncells_i(2), + nbins_i(2)); + auto dn_acc = dn_scatter.access(); + dn_acc(i1_ind, i2_ind, i3_ind, e_ind) += particles.weight(p); + } else { + raise::KernelError(HERE, "invalid dimension"); + } + } + }; + +} // namespace kernel + +#endif // KERNELS_PARTICLE_ENERGY_DISTRIBUTION_HPP diff --git a/src/kernels/particle_moments.hpp b/src/kernels/particle_moments.hpp index 1244ba21d..bb5e837b0 100644 --- a/src/kernels/particle_moments.hpp +++ b/src/kernels/particle_moments.hpp @@ -710,85 +710,6 @@ namespace kernel { } }; - template - class ParticleDistribution_kernel { - const ParticleArrays particles; - const bool is_massive; - const M metric; - - const real_t e_min, e_max; - const bool log_bins; - const size_t n_bins; - - scatter_array_t dn_scatter; - - public: - ParticleDistribution_kernel(const Particles& particles, - const scatter_array_t& dn_scatter, - real_t e_min, - real_t e_max, - bool log_bins, - size_t n_bins, - const M& metric) - : particles { static_cast(particles) } - , is_massive { (particles.mass() != 0.0f) } - , dn_scatter { dn_scatter } - , e_min { e_min } - , e_max { e_max } - , log_bins { log_bins } - , n_bins { n_bins } - , metric { metric } {} - - Inline void operator()(prtlidx_t p) const { - if (particles.tag(p) != ParticleTag::alive) { - return; - } - real_t en; - if constexpr (S == SimEngine::SRPIC) { - if (is_massive) { - en = U2GAMMA(particles.ux1(p), particles.ux2(p), particles.ux3(p)) - ONE; - } else { - en = NORM(particles.ux1(p), particles.ux2(p), particles.ux3(p)); - } - } else if constexpr (S == SimEngine::GRPIC) { - coord_t x_Code { ZERO }; - x_Code[0] = static_cast(particles.i1(p)) + - static_cast(particles.dx1(p)); - x_Code[1] = static_cast(particles.i2(p)) + - static_cast(particles.dx2(p)); - - // raise full covariant 4-vector to get correct contravariant u^0 - // u^i != h^{ij} u_j - const real_t u_0_cov { metric.u_0( - x_Code, - { particles.ux1(p), particles.ux2(p), particles.ux3(p) }, - (is_massive) ? ONE : ZERO) }; - vec_t u_cntrv_4d { ZERO }; - metric.template transform_4d( - x_Code, - { u_0_cov, particles.ux1(p), particles.ux2(p), particles.ux3(p) }, - u_cntrv_4d); - // in GR: u^0 = Gamma/alpha - const real_t Gamma { metric.alpha(x_Code) * u_cntrv_4d[0] }; - en = is_massive ? (Gamma - ONE) : Gamma; - } - if (log_bins) { - en = math::log10(en); - } - size_t e_ind = 0; - if (en <= e_min) { - e_ind = 0; - } else if (en >= e_max) { - e_ind = n_bins; - } else { - e_ind = static_cast( - static_cast(n_bins) * (en - e_min) / (e_max - e_min)); - } - auto dn_acc = dn_scatter.access(); - dn_acc(e_ind) += particles.weight(p); - } - }; - } // namespace kernel #endif // KERNELS_PARTICLE_MOMENTS_HPP diff --git a/src/output/writer.cpp b/src/output/writer.cpp index 50630bbcc..89410acc3 100644 --- a/src/output/writer.cpp +++ b/src/output/writer.cpp @@ -200,14 +200,34 @@ namespace out { } } - void Writer::defineSpectraOutputs(const std::vector& specs) { + void Writer::defineSpectraOutputs(const std::vector& specs, + const std::vector& num_spatial_bins) { m_spectra_writers.clear(); for (const auto& s : specs) { m_spectra_writers.emplace_back(s); } m_io.DefineVariable("sEbn", {}, {}, { adios2::UnknownDim }); + const auto spatial_binning_enabled = std::any_of(num_spatial_bins.begin(), + num_spatial_bins.end(), + [](const auto& n) { + return n != 1u; + }); + const auto nspec_dims = spatial_binning_enabled ? num_spatial_bins.size() + 1u + : 1u; for (const auto& sp : m_spectra_writers) { - m_io.DefineVariable(sp.name(), {}, {}, { adios2::UnknownDim }); + m_io.DefineVariable(sp.name(), + {}, + {}, + adios2::Dims(nspec_dims, adios2::UnknownDim)); + } + if (spatial_binning_enabled) { + const auto dim = num_spatial_bins.size(); + for (auto d { 0u }; d < dim; ++d) { + m_io.DefineVariable("sX" + std::to_string(d + 1) + "bn", + {}, + {}, + { adios2::UnknownDim }); + } } } @@ -383,6 +403,47 @@ namespace out { m_keepalive.emplace_back(array_h); } + template + void PutSpectrumSpatial(adios2::IO& io, + adios2::Engine& writer, + std::vector& keepalive, + const std::string& varname, + const HostView& counts_h) { + auto var = io.InquireVariable(varname); + + adios2::Dims start(N, 0u), count(N, 0u), zeros(N, 0u); + auto ntot { 1ul }; + for (auto d { 0u }; d < N; ++d) { + count[d] = counts_h.extent(d); + ntot *= counts_h.extent(d); + } + +#if defined(MPI_ENABLED) + HostView counts_h_all { "counts_h_all", counts_h.layout() }; + int rank; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Reduce(counts_h.data(), + counts_h_all.data(), + static_cast(ntot), + mpi::get_type(), + MPI_SUM, + MPI_ROOT_RANK, + MPI_COMM_WORLD); + if (rank == MPI_ROOT_RANK) { + var.SetSelection(adios2::Box(start, count)); + writer.Put(var, counts_h_all.data(), adios2::Mode::Deferred); + keepalive.emplace_back(counts_h_all); + } else { + var.SetSelection(adios2::Box(start, zeros)); + writer.Put(var, nullptr, adios2::Mode::Sync); + } +#else + var.SetSelection(adios2::Box(start, count)); + writer.Put(var, counts_h.data(), adios2::Mode::Deferred); + keepalive.emplace_back(counts_h); +#endif + } + void Writer::writeSpectrum(const array_t& counts, const std::string& varname) { auto var = m_io.InquireVariable(varname); @@ -416,6 +477,43 @@ namespace out { #endif } + template + void Writer::writeSpectrumSpatial(const nddata_t& counts, + const std::string& varname) { + static_assert(N >= 2 and N <= 4, "writeSpectrumSpatial: N must be 2, 3 or 4"); + // host-resident copy, layout inherited from `counts` + auto counts_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), + counts); + using counts_h_t = decltype(counts_h); + using layout_t = typename counts_h_t::array_layout; + if constexpr (std::is_same::value) { + PutSpectrumSpatial(m_io, m_writer, m_keepalive, varname, counts_h); + } else { + // ADIOS2 reads the raw buffer as row-major: remap on the host + using counts_rm_t = + Kokkos::View; + counts_rm_t counts_rm {}; + if constexpr (N == 2) { + counts_rm = counts_rm_t { "counts_rm", + counts_h.extent(0), + counts_h.extent(1) }; + } else if constexpr (N == 3) { + counts_rm = counts_rm_t { "counts_rm", + counts_h.extent(0), + counts_h.extent(1), + counts_h.extent(2) }; + } else { + counts_rm = counts_rm_t { "counts_rm", + counts_h.extent(0), + counts_h.extent(1), + counts_h.extent(2), + counts_h.extent(3) }; + } + Kokkos::deep_copy(counts_rm, counts_h); + PutSpectrumSpatial(m_io, m_writer, m_keepalive, varname, counts_rm); + } + } + void Writer::writeSpectrumBins(const array_t& e_bins, const std::string& varname) { auto var = m_io.InquireVariable(varname); @@ -567,4 +665,12 @@ namespace out { WRITE_FIELD(Dim::_3D, 6) #undef WRITE_FIELD +#define WRITE_SPECTRUM_SPATIAL(N) \ + template void Writer::writeSpectrumSpatial(const nddata_t&, \ + const std::string&); + WRITE_SPECTRUM_SPATIAL(2) + WRITE_SPECTRUM_SPATIAL(3) + WRITE_SPECTRUM_SPATIAL(4) +#undef WRITE_SPECTRUM_SPATIAL + } // namespace out diff --git a/src/output/writer.h b/src/output/writer.h index 7231c738b..6e9a139d2 100644 --- a/src/output/writer.h +++ b/src/output/writer.h @@ -121,7 +121,8 @@ namespace out { const std::vector& loc_shape); void defineFieldOutputs(const SimEngine&, const std::vector&); - void defineSpectraOutputs(const std::vector&); + void defineSpectraOutputs(const std::vector&, + const std::vector&); void writeMesh(unsigned short, const array_t&, @@ -138,6 +139,11 @@ namespace out { npart_t, const std::string&); void writeSpectrum(const array_t&, const std::string&); + + template + void writeSpectrumSpatial(const nddata_t&, + const std::string&); // @TODO implement for N = 2, 3, 4 + void writeSpectrumBins(const array_t&, const std::string&); void beginWriting(WriteModeTags, timestep_t, simtime_t); diff --git a/tests/framework/CMakeLists.txt b/tests/framework/CMakeLists.txt index 818ea0f2c..3696c74a9 100644 --- a/tests/framework/CMakeLists.txt +++ b/tests/framework/CMakeLists.txt @@ -66,9 +66,11 @@ if(${output}) if(${mpi}) gen_test(checkpoint-extents true) gen_test(checkpoint-data true) + gen_test(spectra true) else() gen_test(checkpoint-extents false) gen_test(checkpoint-data false) + gen_test(spectra false) endif() endif() diff --git a/tests/framework/spectra.cpp b/tests/framework/spectra.cpp new file mode 100644 index 000000000..0b8c6a40a --- /dev/null +++ b/tests/framework/spectra.cpp @@ -0,0 +1,303 @@ +/** + * Tests the spatially binned spectra output. + * + * Particles with known cell positions and energies are placed on a 16x12 + * Minkowski grid, binned into 4x3 spatial bins x 4 energy bins, written out + * via `Metadomain::WriteSpectra`, and read back from the `.bp` file. + * + * In MPI mode the particles are distributed over the two subdomains, which + * exercises the reduction of the per-rank histograms in `writeSpectrumSpatial`. + */ + +#include "enums.h" +#include "global.h" + +#include "utils/comparators.h" +#include "utils/error.h" +#include "utils/formatting.h" + +#include "metrics/minkowski.h" + +#include "framework/domain/metadomain.h" +#include "framework/parameters/parameters.h" + +#include +#include +#include + +#if defined(MPI_ENABLED) + #include +#endif + +#include +#include +#include +#include +#include + +using namespace ntt; +using namespace metric; + +const std::string SIMNAME = "test_spectra"; + +namespace { + constexpr ncells_t NX1 = 16, NX2 = 12; + constexpr std::size_t NB1 = 4, NB2 = 3, NBE = 4; + constexpr real_t E_MIN = 0.0, E_MAX = 4.0; + + struct TestPrtl { + ncells_t gi1, gi2; // global cell index + real_t gamma; + real_t weight; + bool alive; + std::size_t b1, b2, be; // expected bin + }; + + // dx = 0.5 in both directions, so no particle sits on a bin edge + const std::vector PRTLS { + { 0, 0, 1.5, 1.0, true, 0, 0, 0 }, + { 2, 3, 1.5, 2.0, true, 0, 0, 0 }, // accumulates with the previous one + { 6, 1, 2.5, 1.0, true, 1, 0, 1 }, + { 9, 7, 3.5, 1.0, true, 2, 1, 2 }, + { 15, 11, 4.5, 1.0, true, 3, 2, 3 }, + { 5, 9, 1.5, 3.0, true, 1, 2, 0 }, + { 13, 5, 100.0, 1.0, true, 3, 1, 3 }, // above e_max -> last energy bin + { 11, 11, 2.5, 5.0, false, 0, 0, 0 }, // dead -> not counted + }; + + void cleanup() { + std::filesystem::remove_all(SIMNAME); + } + + auto read_block(adios2::Engine& reader, + adios2::Variable& var, + std::vector& data) -> adios2::Dims { + for (const auto& blk : reader.BlocksInfo(var, reader.CurrentStep())) { + std::size_t ntot { 1 }; + for (const auto& c : blk.Count) { + ntot *= c; + } + if (ntot == 0) { + continue; + } + var.SetBlockSelection(blk.BlockID); + data.resize(ntot); + reader.Get(var, data.data(), adios2::Mode::Sync); + return blk.Count; + } + raise::Error("no non-empty block written for " + var.Name(), HERE); + return {}; + } +} // namespace + +auto main(int argc, char* argv[]) -> int { + GlobalInitialize(argc, argv); + + try { + using M = Minkowski; + + const std::vector res { NX1, NX2 }; + const boundaries_t extent { + { static_cast(0.0), static_cast(NX1) }, + { static_cast(0.0), static_cast(NX2) } + }; + const boundaries_t fldsbc { + { FldsBC::PERIODIC, FldsBC::PERIODIC }, + { FldsBC::PERIODIC, FldsBC::PERIODIC } + }; + const boundaries_t prtlbc { + { PrtlBC::PERIODIC, PrtlBC::PERIODIC }, + { PrtlBC::PERIODIC, PrtlBC::PERIODIC } + }; + const std::vector decomp { -1, -1 }; + + const std::vector species_params { + ParticleSpecies { static_cast(1), + "e-", 1.0f, + -1.0f, + static_cast(100), + timestep_t { 0 }, + timestep_t { 0 }, + ParticlePusher::BORIS, + false, RadiativeDrag::NONE, + EmissionType::NONE, + static_cast(0), + static_cast(0) } + }; + +#if !defined(MPI_ENABLED) + const unsigned int ndomains { 1 }; + adios2::ADIOS adios; +#else + int mpi_size; + MPI_Comm_size(MPI_COMM_WORLD, &mpi_size); + raise::ErrorIf(mpi_size != 2, "this test requires exactly 2 MPI ranks", HERE); + const unsigned int ndomains { static_cast(mpi_size) }; + adios2::ADIOS adios { MPI_COMM_WORLD }; +#endif + + const auto out_step = timestep_t { 5 }; + const auto out_time = simtime_t { 1.25 }; + + Metadomain md { ndomains, decomp, res, extent, + fldsbc, prtlbc, {}, species_params }; + + auto* local = md.subdomain_ptr(md.l_subdomain_indices()[0]); + const auto off = local->offset_ncells(); + const auto nloc = local->mesh.n_active(); + + // place the particles that fall into this subdomain + { + auto& sp = local->species[0]; + + auto i1_h = Kokkos::create_mirror_view(sp.i1); + auto i2_h = Kokkos::create_mirror_view(sp.i2); + auto dx1_h = Kokkos::create_mirror_view(sp.dx1); + auto dx2_h = Kokkos::create_mirror_view(sp.dx2); + auto ux1_h = Kokkos::create_mirror_view(sp.ux1); + auto w_h = Kokkos::create_mirror_view(sp.weight); + auto tag_h = Kokkos::create_mirror_view(sp.tag); + + npart_t np { 0 }; + for (const auto& p : PRTLS) { + if (p.gi1 < off[0] or p.gi1 >= off[0] + nloc[0] or p.gi2 < off[1] or + p.gi2 >= off[1] + nloc[1]) { + continue; + } + i1_h(np) = static_cast(p.gi1 - off[0]); + i2_h(np) = static_cast(p.gi2 - off[1]); + dx1_h(np) = static_cast(0.5); + dx2_h(np) = static_cast(0.5); + ux1_h(np) = math::sqrt(p.gamma * p.gamma - ONE); + w_h(np) = p.weight; + tag_h(np) = p.alive ? ParticleTag::alive : ParticleTag::dead; + ++np; + } + sp.set_npart(np); + + Kokkos::deep_copy(sp.i1, i1_h); + Kokkos::deep_copy(sp.i2, i2_h); + Kokkos::deep_copy(sp.dx1, dx1_h); + Kokkos::deep_copy(sp.dx2, dx2_h); + Kokkos::deep_copy(sp.ux1, ux1_h); + Kokkos::deep_copy(sp.weight, w_h); + Kokkos::deep_copy(sp.tag, tag_h); + } + + SimulationParams params; + params.set("simulation.name", SIMNAME); + params.set("output.format", std::string { "BPFile" }); + params.set("output.debug.ghosts", false); + params.set("output.fields.quantities", std::vector {}); + params.set("output.fields.custom", std::vector {}); + params.set("output.fields.downsampling", std::vector { 1, 1 }); + params.set("output.particles.species", std::vector {}); + params.set("output.spectra.e_min", E_MIN); + params.set("output.spectra.e_max", E_MAX); + params.set("output.spectra.log_bins", false); + params.set("output.spectra.num_energy_bins", NBE); + params.set("output.spectra.num_spatial_bins", + std::vector { NB1, NB2 }); + for (const auto& type : { "fields", "particles", "spectra" }) { + params.set("output." + std::string(type) + ".interval", timestep_t { 1 }); + params.set("output." + std::string(type) + ".interval_time", + simtime_t { -1.0 }); + } + params.setRawData(toml::value { toml::table {} }); + + md.InitWriter(&adios, params); + md.WriteSpectra(params, local, out_step, out_time); + + adios.FlushAll(); + + // ── read back ───────────────────────────────────────────────────────────── + { + adios2::IO io = adios.DeclareIO("read-spectra"); + io.SetEngine("BPFile"); + + namespace fs = std::filesystem; + adios2::Engine reader = io.Open( + fs::path(SIMNAME) / fs::path("spectra") / + fs::path(fmt::format("spectra.%08lu.bp", out_step)), + adios2::Mode::Read); + raise::ErrorIf(reader.BeginStep() != adios2::StepStatus::OK, + "no step in the spectra file", + HERE); + + { + auto var = io.InquireVariable("sN_1"); + raise::ErrorIf(not var, "sN_1 not found", HERE); + std::vector dn; + const auto shape = read_block(reader, var, dn); + + raise::ErrorIf( + shape != adios2::Dims({ NB1, NB2, NBE }), + fmt::format( + "sN_1 shape is %s, expected %s", + fmt::formatVector(shape).c_str(), + fmt::formatVector(std::vector { NB1, NB2, NBE }).c_str()), + HERE); + + std::vector expected(NB1 * NB2 * NBE, ZERO); + for (const auto& p : PRTLS) { + if (p.alive) { + expected[(p.b1 * NB2 + p.b2) * NBE + p.be] += p.weight; + } + } + for (auto i { 0u }; i < expected.size(); ++i) { + raise::ErrorIf(not cmp::AlmostEqual(dn[i], expected[i]), + fmt::format("sN_1[%u] = %f, expected %f", + i, + static_cast(dn[i]), + static_cast(expected[i])), + HERE); + } + } + + // bin edges: energy in [e_min, e_max], spatial in physical units + const std::vector>> bins { + { "sEbn", { 0.0, 1.0, 2.0, 3.0, 4.0 } }, + { "sX1bn", { 0.0, 4.0, 8.0, 12.0, 16.0 } }, + { "sX2bn", { 0.0, 4.0, 8.0, 12.0 } } + }; + for (const auto& [name, expected] : bins) { + auto var = io.InquireVariable(name); + raise::ErrorIf(not var, name + " not found", HERE); + std::vector edges; + const auto shape = read_block(reader, var, edges); + raise::ErrorIf(shape != adios2::Dims({ expected.size() }), + fmt::format("%s has %lu edges, expected %lu", + name.c_str(), + shape.empty() ? 0ul : shape[0], + expected.size()), + HERE); + for (auto i { 0u }; i < expected.size(); ++i) { + raise::ErrorIf(not cmp::AlmostEqual(edges[i], expected[i]), + fmt::format("%s[%u] = %f, expected %f", + name.c_str(), + i, + static_cast(edges[i]), + static_cast(expected[i])), + HERE); + } + } + + reader.EndStep(); + reader.Close(); + } + + } catch (const std::exception& e) { + std::cerr << e.what() << '\n'; + CallOnce([&] { + cleanup(); + }); + GlobalFinalize(); + return 1; + } + + CallOnce([&] { + cleanup(); + }); + GlobalFinalize(); + return 0; +}