From 68768020f45812f6f1e768b5919bd9ed59f2e82e Mon Sep 17 00:00:00 2001 From: Luke Roberts Date: Mon, 11 Aug 2025 13:32:42 -0600 Subject: [PATCH 1/3] Add imex advection example --- example/CMakeLists.txt | 1 + example/imex_advection/CMakeLists.txt | 39 +++ .../advection/advection_package.cpp | 222 ++++++++++++++ .../advection/advection_package.hpp | 274 ++++++++++++++++++ example/imex_advection/driver.cpp | 186 ++++++++++++ example/imex_advection/driver.hpp | 53 ++++ example/imex_advection/main.cpp | 66 +++++ example/imex_advection/parthinput.advection | 77 +++++ example/imex_advection/pgen/pgen.hpp | 27 ++ .../imex_advection/pgen/scalar_advection.cpp | 125 ++++++++ .../utilities/imexrk_integrator.cpp | 53 ++++ .../utilities/imexrk_integrator.hpp | 59 ++++ .../utilities/index_permutation.hpp | 120 ++++++++ .../imex_advection/utilities/reconstruct.hpp | 72 +++++ .../imex_advection/utilities/scratch_pack.hpp | 70 +++++ example/imex_advection/utilities/stokes.hpp | 261 +++++++++++++++++ .../imex_advection/utilities/three_vec.hpp | 101 +++++++ .../utilities/variable_macro.hpp | 29 ++ 18 files changed, 1835 insertions(+) create mode 100644 example/imex_advection/CMakeLists.txt create mode 100644 example/imex_advection/advection/advection_package.cpp create mode 100644 example/imex_advection/advection/advection_package.hpp create mode 100644 example/imex_advection/driver.cpp create mode 100644 example/imex_advection/driver.hpp create mode 100644 example/imex_advection/main.cpp create mode 100644 example/imex_advection/parthinput.advection create mode 100644 example/imex_advection/pgen/pgen.hpp create mode 100644 example/imex_advection/pgen/scalar_advection.cpp create mode 100644 example/imex_advection/utilities/imexrk_integrator.cpp create mode 100644 example/imex_advection/utilities/imexrk_integrator.hpp create mode 100644 example/imex_advection/utilities/index_permutation.hpp create mode 100644 example/imex_advection/utilities/reconstruct.hpp create mode 100644 example/imex_advection/utilities/scratch_pack.hpp create mode 100644 example/imex_advection/utilities/stokes.hpp create mode 100644 example/imex_advection/utilities/three_vec.hpp create mode 100644 example/imex_advection/utilities/variable_macro.hpp diff --git a/example/CMakeLists.txt b/example/CMakeLists.txt index 1475ae8e1db2..bc125231166c 100644 --- a/example/CMakeLists.txt +++ b/example/CMakeLists.txt @@ -15,6 +15,7 @@ add_subdirectory(stochastic_subgrid) add_subdirectory(advection) add_subdirectory(boundary_exchange) add_subdirectory(fine_advection) +add_subdirectory(imex_advection) add_subdirectory(calculate_pi) add_subdirectory(kokkos_pi) add_subdirectory(particles) diff --git a/example/imex_advection/CMakeLists.txt b/example/imex_advection/CMakeLists.txt new file mode 100644 index 000000000000..63b8b6f44358 --- /dev/null +++ b/example/imex_advection/CMakeLists.txt @@ -0,0 +1,39 @@ +#========================================================================================= +# (C) (or copyright) 2025. Triad National Security, LLC. All rights reserved. +# +# This program was produced under U.S. Government contract 89233218CNA000001 for Los +# Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC +# for the U.S. Department of Energy/National Nuclear Security Administration. All rights +# in the program are reserved by Triad National Security, LLC, and the U.S. Department +# of Energy/National Nuclear Security Administration. The Government is granted for +# itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide +# license in this material to reproduce, prepare derivative works, distribute copies to +# the public, perform publicly and display publicly, and to permit others to do so. +#========================================================================================= +get_property(DRIVER_LIST GLOBAL PROPERTY DRIVERS_USED_IN_TESTS) +if( "imex_advection-example" IN_LIST DRIVER_LIST OR NOT PARTHENON_DISABLE_EXAMPLES) + add_executable( + imex_advection-example + driver.cpp + driver.hpp + main.cpp + + advection/advection_package.cpp + advection/advection_package.hpp + + pgen/pgen.hpp + pgen/scalar_advection.cpp + + utilities/imexrk_integrator.cpp + utilities/imexrk_integrator.hpp + utilities/index_permutation.hpp + utilities/reconstruct.hpp + utilities/scratch_pack.hpp + utilities/stokes.hpp + utilities/three_vec.hpp + utilities/variable_macro.hpp + ) + + target_link_libraries(imex_advection-example PRIVATE Parthenon::parthenon) + lint_target(imex_advection-example) +endif() diff --git a/example/imex_advection/advection/advection_package.cpp b/example/imex_advection/advection/advection_package.cpp new file mode 100644 index 000000000000..b30eecc39b85 --- /dev/null +++ b/example/imex_advection/advection/advection_package.cpp @@ -0,0 +1,222 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "advection_package.hpp" +#include "defs.hpp" +#include "kokkos_abstraction.hpp" +#include "reconstruct/dc_inline.hpp" +#include "utils/error_checking.hpp" + +using namespace parthenon::package::prelude; + +// *************************************************// +// define the "physics" package Advect, which *// +// includes defining various functions that control*// +// how parthenon functions and any tasks needed to *// +// implement the "physics" *// +// *************************************************// + +namespace advection_package { + +std::shared_ptr Initialize(ParameterInput *pin) { + auto pkg = std::make_shared("advection_package"); + + Real cfl = pin->GetOrAddReal("advection", "cfl", 0.45); + pkg->AddParam<>("cfl", cfl); + + Real vx = pin->GetOrAddReal("advection", "vx", 1.0); + Real vy = pin->GetOrAddReal("advection", "vy", 0.0); + Real vz = pin->GetOrAddReal("advection", "vz", 0.0); + pkg->AddParam<>("vx", vx); + pkg->AddParam<>("vy", vy); + pkg->AddParam<>("vz", vz); + + Real refine_tol = pin->GetOrAddReal("advection", "refine_tol", 0.3); + pkg->AddParam<>("refine_tol", refine_tol); + Real derefine_tol = pin->GetOrAddReal("advection", "derefine_tol", 0.03); + pkg->AddParam<>("derefine_tol", derefine_tol); + + auto profile_str = pin->GetOrAddString("advection", "profile", "wave"); + if (!((profile_str == "wave") || (profile_str == "smooth_gaussian") || (profile_str == "hard_sphere") || + (profile_str == "block"))) { + PARTHENON_FAIL(("Unknown profile in advection example: " + profile_str).c_str()); + } + pkg->AddParam<>("profile", profile_str); + + bool do_regular_advection = pin->GetOrAddBoolean("advection", "do_regular_advection", true); + pkg->AddParam<>("do_regular_advection", do_regular_advection); + Metadata m({Metadata::Cell, Metadata::Independent, Metadata::WithFluxes, Metadata::FillGhost}); + pkg->AddField(m); + + bool do_CT_advection = pin->GetOrAddBoolean("advection", "do_CT_advection", true); + pkg->AddParam<>("do_CT_advection", do_CT_advection); + if (do_CT_advection) { + auto m = Metadata({Metadata::Face, Metadata::Independent, Metadata::WithFluxes, Metadata::FillGhost}); + m.RegisterRefinementOps(); + pkg->AddField(m); + pkg->AddField(m); + pkg->AddField( + Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{4})); + pkg->AddField( + Metadata({Metadata::Face, Metadata::Derived, Metadata::OneCopy}, std::vector{2})); + pkg->AddField( + Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{3})); + pkg->AddField( + Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{3})); + pkg->AddField(Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy})); + pkg->AddField(Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy})); + } + + pkg->CheckRefinementMesh = CheckRefinementMesh; + pkg->EstimateTimestepMesh = EstimateTimestep; + pkg->FillDerivedMesh = FillDerived; + return pkg; +} + +void CheckRefinementMesh(MeshData *md, parthenon::ParArray1D &amr_tags) { + std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + auto do_regular_advection = pkg->Param("do_regular_advection"); + if (do_regular_advection) { + // refine on advected, for example. could also be a derived quantity + static auto desc = parthenon::MakePackDescriptor(md); + auto pack = desc.GetPack(md); + + const auto &refine_tol = pkg->Param("refine_tol"); + const auto &derefine_tol = pkg->Param("derefine_tol"); + + auto ib = md->GetBoundsI(IndexDomain::entire); + auto jb = md->GetBoundsJ(IndexDomain::entire); + auto kb = md->GetBoundsK(IndexDomain::entire); + auto scatter_tags = amr_tags.ToScatterView(); + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, 0, 0, 0, pack.GetNBlocks() - 1, 0, pack.GetMaxNumberOfVars() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t team_member, const int b, const int n, const int k) { + if (n > pack.GetUpperBound(b)) return; + typename Kokkos::MinMax::value_type minmax; + par_reduce_inner( + parthenon::inner_loop_pattern_ttr_tag, team_member, jb.s, jb.e, ib.s, ib.e, + [&](const int j, const int i, typename Kokkos::MinMax::value_type &lminmax) { + lminmax.min_val = (pack(b, n, k, j, i) < lminmax.min_val ? pack(b, n, k, j, i) : lminmax.min_val); + lminmax.max_val = (pack(b, n, k, j, i) > lminmax.max_val ? pack(b, n, k, j, i) : lminmax.max_val); + }, + Kokkos::MinMax(minmax)); + + auto tags_access = scatter_tags.access(); + auto flag = AmrTag::same; + if (minmax.max_val > refine_tol && minmax.min_val < derefine_tol) flag = AmrTag::refine; + if (minmax.max_val < derefine_tol) flag = AmrTag::derefine; + tags_access(b).update(flag); + }); + amr_tags.ContributeScatter(scatter_tags); + } +} + +Real EstimateTimestep(MeshData *md) { + std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + const auto &cfl = pkg->Param("cfl"); + const auto &vx = pkg->Param("vx"); + const auto &vy = pkg->Param("vy"); + const auto &vz = pkg->Param("vz"); + + static auto desc = parthenon::MakePackDescriptor(md); + auto pack = desc.GetPack(md); + + IndexRange ib = md->GetBoundsI(IndexDomain::interior); + IndexRange jb = md->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md->GetBoundsK(IndexDomain::interior); + + // This is obviously overkill for this constant velocity problem + Real min_dt; + parthenon::par_reduce( + parthenon::loop_pattern_mdrange_tag, PARTHENON_AUTO_LABEL, DevExecSpace(), 0, pack.GetNBlocks() - 1, kb.s, kb.e, + jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int b, const int k, const int j, const int i, Real &lmin_dt) { + auto &coords = pack.GetCoordinates(b); + lmin_dt = std::min(lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vx))); + lmin_dt = std::min(lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vy))); + lmin_dt = std::min(lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vz))); + }, + Kokkos::Min(min_dt)); + + return cfl * min_dt / 2.0; +} + +TaskStatus FillDerived(MeshData *md) { + static auto desc = parthenon::MakePackDescriptor(md); + auto pack = desc.GetPack(md); + + std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + + IndexRange ib = md->GetBoundsI(IndexDomain::interior); + IndexRange jb = md->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md->GetBoundsK(IndexDomain::interior); + const int ndim = md->GetMeshPointer()->ndim; + const int nghost = parthenon::Globals::nghost; + + auto do_CT_advection = pkg->Param("do_CT_advection"); + if (do_CT_advection) { + using TE = parthenon::TopologicalElement; + parthenon::par_for( + PARTHENON_AUTO_LABEL, 0, pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + pack(b, Conserved::C_cc(0), k, j, i) = + 0.5 * (pack(b, TE::F1, Conserved::C(), k, j, i) + pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0))); + pack(b, Conserved::C_cc(1), k, j, i) = + 0.5 * (pack(b, TE::F2, Conserved::C(), k, j, i) + pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i)); + pack(b, Conserved::C_cc(2), k, j, i) = + 0.5 * (pack(b, TE::F3, Conserved::C(), k, j, i) + pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i)); + auto &coords = pack.GetCoordinates(b); + pack(b, Conserved::divC(), k, j, i) = + (pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0)) - pack(b, TE::F1, Conserved::C(), k, j, i)) / + coords.Dxc(k, j, i) + + (pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i) - pack(b, TE::F2, Conserved::C(), k, j, i)) / + coords.Dxc(k, j, i) + + (pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i) - pack(b, TE::F3, Conserved::C(), k, j, i)) / + coords.Dxc(k, j, i); + + pack(b, Conserved::D_cc(0), k, j, i) = + 0.5 * (pack(b, TE::F1, Conserved::D(), k, j, i) + pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0))); + pack(b, Conserved::D_cc(1), k, j, i) = + 0.5 * (pack(b, TE::F2, Conserved::D(), k, j, i) + pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i)); + pack(b, Conserved::D_cc(2), k, j, i) = + 0.5 * (pack(b, TE::F3, Conserved::D(), k, j, i) + pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i)); + pack(b, Conserved::divD(), k, j, i) = + (pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0)) - pack(b, TE::F1, Conserved::D(), k, j, i)) / + coords.Dxc(k, j, i) + + (pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i) - pack(b, TE::F2, Conserved::D(), k, j, i)) / + coords.Dxc(k, j, i) + + (pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i) - pack(b, TE::F3, Conserved::D(), k, j, i)) / + coords.Dxc(k, j, i); + }); + } + return TaskStatus::complete; +} +} // namespace advection_package diff --git a/example/imex_advection/advection/advection_package.hpp b/example/imex_advection/advection/advection_package.hpp new file mode 100644 index 000000000000..62678a8c9158 --- /dev/null +++ b/example/imex_advection/advection/advection_package.hpp @@ -0,0 +1,274 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== +#ifndef SRC_ADVECTION_PACKAGE_HPP_ +#define SRC_ADVECTION_PACKAGE_HPP_ + +#include +#include +#include +#include + +#include +#include +#include + +#include "../utilities/variable_macro.hpp" + +namespace advection_package { +using namespace parthenon::package::prelude; + +namespace Conserved { +VARIABLE(advection, phi); +VARIABLE(advection, C); +VARIABLE(advection, D); +VARIABLE(advection, recon); +VARIABLE(advection, recon_f); +VARIABLE(advection, C_cc); +VARIABLE(advection, D_cc); +VARIABLE(advection, divC); +VARIABLE(advection, divD); +} // namespace Conserved + +std::shared_ptr Initialize(ParameterInput *pin); +AmrTag CheckRefinement(MeshBlockData *rc); +void CheckRefinementMesh(MeshData *md, parthenon::ParArray1D &amr_tags); +Real EstimateTimestep(MeshData *md); +TaskStatus FillDerived(MeshData *md); + +template +TaskStatus CalculateFluxes(pack_desc_t &desc, parthenon::TopologicalElement FACE, parthenon::CellLevel cl, + MeshData *md) { + using TE = parthenon::TopologicalElement; + + std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + + // Pull out velocity and piecewise constant reconstruction offsets + // for the given direction + Real v; + int ioff{0}, joff{0}, koff{0}; + if (FACE == TE::F1) { + v = pkg->Param("vx"); + if (v > 0) ioff = -1; + } else if (FACE == TE::F2) { + v = pkg->Param("vy"); + if (v > 0) joff = -1; + } else if (FACE == TE::F3) { + v = pkg->Param("vz"); + if (v > 0) koff = -1; + } + + auto pack = desc.GetPack(md); + + IndexRange ib = md->GetBoundsI(cl, IndexDomain::interior, FACE); + IndexRange jb = md->GetBoundsJ(cl, IndexDomain::interior, FACE); + IndexRange kb = md->GetBoundsK(cl, IndexDomain::interior, FACE); + constexpr int unused_scratch_size = 0; + constexpr int unused_scratch_level = 1; + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, unused_scratch_size, unused_scratch_level, 0, pack.GetNBlocks() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); + for (int l = pack.GetLowerBound(b); l <= pack.GetUpperBound(b); ++l) { + parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { + const auto [j, i] = idxer(idx); + // Calculate the flux using upwind donor cell reconstruction + pack.flux(b, FACE, l, k, j, i) = v * pack(b, l, k + koff, j + joff, i + ioff); + }); + } + }); + return TaskStatus::complete; +} + +inline TaskStatus ImplicitSourceUpdate(Real dt, MeshData *md_in, MeshData *md_out) { + using TE = parthenon::TopologicalElement; + + auto desc = parthenon::MakePackDescriptor(md_in); + auto pack_in = desc.GetPack(md_in); + auto pack_out = desc.GetPack(md_out); + + const Real tau = 0.5; + const Real phi0 = 0.0; + + IndexRange ib = md_in->GetBoundsI(IndexDomain::interior); + IndexRange jb = md_in->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md_in->GetBoundsK(IndexDomain::interior); + parthenon::par_for( + PARTHENON_AUTO_LABEL, 0, pack_in.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + pack_out(b, Conserved::phi(), k, j, i) = + (pack_in(b, Conserved::phi(), k, j, i) + dt / tau * phi0) / (1.0 + dt / tau); + }); + return TaskStatus::complete; +} + +inline TaskStatus Source(MeshData *md_in, MeshData *md_out) { + using TE = parthenon::TopologicalElement; + + auto desc = parthenon::MakePackDescriptor(md_in); + auto pack_in = desc.GetPack(md_in); + auto pack_out = desc.GetPack(md_out); + + const Real tau = 0.5; + const Real phi0 = 0.0; + + IndexRange ib = md_in->GetBoundsI(IndexDomain::interior); + IndexRange jb = md_in->GetBoundsJ(IndexDomain::interior); + IndexRange kb = md_in->GetBoundsK(IndexDomain::interior); + parthenon::par_for( + PARTHENON_AUTO_LABEL, 0, pack_in.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + pack_out(b, Conserved::phi(), k, j, i) = (phi0 - pack_in(b, Conserved::phi(), k, j, i)) / tau; + }); + return TaskStatus::complete; +} + +template +TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon::CellLevel cl, Real fac, + MeshData *md) { + using TE = parthenon::TopologicalElement; + using recon = Conserved::recon; + using recon_f = Conserved::recon_f; + + int ndim = md->GetParentPointer()->ndim; + static auto desc = + parthenon::MakePackDescriptor(md, {}, {parthenon::PDOpt::WithFluxes}); + auto pack = desc.GetPack(md); + + // 1. Reconstruct the component of the flux field pointing in the direction of + // edge in the quartants of the chosen edge + TE fe; + if (edge == TE::E1) fe = TE::F1; + if (edge == TE::E2) fe = TE::F2; + if (edge == TE::E3) fe = TE::F3; + IndexRange ib = md->GetBoundsI(cl, IndexDomain::interior, TE::CC); + IndexRange jb = md->GetBoundsJ(cl, IndexDomain::interior, TE::CC); + IndexRange kb = md->GetBoundsK(cl, IndexDomain::interior, TE::CC); + int koff = edge == TE::E3 ? ndim > 2 : 0; + int joff = edge == TE::E2 ? ndim > 1 : 0; + int ioff = edge == TE::E1 ? ndim > 0 : 0; + constexpr int scratch_size = 0; + constexpr int scratch_level = 1; + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s - (ndim > 2), kb.e + (ndim > 2), + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s - (ndim > 1), jb.e + (ndim > 1)}, {ib.s - (ndim > 0), ib.e + (ndim > 0)}); + if (pack.GetLowerBound(b, flux_var()) <= pack.GetUpperBound(b, flux_var())) { + parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { + const auto [j, i] = idxer(idx); + // Piecewise linear in the orthogonal directions, could do + // something better here + pack(b, TE::CC, recon(0), k, j, i) = + 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + pack(b, TE::CC, recon(1), k, j, i) = + 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + pack(b, TE::CC, recon(2), k, j, i) = + 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + pack(b, TE::CC, recon(3), k, j, i) = + 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + }); + } + }); + + // 2. Calculate the quartant averaged flux + koff = edge != TE::E3 ? ndim > 2 : 0; + joff = edge != TE::E2 ? ndim > 1 : 0; + ioff = edge != TE::E1 ? ndim > 0 : 0; + Real npoints = (koff + 1) * (joff + 1) * (ioff + 1); + ib = md->GetBoundsI(cl, IndexDomain::interior, edge); + jb = md->GetBoundsJ(cl, IndexDomain::interior, edge); + kb = md->GetBoundsK(cl, IndexDomain::interior, edge); + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); + if (pack.GetLowerBound(b, var()) <= pack.GetUpperBound(b, var())) { + parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { + const auto [j, i] = idxer(idx); + pack.flux(b, edge, var(), k, j, i) = 0.0; + for (int dk = -koff; dk < 1; ++dk) + for (int dj = -joff; dj < 1; ++dj) + for (int di = -ioff; di < 1; ++di) { + // TODO(LFR): Pick out the correct component of the + // reconstructed flux, current version is not an issue for + // piecewise constant flux though. + pack.flux(b, edge, var(), k, j, i) += pack(b, TE::CC, recon(0), k + dk, j + dj, i + di); + } + pack.flux(b, edge, var(), k, j, i) /= npoints; + pack.flux(b, edge, var(), k, j, i) *= fac; + }); + } + }); + + // 3. Reconstruct the transverse components of the advected field at the edge + std::vector faces{TE::F2, TE::F3}; + if (edge == TE::E2) faces = {TE::F3, TE::F1}; + if (edge == TE::E3) faces = {TE::F1, TE::F2}; + for (auto f : faces) { + ib = md->GetBoundsI(cl, IndexDomain::interior, f); + jb = md->GetBoundsJ(cl, IndexDomain::interior, f); + kb = md->GetBoundsK(cl, IndexDomain::interior, f); + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s - (ndim > 2), + kb.e + (ndim > 2), KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s - (ndim > 1), jb.e + (ndim > 1)}, {ib.s - (ndim > 0), ib.e + (ndim > 0)}); + if (pack.GetLowerBound(b, var()) <= pack.GetUpperBound(b, var())) { + parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { + const auto [j, i] = idxer(idx); + // Piecewise linear in the orthogonal directions, could do + // something better here + pack(b, f, recon_f(0), k, j, i) = pack(b, f, var(), k, j, i); + pack(b, f, recon_f(1), k, j, i) = pack(b, f, var(), k, j, i); + }); + } + }); + } + + // 4. Add the diffusive piece to the numerical flux, which is proportional to + // the curl + ib = md->GetBoundsI(cl, IndexDomain::interior, edge); + jb = md->GetBoundsJ(cl, IndexDomain::interior, edge); + kb = md->GetBoundsK(cl, IndexDomain::interior, edge); + TE f1 = faces[0]; + TE f2 = faces[1]; + std::array d1{ndim > 0, ndim > 1, ndim > 2}; + std::array d2 = d1; + d1[static_cast(edge) % 3] = 0; + d1[static_cast(f1) % 3] = 0; + d2[static_cast(edge) % 3] = 0; + d2[static_cast(f2) % 3] = 0; + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); + if (pack.GetLowerBound(b, var()) <= pack.GetUpperBound(b, var())) { + parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { + const auto [j, i] = idxer(idx); + pack.flux(b, edge, var(), k, j, i) += + 0.5 * (pack(b, f1, recon_f(0), k, j, i) - pack(b, f1, recon_f(1), k - d1[2], j - d1[1], i - d1[0])); + pack.flux(b, edge, var(), k, j, i) -= + 0.5 * (pack(b, f2, recon_f(0), k, j, i) - pack(b, f2, recon_f(1), k - d2[2], j - d2[1], i - d2[0])); + }); + } + }); + + // Add in the diffusive component + return TaskStatus::complete; +} + +} // namespace advection_package + +#endif // SRC_ADVECTION_PACKAGE_HPP_ diff --git a/example/imex_advection/driver.cpp b/example/imex_advection/driver.cpp new file mode 100644 index 000000000000..b4810a0641a7 --- /dev/null +++ b/example/imex_advection/driver.cpp @@ -0,0 +1,186 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#include +#include +#include + +// Local Includes +#include "advection/advection_package.hpp" +#include "amr_criteria/refinement_package.hpp" +#include "bvals/comms/bvals_in_one.hpp" +#include "driver.hpp" +#include "interface/metadata.hpp" +#include "interface/state_descriptor.hpp" +#include "interface/update.hpp" +#include "mesh/meshblock_pack.hpp" +#include "parthenon/driver.hpp" +#include "prolong_restrict/prolong_restrict.hpp" +#include "utilities/stokes.hpp" + +using namespace parthenon::driver::prelude; + +namespace scalar_imex { + +// Load the required sub-packages +Packages_t ProcessPackages(std::unique_ptr &pin) { + Packages_t packages; + + if (pin->GetOrAddBoolean("scalar_imex", "advection", true)) packages.Add(advection_package::Initialize(pin.get())); + + auto app = std::make_shared("scalar_imex_app"); + packages.Add(app); + + return packages; +} + +// *************************************************// +// define the application driver. in this case, *// +// that mostly means defining the MakeTaskList *// +// function. *// +// *************************************************// +ScalarIMEXDriver::ScalarIMEXDriver(ParameterInput *pin, ApplicationInput *app_in, Mesh *pm) + : parthenon::MultiStageDriverGeneric(pin, app_in, pm) { + + // fail if these are not specified in the input file + pin->CheckRequired("parthenon/mesh", "ix1_bc"); + pin->CheckRequired("parthenon/mesh", "ox1_bc"); + pin->CheckRequired("parthenon/mesh", "ix2_bc"); + pin->CheckRequired("parthenon/mesh", "ox2_bc"); + + // warn if these fields aren't specified in the input file + pin->CheckDesired("parthenon/mesh", "refinement"); + pin->CheckDesired("parthenon/mesh", "numlevel"); + + // Determine which packages to include in driver, allowing for packages to be loaded but not + // run in the driver + do_advection = pin->GetOrAddBoolean("scalar_imex", "do_advection", pin->GetBoolean("scalar_imex", "advection")); +} + +// See the advection.hpp declaration for a description of how this function gets +// called. +TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int stage) { + using namespace parthenon::Update; + TaskCollection tc; + TaskID none(0); + + const Real dt = tm.dt; + const int ndim = pmesh->ndim; + + auto partitions = pmesh->GetDefaultBlockPartitions(); + TaskRegion &single_tasklist_per_pack_region = tc.AddRegion(partitions.size()); + + // Integration steps: + // for i = 1, ..., nstages + // 1. if i == 1, copy data from base register to first stage register + // 2. Do *local* implicit update from current state of register using U^(i) = + // U^(i)* + dt * \tilde a_{ii} R(U^(i)) + // 3. Communicate and fill derived on U^(i) + // 4. Calculate partial_x F(U^(i)) and R(U^(i)) + // 5. for j = i + 1, ..., nstages + // a. if i == 1, add to base register and store in j register, otherwise + // add and store in j register b. Add -dt * \tilde a_{ji} * partial_x + // F(U^(i)) + dt * a_{ji} * R(U^(i)) + // 6. Update U^base += -dt * \tilde b_{i} * partial_x F(U^(i)) + dt * b_{i} * + // R(U^(i)) + // 7. if i == nstages, communicate and fill derived on base register + + using namespace advection_package::Conserved; + static auto desc_phi = parthenon::MakePackDescriptor( + pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Cell}, + {parthenon::PDOpt::WithFluxes}); + using pack_desc_phi_t = decltype(desc_phi); + + for (int i = 0; i < partitions.size(); i++) { + auto &tl = single_tasklist_per_pack_region[i]; + auto &mbase = pmesh->mesh_data.Add("base", partitions[i]); + auto &mc0 = pmesh->mesh_data.Add(integrator->GetStageName(stage), mbase); + using namespace advection_package::Conserved; + + auto set_stage0 = none; + if (stage == 1) { + // Copy base data into this stage, need to be careful that it is a full + // copy + set_stage0 = tl.AddTask(none, WeightedSumDataAll, mbase.get(), mbase.get(), 1.0, 0.0, mc0.get()); + } + + // Do implicit update here + auto implicit_update = set_stage0; + const auto stage_dt = integrator->a(stage, stage) * dt; + if (do_advection) + implicit_update = tl.AddTask(set_stage0, advection_package::ImplicitSourceUpdate, stage_dt, mc0.get(), mc0.get()); + + // Update state of current stage here, as this is the first place it is + // finalized + auto boundaries_stage = parthenon::AddBoundaryExchangeTasks(implicit_update, tl, mc0, pmesh->multilevel); + + auto fill_derived_stage = tl.AddTask(boundaries_stage, parthenon::Update::FillDerived>, mc0.get()); + + using TT = parthenon::TopologicalType; + using TE = parthenon::TopologicalElement; + std::vector faces{TE::F1}; + if (pmesh->ndim > 1) faces.push_back(TE::F2); + if (pmesh->ndim > 2) faces.push_back(TE::F3); + auto flx = none; + for (auto face : faces) { + if (do_advection) + flx = flx | tl.AddTask(fill_derived_stage, advection_package::CalculateFluxes, desc_phi, face, + parthenon::CellLevel::same, mc0.get()); + } + + auto set_flux = parthenon::AddFluxCorrectionTasks(flx, tl, mc0, pmesh->multilevel); + + auto &mdudt_F = pmesh->mesh_data.Add("dUdtF", mbase); + auto &mdudt_R = pmesh->mesh_data.Add("dUdtR", mbase); + auto flux_div = tl.AddTask(set_flux, StokesAll, mc0.get(), mdudt_F.get()); + auto source = set_flux; + if (do_advection) source = source | tl.AddTask(set_flux, advection_package::Source, mc0.get(), mdudt_R.get()); + + auto update_stages = flux_div; + for (int stage2 = stage + 1; stage2 <= integrator->nstages; ++stage2) { + auto label2 = integrator->GetStageName(stage2); + auto &md_in = stage == 1 ? pmesh->mesh_data.Add("base", partitions[i]) : pmesh->mesh_data.Add(label2, mbase); + auto &md_out = pmesh->mesh_data.Add(label2, mbase); + auto add_flux_div = tl.AddTask(source, WeightedSumDataAll, md_in.get(), mdudt_F.get(), 1.0, + integrator->at(stage2, stage) * dt, md_out.get()); + auto add_source = tl.AddTask(add_flux_div, WeightedSumDataAll, md_out.get(), mdudt_R.get(), 1.0, + integrator->a(stage2, stage) * dt, md_out.get()); + update_stages = update_stages | add_source; + } + + // Add this contribution to the base stage + auto add_flux_div = tl.AddTask(source, WeightedSumDataAll, mbase.get(), mdudt_F.get(), 1.0, + integrator->bt(stage) * dt, mbase.get()); + auto add_source = tl.AddTask(add_flux_div, WeightedSumDataAll, mbase.get(), mdudt_R.get(), 1.0, + integrator->b(stage) * dt, mbase.get()); + + if (stage == integrator->nstages) { + auto boundaries = parthenon::AddBoundaryExchangeTasks(add_source, tl, mbase, pmesh->multilevel); + + auto fill_derived = tl.AddTask(boundaries, parthenon::Update::FillDerived>, mbase.get()); + + auto dealloc = tl.AddTask(fill_derived, SparseDealloc, mbase.get()); + auto new_dt = tl.AddTask(dealloc, EstimateTimestep>, mbase.get()); + if (pmesh->adaptive) { + auto tag_refine = tl.AddTask(new_dt, parthenon::Refinement::Tag>, mbase.get()); + } + } + } + + return tc; +} + +} // namespace scalar_imex diff --git a/example/imex_advection/driver.hpp b/example/imex_advection/driver.hpp new file mode 100644 index 000000000000..56d833a8d60f --- /dev/null +++ b/example/imex_advection/driver.hpp @@ -0,0 +1,53 @@ +//======================================================================================== +// (C) (or copyright) 2020-2025. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#ifndef SRC_XMHD_DRIVER_HPP_ +#define SRC_XMHD_DRIVER_HPP_ + +#include +#include + +#include +#include +#include + +#include "utilities/imexrk_integrator.hpp" + +namespace scalar_imex { +using namespace parthenon::driver::prelude; + +class ScalarIMEXDriver : public parthenon::MultiStageDriverGeneric { + bool do_hydro; + bool do_advection; + bool do_em; + bool do_scalar_imex; + + public: + ScalarIMEXDriver(ParameterInput *pin, ApplicationInput *app_in, Mesh *pm); + // This next function essentially defines the driver. + // Call graph looks like + // main() + // EvolutionDriver::Execute (driver.cpp) + // MultiStageBlockTaskDriver::Step (multistage.cpp) + // DriverUtils::ConstructAndExecuteTaskLists (driver.hpp) + // ScalarIMEXDriver::MakeTaskCollection (advection_driver.cpp) + TaskCollection MakeTaskCollection(BlockList_t &blocks, int stage) final; +}; + +parthenon::Packages_t ProcessPackages(std::unique_ptr &pin); + +} // namespace scalar_imex +#endif // SRC_XMHD_DRIVER_HPP_ diff --git a/example/imex_advection/main.cpp b/example/imex_advection/main.cpp new file mode 100644 index 000000000000..37523db0e40e --- /dev/null +++ b/example/imex_advection/main.cpp @@ -0,0 +1,66 @@ +//======================================================================================== +// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#include "parthenon_manager.hpp" + +#include "driver.hpp" +#include "pgen/pgen.hpp" + +int main(int argc, char *argv[]) { + using parthenon::ParthenonManager; + using parthenon::ParthenonStatus; + ParthenonManager pman; + + // Redefine parthenon defaults + pman.app_input->ProcessPackages = scalar_imex::ProcessPackages; + + // call ParthenonInit to initialize MPI and Kokkos, parse the input deck, and + // set up + auto manager_status = pman.ParthenonInitEnv(argc, argv); + if (manager_status == ParthenonStatus::complete) { + pman.ParthenonFinalize(); + return 0; + } + if (manager_status == ParthenonStatus::error) { + pman.ParthenonFinalize(); + return 1; + } + + const auto problem = pman.pinput->GetOrAddString("scalar_imex", "problem_id", "unset"); + if (problem == "advection") { + pman.app_input->ProblemGenerator = scalar_advection::ProblemGenerator; + } else { + PARTHENON_FAIL("Unknown problem generator."); + } + + // Now that ParthenonInit has been called and setup succeeded, the code can + // now make use of MPI and Kokkos. This needs to be scoped so that the driver + // object is destructed before Finalize + pman.ParthenonInitPackagesAndMesh(); + { + // Initialize the driver + scalar_imex::ScalarIMEXDriver driver(pman.pinput.get(), pman.app_input.get(), pman.pmesh.get()); + + // This line actually runs the simulation + auto driver_status = driver.Execute(); + } + // call MPI_Finalize and Kokkos::finalize if necessary + pman.ParthenonFinalize(); + + // MPI and Kokkos can no longer be used + + return (0); +} diff --git a/example/imex_advection/parthinput.advection b/example/imex_advection/parthinput.advection new file mode 100644 index 000000000000..f80a4cb1b504 --- /dev/null +++ b/example/imex_advection/parthinput.advection @@ -0,0 +1,77 @@ +# ======================================================================================== +# (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved. +# +# This program was produced under U.S. Government contract 89233218CNA000001 for Los +# Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC +# for the U.S. Department of Energy/National Nuclear Security Administration. All rights +# in the program are reserved by Triad National Security, LLC, and the U.S. Department +# of Energy/National Nuclear Security Administration. The Government is granted for +# itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide +# license in this material to reproduce, prepare derivative works, distribute copies to +# the public, perform publicly and display publicly, and to permit others to do so. +# ======================================================================================== + + +problem_id = advection + + +refinement = adaptive +numlevel = 3 + +nx1 = 64 +x1min = -0.5 +x1max = 0.5 +ix1_bc = periodic +ox1_bc = periodic + +nx2 = 64 +x2min = -0.5 +x2max = 0.5 +ix2_bc = periodic +ox2_bc = periodic + +nx3 = 1 +x3min = -0.5 +x3max = 0.5 +ix3_bc = periodic +ox3_bc = periodic + + +nx1 = 16 +nx2 = 16 +nx3 = 1 + + +nlim = -1 +tlim = 1.0 +integrator = SSP2-(2,2,2) +ncycle_out_mesh = -10000 + + +problem_id = advection +advection = true +hydro = false + + +cfl = 0.45 +vx = 1.0 +vy = 1.0 +vz = 1.0 +profile = hard_sphere + +refine_tol = 0.3 # control the package specific refinement tagging function +derefine_tol = 0.03 + +# Include advection equation for regular resolution CC fields +do_regular_advection = true +# Include vector wave equation using CT +do_CT_advection = false + + +file_type = rst +dt = 0.05 + + +file_type = hdf5 +dt = 0.05 +variables = advection.phi diff --git a/example/imex_advection/pgen/pgen.hpp b/example/imex_advection/pgen/pgen.hpp new file mode 100644 index 000000000000..8757c69c335c --- /dev/null +++ b/example/imex_advection/pgen/pgen.hpp @@ -0,0 +1,27 @@ +#ifndef PGEN_PGEN_HPP_ +#define PGEN_PGEN_HPP_ +//======================================================================================== +// (C) (or copyright) 2020-2025. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#include +#include + +namespace scalar_advection { +using namespace parthenon::driver::prelude; +void ProblemGenerator(MeshBlock *pmb, parthenon::ParameterInput *pin); +} // namespace scalar_advection + +#endif // PGEN_PGEN_HPP_ diff --git a/example/imex_advection/pgen/scalar_advection.cpp b/example/imex_advection/pgen/scalar_advection.cpp new file mode 100644 index 000000000000..63033ba2d2f7 --- /dev/null +++ b/example/imex_advection/pgen/scalar_advection.cpp @@ -0,0 +1,125 @@ +//======================================================================================== +// (C) (or copyright) 2020-2025. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#include +#include +#include + +#include + +#include "../advection/advection_package.hpp" + +using namespace parthenon::package::prelude; +using namespace parthenon; + +// *************************************************// +// redefine some weakly linked parthenon functions *// +// *************************************************// +namespace { +KOKKOS_INLINE_FUNCTION +Real sign(Real a, Real b) { + if (a == 0.0) { + return b > 0.0 ? 1.0 : -1.0; + } + return a > 0.0 ? 1.0 : -1.0; +} +} // namespace +namespace scalar_advection { + +void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) { + using parthenon::MetadataFlag; + + auto &data = pmb->meshblock_data.Get(); + + auto pkg = pmb->packages.Get("advection_package"); + const auto &profile = pkg->Param("profile"); + auto do_regular_advection = pkg->Param("do_regular_advection"); + auto do_CT_advection = pkg->Param("do_CT_advection"); + + auto cellbounds = pmb->cellbounds; + IndexRange ib = cellbounds.GetBoundsI(IndexDomain::interior); + IndexRange jb = cellbounds.GetBoundsJ(IndexDomain::interior); + IndexRange kb = cellbounds.GetBoundsK(IndexDomain::interior); + + auto coords = pmb->coords; + + using namespace advection_package::Conserved; + static auto desc = parthenon::MakePackDescriptor(data.get()); + auto pack = desc.GetPack(data.get()); + + int profile_type; + if (profile == "wave") profile_type = 0; + if (profile == "smooth_gaussian") profile_type = 1; + if (profile == "hard_sphere") profile_type = 2; + if (profile == "block") profile_type = 3; + Real amp = 1.0; + const int b = 0; + const int ndim = pmb->pmy_mesh->ndim; + const int nghost = parthenon::Globals::nghost; + + if (do_regular_advection) { + parthenon::par_for( + PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) { + const int kf = ndim > 2 ? (k - nghost) * 2 + nghost : k; + const int jf = ndim > 1 ? (j - nghost) * 2 + nghost : j; + const int fi = ndim > 0 ? (i - nghost) * 2 + nghost : i; + if (profile_type == 0) { + PARTHENON_FAIL("Profile type wave not implemented."); + } else if (profile_type == 1) { + Real rsq = coords.Xc<1>(i) * coords.Xc<1>(i) + coords.Xc<2>(j) * coords.Xc<2>(j) + + coords.Xc<3>(k) * coords.Xc<3>(k); + pack(b, phi(), k, j, i) = 1. + amp * exp(-100.0 * rsq); + } else if (profile_type == 2) { + Real rsq = coords.Xc<1>(i) * coords.Xc<1>(i) + coords.Xc<2>(j) * coords.Xc<2>(j) + + coords.Xc<3>(k) * coords.Xc<3>(k); + pack(b, phi(), k, j, i) = (rsq < 0.15 * 0.15 ? 1.0 : 0.0); + } else { + pack(b, phi(), k, j, i) = 0.0; + } + }); + } + + if (do_CT_advection) { + ib = cellbounds.GetBoundsI(IndexDomain::interior, TE::F2); + jb = cellbounds.GetBoundsJ(IndexDomain::interior, TE::F2); + kb = cellbounds.GetBoundsK(IndexDomain::interior, TE::F2); + const Real x0 = 0.2; + parthenon::par_for( + PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) { + auto &coords = pack.GetCoordinates(b); + Real xlo = coords.X(k, j, i); + Real xhi = coords.X(k, j, i + 1); + Real Alo = std::abs(xlo) < x0 ? sign(xlo, xhi) * (x0 - std::abs(xlo)) : 0.0; + Real Ahi = std::abs(xhi) < x0 ? sign(xhi, xlo) * (x0 - std::abs(xhi)) : 0.0; + pack(b, TE::F2, C(), k, j, i) = (Alo - Ahi) / (xhi - xlo); + }); + + ib = cellbounds.GetBoundsI(IndexDomain::interior, TE::F3); + jb = cellbounds.GetBoundsJ(IndexDomain::interior, TE::F3); + kb = cellbounds.GetBoundsK(IndexDomain::interior, TE::F3); + parthenon::par_for( + PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) { + auto &coords = pack.GetCoordinates(b); + Real xlo = coords.X(k, j, i); + Real xhi = coords.X(k, j, i + 1); + Real Alo = std::abs(xlo) < x0 ? sign(xlo, xhi) * (x0 - std::abs(xlo)) : 0.0; + Real Ahi = std::abs(xhi) < x0 ? sign(xhi, xlo) * (x0 - std::abs(xhi)) : 0.0; + pack(b, TE::F3, D(), k, j, i) = (Alo - Ahi) / (xhi - xlo); + }); + } +} + +} // namespace scalar_advection diff --git a/example/imex_advection/utilities/imexrk_integrator.cpp b/example/imex_advection/utilities/imexrk_integrator.cpp new file mode 100644 index 000000000000..aa77c397e0c1 --- /dev/null +++ b/example/imex_advection/utilities/imexrk_integrator.cpp @@ -0,0 +1,53 @@ +//======================================================================================== +// (C) (or copyright) 2020-2025. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#include +#include +#include + +#include + +#include "../utilities/imexrk_integrator.hpp" + +namespace scalar_imex { + +using namespace parthenon::package::prelude; + +IMEXRKIntegrator::IMEXRKIntegrator(const std::string &name) { + if (name == "SSP2-(2,2,2)") { + const Real gam = 1.0 - 1.0 / sqrt(2.0); + nstages = 2; + a_ = {{gam, 0.0}, {1.0 - 2.0 * gam, gam}}; + b_ = {0.5, 0.5}; + + at_ = {{0, 0}, {1, 0}}; + bt_ = {0.5, 0.5}; + } else if (name == "SSP3-(3,3,2)") { + const Real gam = 1.0 - 1.0 / sqrt(2.0); + nstages = 3; + a_ = {{gam, 0, 0}, {1 - 2 * gam, gam, 0}, {0.5 - gam, 0, gam}}; + b_ = {1.0 / 6.0, 1.0 / 6.0, 2.0 / 3.0}; + + at_ = {{0, 0, 0}, {1.0, 0, 0}, {0.25, 0.25, 0}}; + bt_ = b_; + } else { + PARTHENON_FAIL("Unknown IMEX-RK integrator."); + } + for (int stage = 1; stage <= nstages; ++stage) + stage_names_.push_back("stage" + std::to_string(stage)); +} + +} // namespace scalar_imex diff --git a/example/imex_advection/utilities/imexrk_integrator.hpp b/example/imex_advection/utilities/imexrk_integrator.hpp new file mode 100644 index 000000000000..4fde2bbac2d2 --- /dev/null +++ b/example/imex_advection/utilities/imexrk_integrator.hpp @@ -0,0 +1,59 @@ +//======================================================================================== +// (C) (or copyright) 2020-2025. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#ifndef SRC_IMEXRK_INTEGRATOR_HPP_ +#define SRC_IMEXRK_INTEGRATOR_HPP_ + +#include +#include + +#include + +namespace scalar_imex { +using namespace parthenon::package::prelude; + +class IMEXRKIntegrator { + public: + explicit IMEXRKIntegrator(const std::string &name); + IMEXRKIntegrator() : IMEXRKIntegrator("SSP2-(2,2,2)") {} + explicit IMEXRKIntegrator(ParameterInput *pin) + : IMEXRKIntegrator(pin->GetOrAddString("parthenon/time", "integrator", "SSP2-(2,2,2)")) {} + + // To conform with other integrators + int nstages; + Real dt; + + Real at(int i, int j) const { return at_[i - 1][j - 1]; } + Real a(int i, int j) const { return a_[i - 1][j - 1]; } + + Real bt(int i) const { return bt_[i - 1]; } + Real b(int i) const { return b_[i - 1]; } + + Real ct(int i) const { return ct_[i - 1]; } + Real c(int i) const { return c_[i - 1]; } + + std::string GetStageName(int i) const { return stage_names_[i - 1]; } + + private: + std::vector> at_, a_; + std::vector bt_, b_; + std::vector ct_, c_; + std::vector stage_names_; +}; + +} // namespace scalar_imex + +#endif // SRC_IMEXRK_INTEGRATOR_HPP_ diff --git a/example/imex_advection/utilities/index_permutation.hpp b/example/imex_advection/utilities/index_permutation.hpp new file mode 100644 index 000000000000..c6f93a161e5f --- /dev/null +++ b/example/imex_advection/utilities/index_permutation.hpp @@ -0,0 +1,120 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== +#ifndef SRC_UTILITIES_INDEX_PERMUTATION_HPP_ +#define SRC_UTILITIES_INDEX_PERMUTATION_HPP_ + +#include + +namespace parthenon { +namespace utils { + +struct IndexingData : public IndexShape { + int ndim; + IndexShape cellbounds; + IndexingData(MeshData *md) + : ndim{md->GetMeshPointer()->ndim}, IndexShape(md->GetBlockData(0)->GetBlockPointer()->cellbounds) {} + + KOKKOS_INLINE_FUNCTION + auto Get3DIndexRangeWithHalo(IndexDomain id, std::array halo, + TopologicalElement te = TopologicalElement::CC) const { + IndexRange ib = GetBoundsI(id, te); + IndexRange jb = GetBoundsJ(id, te); + IndexRange kb = GetBoundsK(id, te); + kb.s -= halo[0] * (ndim > 2); + kb.e += halo[0] * (ndim > 2); + jb.s -= halo[1] * (ndim > 1); + jb.e += halo[1] * (ndim > 1); + ib.s -= halo[2] * (ndim > 0); + ib.e += halo[2] * (ndim > 0); + return std::make_tuple(kb, jb, ib); + } + + KOKKOS_INLINE_FUNCTION + auto Get3DIndexRange(IndexDomain id, TopologicalElement te = TopologicalElement::CC) const { + return Get3DIndexRangeWithHalo(id, {0, 0, 0}, te); + } + + KOKKOS_INLINE_FUNCTION + auto GetReconstructionRange(TopologicalElement flux_te) const { + return Get3DIndexRangeWithHalo(IndexDomain::interior, GetOffsetArray(flux_te)); + } + + KOKKOS_INLINE_FUNCTION + std::array GetOffsetArray(TopologicalElement flux_te) const { + return std::array{TopologicalOffsetK(flux_te) * (ndim > 2), TopologicalOffsetJ(flux_te) * (ndim > 1), + TopologicalOffsetI(flux_te) * (ndim > 0)}; + } +}; + +inline auto Get3DIndexRangeWithHalo(MeshData *md, IndexDomain id, std::array halo, + TopologicalElement te = TopologicalElement::CC) { + const int ndim = md->GetMeshPointer()->ndim; + IndexRange ib = md->GetBoundsI(id, te); + IndexRange jb = md->GetBoundsJ(id, te); + IndexRange kb = md->GetBoundsK(id, te); + kb.s -= halo[0] * (ndim > 2); + kb.e += halo[0] * (ndim > 2); + jb.s -= halo[1] * (ndim > 1); + jb.e += halo[1] * (ndim > 1); + ib.s -= halo[2] * (ndim > 0); + ib.e += halo[2] * (ndim > 0); + return std::make_tuple(kb, jb, ib); +} + +inline auto Get3DIndexRange(MeshData *md, IndexDomain id, TopologicalElement te = TopologicalElement::CC) { + return Get3DIndexRangeWithHalo(md, id, {0, 0, 0}, te); +} + +inline CoordinateDirection PermuteDirection(CoordinateDirection base_dir, CoordinateDirection relative_dir) { + return static_cast(((base_dir - 1 + relative_dir - 1) % 3) + 1); +} + +inline auto GetOffsetsForDirection(MeshData *md, CoordinateDirection dir) { + const int ndim = md->GetMeshPointer()->ndim; + int di = (dir == CoordinateDirection::X1DIR) * (ndim > 0); + int dj = (dir == CoordinateDirection::X2DIR) * (ndim > 1); + int dk = (dir == CoordinateDirection::X3DIR) * (ndim > 2); + return std::make_tuple(dk, dj, di); +} + +inline auto GetPermutedOffsetsForRelativeDirection(MeshData *md, CoordinateDirection base_dir, + CoordinateDirection relative_dir) { + auto absolute_dir = PermuteDirection(base_dir, relative_dir); + return GetOffsetsForDirection(md, absolute_dir); +} + +inline auto GetDirection(TopologicalElement te) { + using TE = TopologicalElement; + if (te == TE::F1 || te == TE::E1) + return CoordinateDirection::X1DIR; + else if (te == TE::F2 || te == TE::E2) + return CoordinateDirection::X2DIR; + else if (te == TE::F3 || te == TE::E3) + return CoordinateDirection::X3DIR; + PARTHENON_FAIL("It makes no sense to be asking for the direction of a node or cell " + "topological element."); + return CoordinateDirection::NODIR; +} + +inline auto GetPermutedTE(CoordinateDirection base_dir, TopologicalElement te) { + const auto absolute_dir = PermuteDirection(base_dir, GetDirection(te)); + return GetTopologicalElements(GetTopologicalType(te))[absolute_dir - 1]; +} + +} // namespace utils +} // namespace parthenon + +#endif // SRC_UTILITIES_INDEX_PERMUTATION_HPP_ diff --git a/example/imex_advection/utilities/reconstruct.hpp b/example/imex_advection/utilities/reconstruct.hpp new file mode 100644 index 000000000000..6cf0fe02cc32 --- /dev/null +++ b/example/imex_advection/utilities/reconstruct.hpp @@ -0,0 +1,72 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#ifndef SRC_UTILITIES_RECONSTRUCT_HPP_ +#define SRC_UTILITIES_RECONSTRUCT_HPP_ + +#include +#include + +#include +#include +#include + +#include "utilities/scratch_pack.hpp" + +namespace scalar_imex { +using namespace parthenon::driver::prelude; + +KOKKOS_FORCEINLINE_FUNCTION +Real mc(const Real dm, const Real dp, const Real alpha = 1.99) { + const Real dc = (dm * dp > 0.0) * 0.5 * (dm + dp); + return std::copysign(std::min(std::fabs(dc), alpha * std::min(std::fabs(dm), std::fabs(dp))), dc); +} + +KOKKOS_INLINE_FUNCTION +void PiecewiseLinear(const Real qm, const Real q0, const Real qp, Real &p, Real &m, const Real slope_limit = 1.99) { + Real dq = qp - q0; + // const Real slope_limit = 0.99 / ((wgt <= 0.5)*wgt + (wgt > 0.5)*(1.0-wgt)); + dq = 0.5 * mc(q0 - qm, dq, slope_limit); + p = q0 + dq; // wgt *dq; + m = q0 - dq; //(1.0 - wgt) * dq; +} + +template +KOKKOS_INLINE_FUNCTION void +Reconstruct(parthenon::team_mbr_t member, const int b, const int k, parthenon::TopologicalElement flux_te, + const parthenon::utils::IndexingData &cellbounds, const pack_t &pack, + parthenon::utils::ScratchPack &wm, parthenon::utils::ScratchPack &wp) { + const auto [kb, jb, ib] = cellbounds.GetReconstructionRange(flux_te); + const auto [kbe, jbe, ibe] = cellbounds.Get3DIndexRange(IndexDomain::entire); + const parthenon::Indexer2D idxer_full({jbe.s, jbe.e}, {ibe.s, ibe.e}); + const int npoints = idxer_full.GetFlatIdx(jb.e, ib.e) - idxer_full.GetFlatIdx(jb.s, ib.s) + 1; + parthenon::Indexer2D idxer_recon({jb.s, jb.e}, {ib.s, ib.e}); + const auto [koff, joff, ioff] = cellbounds.GetOffsetArray(flux_te); + // Do L/R reconstruction across this slab + for (int l = pack.GetLowerBound(b); l <= pack.GetUpperBound(b); ++l) { + Real *pl = &pack(b, l, k - koff, jb.s - joff, ib.s - ioff); + Real *pm = &pack(b, l, k, jb.s, ib.s); + Real *pu = &pack(b, l, k + koff, jb.s + joff, ib.s + ioff); + Real *pwp = &wp(l, jb.s, ib.s); + Real *pwm = &wm(l, jb.s, ib.s); + parthenon::par_for_inner(member, 0, npoints - 1, + [&](const int idx) { PiecewiseLinear(pl[idx], pm[idx], pu[idx], pwp[idx], pwm[idx]); }); + } +} + +} // namespace scalar_imex + +#endif // SRC_UTILITIES_RECONSTRUCT_HPP_ diff --git a/example/imex_advection/utilities/scratch_pack.hpp b/example/imex_advection/utilities/scratch_pack.hpp new file mode 100644 index 000000000000..c207499b4229 --- /dev/null +++ b/example/imex_advection/utilities/scratch_pack.hpp @@ -0,0 +1,70 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== +#ifndef SRC_UTILITIES_SCRATCH_PACK_HPP_ +#define SRC_UTILITIES_SCRATCH_PACK_HPP_ + +#include + +namespace parthenon { +namespace utils { + +template +struct ScratchPack { + using scratch_pad_t = ScratchPad3D; + + template + KOKKOS_INLINE_FUNCTION ScratchPack(scratch_space_t &&scratch_space, const pack_t *ppack, int b, int nj, int ni) + : ppack{ppack}, block{b}, + data(std::forward(scratch_space), ppack->GetUpperBound(b) + 1, nj, ni) {} + + template + KOKKOS_FORCEINLINE_FUNCTION Real &operator()(const Tin &t, int j, int i) { + const int vidx = ppack->GetLowerBound(block, t) + t.idx; + return data(vidx, j, i); + } + + KOKKOS_FORCEINLINE_FUNCTION + Real &operator()(int v, int j, int i) { return data(v, j, i); } + + static std::size_t get_size_in_bytes(const pack_t &pack, int nj, int ni) { + const int nvar = pack.GetUpperBoundHost(0) + 1; // Assuming no sparse fields for the time being + return scratch_pad_t::shmem_size(nvar, nj, ni); + } + + const pack_t *ppack; + int block; + scratch_pad_t data; +}; + +template +KOKKOS_INLINE_FUNCTION void swap(ScratchPack &a, ScratchPack &b) { + auto *tmp = a.data.data(); + a.data.assign_data(b.data.data()); + b.data.assign_data(tmp); + + int t_block = a.block; + a.block = b.block; + b.block = t_block; + + auto *t_ppack = a.ppack; + a.ppack = b.ppack; + b.ppack = t_ppack; +} + +} // namespace utils +} // namespace parthenon + +#endif // SRC_UTILITIES_SCRATCH_PACK_HPP_ diff --git a/example/imex_advection/utilities/stokes.hpp b/example/imex_advection/utilities/stokes.hpp new file mode 100644 index 000000000000..f7a4991f0156 --- /dev/null +++ b/example/imex_advection/utilities/stokes.hpp @@ -0,0 +1,261 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== + +#ifndef SRC_STOKES_HPP_ +#define SRC_STOKES_HPP_ + +#include +#include + +#include +#include +#include + +#include "../utilities/index_permutation.hpp" +#include "../utilities/scratch_pack.hpp" + +namespace scalar_imex { +using namespace parthenon::driver::prelude; + +template +TaskStatus WeightedSumDataElement(parthenon::CellLevel cl, parthenon::TopologicalElement te, pack_desc_t pd, + MeshData *in1, MeshData *in2, Real w1, Real w2, MeshData *out) { + auto pack1 = pd.GetPack(in1); + auto pack2 = pd.GetPack(in2); + auto pack_out = pd.GetPack(out); + + IndexRange ib = in1->GetBoundsI(cl, IndexDomain::entire, te); + IndexRange jb = in1->GetBoundsJ(cl, IndexDomain::entire, te); + IndexRange kb = in1->GetBoundsK(cl, IndexDomain::entire, te); + + constexpr int scratch_size = 0; + constexpr int scratch_level = 1; + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack1.GetNBlocks() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); + PARTHENON_REQUIRE(pack1.GetLowerBound(b) == pack2.GetLowerBound(b), "Packs are different size."); + PARTHENON_REQUIRE(pack1.GetLowerBound(b) == pack_out.GetLowerBound(b), "Packs are different size."); + PARTHENON_REQUIRE(pack1.GetUpperBound(b) == pack2.GetUpperBound(b), "Packs are different size."); + PARTHENON_REQUIRE(pack1.GetUpperBound(b) == pack_out.GetUpperBound(b), "Packs are different size."); + for (int l = pack1.GetLowerBound(b); l <= pack1.GetUpperBound(b); ++l) { + const auto [j, i] = idxer(0); + Real *out = &pack_out(b, te, l, k, j, i); + Real const *const one = &pack1(b, te, l, k, j, i); + Real const *const two = &pack2(b, te, l, k, j, i); + parthenon::par_for_inner(member, 0, idxer.size() - 1, + [&](const int idx) { out[idx] = w1 * one[idx] + w2 * two[idx]; }); + } + }); + return TaskStatus::complete; +} + +template +TaskStatus WeightedSumData(parthenon::CellLevel cl, parthenon::TopologicalType tt, pack_desc_t pd, MeshData *in1, + MeshData *in2, Real w1, Real w2, MeshData *out) { + for (auto te : parthenon::GetTopologicalElements(tt)) + WeightedSumDataElement(cl, te, pd, in1, in2, w1, w2, out); + return TaskStatus::complete; +} + +TaskStatus WeightedSumDataAll(MeshData *in1, MeshData *in2, Real w1, Real w2, MeshData *out) { + auto pmesh = in1->GetMeshPointer(); + constexpr auto cl = parthenon::CellLevel::same; + { + static auto desc = parthenon::MakePackDescriptor( + pmesh->resolved_packages.get(), {parthenon::Metadata::Independent, parthenon::Metadata::Cell}, {}); + if (desc.nvar_tot > 0) WeightedSumData(cl, parthenon::TopologicalType::Cell, desc, in1, in2, w1, w2, out); + } + { + static auto desc = parthenon::MakePackDescriptor( + pmesh->resolved_packages.get(), {parthenon::Metadata::Independent, parthenon::Metadata::Face}, {}); + if (desc.nvar_tot > 0) WeightedSumData(cl, parthenon::TopologicalType::Face, desc, in1, in2, w1, w2, out); + } + { + static auto desc = parthenon::MakePackDescriptor( + pmesh->resolved_packages.get(), {parthenon::Metadata::Independent, parthenon::Metadata::Edge}, {}); + if (desc.nvar_tot > 0) WeightedSumData(cl, parthenon::TopologicalType::Edge, desc, in1, in2, w1, w2, out); + } + return TaskStatus::complete; +} + +template +void StokesZero(parthenon::CellLevel cl, parthenon::TopologicalElement TeVar, pack_desc_t &pd, MeshData *out) { + auto pack_out = pd.GetPack(out); + + IndexRange ib = out->GetBoundsI(cl, IndexDomain::interior, TeVar); + IndexRange jb = out->GetBoundsJ(cl, IndexDomain::interior, TeVar); + IndexRange kb = out->GetBoundsK(cl, IndexDomain::interior, TeVar); + + constexpr int scratch_size = 0; + constexpr int scratch_level = 1; + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack_out.GetNBlocks() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); + for (int l = pack_out.GetLowerBound(b); l <= pack_out.GetUpperBound(b); ++l) { + parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { + const auto [j, i] = idxer(idx); + pack_out(b, TeVar, l, k, j, i) = 0.0; + }); + } + }); +} + +template +void StokesComponent(Real fac, parthenon::CellLevel cl, parthenon::TopologicalElement TeVar, + parthenon::TopologicalElement TeFlux, pack_desc_t &pd, int ndim, MeshData *in, + MeshData *out) { + auto pack_in = pd.GetPack(in); + auto pack_out = pd.GetPack(out); + + IndexRange ib = in->GetBoundsI(cl, IndexDomain::interior, TeVar); + IndexRange jb = in->GetBoundsJ(cl, IndexDomain::interior, TeVar); + IndexRange kb = in->GetBoundsK(cl, IndexDomain::interior, TeVar); + int ioff = TopologicalOffsetI(TeFlux) - TopologicalOffsetI(TeVar); + int joff = TopologicalOffsetJ(TeFlux) - TopologicalOffsetJ(TeVar); + int koff = TopologicalOffsetK(TeFlux) - TopologicalOffsetK(TeVar); + PARTHENON_REQUIRE(ioff == 1 || ioff == 0, "Bad combination of TeVar and TeFlux"); + PARTHENON_REQUIRE(joff == 1 || joff == 0, "Bad combination of TeVar and TeFlux"); + PARTHENON_REQUIRE(koff == 1 || koff == 0, "Bad combination of TeVar and TeFlux"); + PARTHENON_REQUIRE((ioff + joff + koff) == 1, "Bad combination of TeVar and TeFlux"); + koff = ndim > 2 ? koff : 0; + joff = ndim > 1 ? joff : 0; + + constexpr size_t scratch_level = 1; + using scratch_pack_t = parthenon::utils::ScratchPack; + const parthenon::utils::IndexingData cellbounds(in); + using TE = parthenon::TopologicalElement; + // Choose the correct memory size for each topological element + const int njtot = cellbounds.ncellsj(IndexDomain::entire, TeVar == TE::CC ? TE::CC : TE::NN); + const int nitot = cellbounds.ncellsi(IndexDomain::entire, TeVar == TE::CC ? TE::CC : TE::NN); + std::size_t scratch_size_in_bytes = parthenon::ScratchPad2D::shmem_size(njtot, nitot) * (2 + koff); + + parthenon::par_for_outer( + PARTHENON_AUTO_LABEL, scratch_size_in_bytes, scratch_level, 0, pack_out.GetNBlocks() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + auto &coords = pack_in.GetCoordinates(b); + parthenon::Indexer2D idxer_entire({0, njtot - 1}, {0, nitot - 1}); + parthenon::ScratchPad2D Vflux(member.team_scratch(scratch_level), njtot, nitot); + parthenon::ScratchPad2D Vvar(member.team_scratch(scratch_level), njtot, nitot); + parthenon::par_for_inner(member, 0, idxer_entire.size() - 1, [&](const int idx) { + const auto [j, i] = idxer_entire(idx); + Vflux(j, i) = coords.Volume(cl, TeFlux, k, j, i); + Vvar(j, i) = coords.Volume(cl, TeVar, k, j, i); + }); + + parthenon::ScratchPad2D Vfluxu(member.team_scratch(scratch_level), koff ? njtot : 0, koff ? nitot : 0); + if (koff) { + parthenon::par_for_inner(member, 0, idxer_entire.size() - 1, [&](const int idx) { + const auto [j, i] = idxer_entire(idx); + Vfluxu(j, i) = coords.Volume(cl, TeFlux, k + koff, j, i); + }); + } + member.team_barrier(); + + PARTHENON_REQUIRE(pack_in.GetLowerBound(b) == pack_out.GetLowerBound(b), "Packs are different size."); + PARTHENON_REQUIRE(pack_in.GetUpperBound(b) == pack_out.GetUpperBound(b), "Packs are different size."); + + const int npoints = idxer_entire.GetFlatIdx(jb.e, ib.e) - idxer_entire.GetFlatIdx(jb.s, ib.s) + 1; + for (int l = pack_out.GetLowerBound(b); l <= pack_out.GetUpperBound(b); ++l) { + Real const *const flxl = &pack_in.flux(b, TeFlux, l, k, jb.s, ib.s); + Real const *const flxu = &pack_in.flux(b, TeFlux, l, k + koff, jb.s + joff, ib.s + ioff); + Real const *const vfl = &Vflux(jb.s, ib.s); + Real const *const vfu = koff ? &Vfluxu(jb.s + joff, ib.s + ioff) : &Vflux(jb.s + joff, ib.s + ioff); + Real const *const vv = &Vvar(jb.s, ib.s); + Real *out = &pack_out(b, TeVar, l, k, jb.s, ib.s); + + parthenon::par_for_inner(member, 0, npoints - 1, [&](const int idx) { + out[idx] += fac * (flxl[idx] * vfl[idx] - flxu[idx] * vfu[idx]) / vv[idx]; + }); + } + }); +} + +template +TaskStatus Stokes(parthenon::CellLevel cl, parthenon::TopologicalType TtVar, pack_desc_t &pd, int ndim, + MeshData *in, MeshData *out) { + using TE = parthenon::TopologicalElement; + using TT = parthenon::TopologicalType; + + // Get the topological type of the generalized flux associated with the + // with variables of topological type TtVar + TT TtFlx = [TtVar] { + if (TtVar == TT::Cell) { + return TT::Face; + } else if (TtVar == TT::Face) { + return TT::Edge; + } else if (TtVar == TT::Edge) { + return TT::Node; + } else { + PARTHENON_FAIL("Stokes does not work for node variables, as they are " + "zero dimensional."); + return TT::Node; + } + }(); + + auto VarTes = GetTopologicalElements(TtVar); + auto FlxTes = GetTopologicalElements(TtFlx); + for (auto vte : VarTes) { + StokesZero(cl, vte, pd, out); + for (auto fte : FlxTes) { + if (IsSubmanifold(fte, vte)) { + Real fac = 1.0; + if (ndim < 3 && fte == TE::F3) continue; + if (ndim < 2 && fte == TE::F2) continue; + if (TtVar == TT::Face) { + // TODO(LFR): This is untested, need to test in parthenon-mhd + // downstream or add a test involving curls Flip the sign if the + // variable is an X1 face and the edge is an X3 edge, or an X2 face + // ... X1 edge, or an X3 face ... X2 edge + const int indicator = ((static_cast(fte) % 3) - (static_cast(vte) % 3) + 3) % 3; + fac = (indicator == 2) ? -1.0 : 1.0; + } + StokesComponent(fac, cl, vte, fte, pd, ndim, in, out); + } + } + } + return TaskStatus::complete; +} + +TaskStatus StokesAll(MeshData *in, MeshData *out) { + auto pmesh = in->GetMeshPointer(); + const int ndim = pmesh->ndim; + constexpr auto cl = parthenon::CellLevel::same; + { + static auto desc = parthenon::MakePackDescriptor( + pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Cell}, + {parthenon::PDOpt::WithFluxes}); + if (desc.nvar_tot > 0) Stokes(cl, parthenon::TopologicalType::Cell, desc, ndim, in, out); + } + { + static auto desc = parthenon::MakePackDescriptor( + pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Face}, + {parthenon::PDOpt::WithFluxes}); + if (desc.nvar_tot > 0) Stokes(cl, parthenon::TopologicalType::Face, desc, ndim, in, out); + } + { + static auto desc = parthenon::MakePackDescriptor( + pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Edge}, + {parthenon::PDOpt::WithFluxes}); + if (desc.nvar_tot > 0) Stokes(cl, parthenon::TopologicalType::Edge, desc, ndim, in, out); + } + return TaskStatus::complete; +} + +} // namespace scalar_imex + +#endif // SRC_STOKES_HPP_ diff --git a/example/imex_advection/utilities/three_vec.hpp b/example/imex_advection/utilities/three_vec.hpp new file mode 100644 index 000000000000..c7d4e71d31d8 --- /dev/null +++ b/example/imex_advection/utilities/three_vec.hpp @@ -0,0 +1,101 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== +#ifndef SRC_UTILITIES_THREE_VEC_HPP_ +#define SRC_UTILITIES_THREE_VEC_HPP_ + +#include + +namespace parthenon { +namespace utils { + +struct ThreeVec { + std::array vals; + + KOKKOS_FORCEINLINE_FUNCTION + ThreeVec() = default; + + KOKKOS_FORCEINLINE_FUNCTION + ThreeVec(Real a1, Real a2, Real a3) : vals{a1, a2, a3} {} + + template + KOKKOS_FORCEINLINE_FUNCTION ThreeVec(const pack_t &pack, int b, var_t, int k, int j, int i) + : vals{pack(b, var_t(0), k, j, i), pack(b, var_t(1), k, j, i), pack(b, var_t(2), k, j, i)} {} + + KOKKOS_FORCEINLINE_FUNCTION + Real &operator[](CoordinateDirection dir) { return vals[dir - 1]; } + + KOKKOS_FORCEINLINE_FUNCTION + const Real &operator[](CoordinateDirection dir) const { return vals[dir - 1]; } +}; + +KOKKOS_FORCEINLINE_FUNCTION +ThreeVec operator+(const ThreeVec &a, const ThreeVec &b) { + ThreeVec out; + for (const auto dir : {X1DIR, X2DIR, X3DIR}) + out[dir] = a[dir] + b[dir]; + return out; +} + +KOKKOS_FORCEINLINE_FUNCTION +ThreeVec operator-(const ThreeVec &a, const ThreeVec &b) { + ThreeVec out; + for (const auto dir : {X1DIR, X2DIR, X3DIR}) + out[dir] = a[dir] - b[dir]; + return out; +} + +KOKKOS_FORCEINLINE_FUNCTION +ThreeVec operator*(Real a, const ThreeVec &b) { + ThreeVec out; + for (const auto dir : {X1DIR, X2DIR, X3DIR}) + out[dir] = a * b[dir]; + return out; +} + +KOKKOS_FORCEINLINE_FUNCTION +ThreeVec operator*(const ThreeVec &b, Real a) { + ThreeVec out; + for (const auto dir : {X1DIR, X2DIR, X3DIR}) + out[dir] = b[dir] * a; + return out; +} + +KOKKOS_FORCEINLINE_FUNCTION +ThreeVec operator/(const ThreeVec &b, Real a) { + ThreeVec out; + for (const auto dir : {X1DIR, X2DIR, X3DIR}) + out[dir] = b[dir] / a; + return out; +} + +KOKKOS_FORCEINLINE_FUNCTION +Real DotProduct(const ThreeVec &a, const ThreeVec &b) { + return a[X1DIR] * b[X1DIR] + a[X2DIR] * b[X2DIR] + a[X3DIR] * b[X3DIR]; +} + +KOKKOS_FORCEINLINE_FUNCTION +ThreeVec CrossProduct(const ThreeVec &a, const ThreeVec &b) { + ThreeVec out; + out[X1DIR] = a[X2DIR] * b[X3DIR] - a[X3DIR] * b[X2DIR]; + out[X2DIR] = a[X3DIR] * b[X1DIR] - a[X1DIR] * b[X3DIR]; + out[X3DIR] = a[X1DIR] * b[X2DIR] - a[X2DIR] * b[X1DIR]; + return out; +} + +} // namespace utils +} // namespace parthenon + +#endif // SRC_UTILITIES_THREE_VEC_HPP_ diff --git a/example/imex_advection/utilities/variable_macro.hpp b/example/imex_advection/utilities/variable_macro.hpp new file mode 100644 index 000000000000..7a9f66df8ed0 --- /dev/null +++ b/example/imex_advection/utilities/variable_macro.hpp @@ -0,0 +1,29 @@ +//======================================================================================== +// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights +// reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 +// for Los Alamos National Laboratory (LANL), which is operated by Triad +// National Security, LLC for the U.S. Department of Energy/National Nuclear +// Security Administration. All rights in the program are reserved by Triad +// National Security, LLC, and the U.S. Department of Energy/National Nuclear +// Security Administration. The Government is granted for itself and others +// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +// in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do +// so. +//======================================================================================== +#ifndef SRC_UTILITIES_VARIABLE_MACRO_HPP_ +#define SRC_UTILITIES_VARIABLE_MACRO_HPP_ + +#include + +#define VARIABLE(ns, varname) \ + struct varname : public parthenon::variable_names::base_t { \ + template \ + KOKKOS_INLINE_FUNCTION varname(Ts &&...args) \ + : parthenon::variable_names::base_t(std::forward(args)...) {} \ + static std::string name() { return #ns "." #varname; } \ + } + +#endif // SRC_UTILITIES_VARIABLE_MACRO_HPP_ From 43b51b7ddd6dc41513745d31518bad15b6ad1558 Mon Sep 17 00:00:00 2001 From: Luke Roberts Date: Mon, 11 Aug 2025 13:38:10 -0600 Subject: [PATCH 2/3] add a couple comments --- example/imex_advection/driver.cpp | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/example/imex_advection/driver.cpp b/example/imex_advection/driver.cpp index b4810a0641a7..2693a18a272c 100644 --- a/example/imex_advection/driver.cpp +++ b/example/imex_advection/driver.cpp @@ -117,7 +117,7 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta set_stage0 = tl.AddTask(none, WeightedSumDataAll, mbase.get(), mbase.get(), 1.0, 0.0, mc0.get()); } - // Do implicit update here + // Do implicit update here to get conserved variables for the current stage auto implicit_update = set_stage0; const auto stage_dt = integrator->a(stage, stage) * dt; if (do_advection) @@ -128,7 +128,8 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta auto boundaries_stage = parthenon::AddBoundaryExchangeTasks(implicit_update, tl, mc0, pmesh->multilevel); auto fill_derived_stage = tl.AddTask(boundaries_stage, parthenon::Update::FillDerived>, mc0.get()); - + + // Calculate fluxes for the current stage using TT = parthenon::TopologicalType; using TE = parthenon::TopologicalElement; std::vector faces{TE::F1}; @@ -146,9 +147,12 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta auto &mdudt_F = pmesh->mesh_data.Add("dUdtF", mbase); auto &mdudt_R = pmesh->mesh_data.Add("dUdtR", mbase); auto flux_div = tl.AddTask(set_flux, StokesAll, mc0.get(), mdudt_F.get()); + + // Calculate source terms for the current stage auto source = set_flux; if (do_advection) source = source | tl.AddTask(set_flux, advection_package::Source, mc0.get(), mdudt_R.get()); - + + // Add the contribution from this stage to the updates for subsequent stages auto update_stages = flux_div; for (int stage2 = stage + 1; stage2 <= integrator->nstages; ++stage2) { auto label2 = integrator->GetStageName(stage2); @@ -166,7 +170,8 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta integrator->bt(stage) * dt, mbase.get()); auto add_source = tl.AddTask(add_flux_div, WeightedSumDataAll, mbase.get(), mdudt_R.get(), 1.0, integrator->b(stage) * dt, mbase.get()); - + + // Perform the last stage cleanup tasks if (stage == integrator->nstages) { auto boundaries = parthenon::AddBoundaryExchangeTasks(add_source, tl, mbase, pmesh->multilevel); From 7924d35b88415386b54b9288d961fce11cedb8f1 Mon Sep 17 00:00:00 2001 From: Luke Roberts Date: Mon, 11 Aug 2025 13:54:23 -0600 Subject: [PATCH 3/3] format --- .../advection/advection_package.cpp | 123 ++++++++++------ .../advection/advection_package.hpp | 74 ++++++---- example/imex_advection/driver.cpp | 76 ++++++---- example/imex_advection/main.cpp | 3 +- .../imex_advection/pgen/scalar_advection.cpp | 15 +- .../utilities/imexrk_integrator.hpp | 3 +- .../utilities/index_permutation.hpp | 21 ++- .../imex_advection/utilities/reconstruct.hpp | 20 ++- .../imex_advection/utilities/scratch_pack.hpp | 10 +- example/imex_advection/utilities/stokes.hpp | 137 +++++++++++------- .../imex_advection/utilities/three_vec.hpp | 6 +- .../utilities/variable_macro.hpp | 12 +- 12 files changed, 314 insertions(+), 186 deletions(-) diff --git a/example/imex_advection/advection/advection_package.cpp b/example/imex_advection/advection/advection_package.cpp index b30eecc39b85..8773acd994a2 100644 --- a/example/imex_advection/advection/advection_package.cpp +++ b/example/imex_advection/advection/advection_package.cpp @@ -62,36 +62,41 @@ std::shared_ptr Initialize(ParameterInput *pin) { pkg->AddParam<>("derefine_tol", derefine_tol); auto profile_str = pin->GetOrAddString("advection", "profile", "wave"); - if (!((profile_str == "wave") || (profile_str == "smooth_gaussian") || (profile_str == "hard_sphere") || - (profile_str == "block"))) { + if (!((profile_str == "wave") || (profile_str == "smooth_gaussian") || + (profile_str == "hard_sphere") || (profile_str == "block"))) { PARTHENON_FAIL(("Unknown profile in advection example: " + profile_str).c_str()); } pkg->AddParam<>("profile", profile_str); - bool do_regular_advection = pin->GetOrAddBoolean("advection", "do_regular_advection", true); + bool do_regular_advection = + pin->GetOrAddBoolean("advection", "do_regular_advection", true); pkg->AddParam<>("do_regular_advection", do_regular_advection); - Metadata m({Metadata::Cell, Metadata::Independent, Metadata::WithFluxes, Metadata::FillGhost}); + Metadata m( + {Metadata::Cell, Metadata::Independent, Metadata::WithFluxes, Metadata::FillGhost}); pkg->AddField(m); bool do_CT_advection = pin->GetOrAddBoolean("advection", "do_CT_advection", true); pkg->AddParam<>("do_CT_advection", do_CT_advection); if (do_CT_advection) { - auto m = Metadata({Metadata::Face, Metadata::Independent, Metadata::WithFluxes, Metadata::FillGhost}); + auto m = Metadata({Metadata::Face, Metadata::Independent, Metadata::WithFluxes, + Metadata::FillGhost}); m.RegisterRefinementOps(); pkg->AddField(m); pkg->AddField(m); - pkg->AddField( - Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{4})); - pkg->AddField( - Metadata({Metadata::Face, Metadata::Derived, Metadata::OneCopy}, std::vector{2})); - pkg->AddField( - Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{3})); - pkg->AddField( - Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{3})); - pkg->AddField(Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy})); - pkg->AddField(Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy})); + pkg->AddField(Metadata( + {Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{4})); + pkg->AddField(Metadata( + {Metadata::Face, Metadata::Derived, Metadata::OneCopy}, std::vector{2})); + pkg->AddField(Metadata( + {Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{3})); + pkg->AddField(Metadata( + {Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector{3})); + pkg->AddField( + Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy})); + pkg->AddField( + Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy})); } pkg->CheckRefinementMesh = CheckRefinementMesh; @@ -101,7 +106,8 @@ std::shared_ptr Initialize(ParameterInput *pin) { } void CheckRefinementMesh(MeshData *md, parthenon::ParArray1D &amr_tags) { - std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + std::shared_ptr pkg = + md->GetMeshPointer()->packages.Get("advection_package"); auto do_regular_advection = pkg->Param("do_regular_advection"); if (do_regular_advection) { // refine on advected, for example. could also be a derived quantity @@ -116,21 +122,29 @@ void CheckRefinementMesh(MeshData *md, parthenon::ParArray1D &amr_ auto kb = md->GetBoundsK(IndexDomain::entire); auto scatter_tags = amr_tags.ToScatterView(); parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, 0, 0, 0, pack.GetNBlocks() - 1, 0, pack.GetMaxNumberOfVars() - 1, kb.s, kb.e, - KOKKOS_LAMBDA(parthenon::team_mbr_t team_member, const int b, const int n, const int k) { + PARTHENON_AUTO_LABEL, 0, 0, 0, pack.GetNBlocks() - 1, 0, + pack.GetMaxNumberOfVars() - 1, kb.s, kb.e, + KOKKOS_LAMBDA(parthenon::team_mbr_t team_member, const int b, const int n, + const int k) { if (n > pack.GetUpperBound(b)) return; typename Kokkos::MinMax::value_type minmax; par_reduce_inner( parthenon::inner_loop_pattern_ttr_tag, team_member, jb.s, jb.e, ib.s, ib.e, - [&](const int j, const int i, typename Kokkos::MinMax::value_type &lminmax) { - lminmax.min_val = (pack(b, n, k, j, i) < lminmax.min_val ? pack(b, n, k, j, i) : lminmax.min_val); - lminmax.max_val = (pack(b, n, k, j, i) > lminmax.max_val ? pack(b, n, k, j, i) : lminmax.max_val); + [&](const int j, const int i, + typename Kokkos::MinMax::value_type &lminmax) { + lminmax.min_val = + (pack(b, n, k, j, i) < lminmax.min_val ? pack(b, n, k, j, i) + : lminmax.min_val); + lminmax.max_val = + (pack(b, n, k, j, i) > lminmax.max_val ? pack(b, n, k, j, i) + : lminmax.max_val); }, Kokkos::MinMax(minmax)); auto tags_access = scatter_tags.access(); auto flag = AmrTag::same; - if (minmax.max_val > refine_tol && minmax.min_val < derefine_tol) flag = AmrTag::refine; + if (minmax.max_val > refine_tol && minmax.min_val < derefine_tol) + flag = AmrTag::refine; if (minmax.max_val < derefine_tol) flag = AmrTag::derefine; tags_access(b).update(flag); }); @@ -139,7 +153,8 @@ void CheckRefinementMesh(MeshData *md, parthenon::ParArray1D &amr_ } Real EstimateTimestep(MeshData *md) { - std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + std::shared_ptr pkg = + md->GetMeshPointer()->packages.Get("advection_package"); const auto &cfl = pkg->Param("cfl"); const auto &vx = pkg->Param("vx"); const auto &vy = pkg->Param("vy"); @@ -155,13 +170,16 @@ Real EstimateTimestep(MeshData *md) { // This is obviously overkill for this constant velocity problem Real min_dt; parthenon::par_reduce( - parthenon::loop_pattern_mdrange_tag, PARTHENON_AUTO_LABEL, DevExecSpace(), 0, pack.GetNBlocks() - 1, kb.s, kb.e, - jb.s, jb.e, ib.s, ib.e, + parthenon::loop_pattern_mdrange_tag, PARTHENON_AUTO_LABEL, DevExecSpace(), 0, + pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int b, const int k, const int j, const int i, Real &lmin_dt) { auto &coords = pack.GetCoordinates(b); - lmin_dt = std::min(lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vx))); - lmin_dt = std::min(lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vy))); - lmin_dt = std::min(lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vz))); + lmin_dt = std::min( + lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vx))); + lmin_dt = std::min( + lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vy))); + lmin_dt = std::min( + lmin_dt, parthenon::robust::ratio(coords.Dxc(k, j, i), std::abs(vz))); }, Kokkos::Min(min_dt)); @@ -169,11 +187,14 @@ Real EstimateTimestep(MeshData *md) { } TaskStatus FillDerived(MeshData *md) { - static auto desc = parthenon::MakePackDescriptor(md); + static auto desc = + parthenon::MakePackDescriptor( + md); auto pack = desc.GetPack(md); - std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + std::shared_ptr pkg = + md->GetMeshPointer()->packages.Get("advection_package"); IndexRange ib = md->GetBoundsI(IndexDomain::interior); IndexRange jb = md->GetBoundsJ(IndexDomain::interior); @@ -185,35 +206,47 @@ TaskStatus FillDerived(MeshData *md) { if (do_CT_advection) { using TE = parthenon::TopologicalElement; parthenon::par_for( - PARTHENON_AUTO_LABEL, 0, pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, - KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + PARTHENON_AUTO_LABEL, 0, pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, + ib.e, KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { pack(b, Conserved::C_cc(0), k, j, i) = - 0.5 * (pack(b, TE::F1, Conserved::C(), k, j, i) + pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0))); + 0.5 * (pack(b, TE::F1, Conserved::C(), k, j, i) + + pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0))); pack(b, Conserved::C_cc(1), k, j, i) = - 0.5 * (pack(b, TE::F2, Conserved::C(), k, j, i) + pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i)); + 0.5 * (pack(b, TE::F2, Conserved::C(), k, j, i) + + pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i)); pack(b, Conserved::C_cc(2), k, j, i) = - 0.5 * (pack(b, TE::F3, Conserved::C(), k, j, i) + pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i)); + 0.5 * (pack(b, TE::F3, Conserved::C(), k, j, i) + + pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i)); auto &coords = pack.GetCoordinates(b); pack(b, Conserved::divC(), k, j, i) = - (pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0)) - pack(b, TE::F1, Conserved::C(), k, j, i)) / + (pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0)) - + pack(b, TE::F1, Conserved::C(), k, j, i)) / coords.Dxc(k, j, i) + - (pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i) - pack(b, TE::F2, Conserved::C(), k, j, i)) / + (pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i) - + pack(b, TE::F2, Conserved::C(), k, j, i)) / coords.Dxc(k, j, i) + - (pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i) - pack(b, TE::F3, Conserved::C(), k, j, i)) / + (pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i) - + pack(b, TE::F3, Conserved::C(), k, j, i)) / coords.Dxc(k, j, i); pack(b, Conserved::D_cc(0), k, j, i) = - 0.5 * (pack(b, TE::F1, Conserved::D(), k, j, i) + pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0))); + 0.5 * (pack(b, TE::F1, Conserved::D(), k, j, i) + + pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0))); pack(b, Conserved::D_cc(1), k, j, i) = - 0.5 * (pack(b, TE::F2, Conserved::D(), k, j, i) + pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i)); + 0.5 * (pack(b, TE::F2, Conserved::D(), k, j, i) + + pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i)); pack(b, Conserved::D_cc(2), k, j, i) = - 0.5 * (pack(b, TE::F3, Conserved::D(), k, j, i) + pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i)); + 0.5 * (pack(b, TE::F3, Conserved::D(), k, j, i) + + pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i)); pack(b, Conserved::divD(), k, j, i) = - (pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0)) - pack(b, TE::F1, Conserved::D(), k, j, i)) / + (pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0)) - + pack(b, TE::F1, Conserved::D(), k, j, i)) / coords.Dxc(k, j, i) + - (pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i) - pack(b, TE::F2, Conserved::D(), k, j, i)) / + (pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i) - + pack(b, TE::F2, Conserved::D(), k, j, i)) / coords.Dxc(k, j, i) + - (pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i) - pack(b, TE::F3, Conserved::D(), k, j, i)) / + (pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i) - + pack(b, TE::F3, Conserved::D(), k, j, i)) / coords.Dxc(k, j, i); }); } diff --git a/example/imex_advection/advection/advection_package.hpp b/example/imex_advection/advection/advection_package.hpp index 62678a8c9158..7755f68125bd 100644 --- a/example/imex_advection/advection/advection_package.hpp +++ b/example/imex_advection/advection/advection_package.hpp @@ -49,11 +49,12 @@ Real EstimateTimestep(MeshData *md); TaskStatus FillDerived(MeshData *md); template -TaskStatus CalculateFluxes(pack_desc_t &desc, parthenon::TopologicalElement FACE, parthenon::CellLevel cl, - MeshData *md) { +TaskStatus CalculateFluxes(pack_desc_t &desc, parthenon::TopologicalElement FACE, + parthenon::CellLevel cl, MeshData *md) { using TE = parthenon::TopologicalElement; - std::shared_ptr pkg = md->GetMeshPointer()->packages.Get("advection_package"); + std::shared_ptr pkg = + md->GetMeshPointer()->packages.Get("advection_package"); // Pull out velocity and piecewise constant reconstruction offsets // for the given direction @@ -78,7 +79,8 @@ TaskStatus CalculateFluxes(pack_desc_t &desc, parthenon::TopologicalElement FACE constexpr int unused_scratch_size = 0; constexpr int unused_scratch_level = 1; parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, unused_scratch_size, unused_scratch_level, 0, pack.GetNBlocks() - 1, kb.s, kb.e, + PARTHENON_AUTO_LABEL, unused_scratch_size, unused_scratch_level, 0, + pack.GetNBlocks() - 1, kb.s, kb.e, KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); for (int l = pack.GetLowerBound(b); l <= pack.GetUpperBound(b); ++l) { @@ -92,7 +94,8 @@ TaskStatus CalculateFluxes(pack_desc_t &desc, parthenon::TopologicalElement FACE return TaskStatus::complete; } -inline TaskStatus ImplicitSourceUpdate(Real dt, MeshData *md_in, MeshData *md_out) { +inline TaskStatus ImplicitSourceUpdate(Real dt, MeshData *md_in, + MeshData *md_out) { using TE = parthenon::TopologicalElement; auto desc = parthenon::MakePackDescriptor(md_in); @@ -106,8 +109,8 @@ inline TaskStatus ImplicitSourceUpdate(Real dt, MeshData *md_in, MeshData< IndexRange jb = md_in->GetBoundsJ(IndexDomain::interior); IndexRange kb = md_in->GetBoundsK(IndexDomain::interior); parthenon::par_for( - PARTHENON_AUTO_LABEL, 0, pack_in.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, - KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + PARTHENON_AUTO_LABEL, 0, pack_in.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, + ib.e, KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { pack_out(b, Conserved::phi(), k, j, i) = (pack_in(b, Conserved::phi(), k, j, i) + dt / tau * phi0) / (1.0 + dt / tau); }); @@ -128,16 +131,17 @@ inline TaskStatus Source(MeshData *md_in, MeshData *md_out) { IndexRange jb = md_in->GetBoundsJ(IndexDomain::interior); IndexRange kb = md_in->GetBoundsK(IndexDomain::interior); parthenon::par_for( - PARTHENON_AUTO_LABEL, 0, pack_in.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, - KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { - pack_out(b, Conserved::phi(), k, j, i) = (phi0 - pack_in(b, Conserved::phi(), k, j, i)) / tau; + PARTHENON_AUTO_LABEL, 0, pack_in.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, + ib.e, KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + pack_out(b, Conserved::phi(), k, j, i) = + (phi0 - pack_in(b, Conserved::phi(), k, j, i)) / tau; }); return TaskStatus::complete; } template -TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon::CellLevel cl, Real fac, - MeshData *md) { +TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, + parthenon::CellLevel cl, Real fac, MeshData *md) { using TE = parthenon::TopologicalElement; using recon = Conserved::recon; using recon_f = Conserved::recon_f; @@ -145,7 +149,8 @@ TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon:: int ndim = md->GetParentPointer()->ndim; static auto desc = parthenon::MakePackDescriptor(md, {}, {parthenon::PDOpt::WithFluxes}); + advection_package::Conserved::recon_f>( + md, {}, {parthenon::PDOpt::WithFluxes}); auto pack = desc.GetPack(md); // 1. Reconstruct the component of the flux field pointing in the direction of @@ -163,22 +168,28 @@ TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon:: constexpr int scratch_size = 0; constexpr int scratch_level = 1; parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s - (ndim > 2), kb.e + (ndim > 2), + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, + kb.s - (ndim > 2), kb.e + (ndim > 2), KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { - parthenon::Indexer2D idxer({jb.s - (ndim > 1), jb.e + (ndim > 1)}, {ib.s - (ndim > 0), ib.e + (ndim > 0)}); + parthenon::Indexer2D idxer({jb.s - (ndim > 1), jb.e + (ndim > 1)}, + {ib.s - (ndim > 0), ib.e + (ndim > 0)}); if (pack.GetLowerBound(b, flux_var()) <= pack.GetUpperBound(b, flux_var())) { parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { const auto [j, i] = idxer(idx); // Piecewise linear in the orthogonal directions, could do // something better here pack(b, TE::CC, recon(0), k, j, i) = - 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + 0.5 * (pack(b, fe, flux_var(), k, j, i) + + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); pack(b, TE::CC, recon(1), k, j, i) = - 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + 0.5 * (pack(b, fe, flux_var(), k, j, i) + + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); pack(b, TE::CC, recon(2), k, j, i) = - 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + 0.5 * (pack(b, fe, flux_var(), k, j, i) + + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); pack(b, TE::CC, recon(3), k, j, i) = - 0.5 * (pack(b, fe, flux_var(), k, j, i) + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); + 0.5 * (pack(b, fe, flux_var(), k, j, i) + + pack(b, fe, flux_var(), k + koff, j + joff, i + ioff)); }); } }); @@ -192,8 +203,8 @@ TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon:: jb = md->GetBoundsJ(cl, IndexDomain::interior, edge); kb = md->GetBoundsK(cl, IndexDomain::interior, edge); parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s, kb.e, - KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s, + kb.e, KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); if (pack.GetLowerBound(b, var()) <= pack.GetUpperBound(b, var())) { parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { @@ -205,7 +216,8 @@ TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon:: // TODO(LFR): Pick out the correct component of the // reconstructed flux, current version is not an issue for // piecewise constant flux though. - pack.flux(b, edge, var(), k, j, i) += pack(b, TE::CC, recon(0), k + dk, j + dj, i + di); + pack.flux(b, edge, var(), k, j, i) += + pack(b, TE::CC, recon(0), k + dk, j + dj, i + di); } pack.flux(b, edge, var(), k, j, i) /= npoints; pack.flux(b, edge, var(), k, j, i) *= fac; @@ -222,9 +234,11 @@ TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon:: jb = md->GetBoundsJ(cl, IndexDomain::interior, f); kb = md->GetBoundsK(cl, IndexDomain::interior, f); parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s - (ndim > 2), - kb.e + (ndim > 2), KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { - parthenon::Indexer2D idxer({jb.s - (ndim > 1), jb.e + (ndim > 1)}, {ib.s - (ndim > 0), ib.e + (ndim > 0)}); + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, + kb.s - (ndim > 2), kb.e + (ndim > 2), + KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + parthenon::Indexer2D idxer({jb.s - (ndim > 1), jb.e + (ndim > 1)}, + {ib.s - (ndim > 0), ib.e + (ndim > 0)}); if (pack.GetLowerBound(b, var()) <= pack.GetUpperBound(b, var())) { parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { const auto [j, i] = idxer(idx); @@ -251,16 +265,18 @@ TaskStatus CalculateVectorFluxes(parthenon::TopologicalElement edge, parthenon:: d2[static_cast(edge) % 3] = 0; d2[static_cast(f2) % 3] = 0; parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s, kb.e, - KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack.GetNBlocks() - 1, kb.s, + kb.e, KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); if (pack.GetLowerBound(b, var()) <= pack.GetUpperBound(b, var())) { parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { const auto [j, i] = idxer(idx); pack.flux(b, edge, var(), k, j, i) += - 0.5 * (pack(b, f1, recon_f(0), k, j, i) - pack(b, f1, recon_f(1), k - d1[2], j - d1[1], i - d1[0])); + 0.5 * (pack(b, f1, recon_f(0), k, j, i) - + pack(b, f1, recon_f(1), k - d1[2], j - d1[1], i - d1[0])); pack.flux(b, edge, var(), k, j, i) -= - 0.5 * (pack(b, f2, recon_f(0), k, j, i) - pack(b, f2, recon_f(1), k - d2[2], j - d2[1], i - d2[0])); + 0.5 * (pack(b, f2, recon_f(0), k, j, i) - + pack(b, f2, recon_f(1), k - d2[2], j - d2[1], i - d2[0])); }); } }); diff --git a/example/imex_advection/driver.cpp b/example/imex_advection/driver.cpp index 2693a18a272c..da36c970cdea 100644 --- a/example/imex_advection/driver.cpp +++ b/example/imex_advection/driver.cpp @@ -39,7 +39,8 @@ namespace scalar_imex { Packages_t ProcessPackages(std::unique_ptr &pin) { Packages_t packages; - if (pin->GetOrAddBoolean("scalar_imex", "advection", true)) packages.Add(advection_package::Initialize(pin.get())); + if (pin->GetOrAddBoolean("scalar_imex", "advection", true)) + packages.Add(advection_package::Initialize(pin.get())); auto app = std::make_shared("scalar_imex_app"); packages.Add(app); @@ -52,7 +53,8 @@ Packages_t ProcessPackages(std::unique_ptr &pin) { // that mostly means defining the MakeTaskList *// // function. *// // *************************************************// -ScalarIMEXDriver::ScalarIMEXDriver(ParameterInput *pin, ApplicationInput *app_in, Mesh *pm) +ScalarIMEXDriver::ScalarIMEXDriver(ParameterInput *pin, ApplicationInput *app_in, + Mesh *pm) : parthenon::MultiStageDriverGeneric(pin, app_in, pm) { // fail if these are not specified in the input file @@ -65,9 +67,10 @@ ScalarIMEXDriver::ScalarIMEXDriver(ParameterInput *pin, ApplicationInput *app_in pin->CheckDesired("parthenon/mesh", "refinement"); pin->CheckDesired("parthenon/mesh", "numlevel"); - // Determine which packages to include in driver, allowing for packages to be loaded but not - // run in the driver - do_advection = pin->GetOrAddBoolean("scalar_imex", "do_advection", pin->GetBoolean("scalar_imex", "advection")); + // Determine which packages to include in driver, allowing for packages to be loaded but + // not run in the driver + do_advection = pin->GetOrAddBoolean("scalar_imex", "do_advection", + pin->GetBoolean("scalar_imex", "advection")); } // See the advection.hpp declaration for a description of how this function gets @@ -100,7 +103,8 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta using namespace advection_package::Conserved; static auto desc_phi = parthenon::MakePackDescriptor( - pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Cell}, + pmesh->resolved_packages.get(), + {parthenon::Metadata::WithFluxes, parthenon::Metadata::Cell}, {parthenon::PDOpt::WithFluxes}); using pack_desc_phi_t = decltype(desc_phi); @@ -114,21 +118,25 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta if (stage == 1) { // Copy base data into this stage, need to be careful that it is a full // copy - set_stage0 = tl.AddTask(none, WeightedSumDataAll, mbase.get(), mbase.get(), 1.0, 0.0, mc0.get()); + set_stage0 = tl.AddTask(none, WeightedSumDataAll, mbase.get(), mbase.get(), 1.0, + 0.0, mc0.get()); } // Do implicit update here to get conserved variables for the current stage auto implicit_update = set_stage0; const auto stage_dt = integrator->a(stage, stage) * dt; if (do_advection) - implicit_update = tl.AddTask(set_stage0, advection_package::ImplicitSourceUpdate, stage_dt, mc0.get(), mc0.get()); + implicit_update = tl.AddTask(set_stage0, advection_package::ImplicitSourceUpdate, + stage_dt, mc0.get(), mc0.get()); // Update state of current stage here, as this is the first place it is // finalized - auto boundaries_stage = parthenon::AddBoundaryExchangeTasks(implicit_update, tl, mc0, pmesh->multilevel); + auto boundaries_stage = + parthenon::AddBoundaryExchangeTasks(implicit_update, tl, mc0, pmesh->multilevel); + + auto fill_derived_stage = tl.AddTask( + boundaries_stage, parthenon::Update::FillDerived>, mc0.get()); - auto fill_derived_stage = tl.AddTask(boundaries_stage, parthenon::Update::FillDerived>, mc0.get()); - // Calculate fluxes for the current stage using TT = parthenon::TopologicalType; using TE = parthenon::TopologicalElement; @@ -138,8 +146,9 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta auto flx = none; for (auto face : faces) { if (do_advection) - flx = flx | tl.AddTask(fill_derived_stage, advection_package::CalculateFluxes, desc_phi, face, - parthenon::CellLevel::same, mc0.get()); + flx = flx | tl.AddTask(fill_derived_stage, + advection_package::CalculateFluxes, + desc_phi, face, parthenon::CellLevel::same, mc0.get()); } auto set_flux = parthenon::AddFluxCorrectionTasks(flx, tl, mc0, pmesh->multilevel); @@ -147,40 +156,49 @@ TaskCollection ScalarIMEXDriver::MakeTaskCollection(BlockList_t &, const int sta auto &mdudt_F = pmesh->mesh_data.Add("dUdtF", mbase); auto &mdudt_R = pmesh->mesh_data.Add("dUdtR", mbase); auto flux_div = tl.AddTask(set_flux, StokesAll, mc0.get(), mdudt_F.get()); - + // Calculate source terms for the current stage auto source = set_flux; - if (do_advection) source = source | tl.AddTask(set_flux, advection_package::Source, mc0.get(), mdudt_R.get()); - + if (do_advection) + source = source | + tl.AddTask(set_flux, advection_package::Source, mc0.get(), mdudt_R.get()); + // Add the contribution from this stage to the updates for subsequent stages auto update_stages = flux_div; for (int stage2 = stage + 1; stage2 <= integrator->nstages; ++stage2) { auto label2 = integrator->GetStageName(stage2); - auto &md_in = stage == 1 ? pmesh->mesh_data.Add("base", partitions[i]) : pmesh->mesh_data.Add(label2, mbase); + auto &md_in = stage == 1 ? pmesh->mesh_data.Add("base", partitions[i]) + : pmesh->mesh_data.Add(label2, mbase); auto &md_out = pmesh->mesh_data.Add(label2, mbase); - auto add_flux_div = tl.AddTask(source, WeightedSumDataAll, md_in.get(), mdudt_F.get(), 1.0, - integrator->at(stage2, stage) * dt, md_out.get()); - auto add_source = tl.AddTask(add_flux_div, WeightedSumDataAll, md_out.get(), mdudt_R.get(), 1.0, - integrator->a(stage2, stage) * dt, md_out.get()); + auto add_flux_div = + tl.AddTask(source, WeightedSumDataAll, md_in.get(), mdudt_F.get(), 1.0, + integrator->at(stage2, stage) * dt, md_out.get()); + auto add_source = + tl.AddTask(add_flux_div, WeightedSumDataAll, md_out.get(), mdudt_R.get(), 1.0, + integrator->a(stage2, stage) * dt, md_out.get()); update_stages = update_stages | add_source; } // Add this contribution to the base stage - auto add_flux_div = tl.AddTask(source, WeightedSumDataAll, mbase.get(), mdudt_F.get(), 1.0, - integrator->bt(stage) * dt, mbase.get()); - auto add_source = tl.AddTask(add_flux_div, WeightedSumDataAll, mbase.get(), mdudt_R.get(), 1.0, - integrator->b(stage) * dt, mbase.get()); - + auto add_flux_div = tl.AddTask(source, WeightedSumDataAll, mbase.get(), mdudt_F.get(), + 1.0, integrator->bt(stage) * dt, mbase.get()); + auto add_source = + tl.AddTask(add_flux_div, WeightedSumDataAll, mbase.get(), mdudt_R.get(), 1.0, + integrator->b(stage) * dt, mbase.get()); + // Perform the last stage cleanup tasks if (stage == integrator->nstages) { - auto boundaries = parthenon::AddBoundaryExchangeTasks(add_source, tl, mbase, pmesh->multilevel); + auto boundaries = + parthenon::AddBoundaryExchangeTasks(add_source, tl, mbase, pmesh->multilevel); - auto fill_derived = tl.AddTask(boundaries, parthenon::Update::FillDerived>, mbase.get()); + auto fill_derived = tl.AddTask( + boundaries, parthenon::Update::FillDerived>, mbase.get()); auto dealloc = tl.AddTask(fill_derived, SparseDealloc, mbase.get()); auto new_dt = tl.AddTask(dealloc, EstimateTimestep>, mbase.get()); if (pmesh->adaptive) { - auto tag_refine = tl.AddTask(new_dt, parthenon::Refinement::Tag>, mbase.get()); + auto tag_refine = + tl.AddTask(new_dt, parthenon::Refinement::Tag>, mbase.get()); } } } diff --git a/example/imex_advection/main.cpp b/example/imex_advection/main.cpp index 37523db0e40e..c95f0356def5 100644 --- a/example/imex_advection/main.cpp +++ b/example/imex_advection/main.cpp @@ -52,7 +52,8 @@ int main(int argc, char *argv[]) { pman.ParthenonInitPackagesAndMesh(); { // Initialize the driver - scalar_imex::ScalarIMEXDriver driver(pman.pinput.get(), pman.app_input.get(), pman.pmesh.get()); + scalar_imex::ScalarIMEXDriver driver(pman.pinput.get(), pman.app_input.get(), + pman.pmesh.get()); // This line actually runs the simulation auto driver_status = driver.Execute(); diff --git a/example/imex_advection/pgen/scalar_advection.cpp b/example/imex_advection/pgen/scalar_advection.cpp index 63033ba2d2f7..b9f72acd5c97 100644 --- a/example/imex_advection/pgen/scalar_advection.cpp +++ b/example/imex_advection/pgen/scalar_advection.cpp @@ -72,18 +72,21 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) { if (do_regular_advection) { parthenon::par_for( - PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) { + PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int k, const int j, const int i) { const int kf = ndim > 2 ? (k - nghost) * 2 + nghost : k; const int jf = ndim > 1 ? (j - nghost) * 2 + nghost : j; const int fi = ndim > 0 ? (i - nghost) * 2 + nghost : i; if (profile_type == 0) { PARTHENON_FAIL("Profile type wave not implemented."); } else if (profile_type == 1) { - Real rsq = coords.Xc<1>(i) * coords.Xc<1>(i) + coords.Xc<2>(j) * coords.Xc<2>(j) + + Real rsq = coords.Xc<1>(i) * coords.Xc<1>(i) + + coords.Xc<2>(j) * coords.Xc<2>(j) + coords.Xc<3>(k) * coords.Xc<3>(k); pack(b, phi(), k, j, i) = 1. + amp * exp(-100.0 * rsq); } else if (profile_type == 2) { - Real rsq = coords.Xc<1>(i) * coords.Xc<1>(i) + coords.Xc<2>(j) * coords.Xc<2>(j) + + Real rsq = coords.Xc<1>(i) * coords.Xc<1>(i) + + coords.Xc<2>(j) * coords.Xc<2>(j) + coords.Xc<3>(k) * coords.Xc<3>(k); pack(b, phi(), k, j, i) = (rsq < 0.15 * 0.15 ? 1.0 : 0.0); } else { @@ -98,7 +101,8 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) { kb = cellbounds.GetBoundsK(IndexDomain::interior, TE::F2); const Real x0 = 0.2; parthenon::par_for( - PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) { + PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int k, const int j, const int i) { auto &coords = pack.GetCoordinates(b); Real xlo = coords.X(k, j, i); Real xhi = coords.X(k, j, i + 1); @@ -111,7 +115,8 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) { jb = cellbounds.GetBoundsJ(IndexDomain::interior, TE::F3); kb = cellbounds.GetBoundsK(IndexDomain::interior, TE::F3); parthenon::par_for( - PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, KOKKOS_LAMBDA(const int k, const int j, const int i) { + PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + KOKKOS_LAMBDA(const int k, const int j, const int i) { auto &coords = pack.GetCoordinates(b); Real xlo = coords.X(k, j, i); Real xhi = coords.X(k, j, i + 1); diff --git a/example/imex_advection/utilities/imexrk_integrator.hpp b/example/imex_advection/utilities/imexrk_integrator.hpp index 4fde2bbac2d2..1cddcde09616 100644 --- a/example/imex_advection/utilities/imexrk_integrator.hpp +++ b/example/imex_advection/utilities/imexrk_integrator.hpp @@ -30,7 +30,8 @@ class IMEXRKIntegrator { explicit IMEXRKIntegrator(const std::string &name); IMEXRKIntegrator() : IMEXRKIntegrator("SSP2-(2,2,2)") {} explicit IMEXRKIntegrator(ParameterInput *pin) - : IMEXRKIntegrator(pin->GetOrAddString("parthenon/time", "integrator", "SSP2-(2,2,2)")) {} + : IMEXRKIntegrator( + pin->GetOrAddString("parthenon/time", "integrator", "SSP2-(2,2,2)")) {} // To conform with other integrators int nstages; diff --git a/example/imex_advection/utilities/index_permutation.hpp b/example/imex_advection/utilities/index_permutation.hpp index c6f93a161e5f..462dde4747e7 100644 --- a/example/imex_advection/utilities/index_permutation.hpp +++ b/example/imex_advection/utilities/index_permutation.hpp @@ -25,7 +25,8 @@ struct IndexingData : public IndexShape { int ndim; IndexShape cellbounds; IndexingData(MeshData *md) - : ndim{md->GetMeshPointer()->ndim}, IndexShape(md->GetBlockData(0)->GetBlockPointer()->cellbounds) {} + : ndim{md->GetMeshPointer()->ndim}, + IndexShape(md->GetBlockData(0)->GetBlockPointer()->cellbounds) {} KOKKOS_INLINE_FUNCTION auto Get3DIndexRangeWithHalo(IndexDomain id, std::array halo, @@ -43,7 +44,8 @@ struct IndexingData : public IndexShape { } KOKKOS_INLINE_FUNCTION - auto Get3DIndexRange(IndexDomain id, TopologicalElement te = TopologicalElement::CC) const { + auto Get3DIndexRange(IndexDomain id, + TopologicalElement te = TopologicalElement::CC) const { return Get3DIndexRangeWithHalo(id, {0, 0, 0}, te); } @@ -54,12 +56,14 @@ struct IndexingData : public IndexShape { KOKKOS_INLINE_FUNCTION std::array GetOffsetArray(TopologicalElement flux_te) const { - return std::array{TopologicalOffsetK(flux_te) * (ndim > 2), TopologicalOffsetJ(flux_te) * (ndim > 1), + return std::array{TopologicalOffsetK(flux_te) * (ndim > 2), + TopologicalOffsetJ(flux_te) * (ndim > 1), TopologicalOffsetI(flux_te) * (ndim > 0)}; } }; -inline auto Get3DIndexRangeWithHalo(MeshData *md, IndexDomain id, std::array halo, +inline auto Get3DIndexRangeWithHalo(MeshData *md, IndexDomain id, + std::array halo, TopologicalElement te = TopologicalElement::CC) { const int ndim = md->GetMeshPointer()->ndim; IndexRange ib = md->GetBoundsI(id, te); @@ -74,11 +78,13 @@ inline auto Get3DIndexRangeWithHalo(MeshData *md, IndexDomain id, std::arr return std::make_tuple(kb, jb, ib); } -inline auto Get3DIndexRange(MeshData *md, IndexDomain id, TopologicalElement te = TopologicalElement::CC) { +inline auto Get3DIndexRange(MeshData *md, IndexDomain id, + TopologicalElement te = TopologicalElement::CC) { return Get3DIndexRangeWithHalo(md, id, {0, 0, 0}, te); } -inline CoordinateDirection PermuteDirection(CoordinateDirection base_dir, CoordinateDirection relative_dir) { +inline CoordinateDirection PermuteDirection(CoordinateDirection base_dir, + CoordinateDirection relative_dir) { return static_cast(((base_dir - 1 + relative_dir - 1) % 3) + 1); } @@ -90,7 +96,8 @@ inline auto GetOffsetsForDirection(MeshData *md, CoordinateDirection dir) return std::make_tuple(dk, dj, di); } -inline auto GetPermutedOffsetsForRelativeDirection(MeshData *md, CoordinateDirection base_dir, +inline auto GetPermutedOffsetsForRelativeDirection(MeshData *md, + CoordinateDirection base_dir, CoordinateDirection relative_dir) { auto absolute_dir = PermuteDirection(base_dir, relative_dir); return GetOffsetsForDirection(md, absolute_dir); diff --git a/example/imex_advection/utilities/reconstruct.hpp b/example/imex_advection/utilities/reconstruct.hpp index 6cf0fe02cc32..e3a4d536bd5a 100644 --- a/example/imex_advection/utilities/reconstruct.hpp +++ b/example/imex_advection/utilities/reconstruct.hpp @@ -32,11 +32,13 @@ using namespace parthenon::driver::prelude; KOKKOS_FORCEINLINE_FUNCTION Real mc(const Real dm, const Real dp, const Real alpha = 1.99) { const Real dc = (dm * dp > 0.0) * 0.5 * (dm + dp); - return std::copysign(std::min(std::fabs(dc), alpha * std::min(std::fabs(dm), std::fabs(dp))), dc); + return std::copysign( + std::min(std::fabs(dc), alpha * std::min(std::fabs(dm), std::fabs(dp))), dc); } KOKKOS_INLINE_FUNCTION -void PiecewiseLinear(const Real qm, const Real q0, const Real qp, Real &p, Real &m, const Real slope_limit = 1.99) { +void PiecewiseLinear(const Real qm, const Real q0, const Real qp, Real &p, Real &m, + const Real slope_limit = 1.99) { Real dq = qp - q0; // const Real slope_limit = 0.99 / ((wgt <= 0.5)*wgt + (wgt > 0.5)*(1.0-wgt)); dq = 0.5 * mc(q0 - qm, dq, slope_limit); @@ -46,13 +48,16 @@ void PiecewiseLinear(const Real qm, const Real q0, const Real qp, Real &p, Real template KOKKOS_INLINE_FUNCTION void -Reconstruct(parthenon::team_mbr_t member, const int b, const int k, parthenon::TopologicalElement flux_te, +Reconstruct(parthenon::team_mbr_t member, const int b, const int k, + parthenon::TopologicalElement flux_te, const parthenon::utils::IndexingData &cellbounds, const pack_t &pack, - parthenon::utils::ScratchPack &wm, parthenon::utils::ScratchPack &wp) { + parthenon::utils::ScratchPack &wm, + parthenon::utils::ScratchPack &wp) { const auto [kb, jb, ib] = cellbounds.GetReconstructionRange(flux_te); const auto [kbe, jbe, ibe] = cellbounds.Get3DIndexRange(IndexDomain::entire); const parthenon::Indexer2D idxer_full({jbe.s, jbe.e}, {ibe.s, ibe.e}); - const int npoints = idxer_full.GetFlatIdx(jb.e, ib.e) - idxer_full.GetFlatIdx(jb.s, ib.s) + 1; + const int npoints = + idxer_full.GetFlatIdx(jb.e, ib.e) - idxer_full.GetFlatIdx(jb.s, ib.s) + 1; parthenon::Indexer2D idxer_recon({jb.s, jb.e}, {ib.s, ib.e}); const auto [koff, joff, ioff] = cellbounds.GetOffsetArray(flux_te); // Do L/R reconstruction across this slab @@ -62,8 +67,9 @@ Reconstruct(parthenon::team_mbr_t member, const int b, const int k, parthenon::T Real *pu = &pack(b, l, k + koff, jb.s + joff, ib.s + ioff); Real *pwp = &wp(l, jb.s, ib.s); Real *pwm = &wm(l, jb.s, ib.s); - parthenon::par_for_inner(member, 0, npoints - 1, - [&](const int idx) { PiecewiseLinear(pl[idx], pm[idx], pu[idx], pwp[idx], pwm[idx]); }); + parthenon::par_for_inner(member, 0, npoints - 1, [&](const int idx) { + PiecewiseLinear(pl[idx], pm[idx], pu[idx], pwp[idx], pwm[idx]); + }); } } diff --git a/example/imex_advection/utilities/scratch_pack.hpp b/example/imex_advection/utilities/scratch_pack.hpp index c207499b4229..3d9df72a2d96 100644 --- a/example/imex_advection/utilities/scratch_pack.hpp +++ b/example/imex_advection/utilities/scratch_pack.hpp @@ -26,9 +26,10 @@ struct ScratchPack { using scratch_pad_t = ScratchPad3D; template - KOKKOS_INLINE_FUNCTION ScratchPack(scratch_space_t &&scratch_space, const pack_t *ppack, int b, int nj, int ni) - : ppack{ppack}, block{b}, - data(std::forward(scratch_space), ppack->GetUpperBound(b) + 1, nj, ni) {} + KOKKOS_INLINE_FUNCTION ScratchPack(scratch_space_t &&scratch_space, const pack_t *ppack, + int b, int nj, int ni) + : ppack{ppack}, block{b}, data(std::forward(scratch_space), + ppack->GetUpperBound(b) + 1, nj, ni) {} template KOKKOS_FORCEINLINE_FUNCTION Real &operator()(const Tin &t, int j, int i) { @@ -40,7 +41,8 @@ struct ScratchPack { Real &operator()(int v, int j, int i) { return data(v, j, i); } static std::size_t get_size_in_bytes(const pack_t &pack, int nj, int ni) { - const int nvar = pack.GetUpperBoundHost(0) + 1; // Assuming no sparse fields for the time being + const int nvar = + pack.GetUpperBoundHost(0) + 1; // Assuming no sparse fields for the time being return scratch_pad_t::shmem_size(nvar, nj, ni); } diff --git a/example/imex_advection/utilities/stokes.hpp b/example/imex_advection/utilities/stokes.hpp index f7a4991f0156..3c6bfa6bbb84 100644 --- a/example/imex_advection/utilities/stokes.hpp +++ b/example/imex_advection/utilities/stokes.hpp @@ -31,8 +31,10 @@ namespace scalar_imex { using namespace parthenon::driver::prelude; template -TaskStatus WeightedSumDataElement(parthenon::CellLevel cl, parthenon::TopologicalElement te, pack_desc_t pd, - MeshData *in1, MeshData *in2, Real w1, Real w2, MeshData *out) { +TaskStatus WeightedSumDataElement(parthenon::CellLevel cl, + parthenon::TopologicalElement te, pack_desc_t pd, + MeshData *in1, MeshData *in2, Real w1, + Real w2, MeshData *out) { auto pack1 = pd.GetPack(in1); auto pack2 = pd.GetPack(in2); auto pack_out = pd.GetPack(out); @@ -44,56 +46,70 @@ TaskStatus WeightedSumDataElement(parthenon::CellLevel cl, parthenon::Topologica constexpr int scratch_size = 0; constexpr int scratch_level = 1; parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack1.GetNBlocks() - 1, kb.s, kb.e, - KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack1.GetNBlocks() - 1, kb.s, + kb.e, KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); - PARTHENON_REQUIRE(pack1.GetLowerBound(b) == pack2.GetLowerBound(b), "Packs are different size."); - PARTHENON_REQUIRE(pack1.GetLowerBound(b) == pack_out.GetLowerBound(b), "Packs are different size."); - PARTHENON_REQUIRE(pack1.GetUpperBound(b) == pack2.GetUpperBound(b), "Packs are different size."); - PARTHENON_REQUIRE(pack1.GetUpperBound(b) == pack_out.GetUpperBound(b), "Packs are different size."); + PARTHENON_REQUIRE(pack1.GetLowerBound(b) == pack2.GetLowerBound(b), + "Packs are different size."); + PARTHENON_REQUIRE(pack1.GetLowerBound(b) == pack_out.GetLowerBound(b), + "Packs are different size."); + PARTHENON_REQUIRE(pack1.GetUpperBound(b) == pack2.GetUpperBound(b), + "Packs are different size."); + PARTHENON_REQUIRE(pack1.GetUpperBound(b) == pack_out.GetUpperBound(b), + "Packs are different size."); for (int l = pack1.GetLowerBound(b); l <= pack1.GetUpperBound(b); ++l) { const auto [j, i] = idxer(0); Real *out = &pack_out(b, te, l, k, j, i); Real const *const one = &pack1(b, te, l, k, j, i); Real const *const two = &pack2(b, te, l, k, j, i); - parthenon::par_for_inner(member, 0, idxer.size() - 1, - [&](const int idx) { out[idx] = w1 * one[idx] + w2 * two[idx]; }); + parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { + out[idx] = w1 * one[idx] + w2 * two[idx]; + }); } }); return TaskStatus::complete; } template -TaskStatus WeightedSumData(parthenon::CellLevel cl, parthenon::TopologicalType tt, pack_desc_t pd, MeshData *in1, - MeshData *in2, Real w1, Real w2, MeshData *out) { +TaskStatus WeightedSumData(parthenon::CellLevel cl, parthenon::TopologicalType tt, + pack_desc_t pd, MeshData *in1, MeshData *in2, + Real w1, Real w2, MeshData *out) { for (auto te : parthenon::GetTopologicalElements(tt)) WeightedSumDataElement(cl, te, pd, in1, in2, w1, w2, out); return TaskStatus::complete; } -TaskStatus WeightedSumDataAll(MeshData *in1, MeshData *in2, Real w1, Real w2, MeshData *out) { +TaskStatus WeightedSumDataAll(MeshData *in1, MeshData *in2, Real w1, Real w2, + MeshData *out) { auto pmesh = in1->GetMeshPointer(); constexpr auto cl = parthenon::CellLevel::same; { static auto desc = parthenon::MakePackDescriptor( - pmesh->resolved_packages.get(), {parthenon::Metadata::Independent, parthenon::Metadata::Cell}, {}); - if (desc.nvar_tot > 0) WeightedSumData(cl, parthenon::TopologicalType::Cell, desc, in1, in2, w1, w2, out); + pmesh->resolved_packages.get(), + {parthenon::Metadata::Independent, parthenon::Metadata::Cell}, {}); + if (desc.nvar_tot > 0) + WeightedSumData(cl, parthenon::TopologicalType::Cell, desc, in1, in2, w1, w2, out); } { static auto desc = parthenon::MakePackDescriptor( - pmesh->resolved_packages.get(), {parthenon::Metadata::Independent, parthenon::Metadata::Face}, {}); - if (desc.nvar_tot > 0) WeightedSumData(cl, parthenon::TopologicalType::Face, desc, in1, in2, w1, w2, out); + pmesh->resolved_packages.get(), + {parthenon::Metadata::Independent, parthenon::Metadata::Face}, {}); + if (desc.nvar_tot > 0) + WeightedSumData(cl, parthenon::TopologicalType::Face, desc, in1, in2, w1, w2, out); } { static auto desc = parthenon::MakePackDescriptor( - pmesh->resolved_packages.get(), {parthenon::Metadata::Independent, parthenon::Metadata::Edge}, {}); - if (desc.nvar_tot > 0) WeightedSumData(cl, parthenon::TopologicalType::Edge, desc, in1, in2, w1, w2, out); + pmesh->resolved_packages.get(), + {parthenon::Metadata::Independent, parthenon::Metadata::Edge}, {}); + if (desc.nvar_tot > 0) + WeightedSumData(cl, parthenon::TopologicalType::Edge, desc, in1, in2, w1, w2, out); } return TaskStatus::complete; } template -void StokesZero(parthenon::CellLevel cl, parthenon::TopologicalElement TeVar, pack_desc_t &pd, MeshData *out) { +void StokesZero(parthenon::CellLevel cl, parthenon::TopologicalElement TeVar, + pack_desc_t &pd, MeshData *out) { auto pack_out = pd.GetPack(out); IndexRange ib = out->GetBoundsI(cl, IndexDomain::interior, TeVar); @@ -103,8 +119,8 @@ void StokesZero(parthenon::CellLevel cl, parthenon::TopologicalElement TeVar, pa constexpr int scratch_size = 0; constexpr int scratch_level = 1; parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack_out.GetNBlocks() - 1, kb.s, kb.e, - KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { + PARTHENON_AUTO_LABEL, scratch_size, scratch_level, 0, pack_out.GetNBlocks() - 1, + kb.s, kb.e, KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { parthenon::Indexer2D idxer({jb.s, jb.e}, {ib.s, ib.e}); for (int l = pack_out.GetLowerBound(b); l <= pack_out.GetUpperBound(b); ++l) { parthenon::par_for_inner(member, 0, idxer.size() - 1, [&](const int idx) { @@ -116,9 +132,10 @@ void StokesZero(parthenon::CellLevel cl, parthenon::TopologicalElement TeVar, pa } template -void StokesComponent(Real fac, parthenon::CellLevel cl, parthenon::TopologicalElement TeVar, - parthenon::TopologicalElement TeFlux, pack_desc_t &pd, int ndim, MeshData *in, - MeshData *out) { +void StokesComponent(Real fac, parthenon::CellLevel cl, + parthenon::TopologicalElement TeVar, + parthenon::TopologicalElement TeFlux, pack_desc_t &pd, int ndim, + MeshData *in, MeshData *out) { auto pack_in = pd.GetPack(in); auto pack_out = pd.GetPack(out); @@ -140,41 +157,54 @@ void StokesComponent(Real fac, parthenon::CellLevel cl, parthenon::TopologicalEl const parthenon::utils::IndexingData cellbounds(in); using TE = parthenon::TopologicalElement; // Choose the correct memory size for each topological element - const int njtot = cellbounds.ncellsj(IndexDomain::entire, TeVar == TE::CC ? TE::CC : TE::NN); - const int nitot = cellbounds.ncellsi(IndexDomain::entire, TeVar == TE::CC ? TE::CC : TE::NN); - std::size_t scratch_size_in_bytes = parthenon::ScratchPad2D::shmem_size(njtot, nitot) * (2 + koff); + const int njtot = + cellbounds.ncellsj(IndexDomain::entire, TeVar == TE::CC ? TE::CC : TE::NN); + const int nitot = + cellbounds.ncellsi(IndexDomain::entire, TeVar == TE::CC ? TE::CC : TE::NN); + std::size_t scratch_size_in_bytes = + parthenon::ScratchPad2D::shmem_size(njtot, nitot) * (2 + koff); parthenon::par_for_outer( - PARTHENON_AUTO_LABEL, scratch_size_in_bytes, scratch_level, 0, pack_out.GetNBlocks() - 1, kb.s, kb.e, + PARTHENON_AUTO_LABEL, scratch_size_in_bytes, scratch_level, 0, + pack_out.GetNBlocks() - 1, kb.s, kb.e, KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b, const int k) { auto &coords = pack_in.GetCoordinates(b); parthenon::Indexer2D idxer_entire({0, njtot - 1}, {0, nitot - 1}); - parthenon::ScratchPad2D Vflux(member.team_scratch(scratch_level), njtot, nitot); - parthenon::ScratchPad2D Vvar(member.team_scratch(scratch_level), njtot, nitot); + parthenon::ScratchPad2D Vflux(member.team_scratch(scratch_level), njtot, + nitot); + parthenon::ScratchPad2D Vvar(member.team_scratch(scratch_level), njtot, + nitot); parthenon::par_for_inner(member, 0, idxer_entire.size() - 1, [&](const int idx) { const auto [j, i] = idxer_entire(idx); Vflux(j, i) = coords.Volume(cl, TeFlux, k, j, i); Vvar(j, i) = coords.Volume(cl, TeVar, k, j, i); }); - parthenon::ScratchPad2D Vfluxu(member.team_scratch(scratch_level), koff ? njtot : 0, koff ? nitot : 0); + parthenon::ScratchPad2D Vfluxu(member.team_scratch(scratch_level), + koff ? njtot : 0, koff ? nitot : 0); if (koff) { - parthenon::par_for_inner(member, 0, idxer_entire.size() - 1, [&](const int idx) { - const auto [j, i] = idxer_entire(idx); - Vfluxu(j, i) = coords.Volume(cl, TeFlux, k + koff, j, i); - }); + parthenon::par_for_inner( + member, 0, idxer_entire.size() - 1, [&](const int idx) { + const auto [j, i] = idxer_entire(idx); + Vfluxu(j, i) = coords.Volume(cl, TeFlux, k + koff, j, i); + }); } member.team_barrier(); - PARTHENON_REQUIRE(pack_in.GetLowerBound(b) == pack_out.GetLowerBound(b), "Packs are different size."); - PARTHENON_REQUIRE(pack_in.GetUpperBound(b) == pack_out.GetUpperBound(b), "Packs are different size."); + PARTHENON_REQUIRE(pack_in.GetLowerBound(b) == pack_out.GetLowerBound(b), + "Packs are different size."); + PARTHENON_REQUIRE(pack_in.GetUpperBound(b) == pack_out.GetUpperBound(b), + "Packs are different size."); - const int npoints = idxer_entire.GetFlatIdx(jb.e, ib.e) - idxer_entire.GetFlatIdx(jb.s, ib.s) + 1; + const int npoints = + idxer_entire.GetFlatIdx(jb.e, ib.e) - idxer_entire.GetFlatIdx(jb.s, ib.s) + 1; for (int l = pack_out.GetLowerBound(b); l <= pack_out.GetUpperBound(b); ++l) { Real const *const flxl = &pack_in.flux(b, TeFlux, l, k, jb.s, ib.s); - Real const *const flxu = &pack_in.flux(b, TeFlux, l, k + koff, jb.s + joff, ib.s + ioff); + Real const *const flxu = + &pack_in.flux(b, TeFlux, l, k + koff, jb.s + joff, ib.s + ioff); Real const *const vfl = &Vflux(jb.s, ib.s); - Real const *const vfu = koff ? &Vfluxu(jb.s + joff, ib.s + ioff) : &Vflux(jb.s + joff, ib.s + ioff); + Real const *const vfu = + koff ? &Vfluxu(jb.s + joff, ib.s + ioff) : &Vflux(jb.s + joff, ib.s + ioff); Real const *const vv = &Vvar(jb.s, ib.s); Real *out = &pack_out(b, TeVar, l, k, jb.s, ib.s); @@ -186,8 +216,8 @@ void StokesComponent(Real fac, parthenon::CellLevel cl, parthenon::TopologicalEl } template -TaskStatus Stokes(parthenon::CellLevel cl, parthenon::TopologicalType TtVar, pack_desc_t &pd, int ndim, - MeshData *in, MeshData *out) { +TaskStatus Stokes(parthenon::CellLevel cl, parthenon::TopologicalType TtVar, + pack_desc_t &pd, int ndim, MeshData *in, MeshData *out) { using TE = parthenon::TopologicalElement; using TT = parthenon::TopologicalType; @@ -221,7 +251,8 @@ TaskStatus Stokes(parthenon::CellLevel cl, parthenon::TopologicalType TtVar, pac // downstream or add a test involving curls Flip the sign if the // variable is an X1 face and the edge is an X3 edge, or an X2 face // ... X1 edge, or an X3 face ... X2 edge - const int indicator = ((static_cast(fte) % 3) - (static_cast(vte) % 3) + 3) % 3; + const int indicator = + ((static_cast(fte) % 3) - (static_cast(vte) % 3) + 3) % 3; fac = (indicator == 2) ? -1.0 : 1.0; } StokesComponent(fac, cl, vte, fte, pd, ndim, in, out); @@ -237,21 +268,27 @@ TaskStatus StokesAll(MeshData *in, MeshData *out) { constexpr auto cl = parthenon::CellLevel::same; { static auto desc = parthenon::MakePackDescriptor( - pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Cell}, + pmesh->resolved_packages.get(), + {parthenon::Metadata::WithFluxes, parthenon::Metadata::Cell}, {parthenon::PDOpt::WithFluxes}); - if (desc.nvar_tot > 0) Stokes(cl, parthenon::TopologicalType::Cell, desc, ndim, in, out); + if (desc.nvar_tot > 0) + Stokes(cl, parthenon::TopologicalType::Cell, desc, ndim, in, out); } { static auto desc = parthenon::MakePackDescriptor( - pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Face}, + pmesh->resolved_packages.get(), + {parthenon::Metadata::WithFluxes, parthenon::Metadata::Face}, {parthenon::PDOpt::WithFluxes}); - if (desc.nvar_tot > 0) Stokes(cl, parthenon::TopologicalType::Face, desc, ndim, in, out); + if (desc.nvar_tot > 0) + Stokes(cl, parthenon::TopologicalType::Face, desc, ndim, in, out); } { static auto desc = parthenon::MakePackDescriptor( - pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes, parthenon::Metadata::Edge}, + pmesh->resolved_packages.get(), + {parthenon::Metadata::WithFluxes, parthenon::Metadata::Edge}, {parthenon::PDOpt::WithFluxes}); - if (desc.nvar_tot > 0) Stokes(cl, parthenon::TopologicalType::Edge, desc, ndim, in, out); + if (desc.nvar_tot > 0) + Stokes(cl, parthenon::TopologicalType::Edge, desc, ndim, in, out); } return TaskStatus::complete; } diff --git a/example/imex_advection/utilities/three_vec.hpp b/example/imex_advection/utilities/three_vec.hpp index c7d4e71d31d8..534aefdeb896 100644 --- a/example/imex_advection/utilities/three_vec.hpp +++ b/example/imex_advection/utilities/three_vec.hpp @@ -31,8 +31,10 @@ struct ThreeVec { ThreeVec(Real a1, Real a2, Real a3) : vals{a1, a2, a3} {} template - KOKKOS_FORCEINLINE_FUNCTION ThreeVec(const pack_t &pack, int b, var_t, int k, int j, int i) - : vals{pack(b, var_t(0), k, j, i), pack(b, var_t(1), k, j, i), pack(b, var_t(2), k, j, i)} {} + KOKKOS_FORCEINLINE_FUNCTION ThreeVec(const pack_t &pack, int b, var_t, int k, int j, + int i) + : vals{pack(b, var_t(0), k, j, i), pack(b, var_t(1), k, j, i), + pack(b, var_t(2), k, j, i)} {} KOKKOS_FORCEINLINE_FUNCTION Real &operator[](CoordinateDirection dir) { return vals[dir - 1]; } diff --git a/example/imex_advection/utilities/variable_macro.hpp b/example/imex_advection/utilities/variable_macro.hpp index 7a9f66df8ed0..f78e231a3d13 100644 --- a/example/imex_advection/utilities/variable_macro.hpp +++ b/example/imex_advection/utilities/variable_macro.hpp @@ -18,12 +18,12 @@ #include -#define VARIABLE(ns, varname) \ - struct varname : public parthenon::variable_names::base_t { \ - template \ - KOKKOS_INLINE_FUNCTION varname(Ts &&...args) \ - : parthenon::variable_names::base_t(std::forward(args)...) {} \ - static std::string name() { return #ns "." #varname; } \ +#define VARIABLE(ns, varname) \ + struct varname : public parthenon::variable_names::base_t { \ + template \ + KOKKOS_INLINE_FUNCTION varname(Ts &&...args) \ + : parthenon::variable_names::base_t(std::forward(args)...) {} \ + static std::string name() { return #ns "." #varname; } \ } #endif // SRC_UTILITIES_VARIABLE_MACRO_HPP_