diff --git a/src/framework/domain/metadomain_io.cpp b/src/framework/domain/metadomain_io.cpp index b9f351a0..db8fc704 100644 --- a/src/framework/domain/metadomain_io.cpp +++ b/src/framework/domain/metadomain_io.cpp @@ -345,8 +345,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; } @@ -788,6 +793,205 @@ 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"); + + + + 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; + + + + 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; + + + + dx2_bin = (static_cast(x2_max) - static_cast(x2_min) )/static_cast(nx2_bins); + + } + + if constexpr (M::PrtlDim == Dim::_3D){ + + x3_min = mesh().extent(in::x3).first; + x3_max = mesh().extent(in::x3).second; + + + dx3_bin = (static_cast(x3_max) - static_cast(x3_min) )/static_cast(nx3_bins); + + + } + + + + //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::cout << "[x1_min, x1_max] =[" << x1_min<< ", " << x1_max << "]" << std::endl; + 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); + } // end shouldWrite("spectra", step, time) + return true; } diff --git a/src/framework/parameters/output.cpp b/src/framework/parameters/output.cpp index 54780097..3562ba8b 100644 --- a/src/framework/parameters/output.cpp +++ b/src/framework/parameters/output.cpp @@ -26,7 +26,7 @@ namespace ntt { "separate_files=false is deprecated", HERE); - 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, @@ -143,6 +143,40 @@ 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_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); /* Stats ---------------------------------------------------------------- */ stats_quantities = toml::find_or(toml_data, @@ -193,6 +227,14 @@ namespace ntt { params->set("output.spectra.log_bins", spectra_log_bins); params->set("output.spectra.n_bins", spectra_n_bins); + params->set("output.spectra3D.e_min", spectra3D_e_min); + params->set("output.spectra3D.e_max", spectra3D_e_max); + params->set("output.spectra3D.log_bins", spectra3D_log_bins); + params->set("output.spectra3D.n_bins", spectra3D_n_bins); + params->set("output.spectra3D.nx1", spectra3D_nx1); + params->set("output.spectra3D.nx2", spectra3D_nx2); + params->set("output.spectra3D.nx3", spectra3D_nx3); + params->set("output.stats.quantities", stats_quantities); params->set("output.stats.custom", stats_custom_quantities); diff --git a/src/framework/parameters/output.h b/src/framework/parameters/output.h index e301fc74..c93c0578 100644 --- a/src/framework/parameters/output.h +++ b/src/framework/parameters/output.h @@ -52,6 +52,14 @@ namespace ntt { bool spectra_log_bins; std::size_t spectra_n_bins; + real_t spectra3D_e_min; + real_t spectra3D_e_max; + bool spectra3D_log_bins; + std::size_t spectra3D_n_bins; + std::size_t spectra3D_nx1; + std::size_t spectra3D_nx2; + std::size_t spectra3D_nx3; + std::vector stats_quantities; std::vector stats_custom_quantities; diff --git a/src/global/defaults.h b/src/global/defaults.h index dbe37ef0..30e4102d 100644 --- a/src/global/defaults.h +++ b/src/global/defaults.h @@ -73,6 +73,14 @@ 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 4aa98f5f..db59b9d9 100644 --- a/src/global/global.h +++ b/src/global/global.h @@ -284,6 +284,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 ff18a036..745df492 100644 --- a/src/output/writer.cpp +++ b/src/output/writer.cpp @@ -361,6 +361,76 @@ 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); @@ -418,6 +488,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 1003e932..345248f1 100644 --- a/src/output/writer.h +++ b/src/output/writer.h @@ -107,6 +107,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);