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..8773acd994a2 --- /dev/null +++ b/example/imex_advection/advection/advection_package.cpp @@ -0,0 +1,255 @@ +//======================================================================================== +// (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..7755f68125bd --- /dev/null +++ b/example/imex_advection/advection/advection_package.hpp @@ -0,0 +1,290 @@ +//======================================================================================== +// (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..da36c970cdea --- /dev/null +++ b/example/imex_advection/driver.cpp @@ -0,0 +1,209 @@ +//======================================================================================== +// (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 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()); + + // 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()); + + // Calculate fluxes for the current stage + 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()); + + // 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); + 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()); + + // Perform the last stage cleanup tasks + 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..c95f0356def5 --- /dev/null +++ b/example/imex_advection/main.cpp @@ -0,0 +1,67 @@ +//======================================================================================== +// (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..b9f72acd5c97 --- /dev/null +++ b/example/imex_advection/pgen/scalar_advection.cpp @@ -0,0 +1,130 @@ +//======================================================================================== +// (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..1cddcde09616 --- /dev/null +++ b/example/imex_advection/utilities/imexrk_integrator.hpp @@ -0,0 +1,60 @@ +//======================================================================================== +// (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..462dde4747e7 --- /dev/null +++ b/example/imex_advection/utilities/index_permutation.hpp @@ -0,0 +1,127 @@ +//======================================================================================== +// (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..e3a4d536bd5a --- /dev/null +++ b/example/imex_advection/utilities/reconstruct.hpp @@ -0,0 +1,78 @@ +//======================================================================================== +// (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..3d9df72a2d96 --- /dev/null +++ b/example/imex_advection/utilities/scratch_pack.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_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..3c6bfa6bbb84 --- /dev/null +++ b/example/imex_advection/utilities/stokes.hpp @@ -0,0 +1,298 @@ +//======================================================================================== +// (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..534aefdeb896 --- /dev/null +++ b/example/imex_advection/utilities/three_vec.hpp @@ -0,0 +1,103 @@ +//======================================================================================== +// (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..f78e231a3d13 --- /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_