Reorganize gk_rz Module Architecture

I’ll use the design, architecture, data, and API/CLI skills to inspect gk_rz and plan a split between interpolation and Gkeyll-specific file handling.

The code already separates projection construction from applying a projection, so there’s a useful boundary to build on. The 3-D path also performs FFT reconstruction with twist-and-shift boundary handling; calling it generic “nodal interpolation” would hide those assumptions. I’ll also account for gk_fluxsurf, which shares the geometry loader.

I recommend keeping explicit geometry transformations in operations/ and moving Gkeyll geometry discovery and loading into diagnostics/gk/. That gives you pg.gk.rz(data) for the convenient simulation-aware workflow, with a reusable core underneath.

I would adjust the argument slightly: scope and the gk prefix aren’t reliable boundaries. Mana is right that a coordinate transformation can apply to multiple equation systems. Your stronger argument is that a core transformation should receive its geometry explicitly; it shouldn’t need to interpret a simulation filename to discover its inputs.

The current implementation already provides much of this separation through resolve_rz_projection() and map_to_rz() in rz.py. The refactor should build on that.

1. State the architectural rule explicitly.

   Define operations as equation-blind transformations whose mathematical inputs are explicit. Diagnostics own model-specific compositions, auxiliary-file selection, and interpretation of simulation conventions.

   This does not require banning every operations → io import: saving data is an explicit effect. The boundary here is implicitly discovering additional scientific inputs from the source filename.

   Update the architecture contract and guidance together. Currently, the import-contract commentary explicitly permits domain-specific operation subpackages, while the architecture skill describes flat, equation-blind verbs.

2. Split geometry loading from geometry data and mathematics.

   geometry.py currently mixes these responsibilities. Separate them into:

   Owner	Responsibility
   diagnostics/gk/geometry.py	Choose auxiliary files, honor overrides, resolve blocks, interpret Gkeyll geometry metadata and layouts
   operations/geometry.py	In-memory geometry records and shared transformation support
   numerics/	Array-only resampling and reconstruction kernels
   Existing io naming machinery	Parse Gkeyll output filenames

   Keep one implementation of filename parsing. Use diagnostic discovery for shared output discovery; keep GK auxiliary suffixes and nodes-versus-mapc2p preference in the GK loader.

   Geometry passed into the core should contain coordinates and physical components, including optional boundary geometry, with no filename or discovery policy.

3. Retain a narrowly defined R–Z operation.

   Move projection construction and application into a flat operations/rz.py. Keep map_to_rz as the core verb and retain reusable projection construction.

   The intended API would be:

   Make projection keyword-only to match the operation signature convention. The functional and fluent forms should alias the same callable.

   Avoid interpolate_from_file: it describes an input mechanism rather than the mathematical transformation. Also avoid promising arbitrary nodal interpolation. The existing 3-D algorithm includes toroidal periodicity, FFT reconstruction, and twist-and-shift handling. Document those coordinate assumptions as part of the operation’s contract.

   Preserve the existing raw modal input requirement, component selection, output grids, and _result behavior. Expanding supported representations should be separate work.

4. Make pg.gk.rz the canonical convenience composition.

   Implement rz() in diagnostics/gk/rz.py as a short sequence:

   • Resolve and load geometry.
   • Construct the projection.
   • Apply map_to_rz.

   Preserve current overrides, geometry preference, optional corner loading, and numerical defaults.

   Move per-block projection discovery and lookup here as well. There are currently duplicate batch implementations in operations and diagnostics; consolidate them. Projection reuse must check grid compatibility, even when datasets share a filename prefix.

   Account for gk_fluxsurf in the same refactor: it shares this geometry machinery. Apply the same ownership split so the new core never needs to import the diagnostic loader.

5. Migrate the public API and generated CLI deliberately.

   Expose rz from pg.gk, giving the generated CLI the existing spelling gk_rz. Classify it as Section.DIAGNOSTICS while retaining Execution.MAP_REPLACE, so existing CLI pipelines continue to work.

   Expose the explicit core operation through pg.map_to_rz and GData.map_to_rz. Keep geometry-object APIs explicitly hidden from CLI generation using the existing mechanism; building CLI geometry serialization is unnecessary for this refactor.

   Removing data.gk_rz() is a breaking Python API change. Schedule that removal for a breaking release. A temporary pg.gk_rz facade alias can point directly to pg.gk.rz, but an operations-to-diagnostics forwarding shim would violate the intended dependency direction. If fluent compatibility is mandatory for the current release, defer the final removal rather than adding that reverse dependency.

6. Verify the boundary and preserve scientific behavior.

   Split existing tests by responsibility, retaining their numerical references:

   • Core tests use explicit in-memory geometry and data without a filename.
   • Diagnostic tests cover file preference, overrides, corner geometry, and per-block resolution.
   • Numerical tests preserve 2-D mapping, 3-D Fourier phase behavior, component selection, and projection compatibility.
   • Surface tests verify canonical aliases, group broadcasting, CLI replacement behavior, and help classification.

   Run the architecture, API/CLI, R–Z, flux-surface, and multiblock checks, followed by the full suite. Update the R–Z example and documentation to show both convenience and explicit-geometry workflows.

The acceptance criterion is concrete: the core can transform a field using supplied geometry without knowing which simulation produced either one; pg.gk.rz supplies that simulation knowledge in one convenient call.

Yes—operations/map.py is a reasonable home for coordinate-mapping operations. I’d keep the function named map_to_rz: the module describes the family of transformations, while the function states its output.

Here, “universal” should mean independent of equation systems and file conventions. R–Z reconstruction still assumes particular coordinate geometry and, in 3-D, periodic field-aligned structure. I wouldn’t introduce a universal map() interface until another mapping establishes what it should share.

I propose this file layout:

File	Purpose and design
operations/map.py — new	Owns map_to_rz(data, *, projection, ...) and resolve_rz_projection(data, geometry, ...). Checks compatibility with the input field, calls existing DG interpolation and numerical kernels, and returns results through _result. Contains no file discovery, loading, species names, or GK metadata interpretation.
operations/geometry.py — new	Owns the frozen, in-memory Geometry, RzProjection, and FluxSurfaceGrid records and their structural invariants. Geometry describes actual coordinates and components, independent of their source. Shared geometry/grid support lives here only where multiple transformations need it.
numerics/rz.py — new	Owns the array-only R–Z reconstruction kernel, including the existing Fourier phase and twist-and-shift calculation. Receives all numerical inputs explicitly; knows nothing about GDataState, files, diagnostics, or CLI behavior.
numerics/resample.py — new	Owns the tensor-grid resampling calculation already shared by R–Z and flux-surface mapping. Accepts arrays and source/target coordinates. Keeps interpolation and extrapolation behavior explicit.
diagnostics/gk/geometry.py — new	Owns Gkeyll geometry resolution and loading: nodes-versus-mapc2p preference, explicit overrides, optional corner files, block substitution, and interpretation of geometry metadata and component layouts. Produces the core Geometry record. Uses existing readers and filename parsing.
diagnostics/gk/rz.py — replace compatibility implementation	Owns canonical rz(data, ...), exposed as pg.gk.rz. Composes geometry loading, projection construction, and mapping. Also owns simulation-aware batch projection resolution and per-block lookup, eliminating the current duplicate implementation.
diagnostics/gk/fluxsurf.py — update	Owns the corresponding GK convenience composition and per-block geometry resolution for flux surfaces. Its numerical transformation moves into operations/map.py, sharing the explicit geometry contract. This prevents the shared loader from pulling operations back into diagnostics.
io/naming.py — retain	Remains the authoritative filename parser. Reuse it; do not copy naming expressions into the new loader.
diagnostics/discovery.py — retain	Continues to own shared output discovery. GK-specific auxiliary-file selection remains in diagnostics/gk/geometry.py.

The public surfaces and verification need these accompanying changes:

File	Change
operations/__init__.py	Export the canonical mapping operations and geometry API.
gdata/gdata.py	Alias map_to_rz directly as a fluent operation. Remove gk_rz at the agreed breaking-release boundary.
diagnostics/gk/__init__.py	Export rz and attach diagnostic CLI metadata so its command remains gk_rz.
__init__.py	Re-export the core API; optionally retain pg.gk_rz as a transitional alias to pg.gk.rz.
tests/test_operations_map.py — new	Test mappings using explicit geometry and filename-free inputs, including result semantics and projection compatibility.
tests/test_diagnostics_gk_geometry.py — new	Test auxiliary-file discovery, geometry decoding, overrides, and blocks.
tests/test_diagnostics_gk_rz.py — new	Test the convenience composition and projection reuse.
Existing architecture, fluent, CLI, and multiblock tests	Update ownership and callable expectations; preserve numerical and pipeline behavior.
Existing R–Z examples and documentation	Demonstrate pg.gk.rz(data) and explicit projection reuse through data.map_to_rz(...).
Architecture skill and contract commentary	Record the agreed ownership rule consistently.

Finally, retire operations/gyrokinetics/ once both R–Z and flux-surface code have moved. It should not remain as a forwarding package that imports diagnostics.

The dependency flow becomes:

I’ll implement the split, remove operations/gyrokinetics and the legacy API names, and update the tests and documentation. I’m using the design, architecture, data, API/CLI, development, and testing skills for this refactor.

operations/map.py already contains a coordinate-grid mapping verb. I’ll add the explicit R–Z and flux-surface transformations there, with their shared geometry records and numerical kernels in separate modules. The GK wrappers will expose pg.gk.rz and pg.gk.fluxsurf; the CLI names will remain gk_rz and gk_fluxsurf.

The package split is in place, and the generated CLI still exposes both commands under their existing names. The first focused run passed 205 tests and found two test-migration issues: a positional argument that is now keyword-only, and an error test that needs explicit geometry to reach input validation. I’ll fix those and add filename-free mapping tests before running the full suite.

The filename-free tests now pass for both 2-D and 3-D fields, including flux-surface extraction. Cached mappings also reject a changed computational grid even when the filename prefix is unchanged. The full suite is running; I’m reviewing the diff for stale exports, documentation, and formatting.

The focused checks passed: 152 tests, with one existing skip for a node-plot fixture limitation. Lint and formatting checks pass too. The full suite is still running its documentation and downloadable-example checks, which compare the Python and CLI outputs.

Implemented the clean refactor.

• Removed operations/gyrokinetics and legacy Python APIs.
• Core mapping lives in operations/map.py, with explicit geometry and reusable projections.
• GK discovery and loading live in diagnostics/gk, exposed through pg.gk.rz() and pg.gk.fluxsurf().
• CLI names remain gk_rz and gk_fluxsurf, classified as diagnostics.
• Updated examples, documentation, architecture rules, and tests.

Validation: 2,000 passed, 6 skipped. Lint and formatting checks passed.