diff --git a/CMakeLists.txt b/CMakeLists.txt index 340369e4..6b470bf7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -120,6 +120,7 @@ set(SINGULARITY_USE_KOKKOS if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/external/singularity-eos/CMakeLists.txt) # this forces singularity to use our version of kokkos-kernels to be # compatible with parthenon's kokkos 5.1 + set(KokkosKernels_ENABLE_SUPERNODAL_SPTRSV OFF) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/kokkos-kernels) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/singularity-eos) else() diff --git a/docs/grid.md b/docs/grid.md index 2bf7e6ee..a4b0274f 100644 --- a/docs/grid.md +++ b/docs/grid.md @@ -126,5 +126,90 @@ The alias `FirstDer` is pulled out of the `ScratchVariableList` can then be used `scratch_firstder` will be registered. +## Coordinates + +Kamayan builds on Parthenon's mesh and coordinate layout (currently +`parthenon::UniformCartesian`), but wraps it behind a small geometry layer so that +code can be written against a consistent coordinate/metric API while the geometry +is selected at runtime. + +At the moment the grid unit supports: + +- `Geometry::cartesian`: standard 1D/2D/3D Cartesian +- `Geometry::cylindrical`: 2D axisymmetric r-z with an implicit azimuthal direction + (the third direction uses `dphi = 2*pi`) + +### `CoordinateSystem` + +The coordinate wrappers implement the `CoordinateSystem` concept +(`src/grid/geometry.hpp`), which is the practical "contract" used throughout the code. +It includes both compile-time and runtime axis dispatch: + +- Compile-time: `template Dx()`, `Xc(idx)`, `Xf(idx)`, `FaceArea(k,j,i)`, ... +- Runtime: `Dx(Axis)`, `Xc(Axis, idx)`, `Xf(Axis, idx)`, `FaceArea(Axis,k,j,i)`, ... + +The main coordinate wrapper is `grid::Coordinates` (`src/grid/geometry.hpp`). It is +a thin layer over Parthenon's coordinates that provides geometry-specific fixes for +metric factors: + +- `X`/`Xi`: coordinate at the Parthenon cell-center location +- `Xc`: coordinate at the cell centroid (differs from `Xi` for cylindrical r) +- `Xf`: coordinate at the face location + +- In Cartesian geometry, methods delegate to Parthenon. +- In cylindrical geometry: + - `Dx()` returns `2*pi` (implicit azimuthal direction) + - `Xc(idx)` computes the radial *centroid* of the cell (not just the + midpoint); `Xi(idx)` remains the "cell center" value from Parthenon + - `FaceArea`, `EdgeLength`, and `CellVolume` use r-z metric factors (e.g. + `V = 0.5 * dphi * (r_{i+1/2}^2 - r_{i-1/2}^2) * dz`) + +### CoordinatePacks + +Parthenon's coordinate methods are largely inline calculations. For geometry-dependent +metric factors (areas, volumes, centroids, etc.) it is often cheaper and simpler to +precompute them once per `MeshBlock` and store them as normal fields. + +The grid unit registers and fills these coordinate fields: + +- Field types live in `src/grid/coordinates.hpp` under `kamayan::grid::coords::*` and + the full list is `grid::CoordFields`. +- The fields are allocated with geometry-aware *degenerate* shapes via + `grid::CoordinateShape`: + - Cartesian: `Dx*`, `FaceArea*`, `EdgeLength*`, `Volume` are scalars; `X*`, `Xc*`, `Xf*` + are 1D arrays per-axis. + - Cylindrical: `Dx*` are scalars; r-dependent quantities (`Volume`, `FaceArea*`, + `EdgeLength*`, and `X*`/`Xc*`/`Xf*` for r) are stored as 1D arrays in the radial direction. +- `grid::CalculateCoordinates` (`src/grid/coordinates.cpp`) fills all `CoordFields` for each + block (using `grid::CoordinateIndexRanges` to respect degenerate/face extents), and is + called from the grid unit's `InitMeshBlockData` callback. + +`grid::CoordinatePack` (`src/grid/coordinates.hpp`) wraps a `SparsePack` holding +these `geom.*` fields and exposes the same API as `grid::Coordinates` but indexed by +`(k,j,i)`. Internally it maps `(k,j,i)` onto the degenerate storage layout (scalar/1D), so +call sites do not need to care about how each metric is stored. + +Example usage (runtime geometry): + +```cpp title="problems/isentropic_vortex.cpp:isen_coords_pack" +--8<-- "problems/isentropic_vortex.cpp:isen_coords_pack" +``` + +### Generic Coordinates (Packs) + +When the geometry is only known at runtime, Kamayan provides variant-based wrappers: + +- `grid::GenericCoordinate` (`src/grid/geometry.hpp`) wraps `grid::Coordinates` in a + `variant` and forwards calls via `visit`. +- `grid::GenericCoordinatePack` (`src/grid/coordinates.hpp`) does the same for + `grid::CoordinatePack`. + +In practice `GenericCoordinatePack` is constructed from a coordinate-field pack (e.g. +`grid::GetPack(grid::CoordFields(), ...)`) and then used like a normal coordinate object. + +These are convenient in code like problem generators, but they do add a small overhead +compared to templating on `Geometry`. For performance-critical kernels prefer templating +on `geom` and constructing the matching `grid::Coordinates`/`grid::CoordinatePack`. + ## Parameters {!assets/generated/grid_parms.md!} diff --git a/docs/kamayan_docs.cpp b/docs/kamayan_docs.cpp index 30e001c4..ec84dff2 100644 --- a/docs/kamayan_docs.cpp +++ b/docs/kamayan_docs.cpp @@ -61,7 +61,7 @@ int main(int argc_in, char *argv_in[]) { // generate all the graphs for the task collection/lists // for the driver as well as all the units auto driver = kamayan::InitPackages(pman, units); - auto pkg = std::make_shared("Test Package"); + auto pkg = std::make_shared("Test Package"); auto block_list = kamayan::MakeTestBlockList(pkg, 1, 8, 3); auto tc = driver.MakeTaskCollection(block_list, 1); std::cout << tc; diff --git a/python/kamayan/cli/app.py b/python/kamayan/cli/app.py index 57d97d72..fb454e3f 100644 --- a/python/kamayan/cli/app.py +++ b/python/kamayan/cli/app.py @@ -57,11 +57,17 @@ def run( input_file: Optional[str] = None, dry_run: bool = False, info: bool = False, + debug: bool = False, **kwargs, ): parthenon_args: typer.Context = kwargs.pop("ctx") km = self.func(*args, **kwargs) + if debug: + import os + + typer.echo(f"PID: {os.getpid()}") + breakpoint() if input_file: km.input_file = input_file @@ -139,6 +145,18 @@ def run( ) func_params += [info_param] + debug_param = inspect.Parameter( + "debug", + inspect.Parameter.KEYWORD_ONLY, + default=typer.Option( + False, + "--debug", + help="Pause execution and wait for debugger to attach", + ), + annotation=bool, + ) + func_params += [debug_param] + ctx_param = inspect.Parameter( "ctx", inspect.Parameter.KEYWORD_ONLY, annotation=typer.Context ) diff --git a/python/kamayan/code_units/Grid.py b/python/kamayan/code_units/Grid.py index 23f25030..574cbc08 100644 --- a/python/kamayan/code_units/Grid.py +++ b/python/kamayan/code_units/Grid.py @@ -12,6 +12,7 @@ _resolution = None | int _dx = None | float refinement_strategy = Literal["none", "static", "adaptive"] +GEOMETRY = Literal["cartesian", "cylindrical"] def _get_N(N: _resolution, dx: _dx, L: float) -> int: @@ -56,7 +57,7 @@ def _set_mesh_block(params: KamayanParams, nxb1: int, nxb2: int, nxb3: int) -> N } -_boundary_conditions = Literal["outflow", "periodic", "user", "reflect"] +_boundary_conditions = Literal["outflow", "periodic", "user", "reflect", "axisymmetric"] @dataclass @@ -125,6 +126,7 @@ class KamayanGrid(Node): xbnd2: tuple[float, float] | None = None xbnd3: tuple[float, float] | None = None numlevel: int = 1 + geometry: GEOMETRY = "cartesian" boundary_conditions = auto_property_node(BoundaryConditions, "boundary_conditions") @@ -144,6 +146,7 @@ def set_params(self, params) -> None: ) _set_mesh_block(params, self.nxb1, self.nxb2, self.nxb3) + params["geometry"] = {"geometry": self.geometry} class UniformGrid(KamayanGrid): @@ -160,6 +163,7 @@ def __init__( dx1: _dx = None, dx2: _dx = None, dx3: _dx = None, + **kwargs, ) -> None: """Set parameters for a uniform grid. @@ -176,6 +180,7 @@ def __init__( dx1: zone size along dimensions of the domain. dx2: zone size along dimensions of the domain. dx3: zone size along dimensions of the domain. + kwargs: kwarguments to init KamayanGrid. Raises: ValueError: When a dimension is over-constrained @@ -224,6 +229,7 @@ def __init__( xbnd1=xbnd1, xbnd2=xbnd2, xbnd3=xbnd3, + **kwargs, ) @@ -333,6 +339,7 @@ def __init__( dx2: _dx = None, dx3: _dx = None, num_levels: int = 1, + **kwargs, ) -> None: """Set parameters for a uniform grid. @@ -353,6 +360,7 @@ def __init__( dx2: zone size along dimensions of the domain at highest resoltution. dx3: zone size along dimensions of the domain at highest resoltution. num_levels: number of total amr refinement levels + kwargs: kwarguments to init KamayanGrid. Raises: ValueError: When a dimension is over-constrained @@ -394,6 +402,7 @@ def __init__( xbnd1=xbnd1, xbnd2=xbnd2, xbnd3=xbnd3, + **kwargs, ) self.refinement_fields = RefinementCollection() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5d5386a8..319a0960 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1,5 +1,7 @@ set(_sources driver/kamayan_driver.cpp + grid/boundary_conditions.cpp + grid/coordinates.cpp grid/grid.cpp grid/grid_refinement.cpp kamayan/callback_dag.cpp @@ -28,6 +30,8 @@ set(_pybind_sources # maybe I should just glob for tests... set(_test_sources grid/tests/test_grid.cpp + grid/tests/test_geometry.cpp + grid/tests/test_coordinates_pack.cpp kamayan_utils/tests/test_strings.cpp kamayan_utils/tests/test_type_list.cpp kamayan_utils/tests/test_type_list_array.cpp diff --git a/src/dispatcher/options.hpp b/src/dispatcher/options.hpp index 68d2a328..8409f1fe 100644 --- a/src/dispatcher/options.hpp +++ b/src/dispatcher/options.hpp @@ -17,19 +17,41 @@ namespace kamayan { -template -struct OptList {}; +// checks that a functor can be instantiated for EnumV +// used to describe functors that can be used with OptList::dispatch below +template +concept OptionFunctor = requires { + requires std::is_void_v().template operator()( + std::declval()...))>; +}; // used to enumerate the allowed values of a PolyOpt in a given // dispatch functor. By default this list is exactly what is written // unless cmake has been configured with -DOPT_enum_opt="enum_v,..." // in that case OptInf::isdef == true and the ParmList() // will just be those defined at configure time -template +template requires(PolyOpt) -struct OptList { +struct OptList { using type = enum_opt; static constexpr auto value = OptInfo::ParmList(); + + // used to dispatch to a runtime template specialization of a type's operator() + template + requires(OptionFunctor && ...) + static bool dispatch(const Functor &functor, const enum_opt parm, Args &&...args) { + bool found_parm = false; + (void)(( + [&]() { + if (parm == enum_vs) { + found_parm = true; + functor.template operator()(std::forward(args)...); + } + }(), + !found_parm) && + ...); + return found_parm; + } }; template diff --git a/src/driver/kamayan_driver.cpp b/src/driver/kamayan_driver.cpp index 0fb41179..c36e9fc6 100644 --- a/src/driver/kamayan_driver.cpp +++ b/src/driver/kamayan_driver.cpp @@ -206,15 +206,6 @@ TaskID KamayanDriver::BuildTaskListRKStage(TaskList &task_list, const Real &dt, auto set_fluxes = parthenon::AddFluxCorrectionTasks( calc_fluxes, task_list, md0, md0->GetMeshPointer()->multilevel); - units_->AddTasksDAG([](KamayanUnit *u) -> auto & { return u->PrepareConserved; }, - [&](KamayanUnit *unit) { - std::string task_label = unit->Name() + "::PrepareConserved"; - prepare = task_list.AddTask(calc_fluxes, task_label, - unit->PrepareConserved.callback, - md1.get()); - }, - "PrepareConserved"); - // now set dudt using flux-divergence / discrete stokes theorem build_dudt = task_list.AddTask(set_fluxes, "grid::FluxesToDuDt", grid::FluxesToDuDt, md0.get(), mdudt.get()); @@ -229,8 +220,18 @@ TaskID KamayanDriver::BuildTaskListRKStage(TaskList &task_list, const Real &dt, next = unit->AddTasksOneStep(next, task_list, md0.get(), mdudt.get()); }); + prepare = next; + units_->AddTasksDAG([](KamayanUnit *u) -> auto & { return u->PrepareConserved; }, + [&](KamayanUnit *unit) { + std::string task_label = unit->Name() + "::PrepareConserved"; + prepare = + task_list.AddTask(prepare, task_label, + unit->PrepareConserved.callback, md0.get()); + }, + "PrepareConserved"); + if (flux_callbacks.size() + one_step_callbacks.size() > 0) { - next = grid::ApplyDuDt(next | prepare, task_list, mbase.get(), md0.get(), md1.get(), + next = grid::ApplyDuDt(prepare, task_list, mbase.get(), md0.get(), md1.get(), mdudt.get(), beta, dt); // now we might need to prepare the conserved vars for the next step diff --git a/src/grid/boundary_conditions.cpp b/src/grid/boundary_conditions.cpp new file mode 100644 index 00000000..f98b393f --- /dev/null +++ b/src/grid/boundary_conditions.cpp @@ -0,0 +1,84 @@ +#include "bvals/boundary_conditions.hpp" +#include +#include +#include +#include + +#include + +#include "basic_types.hpp" +#include "grid/grid.hpp" +#include "grid/grid_types.hpp" +#include "kamayan/fields.hpp" +#include "pack/pack_utils.hpp" + +namespace kamayan::grid { +using TE = TopologicalElement; +// use this as an interface to write boundary conditions per variable topological element +// types that can be registered through parthenon's appliciation input interface at the +// bottom of this file +using BoundaryConditionInterface = + std::function &, bool)>; + +void Axisymmetric(const TE el, std::shared_ptr &mbd, bool coarse) { + const auto ttFlag = [el] { + const auto tt = GetTopologicalType(el); + switch (tt) { + case (TopologicalType::Cell): + return Metadata::Cell; + case (TopologicalType::Face): + return Metadata::Face; + case (TopologicalType::Edge): + return Metadata::Edge; + case (TopologicalType::Node): + return Metadata::Node; + default: + PARTHENON_FAIL("Unknown topological type") + } + }(); + + std::vector flags{Metadata::FillGhost, ttFlag}; + std::set opts = coarse ? std::set{PDOpt::Coarse} : std::set{}; + const auto desc = + MakePackDescriptor(mbd.get(), flags, opts); + + auto q = desc.GetPack(mbd.get()); + const int b = 0; + const int lstart = q.GetLowerBoundHost(b); + const int lend = q.GetUpperBoundHost(b); + // early return if we don't get anything + if (lend < lstart) return; + auto nb = parthenon::IndexRange{lstart, lend}; + + MeshBlock *pmb = mbd->GetBlockPointer(); + const auto &bounds = (coarse ? pmb->c_cellbounds : pmb->cellbounds); + + const auto &range = bounds.GetBoundsI(IndexDomain::interior, el); + const int ref = range.s; + + constexpr IndexDomain domain = IndexDomain::inner_x1; + // used for reflections + const int offset = 2 * ref - 1; + + pmb->par_for_bndry( + PARTHENON_AUTO_LABEL, nb, domain, el, coarse, false, + KOKKOS_LAMBDA(const int &l, const int &k, const int &j, const int &i) { + const auto vec_component = q(b, el, l).vector_component; + const Real factor = + (vec_component == 1 || vec_component == 3 || el == TE::F1) ? -1.0 : 1.0; + q(b, el, l, k, j, i) = factor * q(b, el, l, k, j, offset - i); + }); +} + +parthenon::BValFunc GenericBC(BoundaryConditionInterface bndry_func) { + return [=](std::shared_ptr &mbd, bool coarse) { + for (auto el : {TE::CC, TE::F1, TE::F2, TE::F3, TE::E1, TE::E2, TE::E3, TE::NN}) + bndry_func(el, mbd, coarse); + }; +} + +void RegisterBoundaryConditions(parthenon::ApplicationInput *app) { + app->RegisterBoundaryCondition(parthenon::BoundaryFace::inner_x1, "axisymmetric", + GenericBC(Axisymmetric)); +} +} // namespace kamayan::grid diff --git a/src/grid/coordinates.cpp b/src/grid/coordinates.cpp new file mode 100644 index 00000000..52550b2e --- /dev/null +++ b/src/grid/coordinates.cpp @@ -0,0 +1,86 @@ +#include "grid/coordinates.hpp" + +#include "basic_types.hpp" +#include "grid/geometry.hpp" +#include "grid/grid.hpp" +#include "grid/grid_types.hpp" +#include "kamayan_utils/parallel.hpp" + +namespace kamayan::grid { + +template +void FillCoords(const Functor &functor, const parthenon::IndexShape cellbounds) { + auto [kb, jb, ib] = CoordinateIndexRanges(cellbounds); + par_for(PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, functor); +} + +template +void CalculateCoordinates(const Coordinates &coords, auto &pack, + const parthenon::IndexShape &cellbounds) { + FillCoords( + KOKKOS_LAMBDA(const int k, const int j, const int i) { + pack(0, coords::Volume(), k, j, i) = coords.CellVolume(k, j, i); + }, + cellbounds); + + [&]() { + (FillCoords>( + KOKKOS_LAMBDA(const int k, const int j, const int i) { + pack(0, coords::Dx(), k, j, i) = coords.template Dx(); + }, + cellbounds), + ...); + (FillCoords>( + KOKKOS_LAMBDA(const int k, const int j, const int i) { + pack(0, coords::X(), k, j, i) = coords.template Xi(k, j, i); + }, + cellbounds), + ...); + (FillCoords>( + KOKKOS_LAMBDA(const int k, const int j, const int i) { + pack(0, coords::Xc(), k, j, i) = coords.template Xc(k, j, i); + }, + cellbounds), + ...); + (FillCoords>( + KOKKOS_LAMBDA(const int k, const int j, const int i) { + pack(0, coords::Xf(), k, j, i) = coords.template Xf(k, j, i); + }, + cellbounds), + ...); + (FillCoords>( + KOKKOS_LAMBDA(const int k, const int j, const int i) { + pack(0, coords::FaceArea(), k, j, i) = + coords.template FaceArea(k, j, i); + }, + cellbounds), + ...); + (FillCoords>( + KOKKOS_LAMBDA(const int k, const int j, const int i) { + pack(0, coords::EdgeLength(), k, j, i) = + coords.template EdgeLength(k, j, i); + }, + cellbounds), + ...); + }.template operator()(); +} + +void CalculateCoordinates(MeshBlock *mb) { + const auto geometry = GetConfig(mb)->Get(); + CalculateCoordinates(mb, geometry); +} + +void CalculateCoordinates(MeshBlock *mb, Geometry geom) { + GeometryOptions::dispatch( + [&]() { + auto cellbounds = mb->cellbounds; + + auto pack = GetPack(CoordFields(), mb); + auto coords = Coordinates(pack, 0); + + CalculateCoordinates(coords, pack, cellbounds); + }, + geom); +} + +} // namespace kamayan::grid diff --git a/src/grid/coordinates.hpp b/src/grid/coordinates.hpp new file mode 100644 index 00000000..ca6d582f --- /dev/null +++ b/src/grid/coordinates.hpp @@ -0,0 +1,660 @@ +#ifndef GRID_COORDINATES_HPP_ +#define GRID_COORDINATES_HPP_ +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include "grid/geometry.hpp" +#include "grid/geometry_types.hpp" +#include "grid/grid.hpp" +#include "grid/grid_types.hpp" +#include "interface/variable_state.hpp" +#include "kamayan_utils/parallel.hpp" +#include "kamayan_utils/strings.hpp" +#include "kamayan_utils/type_abstractions.hpp" +#include "kamayan_utils/type_list.hpp" +#include "kokkos_types.hpp" +#include "utils/concepts_lite.hpp" + +namespace kamayan::grid { + +template +struct CoordVar : parthenon::variable_names::base_t { + template + KOKKOS_INLINE_FUNCTION CoordVar(Args &&...args) + : parthenon::variable_names::base_t(std::forward(args)...) {} + + static std::string name() { return std::string(var_name.value); } + static constexpr auto label = var_name; + static constexpr auto element = el; +}; + +namespace impl { +template +constexpr auto AxisNamer(const strings::CompileTimeString base, Axis ax) { + if (ax == Axis::IAXIS) { + return strings::concat_cts(base, strings::make_cts("1")); + } else if (ax == Axis::JAXIS) { + return strings::concat_cts(base, strings::make_cts("2")); + } else if (ax == Axis::KAXIS) { + return strings::concat_cts(base, strings::make_cts("3")); + } + return strings::concat_cts(base, strings::make_cts("1")); +} +} // namespace impl + +constexpr TopologicalElement AxisToTE(const TopologicalElement el, Axis ax) { + return static_cast(static_cast(el) + AxisToInt(ax) - 1); +} + +template +using AxisCoord = CoordVar; + +namespace coords { +template +using Dx = AxisCoord<"geom.Dx", ax>; + +template +using X = AxisCoord<"geom.X", ax>; + +template +using Xc = AxisCoord<"geom.Xc", ax>; + +template +using Xf = AxisCoord<"geom.Xf", ax, AxisToTE(TopologicalElement::F1, ax)>; + +template +using FaceArea = AxisCoord<"geom.FaceArea", ax, AxisToTE(TopologicalElement::F1, ax)>; + +template +using EdgeLength = AxisCoord<"geom.EdgeLength", ax, AxisToTE(TopologicalElement::E1, ax)>; + +using Volume = CoordVar<"geom.Volume">; +} // namespace coords + +template