From 2d0a157db9d0ccb36f90378765fee8740506d404 Mon Sep 17 00:00:00 2001 From: A Sullivan Date: Fri, 31 Jul 2026 10:28:19 -0400 Subject: [PATCH 01/17] added spectra3D components to entity_1.5rc --- src/framework/domain/metadomain_io.cpp | 249 ++++++++++++++++++++++++- src/framework/parameters/output.cpp | 45 ++++- src/framework/parameters/output.h | 8 + src/global/defaults.h | 7 + src/global/global.h | 1 + src/output/writer.cpp | 70 +++++++ src/output/writer.h | 1 + 7 files changed, 378 insertions(+), 3 deletions(-) diff --git a/src/framework/domain/metadomain_io.cpp b/src/framework/domain/metadomain_io.cpp index 3fc0f5f57..b42247634 100644 --- a/src/framework/domain/metadomain_io.cpp +++ b/src/framework/domain/metadomain_io.cpp @@ -109,7 +109,7 @@ namespace ntt { spectra_species.push_back(sp.index()); } g_writer.defineSpectraOutputs(spectra_species); - for (const auto& type : { "fields", "particles", "spectra" }) { + for (const auto& type : { "fields", "particles", "spectra", "spectra3D" }) { g_writer.addTracker(type, params.template get( "output." + std::string(type) + ".interval"), @@ -382,8 +382,13 @@ namespace ntt { g_writer.shouldWrite("spectra", finished_step, finished_time); + const auto write_spectra3D = params.template get( + "output.spectra3D.enable") and + g_writer.shouldWrite("spectra3D", + finished_step, + finished_time); const auto extension = params.template get("output.format"); - if (not(write_fields or write_particles or write_spectra) and + if (not(write_fields or write_particles or write_spectra or write_spectra3D) and extension != "disabled") { return false; } @@ -842,6 +847,246 @@ namespace ntt { g_writer.endWriting(WriteMode::Spectra); } // end shouldWrite("spectra", step, time) + if (write_spectra3D) { + g_writer.beginWriting(WriteMode::Spectra3D, current_step, current_time); + const auto log_bins = params.template get( + "output.spectra3D.log_bins"); + const auto n_bins = params.template get( + "output.spectra3D.n_bins"); + const auto& metric = local_domain->mesh.metric; + // extract the number of bins globally in each direction + const auto nx1_bins = params.template get("output.spectra3D.nx1"); + const auto nx2_bins = params.template get("output.spectra3D.nx2"); + const auto nx3_bins = params.template get("output.spectra3D.nx3"); + + // select the min and max energy for the spectra + auto e_min = params.template get("output.spectra3D.e_min"); + auto e_max = params.template get("output.spectra3D.e_max"); + + + const auto x1_size_local = local_domain->mesh.n_active()[0]; + const auto x1_offset_local = local_domain->offset_ncells()[0]; + const auto x1_size_global = mesh().n_active()[0]; + + // divide based on physical coordinates maybe rather than cells + + std::size_t x2_size_local = 0; + std::size_t x2_offset_local = 0; + std::size_t x2_size_global = 0; + + std::size_t x3_size_local = 0; + std::size_t x3_offset_local = 0; + std::size_t x3_size_global = 0; + + auto x1_min = mesh().extent(in::x1).first; // extent is in physical units, not code + decltype(x1_min) x2_min = 0; + decltype(x1_min) x3_min = 0; + + auto x1_max = mesh().extent(in::x1).second; // pairs in c++ are addressed by first and secocnd + decltype(x1_max) x2_max = 0; + decltype(x1_max) x3_max = 0; + + //auto x1_extent_local = mesh().extent(in::x1) + //const auto nx1 = local_domain->mesh.n_active(in::x1); + //const auto nx2 = local_domain->mesh.n_active(in::x2); + + //auto dx1 = (x1_extent_local.second - x1_extent_local.first) / (real_t)(nx1_bin - 1); + auto x1_extent_local = local_domain->mesh.extent(in::x1); + auto x1_min_local = x1_extent_local.first; + auto x1_max_local = x1_extent_local.second; + + auto x1_ind_rank_min = static_cast(static_cast(nx1_bins) * (x1_min_local - x1_min) / (x1_max - x1_min)); + auto x1_ind_rank_max = static_cast(static_cast(nx1_bins) * (x1_max_local - x1_min) / (x1_max - x1_min)); + + decltype(x1_min_local) x2_min_local=0; + decltype(x1_max_local) x2_max_local=0; + decltype(x1_ind_rank_min) x2_ind_rank_min=0; + decltype(x1_ind_rank_max) x2_ind_rank_max=0; + + decltype(x1_min_local) x3_min_local=0; + decltype(x1_max_local) x3_max_local=0; + decltype(x1_ind_rank_min) x3_ind_rank_min=0; + decltype(x1_ind_rank_max) x3_ind_rank_max=0; + + real_t dx1_bin = (x1_max - x1_min )/nx1_bins; + real_t dx2_bin = 1.; + real_t dx3_bin = 1.; + + + + if constexpr (M::PrtlDim == Dim::_2D or M::PrtlDim == Dim::_3D){ // only pick x2 if simulation is in 2D + + x2_min = mesh().extent(in::x2).first; + x2_max = mesh().extent(in::x2).second; + + auto x2_extent_local = local_domain->mesh.extent(in::x2); + x2_min_local = x2_extent_local.first; + x2_max_local = x2_extent_local.second; + + dx2_bin = (static_cast(x2_max) - static_cast(x2_min) )/static_cast(nx2_bins); + + x2_ind_rank_min = static_cast(static_cast(nx2_bins) * (x2_min_local - x1_min) / (x2_max - x2_min)); + x2_ind_rank_max = static_cast(static_cast(nx2_bins) * (x2_max_local - x1_min) / (x2_max - x2_min)); + } + + if constexpr (M::PrtlDim == Dim::_3D){ + + x3_min = mesh().extent(in::x3).first; + x3_max = mesh().extent(in::x3).second; + + auto x3_extent_local = local_domain->mesh.extent(in::x3); + x3_min_local = x3_extent_local.first; + x3_max_local = x3_extent_local.second; + + dx3_bin = (static_cast(x3_max) - static_cast(x3_min) )/static_cast(nx3_bins); + + x3_ind_rank_min = static_cast(static_cast(nx3_bins) * (x3_min_local - x3_min) / (x1_max - x1_min)); + x3_ind_rank_max = static_cast(static_cast(nx3_bins) * (x3_max_local - x3_min) / (x1_max - x1_min)); + } + + + + //auto dx_slice = mesh().dx1; + + + //auto dxslice = (x1_max - x1_min) / x_bins + + if (log_bins) { + e_min = math::log10(e_min); + e_max = math::log10(e_max); + } + array_t energy { "energy", n_bins + 1 }; + // generating the energy bins + Kokkos::parallel_for( + "GenerateEnergyBins", + n_bins + 1, + Lambda(index_t e) { + if (log_bins) { + energy(e) = math::pow(10.0, e_min + (e_max - e_min) * e / n_bins); + } else { + energy(e) = e_min + (e_max - e_min) * e / n_bins; + } + }); + + for (const auto& spec : g_writer.spectraWriters()) { + auto& species = local_domain->species[spec.species() - 1]; + array_t dn3d { "dn3d", nx1_bins, nx2_bins, nx3_bins, n_bins }; + auto dn3d_scatter = Kokkos::Experimental::create_scatter_view(dn3d); + auto ux1 = species.ux1; + auto ux2 = species.ux2; + auto ux3 = species.ux3; + auto i1 = species.i1; + auto dx1 = species.dx1; + //adeep_copy; + decltype(i1) i2; + decltype(i1) i3; + decltype(dx1) dx2; + decltype(dx1) dx3; + if constexpr (M::PrtlDim == Dim::_2D or M::PrtlDim == Dim::_3D){ // only pick x2 if simulation is in 2D + i2 = species.i2; + dx2 = species.dx2; + } + if constexpr (M::PrtlDim == Dim::_3D){ + i3 = species.i3; + dx3 = species.dx3; + } + auto weight = species.weight; + auto tag = species.tag; + const auto is_massive = species.mass() > 0.0f; + Kokkos::parallel_for( + "ComputeSpectra", + species.rangeActiveParticles(), + Lambda(index_t p) { + if (tag(p) != ParticleTag::alive) { + return; + } + + coord_t x_Cd {ZERO}; + if (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { + x_Cd[0] = static_cast(i1(p)) + static_cast(dx1(p)); + } + if (D == Dim::_2D or D == Dim::_3D) { + x_Cd[1] = static_cast(i2(p)) + static_cast(dx2(p)); + } + if (D == Dim::_3D) { + x_Cd[2] = static_cast(i3(p)) + static_cast(dx3(p)); + } + coord_t x_Ph { ZERO }; + metric.template convert(x_Cd, x_Ph); + + real_t en; + if (is_massive) { + en = U2GAMMA(ux1(p), ux2(p), ux3(p)) - ONE; + } else { + en = NORM(ux1(p), ux2(p), ux3(p)); + } + if (log_bins) { + en = math::log10(en); + } + std::size_t e_ind = 0; + if (en <= e_min) { + e_ind = 0; + } else if (en >= e_max) { + e_ind = n_bins-1; + } else { + e_ind = static_cast( + static_cast(n_bins) * (en - e_min) / (e_max - e_min)); + } + + std::size_t x1_ind = 0; + if (x_Ph[0]<= x1_min) { + x1_ind = 0; + } else if (x_Ph[0] >= x1_max) { + x1_ind = nx1_bins-1; + } else { + x1_ind = static_cast( + static_cast(nx1_bins) * (x_Ph[0] - x1_min) / (x1_max - x1_min)); + } + + + std::size_t x2_ind = 0; + if (M::PrtlDim == Dim::_2D or M::PrtlDim == Dim::_3D){ // only pick x2 if simulation is in 2D + if (x_Ph[1] <= x2_min) { + x2_ind = 0; + } else if (x_Ph[1] >= x2_max) { + x2_ind = nx2_bins-1; + } else { + x2_ind = static_cast( + static_cast(nx2_bins) * (x_Ph[1] - x2_min) / (x2_max - x2_min)); + } + //real_t x2_center = x2_min + (x2_ind+0.5)*dx2_bin; + //owns_x2 = (x2_center >= x2_min_local) && (x2_center <= x2_max_local); + } + std::size_t x3_ind = 0; + if (M::PrtlDim == Dim::_3D){ // only pick x3 if simulation is in 3D + + if (x_Ph[2] <= x3_min) { + x3_ind = 0; + } else if (x_Ph[2] >= x3_max) { + x3_ind = nx3_bins-1; + } else { + x3_ind = static_cast( + static_cast(nx3_bins) * (x_Ph[2] - x3_min) / (x3_max - x3_min)); + } + + + } + + // now I want to save the nx_bins contained in each rank + // can save array of x1_inds which are saved? + // maybe I can ask what local x_min and x_max are for this rank, pass that to writeSpectrum3D + + + auto dn3d_acc = dn3d_scatter.access(); + dn3d_acc(x1_ind, x2_ind, x3_ind, e_ind) += weight(p); + }); + Kokkos::Experimental::contribute(dn3d, dn3d_scatter); + g_writer.writeSpectrum3D(dn3d, spec.name()); + } + g_writer.writeSpectrumBins(energy, "sEbn"); + g_writer.endWriting(WriteMode::Spectra3D); + } + return true; } diff --git a/src/framework/parameters/output.cpp b/src/framework/parameters/output.cpp index 6c1858e9d..13e1a5463 100644 --- a/src/framework/parameters/output.cpp +++ b/src/framework/parameters/output.cpp @@ -33,7 +33,7 @@ namespace ntt { HERE); categories.emplace(); - for (const auto& category : { "fields", "particles", "spectra", "stats" }) { + for (const auto& category : { "fields", "particles", "spectra", "spectra3D", "stats" }) { const auto q_int = toml::find_or(toml_data, "output", category, @@ -158,6 +158,41 @@ namespace ntt { "spectra", "n_bins", defaults::output::spec_nbins); + + /* Spectra3D ------------------------------------------------------------ */ + spectra3d_e_min = toml::find_or(toml_data, "output", "spectra3D", "e_min", defaults::output::spec3d_emin); + + spectra3d_e_max = toml::find_or(toml_data, "output", "spectra3D", "e_max", defaults::output::spec3d_emin); + + spectra3d_log_bins = toml::find_or(toml_data, + "output", + "spectra3D", + "log_bins", + defaults::output::spec3d_log); + + spectra3d_n_bins = toml::find_or(toml_data, + "output", + "spectra3D", + "n_bins", + defaults::output::spec3d_nbins); + + spectra3d_nx1 = toml::find_or(toml_data, + "output", + "spectra3D", + "nx1", + defaults::output::spec3d_nx1); + + spectra3d_nx2 = toml::find_or(toml_data, + "output", + "spectra3D", + "nx2", + defaults::output::spec3d_nx1); + + spectra3d_nx3 = toml::find_or(toml_data, + "output", + "spectra3D", + "nx3", + defaults::output::spec3d_nx1); /* Stats ---------------------------------------------------------------- */ stats_quantities = toml::find_or(toml_data, @@ -209,6 +244,14 @@ namespace ntt { params->set("output.spectra.log_bins", spectra_log_bins.value()); params->set("output.spectra.n_bins", spectra_n_bins.value()); + params->set("output.spectra3D.e_min", spectra3d_e_min.value()); + params->set("output.spectra3D.e_max", spectra3d_e_max.value()); + params->set("output.spectra3D.log_bins", spectra3d_log_bins.value()); + params->set("output.spectra3D.n_bins", spectra3d_n_bins.value()); + params->set("output.spectra3D.nx1", spectra3d_nx1.value()); + params->set("output.spectra3D.nx2", spectra3d_nx2.value()); + params->set("output.spectra3D.nx3", spectra3d_nx3.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..680751514 100644 --- a/src/framework/parameters/output.h +++ b/src/framework/parameters/output.h @@ -55,6 +55,14 @@ namespace ntt { std::optional spectra_log_bins; std::optional spectra_n_bins; + std::optional spectra3d_e_min; + std::optional spectra3d_e_max; + std::optional spectra3d_log_bins; + std::optional spectra3d_n_bins; + std::optional spectra3d_nx1; + std::optional spectra3d_nx2; + std::optional spectra3d_nx3; + std::optional> stats_quantities; std::optional> stats_custom_quantities; diff --git a/src/global/defaults.h b/src/global/defaults.h index e1387677e..0b7e30355 100644 --- a/src/global/defaults.h +++ b/src/global/defaults.h @@ -74,6 +74,13 @@ namespace ntt::defaults { const real_t spec_emax = 1e3; const bool spec_log = true; const std::size_t spec_nbins = 200; + const real_t spec3d_emin = 1e-3; + const real_t spec3d_emax = 1e3; + const bool spec3d_log = true; + const std::size_t spec3d_nbins = 200; + const std::size_t spec3d_nx1 = 1; + const std::size_t spec3d_nx2 = 1; + const std::size_t spec3d_nx3 = 1; const std::vector stats_quantities = { "B^2", "E^2", "ExB", diff --git a/src/global/global.h b/src/global/global.h index c3fbc5061..00047359a 100644 --- a/src/global/global.h +++ b/src/global/global.h @@ -298,6 +298,7 @@ namespace WriteMode { Particles = 1 << 1, Spectra = 1 << 2, Stats = 1 << 3, + Spectra3D = 1 << 4, }; } // namespace WriteMode diff --git a/src/output/writer.cpp b/src/output/writer.cpp index a84745c5b..317d95af9 100644 --- a/src/output/writer.cpp +++ b/src/output/writer.cpp @@ -383,6 +383,74 @@ namespace out { #endif } +// spectrum3D, broken up into separate sub-domains + void Writer::writeSpectrum3D(const array_t& counts3D, + const std::string& varname) { + // need to include the rank specific coordinates contained here + std::string varname3d = varname + "_3D"; + //auto var = m_io.InquireVariable(varname); + auto counts3D_h = Kokkos::create_mirror_view(counts3D); + // copy to host + Kokkos::deep_copy(counts3D_h, counts3D); +#if defined(MPI_ENABLED) + array_t counts3D_all { "counts3D_all", counts3D.extent(0), counts3D.extent(1), counts3D.extent(2), counts3D.extent(3) }; + //auto counts_h_all = Kokkos::create_mirror_view(counts_all); + int rank; + int size; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + + auto counts3D_h_all = Kokkos::create_mirror_view(counts3D_all); + + + MPI_Allreduce(counts3D_h.data(), + counts3D_h_all.data(), + counts3D_h.extent(0)*counts3D_h.extent(1)*counts3D_h.extent(2)*counts3D_h.extent(3),// size so probably need to multiply counts_h.extent(0)*counts_h.extent(1)*counts_h.extent(2) + mpi::get_type(), + MPI_SUM, + MPI_COMM_WORLD); // size + //Kokkos::deep_copy() +#else + int rank = 0; + int size = 1; +#endif + using Shape = std::vector; + + Shape shape = {counts3D.extent(0), counts3D.extent(1), counts3D.extent(2), counts3D.extent(3)}; // global size of counts3D_all + Shape start, count; // per-rank selection + + // split the domain along the x1 direction across the different ranks + auto split_x1 = [&](std::size_t n, int this_rank, int nranks) + { + // base number of cells to allocate to each rank + std::size_t base = n / nranks; + // remainder that do not fit neatly in one rank + std::size_t rem = n % nranks; + // number of cells to allocate to the rank including the remainder + std::size_t nloc = base + (this_rank <(int)rem ? 1 : 0); // allocate one more if this rank is less that the remainder + //offset to allocate + std::size_t off = base * this_rank + std::min(this_rank, rem); + + return std::pair{nloc, off}; + }; + + auto [nloc0, off0] = split_x1(counts3D.extent(0), rank, size); + start = {off0, 0, 0, 0}; + count = {nloc0, counts3D.extent(1), counts3D.extent(2), counts3D.extent(3)}; + + auto var = m_io.InquireVariable(varname3d); + if (!var) + { + var = m_io.DefineVariable(varname3d, shape, start, count, adios2::ConstantDims); + } + + //auto var = m_io.DefineVariable(varname3d, shape, start, count, adios2::ConstantDims); + + auto counts3D_h_all_slab = Kokkos::subview(counts3D_h_all, Kokkos::make_pair(off0, off0+nloc0), Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()); + + m_writer.Put(var, counts3D_h_all_slab, adios2::Mode::Sync); + } + void Writer::writeSpectrumBins(const array_t& e_bins, const std::string& varname) { auto var = m_io.InquireVariable(varname); @@ -448,6 +516,8 @@ namespace out { mode_str = "particles"; } else if (write_mode == WriteMode::Spectra) { mode_str = "spectra"; + } else if (write_mode == WriteMode::Spectra3D) { + mode_str = "spectra3D"; } else { raise::Fatal("Unknown write mode", HERE); } diff --git a/src/output/writer.h b/src/output/writer.h index 055cb215e..b2516b33f 100644 --- a/src/output/writer.h +++ b/src/output/writer.h @@ -127,6 +127,7 @@ namespace out { npart_t, const std::string&); void writeSpectrum(const array_t&, const std::string&); + void writeSpectrum3D(const array_t&, const std::string&); void writeSpectrumBins(const array_t&, const std::string&); void beginWriting(WriteModeTags, timestep_t, simtime_t); From 7c828be35e4d6d8afd9b1bcab64fd75351927260 Mon Sep 17 00:00:00 2001 From: A Sullivan Date: Fri, 31 Jul 2026 12:56:13 -0400 Subject: [PATCH 02/17] fixed bug in metadomain_io.cpp and added template g_checkpoint_writer.io().template in metadomain_chckpt.cpp --- pgens/shock/shock.toml | 14 +++++++++++++- src/framework/domain/metadomain_chckpt.cpp | 6 +++--- src/framework/domain/metadomain_io.cpp | 4 ++-- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/pgens/shock/shock.toml b/pgens/shock/shock.toml index 8cde86990..37515078f 100644 --- a/pgens/shock/shock.toml +++ b/pgens/shock/shock.toml @@ -63,7 +63,19 @@ stride = 10 [output.spectra] - enable = false + enable = true + e_min = 1e-2 + e_max = 1e1 + log_bins = true + + [output.spectra3D] + enable = true + e_min = 1e-2 + e_max = 1e1 + log_bins = true + nx1 = 40 + nx2 = 3 + nx3 = 1 [diagnostics] log_level = "WARNING" diff --git a/src/framework/domain/metadomain_chckpt.cpp b/src/framework/domain/metadomain_chckpt.cpp index 1116df85e..8ed0f7e99 100644 --- a/src/framework/domain/metadomain_chckpt.cpp +++ b/src/framework/domain/metadomain_chckpt.cpp @@ -80,17 +80,17 @@ namespace ntt { species.CheckpointDeclare(g_checkpoint_writer.io()); } for (auto d { 0u }; d < M::Dim; ++d) { - g_checkpoint_writer.io().DefineVariable( + g_checkpoint_writer.io().template DefineVariable( fmt::format("subdomain_x%d_min", d + 1), { adios2::UnknownDim }, { adios2::UnknownDim }, { adios2::UnknownDim }); - g_checkpoint_writer.io().DefineVariable( + g_checkpoint_writer.io().template DefineVariable( fmt::format("subdomain_x%d_max", d + 1), { adios2::UnknownDim }, { adios2::UnknownDim }, { adios2::UnknownDim }); - g_checkpoint_writer.io().DefineVariable( + g_checkpoint_writer.io().template DefineVariable( fmt::format("subdomain_nx%d", d + 1), { adios2::UnknownDim }, { adios2::UnknownDim }, diff --git a/src/framework/domain/metadomain_io.cpp b/src/framework/domain/metadomain_io.cpp index b42247634..8669efe10 100644 --- a/src/framework/domain/metadomain_io.cpp +++ b/src/framework/domain/metadomain_io.cpp @@ -960,7 +960,7 @@ namespace ntt { Kokkos::parallel_for( "GenerateEnergyBins", n_bins + 1, - Lambda(index_t e) { + Lambda(uint32_t e) { if (log_bins) { energy(e) = math::pow(10.0, e_min + (e_max - e_min) * e / n_bins); } else { @@ -996,7 +996,7 @@ namespace ntt { Kokkos::parallel_for( "ComputeSpectra", species.rangeActiveParticles(), - Lambda(index_t p) { + Lambda(prtlidx_t p) { if (tag(p) != ParticleTag::alive) { return; } From 23c3128d1e40b61273366473c4c9b63f4bde8624 Mon Sep 17 00:00:00 2001 From: A Sullivan Date: Mon, 14 Sep 2026 13:04:58 -0400 Subject: [PATCH 03/17] fixed some typos in the spectra3D additions (GPU still not fixed) --- src/framework/domain/metadomain_io.cpp | 77 +++++--------------------- src/framework/parameters/output.cpp | 6 +- src/output/writer.cpp | 1 + 3 files changed, 17 insertions(+), 67 deletions(-) diff --git a/src/framework/domain/metadomain_io.cpp b/src/framework/domain/metadomain_io.cpp index 8669efe10..78efe38c6 100644 --- a/src/framework/domain/metadomain_io.cpp +++ b/src/framework/domain/metadomain_io.cpp @@ -864,19 +864,6 @@ namespace ntt { auto e_max = params.template get("output.spectra3D.e_max"); - const auto x1_size_local = local_domain->mesh.n_active()[0]; - const auto x1_offset_local = local_domain->offset_ncells()[0]; - const auto x1_size_global = mesh().n_active()[0]; - - // divide based on physical coordinates maybe rather than cells - - std::size_t x2_size_local = 0; - std::size_t x2_offset_local = 0; - std::size_t x2_size_global = 0; - - std::size_t x3_size_local = 0; - std::size_t x3_offset_local = 0; - std::size_t x3_size_global = 0; auto x1_min = mesh().extent(in::x1).first; // extent is in physical units, not code decltype(x1_min) x2_min = 0; @@ -886,62 +873,25 @@ namespace ntt { decltype(x1_max) x2_max = 0; decltype(x1_max) x3_max = 0; - //auto x1_extent_local = mesh().extent(in::x1) - //const auto nx1 = local_domain->mesh.n_active(in::x1); - //const auto nx2 = local_domain->mesh.n_active(in::x2); - - //auto dx1 = (x1_extent_local.second - x1_extent_local.first) / (real_t)(nx1_bin - 1); - auto x1_extent_local = local_domain->mesh.extent(in::x1); - auto x1_min_local = x1_extent_local.first; - auto x1_max_local = x1_extent_local.second; - - auto x1_ind_rank_min = static_cast(static_cast(nx1_bins) * (x1_min_local - x1_min) / (x1_max - x1_min)); - auto x1_ind_rank_max = static_cast(static_cast(nx1_bins) * (x1_max_local - x1_min) / (x1_max - x1_min)); - - decltype(x1_min_local) x2_min_local=0; - decltype(x1_max_local) x2_max_local=0; - decltype(x1_ind_rank_min) x2_ind_rank_min=0; - decltype(x1_ind_rank_max) x2_ind_rank_max=0; - - decltype(x1_min_local) x3_min_local=0; - decltype(x1_max_local) x3_max_local=0; - decltype(x1_ind_rank_min) x3_ind_rank_min=0; - decltype(x1_ind_rank_max) x3_ind_rank_max=0; + - real_t dx1_bin = (x1_max - x1_min )/nx1_bins; - real_t dx2_bin = 1.; - real_t dx3_bin = 1.; - if constexpr (M::PrtlDim == Dim::_2D or M::PrtlDim == Dim::_3D){ // only pick x2 if simulation is in 2D + if constexpr (D == Dim::_2D or D == Dim::_3D){ // only pick x2 if simulation is in 2D x2_min = mesh().extent(in::x2).first; x2_max = mesh().extent(in::x2).second; - - auto x2_extent_local = local_domain->mesh.extent(in::x2); - x2_min_local = x2_extent_local.first; - x2_max_local = x2_extent_local.second; - - dx2_bin = (static_cast(x2_max) - static_cast(x2_min) )/static_cast(nx2_bins); - - x2_ind_rank_min = static_cast(static_cast(nx2_bins) * (x2_min_local - x1_min) / (x2_max - x2_min)); - x2_ind_rank_max = static_cast(static_cast(nx2_bins) * (x2_max_local - x1_min) / (x2_max - x2_min)); + } - if constexpr (M::PrtlDim == Dim::_3D){ + if constexpr (D == Dim::_3D){ x3_min = mesh().extent(in::x3).first; x3_max = mesh().extent(in::x3).second; - auto x3_extent_local = local_domain->mesh.extent(in::x3); - x3_min_local = x3_extent_local.first; - x3_max_local = x3_extent_local.second; - dx3_bin = (static_cast(x3_max) - static_cast(x3_min) )/static_cast(nx3_bins); - - x3_ind_rank_min = static_cast(static_cast(nx3_bins) * (x3_min_local - x3_min) / (x1_max - x1_min)); - x3_ind_rank_max = static_cast(static_cast(nx3_bins) * (x3_max_local - x3_min) / (x1_max - x1_min)); + } @@ -982,11 +932,11 @@ namespace ntt { decltype(i1) i3; decltype(dx1) dx2; decltype(dx1) dx3; - if constexpr (M::PrtlDim == Dim::_2D or M::PrtlDim == Dim::_3D){ // only pick x2 if simulation is in 2D + if constexpr (D == Dim::_2D or D == Dim::_3D){ // only pick x2 if simulation is in 2D i2 = species.i2; dx2 = species.dx2; } - if constexpr (M::PrtlDim == Dim::_3D){ + if constexpr (D == Dim::_3D){ i3 = species.i3; dx3 = species.dx3; } @@ -1002,13 +952,13 @@ namespace ntt { } coord_t x_Cd {ZERO}; - if (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { + if constexpr (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { x_Cd[0] = static_cast(i1(p)) + static_cast(dx1(p)); } - if (D == Dim::_2D or D == Dim::_3D) { + if constexpr (D == Dim::_2D or D == Dim::_3D) { x_Cd[1] = static_cast(i2(p)) + static_cast(dx2(p)); } - if (D == Dim::_3D) { + if constexpr (D == Dim::_3D) { x_Cd[2] = static_cast(i3(p)) + static_cast(dx3(p)); } coord_t x_Ph { ZERO }; @@ -1045,7 +995,7 @@ namespace ntt { std::size_t x2_ind = 0; - if (M::PrtlDim == Dim::_2D or M::PrtlDim == Dim::_3D){ // only pick x2 if simulation is in 2D + if constexpr (D == Dim::_2D or D == Dim::_3D){ // only pick x2 if simulation is in 2D if (x_Ph[1] <= x2_min) { x2_ind = 0; } else if (x_Ph[1] >= x2_max) { @@ -1054,11 +1004,10 @@ namespace ntt { x2_ind = static_cast( static_cast(nx2_bins) * (x_Ph[1] - x2_min) / (x2_max - x2_min)); } - //real_t x2_center = x2_min + (x2_ind+0.5)*dx2_bin; - //owns_x2 = (x2_center >= x2_min_local) && (x2_center <= x2_max_local); + } std::size_t x3_ind = 0; - if (M::PrtlDim == Dim::_3D){ // only pick x3 if simulation is in 3D + if constexpr (D == Dim::_3D){ // only pick x3 if simulation is in 3D if (x_Ph[2] <= x3_min) { x3_ind = 0; diff --git a/src/framework/parameters/output.cpp b/src/framework/parameters/output.cpp index 13e1a5463..570d73006 100644 --- a/src/framework/parameters/output.cpp +++ b/src/framework/parameters/output.cpp @@ -162,7 +162,7 @@ namespace ntt { /* Spectra3D ------------------------------------------------------------ */ spectra3d_e_min = toml::find_or(toml_data, "output", "spectra3D", "e_min", defaults::output::spec3d_emin); - spectra3d_e_max = toml::find_or(toml_data, "output", "spectra3D", "e_max", defaults::output::spec3d_emin); + spectra3d_e_max = toml::find_or(toml_data, "output", "spectra3D", "e_max", defaults::output::spec3d_emax); spectra3d_log_bins = toml::find_or(toml_data, "output", @@ -186,13 +186,13 @@ namespace ntt { "output", "spectra3D", "nx2", - defaults::output::spec3d_nx1); + defaults::output::spec3d_nx2); spectra3d_nx3 = toml::find_or(toml_data, "output", "spectra3D", "nx3", - defaults::output::spec3d_nx1); + defaults::output::spec3d_nx3); /* Stats ---------------------------------------------------------------- */ stats_quantities = toml::find_or(toml_data, diff --git a/src/output/writer.cpp b/src/output/writer.cpp index 317d95af9..80bfe360f 100644 --- a/src/output/writer.cpp +++ b/src/output/writer.cpp @@ -413,6 +413,7 @@ namespace out { #else int rank = 0; int size = 1; + auto counts3D_h_all = counts3D_h; #endif using Shape = std::vector; From 7bae89df3f7204789b215bcc82951e43dd5b27c2 Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 14 Sep 2026 14:15:36 -0400 Subject: [PATCH 04/17] gitignore devenv --- .gitignore | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index e683ed8c5..cd508a596 100644 --- a/.gitignore +++ b/.gitignore @@ -64,4 +64,8 @@ action-token ignore-* tombi/ tidy/ -.claude \ No newline at end of file +.claude + +# devenv +.devenv* +devenv.local.nix From 50040dcbd3d69d99cac6179fd7c93d52027d3013 Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 14 Sep 2026 15:06:22 -0400 Subject: [PATCH 05/17] spectra3d merged with regular spectra output --- input.example.toml | 8 +- src/framework/domain/metadomain_io.cpp | 132 ++++++++++++------------- src/framework/parameters/output.cpp | 86 +++++++--------- src/framework/parameters/output.h | 17 +--- src/global/defaults.h | 9 +- src/global/global.h | 1 - 6 files changed, 109 insertions(+), 144 deletions(-) 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/src/framework/domain/metadomain_io.cpp b/src/framework/domain/metadomain_io.cpp index a1bc13f00..793955959 100644 --- a/src/framework/domain/metadomain_io.cpp +++ b/src/framework/domain/metadomain_io.cpp @@ -383,10 +383,10 @@ namespace ntt { finished_step, finished_time); const auto write_spectra3D = params.template get( - "output.spectra3D.enable") and - g_writer.shouldWrite("spectra3D", - finished_step, - finished_time); + "output.spectra3D.enable") and + g_writer.shouldWrite("spectra3D", + finished_step, + finished_time); const auto extension = params.template get("output.format"); if (not(write_fields or write_particles or write_spectra or write_spectra3D) and extension != "disabled") { @@ -856,18 +856,19 @@ namespace ntt { "output.spectra3D.log_bins"); const auto n_bins = params.template get( "output.spectra3D.n_bins"); - const auto& metric = local_domain->mesh.metric; + const auto& metric = local_domain->mesh.metric; // extract the number of bins globally in each direction - const auto nx1_bins = params.template get("output.spectra3D.nx1"); - const auto nx2_bins = params.template get("output.spectra3D.nx2"); - const auto nx3_bins = params.template get("output.spectra3D.nx3"); + const auto nx1_bins = params.template get( + "output.spectra3D.nx1"); + const auto nx2_bins = params.template get( + "output.spectra3D.nx2"); + const auto nx3_bins = params.template get( + "output.spectra3D.nx3"); // select the min and max energy for the spectra auto e_min = params.template get("output.spectra3D.e_min"); auto e_max = params.template get("output.spectra3D.e_max"); - - auto x1_min = mesh().extent(in::x1).first; // extent is in physical units, not code decltype(x1_min) x2_min = 0; decltype(x1_min) x3_min = 0; @@ -875,34 +876,23 @@ namespace ntt { auto x1_max = mesh().extent(in::x1).second; // pairs in c++ are addressed by first and secocnd decltype(x1_max) x2_max = 0; decltype(x1_max) x3_max = 0; - - + if constexpr (D == Dim::_2D or + D == Dim::_3D) { // only pick x2 if simulation is in 2D - - - if constexpr (D == Dim::_2D or D == Dim::_3D){ // only pick x2 if simulation is in 2D - - x2_min = mesh().extent(in::x2).first; - x2_max = mesh().extent(in::x2).second; - + x2_min = mesh().extent(in::x2).first; + x2_max = mesh().extent(in::x2).second; } - if constexpr (D == Dim::_3D){ - - x3_min = mesh().extent(in::x3).first; - x3_max = mesh().extent(in::x3).second; + if constexpr (D == Dim::_3D) { - - + x3_min = mesh().extent(in::x3).first; + x3_max = mesh().extent(in::x3).second; } - - - //auto dx_slice = mesh().dx1; - + // auto dx_slice = mesh().dx1; - //auto dxslice = (x1_max - x1_min) / x_bins + // auto dxslice = (x1_max - x1_min) / x_bins if (log_bins) { e_min = math::log10(e_min); @@ -922,25 +912,26 @@ namespace ntt { }); for (const auto& spec : g_writer.spectraWriters()) { - auto& species = local_domain->species[spec.species() - 1]; + auto& species = local_domain->species[spec.species() - 1]; array_t dn3d { "dn3d", nx1_bins, nx2_bins, nx3_bins, n_bins }; - auto dn3d_scatter = Kokkos::Experimental::create_scatter_view(dn3d); - auto ux1 = species.ux1; - auto ux2 = species.ux2; - auto ux3 = species.ux3; - auto i1 = species.i1; - auto dx1 = species.dx1; - //adeep_copy; - decltype(i1) i2; - decltype(i1) i3; + auto dn3d_scatter = Kokkos::Experimental::create_scatter_view(dn3d); + auto ux1 = species.ux1; + auto ux2 = species.ux2; + auto ux3 = species.ux3; + auto i1 = species.i1; + auto dx1 = species.dx1; + // adeep_copy; + decltype(i1) i2; + decltype(i1) i3; decltype(dx1) dx2; decltype(dx1) dx3; - if constexpr (D == Dim::_2D or D == Dim::_3D){ // only pick x2 if simulation is in 2D - i2 = species.i2; + if constexpr (D == Dim::_2D or + D == Dim::_3D) { // only pick x2 if simulation is in 2D + i2 = species.i2; dx2 = species.dx2; } - if constexpr (D == Dim::_3D){ - i3 = species.i3; + if constexpr (D == Dim::_3D) { + i3 = species.i3; dx3 = species.dx3; } auto weight = species.weight; @@ -953,15 +944,15 @@ namespace ntt { if (tag(p) != ParticleTag::alive) { return; } - - coord_t x_Cd {ZERO}; - if constexpr (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { + + coord_t x_Cd { ZERO }; + if constexpr (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { x_Cd[0] = static_cast(i1(p)) + static_cast(dx1(p)); } - if constexpr (D == Dim::_2D or D == Dim::_3D) { + if constexpr (D == Dim::_2D or D == Dim::_3D) { x_Cd[1] = static_cast(i2(p)) + static_cast(dx2(p)); } - if constexpr (D == Dim::_3D) { + if constexpr (D == Dim::_3D) { x_Cd[2] = static_cast(i3(p)) + static_cast(dx3(p)); } coord_t x_Ph { ZERO }; @@ -980,56 +971,55 @@ namespace ntt { if (en <= e_min) { e_ind = 0; } else if (en >= e_max) { - e_ind = n_bins-1; + e_ind = n_bins - 1; } else { e_ind = static_cast( static_cast(n_bins) * (en - e_min) / (e_max - e_min)); } std::size_t x1_ind = 0; - if (x_Ph[0]<= x1_min) { + if (x_Ph[0] <= x1_min) { x1_ind = 0; } else if (x_Ph[0] >= x1_max) { - x1_ind = nx1_bins-1; + x1_ind = nx1_bins - 1; } else { - x1_ind = static_cast( - static_cast(nx1_bins) * (x_Ph[0] - x1_min) / (x1_max - x1_min)); + x1_ind = static_cast(static_cast(nx1_bins) * + (x_Ph[0] - x1_min) / + (x1_max - x1_min)); } - std::size_t x2_ind = 0; - if constexpr (D == Dim::_2D or D == Dim::_3D){ // only pick x2 if simulation is in 2D + if constexpr (D == Dim::_2D or + D == Dim::_3D) { // only pick x2 if simulation is in 2D if (x_Ph[1] <= x2_min) { x2_ind = 0; } else if (x_Ph[1] >= x2_max) { - x2_ind = nx2_bins-1; + x2_ind = nx2_bins - 1; } else { - x2_ind = static_cast( - static_cast(nx2_bins) * (x_Ph[1] - x2_min) / (x2_max - x2_min)); + x2_ind = static_cast(static_cast(nx2_bins) * + (x_Ph[1] - x2_min) / + (x2_max - x2_min)); } - } std::size_t x3_ind = 0; - if constexpr (D == Dim::_3D){ // only pick x3 if simulation is in 3D - + if constexpr (D == Dim::_3D) { // only pick x3 if simulation is in 3D + if (x_Ph[2] <= x3_min) { x3_ind = 0; } else if (x_Ph[2] >= x3_max) { - x3_ind = nx3_bins-1; + x3_ind = nx3_bins - 1; } else { - x3_ind = static_cast( - static_cast(nx3_bins) * (x_Ph[2] - x3_min) / (x3_max - x3_min)); + x3_ind = static_cast(static_cast(nx3_bins) * + (x_Ph[2] - x3_min) / + (x3_max - x3_min)); } - - - } + } // now I want to save the nx_bins contained in each rank // can save array of x1_inds which are saved? // maybe I can ask what local x_min and x_max are for this rank, pass that to writeSpectrum3D - - auto dn3d_acc = dn3d_scatter.access(); + auto dn3d_acc = dn3d_scatter.access(); dn3d_acc(x1_ind, x2_ind, x3_ind, e_ind) += weight(p); }); Kokkos::Experimental::contribute(dn3d, dn3d_scatter); @@ -1037,7 +1027,7 @@ namespace ntt { } g_writer.writeSpectrumBins(energy, "sEbn"); g_writer.endWriting(WriteMode::Spectra3D); - } + } return true; } diff --git a/src/framework/parameters/output.cpp b/src/framework/parameters/output.cpp index 570d73006..7591895a1 100644 --- a/src/framework/parameters/output.cpp +++ b/src/framework/parameters/output.cpp @@ -33,7 +33,7 @@ namespace ntt { HERE); categories.emplace(); - for (const auto& category : { "fields", "particles", "spectra", "spectra3D", "stats" }) { + for (const auto& category : { "fields", "particles", "spectra", "stats" }) { const auto q_int = toml::find_or(toml_data, "output", category, @@ -153,46 +153,37 @@ 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); - - /* Spectra3D ------------------------------------------------------------ */ - spectra3d_e_min = toml::find_or(toml_data, "output", "spectra3D", "e_min", defaults::output::spec3d_emin); - - spectra3d_e_max = toml::find_or(toml_data, "output", "spectra3D", "e_max", defaults::output::spec3d_emax); - - spectra3d_log_bins = toml::find_or(toml_data, - "output", - "spectra3D", - "log_bins", - defaults::output::spec3d_log); - - spectra3d_n_bins = toml::find_or(toml_data, - "output", - "spectra3D", - "n_bins", - defaults::output::spec3d_nbins); - - spectra3d_nx1 = toml::find_or(toml_data, - "output", - "spectra3D", - "nx1", - defaults::output::spec3d_nx1); - - spectra3d_nx2 = toml::find_or(toml_data, - "output", - "spectra3D", - "nx2", - defaults::output::spec3d_nx2); - - spectra3d_nx3 = toml::find_or(toml_data, - "output", - "spectra3D", - "nx3", - defaults::output::spec3d_nx3); + 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, @@ -242,15 +233,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.spectra3D.e_min", spectra3d_e_min.value()); - params->set("output.spectra3D.e_max", spectra3d_e_max.value()); - params->set("output.spectra3D.log_bins", spectra3d_log_bins.value()); - params->set("output.spectra3D.n_bins", spectra3d_n_bins.value()); - params->set("output.spectra3D.nx1", spectra3d_nx1.value()); - params->set("output.spectra3D.nx2", spectra3d_nx2.value()); - params->set("output.spectra3D.nx3", spectra3d_nx3.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 680751514..a51bf9290 100644 --- a/src/framework/parameters/output.h +++ b/src/framework/parameters/output.h @@ -50,18 +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 spectra3d_e_min; - std::optional spectra3d_e_max; - std::optional spectra3d_log_bins; - std::optional spectra3d_n_bins; - std::optional spectra3d_nx1; - std::optional spectra3d_nx2; - std::optional spectra3d_nx3; + 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/defaults.h b/src/global/defaults.h index 919b11fde..82b4364cc 100644 --- a/src/global/defaults.h +++ b/src/global/defaults.h @@ -75,14 +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 real_t spec3d_emin = 1e-3; - const real_t spec3d_emax = 1e3; - const bool spec3d_log = true; - const std::size_t spec3d_nbins = 200; - const std::size_t spec3d_nx1 = 1; - const std::size_t spec3d_nx2 = 1; - const std::size_t spec3d_nx3 = 1; + const std::size_t spec_num_e_bins = 200; const std::vector stats_quantities = { "B^2", "E^2", "ExB", diff --git a/src/global/global.h b/src/global/global.h index 00047359a..c3fbc5061 100644 --- a/src/global/global.h +++ b/src/global/global.h @@ -298,7 +298,6 @@ namespace WriteMode { Particles = 1 << 1, Spectra = 1 << 2, Stats = 1 << 3, - Spectra3D = 1 << 4, }; } // namespace WriteMode From 2076f0f6ccd16223ce9ac00dff152910e9b017e3 Mon Sep 17 00:00:00 2001 From: haykh Date: Tue, 15 Sep 2026 11:26:06 -0400 Subject: [PATCH 06/17] particle energy distribution calc refactored (incl. spatial) --- src/framework/domain/io/spectra.cpp | 157 ++++++++++--- src/global/arch/kokkos_aliases.h | 34 ++- src/kernels/particle_energy_distribution.hpp | 218 +++++++++++++++++++ src/kernels/particle_moments.hpp | 79 ------- 4 files changed, 375 insertions(+), 113 deletions(-) create mode 100644 src/kernels/particle_energy_distribution.hpp diff --git a/src/framework/domain/io/spectra.cpp b/src/framework/domain/io/spectra.cpp index 317c76b6d..ba54c6a65 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,143 @@ 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 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); + } + } + + { + // 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"); } - g_writer.writeSpectrumBins(energy, "sEbn"); g_writer.endWriting(WriteMode::Spectra); } 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/kernels/particle_energy_distribution.hpp b/src/kernels/particle_energy_distribution.hpp new file mode 100644 index 000000000..d1590480b --- /dev/null +++ b/src/kernels/particle_energy_distribution.hpp @@ -0,0 +1,218 @@ +/** + * @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; + } 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(static_cast(nbins) * 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)); + 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)); + 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 From 69899388b92ff2fc96010afc135e0722b8a0f4c3 Mon Sep 17 00:00:00 2001 From: haykh Date: Tue, 15 Sep 2026 11:47:39 -0400 Subject: [PATCH 07/17] declare and call spatial bins in writing --- src/framework/domain/io/init.cpp | 4 +++- src/framework/domain/io/spectra.cpp | 8 ++++++++ src/output/writer.h | 3 ++- 3 files changed, 13 insertions(+), 2 deletions(-) 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 ba54c6a65..b053fca54 100644 --- a/src/framework/domain/io/spectra.cpp +++ b/src/framework/domain/io/spectra.cpp @@ -143,6 +143,7 @@ namespace ntt { nbins_i, { e_min, e_max, log_bins, num_energy_bins }, local_domain->mesh.metric); + // @TODO: write } } @@ -165,6 +166,13 @@ namespace ntt { g_writer.writeSpectrumBins(energy, "sEbn"); } + if (spatial_binning_enabled) { + for (auto d = 0u; d < dim; ++d) { + array_t xi { "xi", num_spatial_bins[d] + 1 }; + // @TODO: fill spatial bins (in physical units) + g_writer.writeSpectrumBins(xi, "sX" + std::to_string(d + 1) + "bn"); + } + } g_writer.endWriting(WriteMode::Spectra); } diff --git a/src/output/writer.h b/src/output/writer.h index 9fee71d80..18193eaad 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&, From 7c88d0fa4bb2580f0b127f051339af80959319ca Mon Sep 17 00:00:00 2001 From: haykh Date: Tue, 15 Sep 2026 13:23:19 -0400 Subject: [PATCH 08/17] spatially binned spectra done (probably) --- src/framework/domain/io/spectra.cpp | 215 ++++------------------- src/output/writer.cpp | 262 ++++++++++++++++++++-------- src/output/writer.h | 6 +- 3 files changed, 231 insertions(+), 252 deletions(-) diff --git a/src/framework/domain/io/spectra.cpp b/src/framework/domain/io/spectra.cpp index b053fca54..f628ff05b 100644 --- a/src/framework/domain/io/spectra.cpp +++ b/src/framework/domain/io/spectra.cpp @@ -66,6 +66,20 @@ namespace ntt { 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, @@ -143,7 +157,8 @@ namespace ntt { nbins_i, { e_min, e_max, log_bins, num_energy_bins }, local_domain->mesh.metric); - // @TODO: write + g_writer.writeSpectrumSpatial(dim + 1u)>(dn, + spec.name()); } } @@ -167,9 +182,26 @@ namespace ntt { 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 }; - // @TODO: fill spatial bins (in physical units) + // # 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"); } } @@ -189,182 +221,3 @@ namespace ntt { // NOLINTEND(bugprone-macro-parentheses) } // namespace ntt - -// if (write_spectra3D) { -// g_writer.beginWriting(WriteMode::Spectra3D, current_step, current_time); -// const auto log_bins = params.template get( -// "output.spectra3D.log_bins"); -// const auto n_bins = params.template get( -// "output.spectra3D.n_bins"); -// const auto& metric = local_domain->mesh.metric; -// // extract the number of bins globally in each direction -// const auto nx1_bins = params.template get( -// "output.spectra3D.nx1"); -// const auto nx2_bins = params.template get( -// "output.spectra3D.nx2"); -// const auto nx3_bins = params.template get( -// "output.spectra3D.nx3"); -// -// // select the min and max energy for the spectra -// auto e_min = params.template get("output.spectra3D.e_min"); -// auto e_max = params.template get("output.spectra3D.e_max"); -// -// auto x1_min = mesh().extent(in::x1).first; // extent is in physical units, not code -// decltype(x1_min) x2_min = 0; -// decltype(x1_min) x3_min = 0; -// -// auto x1_max = mesh().extent(in::x1).second; // pairs in c++ are addressed by first and secocnd -// decltype(x1_max) x2_max = 0; -// decltype(x1_max) x3_max = 0; -// -// if constexpr (D == Dim::_2D or -// D == Dim::_3D) { // only pick x2 if simulation is in 2D -// -// x2_min = mesh().extent(in::x2).first; -// x2_max = mesh().extent(in::x2).second; -// } -// -// if constexpr (D == Dim::_3D) { -// -// x3_min = mesh().extent(in::x3).first; -// x3_max = mesh().extent(in::x3).second; -// } -// -// // auto dx_slice = mesh().dx1; -// -// // auto dxslice = (x1_max - x1_min) / x_bins -// -// if (log_bins) { -// e_min = math::log10(e_min); -// e_max = math::log10(e_max); -// } -// array_t energy { "energy", n_bins + 1 }; -// // generating the energy bins -// Kokkos::parallel_for( -// "GenerateEnergyBins", -// n_bins + 1, -// Lambda(uint32_t e) { -// if (log_bins) { -// energy(e) = math::pow(10.0, e_min + (e_max - e_min) * e / n_bins); -// } else { -// energy(e) = e_min + (e_max - e_min) * e / n_bins; -// } -// }); -// -// for (const auto& spec : g_writer.spectraWriters()) { -// auto& species = local_domain->species[spec.species() - 1]; -// array_t dn3d { "dn3d", nx1_bins, nx2_bins, nx3_bins, n_bins }; -// auto dn3d_scatter = Kokkos::Experimental::create_scatter_view(dn3d); -// auto ux1 = species.ux1; -// auto ux2 = species.ux2; -// auto ux3 = species.ux3; -// auto i1 = species.i1; -// auto dx1 = species.dx1; -// // adeep_copy; -// decltype(i1) i2; -// decltype(i1) i3; -// decltype(dx1) dx2; -// decltype(dx1) dx3; -// if constexpr (D == Dim::_2D or -// D == Dim::_3D) { // only pick x2 if simulation is in 2D -// i2 = species.i2; -// dx2 = species.dx2; -// } -// if constexpr (D == Dim::_3D) { -// i3 = species.i3; -// dx3 = species.dx3; -// } -// auto weight = species.weight; -// auto tag = species.tag; -// const auto is_massive = species.mass() > 0.0f; -// Kokkos::parallel_for( -// "ComputeSpectra", -// species.rangeActiveParticles(), -// Lambda(prtlidx_t p) { -// if (tag(p) != ParticleTag::alive) { -// return; -// } -// -// coord_t x_Cd { ZERO }; -// if constexpr (D == Dim::_1D or D == Dim::_2D or D == Dim::_3D) { -// x_Cd[0] = static_cast(i1(p)) + static_cast(dx1(p)); -// } -// if constexpr (D == Dim::_2D or D == Dim::_3D) { -// x_Cd[1] = static_cast(i2(p)) + static_cast(dx2(p)); -// } -// if constexpr (D == Dim::_3D) { -// x_Cd[2] = static_cast(i3(p)) + static_cast(dx3(p)); -// } -// coord_t x_Ph { ZERO }; -// metric.template convert(x_Cd, x_Ph); -// -// real_t en; -// if (is_massive) { -// en = U2GAMMA(ux1(p), ux2(p), ux3(p)) - ONE; -// } else { -// en = NORM(ux1(p), ux2(p), ux3(p)); -// } -// if (log_bins) { -// en = math::log10(en); -// } -// std::size_t e_ind = 0; -// if (en <= e_min) { -// e_ind = 0; -// } else if (en >= e_max) { -// e_ind = n_bins - 1; -// } else { -// e_ind = static_cast( -// static_cast(n_bins) * (en - e_min) / (e_max - e_min)); -// } -// -// std::size_t x1_ind = 0; -// if (x_Ph[0] <= x1_min) { -// x1_ind = 0; -// } else if (x_Ph[0] >= x1_max) { -// x1_ind = nx1_bins - 1; -// } else { -// x1_ind = static_cast(static_cast(nx1_bins) * -// (x_Ph[0] - x1_min) / -// (x1_max - x1_min)); -// } -// -// std::size_t x2_ind = 0; -// if constexpr (D == Dim::_2D or -// D == Dim::_3D) { // only pick x2 if simulation is in 2D -// if (x_Ph[1] <= x2_min) { -// x2_ind = 0; -// } else if (x_Ph[1] >= x2_max) { -// x2_ind = nx2_bins - 1; -// } else { -// x2_ind = static_cast(static_cast(nx2_bins) * -// (x_Ph[1] - x2_min) / -// (x2_max - x2_min)); -// } -// } -// std::size_t x3_ind = 0; -// if constexpr (D == Dim::_3D) { // only pick x3 if simulation is in 3D -// -// if (x_Ph[2] <= x3_min) { -// x3_ind = 0; -// } else if (x_Ph[2] >= x3_max) { -// x3_ind = nx3_bins - 1; -// } else { -// x3_ind = static_cast(static_cast(nx3_bins) * -// (x_Ph[2] - x3_min) / -// (x3_max - x3_min)); -// } -// } -// -// // now I want to save the nx_bins contained in each rank -// // can save array of x1_inds which are saved? -// // maybe I can ask what local x_min and x_max are for this rank, pass that to writeSpectrum3D -// -// auto dn3d_acc = dn3d_scatter.access(); -// dn3d_acc(x1_ind, x2_ind, x3_ind, e_ind) += weight(p); -// }); -// Kokkos::Experimental::contribute(dn3d, dn3d_scatter); -// g_writer.writeSpectrum3D(dn3d, spec.name()); -// } -// g_writer.writeSpectrumBins(energy, "sEbn"); -// g_writer.endWriting(WriteMode::Spectra3D); -// } diff --git a/src/output/writer.cpp b/src/output/writer.cpp index 69c509fb5..568143f9d 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,75 +477,130 @@ namespace out { #endif } -// spectrum3D, broken up into separate sub-domains - void Writer::writeSpectrum3D(const array_t& counts3D, - const std::string& varname) { - // need to include the rank specific coordinates contained here - std::string varname3d = varname + "_3D"; - //auto var = m_io.InquireVariable(varname); - auto counts3D_h = Kokkos::create_mirror_view(counts3D); - // copy to host - Kokkos::deep_copy(counts3D_h, counts3D); -#if defined(MPI_ENABLED) - array_t counts3D_all { "counts3D_all", counts3D.extent(0), counts3D.extent(1), counts3D.extent(2), counts3D.extent(3) }; - //auto counts_h_all = Kokkos::create_mirror_view(counts_all); - int rank; - int size; - MPI_Comm_rank(MPI_COMM_WORLD, &rank); - MPI_Comm_size(MPI_COMM_WORLD, &size); - - auto counts3D_h_all = Kokkos::create_mirror_view(counts3D_all); - - - MPI_Allreduce(counts3D_h.data(), - counts3D_h_all.data(), - counts3D_h.extent(0)*counts3D_h.extent(1)*counts3D_h.extent(2)*counts3D_h.extent(3),// size so probably need to multiply counts_h.extent(0)*counts_h.extent(1)*counts_h.extent(2) - mpi::get_type(), - MPI_SUM, - MPI_COMM_WORLD); // size - //Kokkos::deep_copy() -#else - int rank = 0; - int size = 1; - auto counts3D_h_all = counts3D_h; -#endif - using Shape = std::vector; - - Shape shape = {counts3D.extent(0), counts3D.extent(1), counts3D.extent(2), counts3D.extent(3)}; // global size of counts3D_all - Shape start, count; // per-rank selection - - // split the domain along the x1 direction across the different ranks - auto split_x1 = [&](std::size_t n, int this_rank, int nranks) - { - // base number of cells to allocate to each rank - std::size_t base = n / nranks; - // remainder that do not fit neatly in one rank - std::size_t rem = n % nranks; - // number of cells to allocate to the rank including the remainder - std::size_t nloc = base + (this_rank <(int)rem ? 1 : 0); // allocate one more if this rank is less that the remainder - //offset to allocate - std::size_t off = base * this_rank + std::min(this_rank, rem); - - return std::pair{nloc, off}; - }; - - auto [nloc0, off0] = split_x1(counts3D.extent(0), rank, size); - start = {off0, 0, 0, 0}; - count = {nloc0, counts3D.extent(1), counts3D.extent(2), counts3D.extent(3)}; - - auto var = m_io.InquireVariable(varname3d); - if (!var) - { - var = m_io.DefineVariable(varname3d, shape, start, count, adios2::ConstantDims); + 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); } - - //auto var = m_io.DefineVariable(varname3d, shape, start, count, adios2::ConstantDims); - - auto counts3D_h_all_slab = Kokkos::subview(counts3D_h_all, Kokkos::make_pair(off0, off0+nloc0), Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()); - - m_writer.Put(var, counts3D_h_all_slab, adios2::Mode::Sync); } + // // spectrum3D, broken up into separate sub-domains + // void Writer::writeSpectrum3D(const array_t& counts3D, + // const std::string& varname) { + // // need to include the rank specific coordinates contained here + // std::string varname3d = varname + "_3D"; + // // auto var = m_io.InquireVariable(varname); + // auto counts3D_h = Kokkos::create_mirror_view(counts3D); + // // copy to host + // Kokkos::deep_copy(counts3D_h, counts3D); + // #if defined(MPI_ENABLED) + // array_t counts3D_all { "counts3D_all", + // counts3D.extent(0), + // counts3D.extent(1), + // counts3D.extent(2), + // counts3D.extent(3) }; + // // auto counts_h_all = Kokkos::create_mirror_view(counts_all); + // int rank; + // int size; + // MPI_Comm_rank(MPI_COMM_WORLD, &rank); + // MPI_Comm_size(MPI_COMM_WORLD, &size); + // + // auto counts3D_h_all = Kokkos::create_mirror_view(counts3D_all); + // + // MPI_Allreduce( + // counts3D_h.data(), + // counts3D_h_all.data(), + // counts3D_h.extent(0) * counts3D_h.extent(1) * counts3D_h.extent(2) * + // counts3D_h.extent( + // 3), // size so probably need to multiply counts_h.extent(0)*counts_h.extent(1)*counts_h.extent(2) + // mpi::get_type(), + // MPI_SUM, + // MPI_COMM_WORLD); // size + // // Kokkos::deep_copy() + // #else + // int rank = 0; + // int size = 1; + // auto counts3D_h_all = counts3D_h; + // #endif + // using Shape = std::vector; + // + // Shape shape = { counts3D.extent(0), + // counts3D.extent(1), + // counts3D.extent(2), + // counts3D.extent(3) }; // global size of counts3D_all + // Shape start, count; // per-rank selection + // + // // split the domain along the x1 direction across the different ranks + // auto split_x1 = [&](std::size_t n, int this_rank, int nranks) { + // // base number of cells to allocate to each rank + // std::size_t base = n / nranks; + // // remainder that do not fit neatly in one rank + // std::size_t rem = n % nranks; + // // number of cells to allocate to the rank including the remainder + // std::size_t nloc = base + + // (this_rank < (int)rem + // ? 1 + // : 0); // allocate one more if this rank is less that the remainder + // // offset to allocate + // std::size_t off = base * this_rank + std::min(this_rank, rem); + // + // return std::pair { nloc, off }; + // }; + // + // auto [nloc0, off0] = split_x1(counts3D.extent(0), rank, size); + // start = { off0, 0, 0, 0 }; + // count = { nloc0, counts3D.extent(1), counts3D.extent(2), counts3D.extent(3) }; + // + // auto var = m_io.InquireVariable(varname3d); + // if (!var) { + // var = m_io.DefineVariable(varname3d, + // shape, + // start, + // count, + // adios2::ConstantDims); + // } + // + // // auto var = m_io.DefineVariable(varname3d, shape, start, count, adios2::ConstantDims); + // + // auto counts3D_h_all_slab = Kokkos::subview(counts3D_h_all, + // Kokkos::make_pair(off0, off0 + nloc0), + // Kokkos::ALL(), + // Kokkos::ALL(), + // Kokkos::ALL()); + // + // m_writer.Put(var, counts3D_h_all_slab, adios2::Mode::Sync); + // } + void Writer::writeSpectrumBins(const array_t& e_bins, const std::string& varname) { auto var = m_io.InquireVariable(varname); @@ -567,8 +683,6 @@ namespace out { mode_str = "particles"; } else if (write_mode == WriteMode::Spectra) { mode_str = "spectra"; - } else if (write_mode == WriteMode::Spectra3D) { - mode_str = "spectra3D"; } else { raise::Fatal("Unknown write mode", HERE); } @@ -638,4 +752,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 18193eaad..6e9a139d2 100644 --- a/src/output/writer.h +++ b/src/output/writer.h @@ -139,7 +139,11 @@ namespace out { npart_t, const std::string&); void writeSpectrum(const array_t&, const std::string&); - void writeSpectrum3D(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); From a94afa584600dbf69ec21caa9e7f5692426ab49a Mon Sep 17 00:00:00 2001 From: haykh Date: Tue, 15 Sep 2026 13:52:22 -0400 Subject: [PATCH 09/17] minor bugs --- src/framework/parameters/output.cpp | 11 ++++++----- src/kernels/particle_energy_distribution.hpp | 2 ++ 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/framework/parameters/output.cpp b/src/framework/parameters/output.cpp index 7591895a1..d2a6a9c44 100644 --- a/src/framework/parameters/output.cpp +++ b/src/framework/parameters/output.cpp @@ -171,11 +171,12 @@ namespace ntt { "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 }); + 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", diff --git a/src/kernels/particle_energy_distribution.hpp b/src/kernels/particle_energy_distribution.hpp index d1590480b..29e53e60e 100644 --- a/src/kernels/particle_energy_distribution.hpp +++ b/src/kernels/particle_energy_distribution.hpp @@ -189,6 +189,7 @@ namespace kernel { 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)) + @@ -206,6 +207,7 @@ namespace kernel { 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"); From aebfdcc3e9c92aa7ab388dab7f3abb03f4611880 Mon Sep 17 00:00:00 2001 From: haykh Date: Tue, 15 Sep 2026 16:36:44 -0400 Subject: [PATCH 10/17] minor bugfix --- src/kernels/particle_energy_distribution.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/kernels/particle_energy_distribution.hpp b/src/kernels/particle_energy_distribution.hpp index 29e53e60e..eba88bac6 100644 --- a/src/kernels/particle_energy_distribution.hpp +++ b/src/kernels/particle_energy_distribution.hpp @@ -90,7 +90,7 @@ namespace kernel { if (en <= energy_binning.e_min) { return 0u; } else if (en >= energy_binning.e_max) { - return energy_binning.n_bins; + return energy_binning.n_bins - 1u; } else { return static_cast(static_cast(energy_binning.n_bins) * (en - energy_binning.e_min) / @@ -161,7 +161,7 @@ namespace kernel { } else if (ni >= dncells * static_cast(nbins)) { return nbins - 1u; } else { - return static_cast(static_cast(nbins) * ni / dncells); + return static_cast(ni / dncells); } } From 56b174c9ffe817ec26dd9a3cc8c728e4546db767 Mon Sep 17 00:00:00 2001 From: haykh Date: Tue, 15 Sep 2026 17:59:50 -0400 Subject: [PATCH 11/17] rm old spec3d --- src/output/writer.cpp | 87 ------------------------------------------- 1 file changed, 87 deletions(-) diff --git a/src/output/writer.cpp b/src/output/writer.cpp index 568143f9d..89410acc3 100644 --- a/src/output/writer.cpp +++ b/src/output/writer.cpp @@ -514,93 +514,6 @@ namespace out { } } - // // spectrum3D, broken up into separate sub-domains - // void Writer::writeSpectrum3D(const array_t& counts3D, - // const std::string& varname) { - // // need to include the rank specific coordinates contained here - // std::string varname3d = varname + "_3D"; - // // auto var = m_io.InquireVariable(varname); - // auto counts3D_h = Kokkos::create_mirror_view(counts3D); - // // copy to host - // Kokkos::deep_copy(counts3D_h, counts3D); - // #if defined(MPI_ENABLED) - // array_t counts3D_all { "counts3D_all", - // counts3D.extent(0), - // counts3D.extent(1), - // counts3D.extent(2), - // counts3D.extent(3) }; - // // auto counts_h_all = Kokkos::create_mirror_view(counts_all); - // int rank; - // int size; - // MPI_Comm_rank(MPI_COMM_WORLD, &rank); - // MPI_Comm_size(MPI_COMM_WORLD, &size); - // - // auto counts3D_h_all = Kokkos::create_mirror_view(counts3D_all); - // - // MPI_Allreduce( - // counts3D_h.data(), - // counts3D_h_all.data(), - // counts3D_h.extent(0) * counts3D_h.extent(1) * counts3D_h.extent(2) * - // counts3D_h.extent( - // 3), // size so probably need to multiply counts_h.extent(0)*counts_h.extent(1)*counts_h.extent(2) - // mpi::get_type(), - // MPI_SUM, - // MPI_COMM_WORLD); // size - // // Kokkos::deep_copy() - // #else - // int rank = 0; - // int size = 1; - // auto counts3D_h_all = counts3D_h; - // #endif - // using Shape = std::vector; - // - // Shape shape = { counts3D.extent(0), - // counts3D.extent(1), - // counts3D.extent(2), - // counts3D.extent(3) }; // global size of counts3D_all - // Shape start, count; // per-rank selection - // - // // split the domain along the x1 direction across the different ranks - // auto split_x1 = [&](std::size_t n, int this_rank, int nranks) { - // // base number of cells to allocate to each rank - // std::size_t base = n / nranks; - // // remainder that do not fit neatly in one rank - // std::size_t rem = n % nranks; - // // number of cells to allocate to the rank including the remainder - // std::size_t nloc = base + - // (this_rank < (int)rem - // ? 1 - // : 0); // allocate one more if this rank is less that the remainder - // // offset to allocate - // std::size_t off = base * this_rank + std::min(this_rank, rem); - // - // return std::pair { nloc, off }; - // }; - // - // auto [nloc0, off0] = split_x1(counts3D.extent(0), rank, size); - // start = { off0, 0, 0, 0 }; - // count = { nloc0, counts3D.extent(1), counts3D.extent(2), counts3D.extent(3) }; - // - // auto var = m_io.InquireVariable(varname3d); - // if (!var) { - // var = m_io.DefineVariable(varname3d, - // shape, - // start, - // count, - // adios2::ConstantDims); - // } - // - // // auto var = m_io.DefineVariable(varname3d, shape, start, count, adios2::ConstantDims); - // - // auto counts3D_h_all_slab = Kokkos::subview(counts3D_h_all, - // Kokkos::make_pair(off0, off0 + nloc0), - // Kokkos::ALL(), - // Kokkos::ALL(), - // Kokkos::ALL()); - // - // m_writer.Put(var, counts3D_h_all_slab, adios2::Mode::Sync); - // } - void Writer::writeSpectrumBins(const array_t& e_bins, const std::string& varname) { auto var = m_io.InquireVariable(varname); From 6697b95e6c3eba98dbeaedb933d696891db193e4 Mon Sep 17 00:00:00 2001 From: haykh Date: Tue, 15 Sep 2026 18:03:58 -0400 Subject: [PATCH 12/17] input for shock fixed back --- pgens/shock/shock.toml | 43 +++++++++++++++++------------------------- 1 file changed, 17 insertions(+), 26 deletions(-) diff --git a/pgens/shock/shock.toml b/pgens/shock/shock.toml index 37515078f..8d3270a74 100644 --- a/pgens/shock/shock.toml +++ b/pgens/shock/shock.toml @@ -40,47 +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 = true - e_min = 1e-2 - e_max = 1e1 - log_bins = true - - [output.spectra3D] - enable = true - e_min = 1e-2 - e_max = 1e1 - log_bins = true - nx1 = 40 - nx2 = 3 - nx3 = 1 + 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 From 7d6bb4590ae7dd86e69aa909a22fd403884eec8e Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 21 Sep 2026 10:33:38 -0400 Subject: [PATCH 13/17] vendor report moved --- cmake/report.cmake | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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( From bdc3172053fcc7aff8e2f20a71e9286c8cbbd28d Mon Sep 17 00:00:00 2001 From: hayk Date: Mon, 21 Sep 2026 08:16:49 -0700 Subject: [PATCH 14/17] perlmutter dep fixed --- dependencies.py | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) 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]]: From 940019dd7d16924fcede9b8f5a3c21c75232949a Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 21 Sep 2026 12:17:05 -0400 Subject: [PATCH 15/17] spec test --- tests/framework/CMakeLists.txt | 2 + tests/framework/spectra.cpp | 303 +++++++++++++++++++++++++++++++++ 2 files changed, 305 insertions(+) create mode 100644 tests/framework/spectra.cpp 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; +} From 7258dcae91fc9fb3c6caf4041cb281c87ab15507 Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 21 Sep 2026 12:17:11 -0400 Subject: [PATCH 16/17] RUNTESTS From 881a9e34e78b2596ae48862c72f25c2b0082137d Mon Sep 17 00:00:00 2001 From: haykh Date: Mon, 21 Sep 2026 12:18:02 -0400 Subject: [PATCH 17/17] RUNTESTS