Add cylindrical geometry

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()

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 <Axis ax> 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<geom>` (`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<Axis::KAXIS>()` returns `2*pi` (implicit azimuthal direction)
+  - `Xc<Axis::IAXIS>(idx)` computes the radial *centroid* of the cell (not just the
+    midpoint); `Xi<Axis::IAXIS>(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<geom, ...>` (`src/grid/coordinates.hpp`) wraps a `SparsePack` holding
+these `geom.*` fields and exposes the same API as `grid::Coordinates<geom>` 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<geom>` in a
+  `variant` and forwards calls via `visit`.
+- `grid::GenericCoordinatePack` (`src/grid/coordinates.hpp`) does the same for
+  `grid::CoordinatePack<geom, ...>`.
+
+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<geom>`/`grid::CoordinatePack<geom>`.
+
 ## Parameters
 {!assets/generated/grid_parms.md!}

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<kamayan::StateDescriptor>("Test Package");
+    auto pkg = std::make_shared<kamayan::KamayanUnit>("Test Package");
     auto block_list = kamayan::MakeTestBlockList(pkg, 1, 8, 3);
     auto tc = driver.MakeTaskCollection(block_list, 1);
     std::cout << tc;

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
         )

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()
 

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

src/dispatcher/options.hpp

@@ -17,19 +17,41 @@
 
 namespace kamayan {
 
-template <typename, auto...>
-struct OptList {};
+// checks that a functor can be instantiated for EnumV
+// used to describe functors that can be used with OptList::dispatch below
+template <typename F, typename EnumT, EnumT EnumV, typename... Args>
+concept OptionFunctor = requires {
+  requires std::is_void_v<decltype(std::declval<F>().template operator()<EnumV>(
+      std::declval<Args>()...))>;
+};
 
 // 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<enum_opt>::isdef == true and the ParmList()
 // will just be those defined at configure time
-template <typename enum_opt, auto enum_v, auto... enum_vs>
+template <typename enum_opt, enum_opt... enum_vs>
 requires(PolyOpt<enum_opt>)
-struct OptList<enum_opt, enum_v, enum_vs...> {
+struct OptList {
   using type = enum_opt;
   static constexpr auto value = OptInfo<enum_opt>::ParmList();
+
+  // used to dispatch to a runtime template specialization of a type's operator()
+  template <typename Functor, typename... Args>
+  requires(OptionFunctor<Functor, enum_opt, enum_vs, Args...> && ...)
+  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()<enum_vs>(std::forward<Args>(args)...);
+                 }
+               }(),
+               !found_parm) &&
+           ...);
+    return found_parm;
+  }
 };
 
 template <typename OL>

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