From f45fb37ce4ca9e0a93ab8c19e960264f77f93cb2 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Thu, 10 Sep 2026 16:49:25 -0400 Subject: [PATCH] Refactor gk-rz into the naming conventions which are used by GK and GK specific, and refactor out the core methods of interpolating on geometry into map.py in operations. Explicit geometry construction is kept general, but the GK-specific nature is kept in diagnostics. Examples are updated. CLI names are still gk_rz. Script pg.gk.rz() is a short hand for chaining a few core commands together, including constructing the geometry object and the projection. Entire-Checkpoint: 01M26HA71FS8K7HRX1BNFMXXQ5 --- .agents/skills/postgkyl-architecture/SKILL.md | 5 +- docs/source/physical-rz.rst | 9 +- examples/README.md | 4 +- examples/cli_tutorial.md | 9 +- examples/scripts/05_gk_rz.py | 18 +- src/postgkyl/__init__.py | 11 +- src/postgkyl/cli/discovery.py | 5 +- src/postgkyl/diagnostics/__init__.py | 9 +- src/postgkyl/diagnostics/gk/__init__.py | 27 +- src/postgkyl/diagnostics/gk/fluxsurf.py | 106 ++++-- src/postgkyl/diagnostics/gk/geometry.py | 184 +++++++++ src/postgkyl/diagnostics/gk/nodes.py | 5 +- src/postgkyl/diagnostics/gk/rz.py | 117 ++++-- src/postgkyl/gdata/gdata.py | 4 +- src/postgkyl/gdatastate/collection.py | 2 +- src/postgkyl/gdatastate/gdatastate.py | 2 +- src/postgkyl/numerics/resample.py | 13 + src/postgkyl/numerics/rz.py | 39 ++ src/postgkyl/operations/__init__.py | 27 +- src/postgkyl/operations/geometry.py | 204 ++++++++++ .../operations/gyrokinetics/__init__.py | 47 --- .../operations/gyrokinetics/fluxsurf.py | 218 ----------- .../operations/gyrokinetics/geometry.py | 355 ----------------- src/postgkyl/operations/gyrokinetics/rz.py | 359 ------------------ src/postgkyl/operations/map.py | 286 +++++++++++++- tests/test_diagnostics_gk_fluxsurf.py | 80 ++++ ...try.py => test_diagnostics_gk_geometry.py} | 100 +---- tests/test_diagnostics_gk_rz.py | 184 +++++++++ tests/test_diagnostics_programs_nodes.py | 9 +- tests/test_gdata_fluent.py | 6 +- tests/test_multiblock.py | 14 +- tests/test_numerics_rz.py | 17 + tests/test_operations_geometry.py | 106 ++++++ tests/test_operations_gk_fluxsurf.py | 173 --------- tests/test_operations_gk_rz.py | 308 --------------- tests/test_operations_map_fluxsurf.py | 105 +++++ tests/test_operations_map_rz.py | 223 +++++++++++ tests/test_postgkyl.py | 13 +- 38 files changed, 1701 insertions(+), 1702 deletions(-) create mode 100644 src/postgkyl/diagnostics/gk/geometry.py create mode 100644 src/postgkyl/numerics/resample.py create mode 100644 src/postgkyl/numerics/rz.py create mode 100644 src/postgkyl/operations/geometry.py delete mode 100644 src/postgkyl/operations/gyrokinetics/__init__.py delete mode 100644 src/postgkyl/operations/gyrokinetics/fluxsurf.py delete mode 100644 src/postgkyl/operations/gyrokinetics/geometry.py delete mode 100644 src/postgkyl/operations/gyrokinetics/rz.py create mode 100644 tests/test_diagnostics_gk_fluxsurf.py rename tests/{test_operations_gk_geometry.py => test_diagnostics_gk_geometry.py} (53%) create mode 100644 tests/test_diagnostics_gk_rz.py create mode 100644 tests/test_numerics_rz.py create mode 100644 tests/test_operations_geometry.py delete mode 100644 tests/test_operations_gk_fluxsurf.py delete mode 100644 tests/test_operations_gk_rz.py create mode 100644 tests/test_operations_map_fluxsurf.py create mode 100644 tests/test_operations_map_rz.py diff --git a/.agents/skills/postgkyl-architecture/SKILL.md b/.agents/skills/postgkyl-architecture/SKILL.md index 17462d09..93a84445 100644 --- a/.agents/skills/postgkyl-architecture/SKILL.md +++ b/.agents/skills/postgkyl-architecture/SKILL.md @@ -29,11 +29,12 @@ remain GData without importing gdata. Reuse `gdatastate/guards.py` and Readers return `(grid, values)` and fill plain metadata; they never import state. Diagnostics are free functions under model families `gk`, `vm`, `pkpm`, or `mom`, -not GData methods. Model-specific loading belongs beside its physics. Resolve +not GData methods. Model-specific loading and auxiliary-file discovery belong +beside their compositions. Coordinate operations receive explicit geometry or projections; +operations has no equation-family subpackages. Resolve output stems/frames through `diagnostics/discovery.py`; keep quantity vocabulary in the equation module's `VARIABLES` table. Use public functions from lower layers and return a state via `_result`, or a Figure for program diagnostics. -for the current major version. Run the import, foreign-floor, facade, and canonical-callable contracts in `tests/test_postgkyl.py` after structural changes. diff --git a/docs/source/physical-rz.rst b/docs/source/physical-rz.rst index e63388c2..8c50c12a 100644 --- a/docs/source/physical-rz.rst +++ b/docs/source/physical-rz.rst @@ -9,6 +9,13 @@ The Python and CLI outputs below are both generated and compared when this websi Computational coordinates are useful for analysis, but a poloidal view places the field in its physical geometry. Keep the companion ``rt_gk_tcv_nt_iwl_3x2v_p1-geo_int_mapc2p.gkyl`` beside the electron density -file so the operation can resolve it by the simulation prefix. +file so ``pg.gk.rz`` can resolve it by the simulation prefix. This diagnostic +loads Gkeyll geometry and composes equation-independent coordinate operations. + +For explicit geometry, construct ``pg.Geometry`` from coordinate arrays, build +``pg.resolve_rz_projection(data, geometry)``, and apply it with +``data.map_to_rz(projection=projection)``. The projection can be reused for +fields on the same computational grid. Three-dimensional reconstruction +assumes periodic field-aligned coordinates and twist-and-shift boundaries. .. include:: _pairs/05_gk_rz.inc diff --git a/examples/README.md b/examples/README.md index ccd49e95..f475cb73 100644 --- a/examples/README.md +++ b/examples/README.md @@ -20,7 +20,7 @@ of truth for "does the tutorial still work," not just this README. | [`02_arithmetic_and_numpy.py`](scripts/02_arithmetic_and_numpy.py) | Weak DG algebra on raw modal data (`*`, `/`, `+`, `.integrate()`) vs. plain NumPy math after `.interpolate()`, and the guardrail between them | | [`03_diagnostics_five_moment.py`](scripts/03_diagnostics_five_moment.py) | The `diagnostics` layer: equation-specific physics (`postgkyl.diagnostics.mom.five_moment`) on top of a `GData`, on a generated shock-tube initial state | | [`04_gyrokinetics.py`](scripts/04_gyrokinetics.py) | The gyrokinetic diagnostics: `pg.gk.load_quantity` (named moments/geometry, resolved by naming convention) and `pg.gk.load_distf` (full distribution function), on the `rt_gk_tcv_iwl*` fixtures | -| [`05_gk_rz.py`](scripts/05_gk_rz.py) | The gyrokinetic R-Z operation: one-line fluent and functional calls plus projection reuse over multiple toroidal angles | +| [`05_gk_rz.py`](scripts/05_gk_rz.py) | Gkeyll geometry discovery with `pg.gk.rz`, plus explicit projection reuse with `map_to_rz` | | [`06_growth.py`](scripts/06_growth.py) | Recover a known energy growth rate and inspect log-space residuals | | [`07_collect_animate.py`](scripts/07_collect_animate.py) | Load many frames, collect a space–time diagram, and animate 1D/2D travelling waves | | [`08_plotly.py`](scripts/08_plotly.py) | Interactive 2D height surfaces and 3D volume isosurfaces | @@ -49,7 +49,7 @@ See [`cli_tutorial.md`](cli_tutorial.md) -- inspecting a file, the `interpolate`/`select`/`plot` chain, discontinuity-preserving plots with `local_poly`, DynVector `info`/`fit`, the gyrokinetic loaders (`gk_load_quantity`, `gk_load_distf`), the `gk_rz` -transformation, `save`, and the generated command inventory. +diagnostic, `save`, and the generated command inventory. ## Running the tests diff --git a/examples/cli_tutorial.md b/examples/cli_tutorial.md index b9551690..480264a3 100644 --- a/examples/cli_tutorial.md +++ b/examples/cli_tutorial.md @@ -125,10 +125,10 @@ pgkyl gk_load_distf --name tests/test_data/rt_gk_tcv_iwl_1x2v_p1 \ ## 7. Map a gyrokinetic field to R-Z -`gk_rz` is a data transformation: it interpolates one raw DG component and +`gk_rz` is a diagnostic composition: it interpolates one raw DG component and maps it onto the physical poloidal plane. Geometry is inferred from the field's filename, preferring nodal geometry and falling back to modal -`mapc2p` geometry. The CLI and Python calls below use the same operation and +`mapc2p` geometry. The CLI and Python calls below use the same diagnostic and defaults: ```bash @@ -139,9 +139,10 @@ pgkyl tests/test_data/rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl \ ```python import postgkyl as pg -mapped = pg.load( +data = pg.load( "tests/test_data/rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl" -).gk_rz(nz_interp=2) +) +mapped = pg.gk.rz(data, nz_interp=2) ``` ## 8. Saving to another format diff --git a/examples/scripts/05_gk_rz.py b/examples/scripts/05_gk_rz.py index 40927921..fdf24388 100644 --- a/examples/scripts/05_gk_rz.py +++ b/examples/scripts/05_gk_rz.py @@ -12,7 +12,6 @@ import numpy as np import postgkyl as pg -from postgkyl.operations import gyrokinetics as gk_ops from _example_paths import TEST_DATA, prepare_output_dir @@ -21,7 +20,7 @@ data = pg.load(TEST_DATA / "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") # The common path: geometry is inferred from FIELD's simulation prefix. -mapped = data.gk_rz(z_axis=0.0, phi_tor=0.0, nz_interp=2) +mapped = pg.gk.rz(data, z_axis=0.0, phi_tor=0.0, nz_interp=2) fig = mapped.plot(title="Electron density in the poloidal plane", xlabel="R [m]", ylabel="Z [m]", @@ -30,19 +29,12 @@ no_show=True) fig.savefig(OUTPUT_DIR / "05_gk_rz.png") -# The functional spelling is the identical operation. -functional = pg.gk_rz(data, z_axis=0.0, phi_tor=0.0, nz_interp=2) -np.testing.assert_allclose(functional.values, mapped.values) - # Reuse geometry and projection when several fields/frames or toroidal angles # share one computational grid. -geometry = gk_ops.resolve_geometry(data.file_name) -projection = gk_ops.resolve_rz_projection(data, - geometry, - z_axis=0.0, - nz_interp=2) -at_zero = gk_ops.map_to_rz(data, projection, phi_tor=0.0) -at_quarter_turn = gk_ops.map_to_rz(data, projection, phi_tor=np.pi / 2) +geometry = pg.gk.resolve_geometry(data.file_name) +projection = pg.resolve_rz_projection(data, geometry, z_axis=0.0, nz_interp=2) +at_zero = data.map_to_rz(projection=projection, phi_tor=0.0) +at_quarter_turn = pg.map_to_rz(data, projection=projection, phi_tor=np.pi / 2) np.testing.assert_allclose(at_zero.values, mapped.values) assert not np.allclose( at_zero.values, at_quarter_turn.values, rtol=1e-12, atol=0.0) diff --git a/src/postgkyl/__init__.py b/src/postgkyl/__init__.py index e68778a4..a8c463d0 100644 --- a/src/postgkyl/__init__.py +++ b/src/postgkyl/__init__.py @@ -17,7 +17,7 @@ info <- operations/ (the info verb, one-or-many) integrate <- operations/ (grid integral, via Gkeyll) interpolate, select <- operations/ (functional verb spellings) - gk_rz <- operations/gyrokinetics/ (domain operation) + map_to_rz <- operations/map.py represent, apply <- operations/ (value_form verbs) available_evaluate_operators <- operations/ (``evaluate``'s RPN token vocabulary) save <- io/ (file output) @@ -75,7 +75,10 @@ select, val2coord, ) -from postgkyl.operations.gyrokinetics import gk_fluxsurf, gk_rz +from postgkyl.operations import (Geometry, RzProjection, FluxSurfaceGrid, + map_to_rz, resolve_rz_projection, + extract_flux_surface, + resolve_flux_surface_grid) from postgkyl.render import animate, plot, plotly, plotly_animate, pyvista from postgkyl.gdatastate import group_blocks from postgkyl.cli_spec import hidden @@ -93,7 +96,9 @@ "integrate", "interpolate", "local_poly", "select", "average", "eval_at_coord_proj", "fft", "magsq", "mask", "grid", "val2coord", "extract_input", "fit", "growth", "differentiate", "map", "represent", - "apply", "gk_rz", "gk_fluxsurf", "save", "collect", "evaluate", "relchange", + "apply", "Geometry", "RzProjection", "FluxSurfaceGrid", "map_to_rz", + "resolve_rz_projection", "extract_flux_surface", + "resolve_flux_surface_grid", "save", "collect", "evaluate", "relchange", "animate", "plotly_animate", "sort", "available_evaluate_operators", "plotly", "pyvista", "gk", "__version__", "version_report" ] diff --git a/src/postgkyl/cli/discovery.py b/src/postgkyl/cli/discovery.py index e8c1670c..d3d42d51 100644 --- a/src/postgkyl/cli/discovery.py +++ b/src/postgkyl/cli/discovery.py @@ -102,10 +102,7 @@ def consider(obj, path: str, name: str) -> None: namespace = relative.rsplit(".", 1)[-1] for name, value in _functions(module): if not value.__module__.startswith("postgkyl.diagnostics"): - # Compatibility re-exports owned by a lower layer keep that owner's - # canonical command (for example operations.gyrokinetics.gk_rz -> - # ``gk_rz``); the diagnostic alias is classified but does not invent - # a second command or move it into the wrong help section. + # Lower-layer re-exports retain their canonical command vocabulary. _classify(value, f"{module.__name__}.{name}") continue consider(value, f"{module.__name__}.{name}", f"{namespace}_{name}") diff --git a/src/postgkyl/diagnostics/__init__.py b/src/postgkyl/diagnostics/__init__.py index 21a1a857..4a14e78a 100644 --- a/src/postgkyl/diagnostics/__init__.py +++ b/src/postgkyl/diagnostics/__init__.py @@ -5,9 +5,8 @@ here take loaded ``GData``/``GDataState`` (one or several) plus physical scalars as keyword-only options, and return a ``GDataState`` (via ``_result``) or, in later layers, a ``Figure``. Equation-blind core verbs -stay in flat ``operations`` modules. Domain-specific transformations live in -operation subpackages (for example ``operations.gyrokinetics``); this layer -is reserved for code that knows what field components physically mean. +stay in flat ``operations`` modules. Model-specific auxiliary discovery and +interpretation live beside their diagnostic compositions. The four public packages mirror Gkeyll's model families: ``gk``, ``vm``, ``pkpm``, and ``mom``. The equation-blind ``discovery`` module @@ -150,11 +149,7 @@ def _combine(function, dataset_names: tuple[str, ...]) -> None: "is_geo_mapc2p", "multib_tag", "nodes_to_RZ", - "map_to_rz", "resolve_geometry", - "resolve_rz_projection", - "extract_flux_surface", - "resolve_flux_surface_grid", ): _function = getattr(gk, _name) if command_spec(_function) is None and hidden_spec(_function) is None: diff --git a/src/postgkyl/diagnostics/gk/__init__.py b/src/postgkyl/diagnostics/gk/__init__.py index aa528cf7..78991429 100644 --- a/src/postgkyl/diagnostics/gk/__init__.py +++ b/src/postgkyl/diagnostics/gk/__init__.py @@ -7,9 +7,8 @@ instruction file's decision record): splitting resolution from physics would give gyrokinetics two homes for one piece of equation knowledge. Only the equation-blind stem/frame discovery is shared, via -``postgkyl.diagnostics.discovery``. Geometry-only transformations live below -this physics layer in ``postgkyl.operations.gyrokinetics``; the R-Z and -flux-surface names exported here are compatibility aliases. +``postgkyl.diagnostics.discovery``. Geometry discovery lives here and composes +explicit coordinate transformations from ``operations.map``. """ from __future__ import annotations @@ -37,12 +36,11 @@ # Layer 13: program-scale diagnostics ported from src_bak's apps/gk_*.py. from .energy_balance import EnergyBalanceTraces, energy_balance_error, energy_balance from .particle_balance import ParticleBalanceTraces, particle_balance, particle_balance_error -from .nodes import GKYL_GEOMETRY_ID, nodes, is_geo_mapc2p, multib_tag, nodes_to_RZ +from .nodes import nodes, multib_tag, nodes_to_RZ -# Compatibility exports: canonical transformation APIs now live under -# postgkyl.operations.gyrokinetics. These imports are exact aliases. -from .rz import Geometry, RzProjection, gk_rz, map_to_rz, resolve_geometry, resolve_rz_projection -from .fluxsurf import FluxSurfaceGrid, extract_flux_surface, resolve_flux_surface_grid +from .geometry import GKYL_GEOMETRY_ID, is_geo_mapc2p, resolve_geometry +from .rz import rz +from .fluxsurf import fluxsurf from typing import Annotated @@ -69,6 +67,9 @@ for _function in (energy_balance, particle_balance, nodes): command(_REPORT_SPEC)(_function) +for _function in (rz, fluxsurf): + command(CommandSpec(Section.DIAGNOSTICS, Execution.MAP_REPLACE))(_function) + __all__ = [ "load_distf", "resolve_frames", @@ -100,13 +101,7 @@ "is_geo_mapc2p", "multib_tag", "nodes_to_RZ", - "Geometry", - "RzProjection", - "gk_rz", - "map_to_rz", "resolve_geometry", - "resolve_rz_projection", - "FluxSurfaceGrid", - "extract_flux_surface", - "resolve_flux_surface_grid", + "rz", + "fluxsurf", ] diff --git a/src/postgkyl/diagnostics/gk/fluxsurf.py b/src/postgkyl/diagnostics/gk/fluxsurf.py index 64bcc7bc..45194182 100644 --- a/src/postgkyl/diagnostics/gk/fluxsurf.py +++ b/src/postgkyl/diagnostics/gk/fluxsurf.py @@ -1,23 +1,83 @@ -"""Compatibility aliases for gyrokinetic flux-surface operations. - -Canonical imports live in :mod:`postgkyl.operations.gyrokinetics`; this path -is scheduled for removal in the next major version. -""" - -from postgkyl.operations.gyrokinetics.fluxsurf import ( - FluxSurfaceGrid, - Geometry, - extract_flux_surface, - flux_surface_grids, - grid_for, - resolve_flux_surface_grid, -) - -__all__ = [ - "Geometry", - "FluxSurfaceGrid", - "extract_flux_surface", - "flux_surface_grids", - "grid_for", - "resolve_flux_surface_grid", -] +"""Gkeyll geometry discovery composed with explicit fluxsurf mapping.""" +from __future__ import annotations + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.operations.map import resolve_flux_surface_grid, extract_flux_surface +from postgkyl.operations.geometry import FluxSurfaceGrid, validate_mapping_grid +from .geometry import geometry_prefix, per_block_path, resolve_geometry + + +def flux_surface_grids(datasets, + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + x_idx: int = 0, + nphi: int = 128, + nz_interp: int = 8) -> dict[str | None, FluxSurfaceGrid]: + """Build one reusable flux-surface grid per block geometry.""" + grids: dict[str | None, FluxSurfaceGrid] = {} + for data in datasets: + key = geometry_prefix(data.file_name) + if key in grids: + validate_mapping_grid(data, grids[key].computational_grid) + continue + block = data.ctx.get("block") + geometry = resolve_geometry(data.file_name, + mapc2p=per_block_path(mapc2p, block), + nodes_file=per_block_path(nodes_file, block)) + grids[key] = resolve_flux_surface_grid(data, + geometry, + x_idx=x_idx, + nphi=nphi, + nz_interp=nz_interp) + return grids + + +def grid_for(grids: dict[str | None, FluxSurfaceGrid], + data: "GDataState") -> FluxSurfaceGrid: + """Return the sampling grid belonging to ``data``'s block.""" + return grids[geometry_prefix(data.file_name)] + + +def fluxsurf(data: "GDataState", + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + x_idx: int = 0, + nphi: int = 128, + nz_interp: int = 8, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Extract one field component on a toroidal flux surface. + + Args: + data: Three-dimensional field-aligned modal dataset. + mapc2p: Explicit modal geometry path. + nodes_file: Explicit nodal geometry path. + x_idx: Radial cell index identifying the surface. + nphi: Number of toroidal-angle slices. + nz_interp: Parallel-direction interpolation factor. + comp: Physical field component to extract. + inplace: Mutate and return ``data`` instead of creating a dataset. + tag: Optional tag for the returned dataset. + label: Optional label for the returned dataset. + """ + geometry = resolve_geometry(data.file_name, + mapc2p=mapc2p, + nodes_file=nodes_file) + fs_grid = resolve_flux_surface_grid(data, + geometry, + x_idx=x_idx, + nphi=nphi, + nz_interp=nz_interp) + return extract_flux_surface(data, + fs_grid=fs_grid, + comp=comp, + inplace=inplace, + tag=tag, + label=label) + + +__all__ = ["fluxsurf", "flux_surface_grids", "grid_for"] diff --git a/src/postgkyl/diagnostics/gk/geometry.py b/src/postgkyl/diagnostics/gk/geometry.py new file mode 100644 index 00000000..fef332f4 --- /dev/null +++ b/src/postgkyl/diagnostics/gk/geometry.py @@ -0,0 +1,184 @@ +"""Resolve and decode Gkeyll's auxiliary geometry files. + +This module owns geometry suffixes, representation selection, multiblock +paths, and file-layout interpretation. Filename parsing belongs to io.naming. +""" +from __future__ import annotations + +import os +import numpy as np + +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.io import parse_output_name +from postgkyl.numerics import nodal_to_cell_centered_grid +from postgkyl.operations import interpolate +from postgkyl.operations.geometry import Geometry, _validate_geometry + +# Mirrors ``enum gkyl_geometry_id`` in gkeyll/core/zero/gkyl_eqn_type.h. +# This foreign-format fact is shared with the grid-node diagnostic, which +# imports it from here instead of maintaining a second copy. +GKYL_GEOMETRY_ID = [ + "GKYL_GEOMETRY_NONE", + "GKYL_GEOMETRY_TOKAMAK", + "GKYL_GEOMETRY_MIRROR", + "GKYL_GEOMETRY_MAPC2P", + "GKYL_GEOMETRY_FROMFILE", +] +_MAPC2P_IDX = GKYL_GEOMETRY_ID.index("GKYL_GEOMETRY_MAPC2P") + + +def is_geo_mapc2p(ctx: dict) -> bool: + """Whether ``ctx`` identifies user-supplied Cartesian MAPC2P geometry. + + Files without ``geometry_type`` retain the historical MAPC2P default. + """ + return ctx.get("geometry_type", _MAPC2P_IDX) == _MAPC2P_IDX + + +def geometry_prefix(file_name: str | None) -> str | None: + """Return the per-block simulation prefix for ``file_name``. + + Parsing is delegated to :mod:`postgkyl.io.naming`, the authoritative home + of Gkeyll's output-name convention. + """ + name = parse_output_name(file_name) + return name.prefix if name is not None else None + + +def per_block_path(path: str | None, block: int | None) -> str | None: + """Substitute a multiblock index for ``'*'`` in a geometry override.""" + if path is None or block is None or "*" not in path: + return path + return path.replace("*", str(block)) + + +def _gauss_nodes(edges: np.ndarray) -> np.ndarray: + """Physical p1 Gauss-node coordinates for a one-dimensional edge grid.""" + centers = 0.5 * (edges[:-1] + edges[1:]) + offsets = np.diff(edges) / (2.0 * np.sqrt(3.0)) + return np.ravel(np.column_stack([centers - offsets, centers + offsets])) + + +def _pointwise_file( + path: str) -> tuple[list[np.ndarray], np.ndarray, GDataState]: + """Read a point-value geometry file and squeeze singleton dimensions.""" + data = GDataState(path) + grid = [np.squeeze(axis) for axis in data.grid] + return grid, np.squeeze(data.values), data + + +def _geometry_components( + values: np.ndarray, data: GDataState, + path: str) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: + """Interpret one geometry value array as ``R``, ``Z``, and optional phi.""" + required = 3 if is_geo_mapc2p(data.ctx) else 2 + if values.ndim < 2 or values.shape[-1] < required: + kind = "Cartesian X/Y/Z" if required == 3 else "R/Z" + raise ValueError( + f"Geometry file '{path}' must contain at least {required} {kind} components." + ) + + if is_geo_mapc2p(data.ctx): + x, y, z = values[..., 0], values[..., 1], values[..., 2] + return np.sqrt(x**2 + y**2), z, np.arctan2(y, x) + r, z = values[..., 0], values[..., 1] + phi = values[..., 2] if r.ndim == 3 and values.shape[-1] >= 3 else None + return r, z, phi + + +def _read_mapc2p_geometry(path: str): + """Interpolate a modal geometry file to physical ``R``, ``Z``, and phi.""" + source = GDataState(path) + field = interpolate(source) + cells = field.values.shape[:-1] + coords = nodal_to_cell_centered_grid(field.grid, cells) + major_r, vert_z, phi = _geometry_components(field.values, field, path) + return coords, major_r, vert_z, phi + + +def _read_nodes_geometry(path: str): + """Read a p1 pointwise nodal geometry file.""" + grid, values, data = _pointwise_file(path) + coords = [] + for dim, axis in enumerate(grid): + if axis.ndim != 1 or axis.shape[0] != values.shape[dim] + 1 \ + or values.shape[dim] % 2: + raise ValueError(f"Unrecognized nodal geometry layout in '{path}'.") + coords.append(_gauss_nodes(axis[::2])) + major_r, vert_z, phi = _geometry_components(values, data, path) + return coords, major_r, vert_z, phi + + +def _read_corner_rz(path: str): + """Read R and Z from a pointwise ``'-geo_corn_nodes.gkyl'`` file.""" + grid, values, data = _pointwise_file(path) + coords = [ + np.linspace(axis[0], axis[-1], n) + for axis, n in zip(grid, values.shape[:-1]) + ] + major_r, vert_z, _ = _geometry_components(values, data, path) + return coords, major_r, vert_z + + +def resolve_geometry(file_name: str | None, + *, + mapc2p: str | None = None, + nodes_file: str | None = None) -> Geometry: + """Resolve and load the geometry belonging to ``file_name``. + + The exact pointwise ``'-geo_int_nodes.gkyl'`` representation is + preferred, with ``'-geo_int_mapc2p.gkyl'`` as the modal fallback. + ``nodes_file`` and ``mapc2p`` override that lookup and are mutually + exclusive. Passing ``mapc2p=''`` explicitly requests the inferred modal + filename. + + Raises: + ValueError: If both overrides are supplied or no geometry can be found. + """ + if mapc2p is not None and nodes_file is not None: + raise ValueError("Pass either mapc2p= or nodes_file=, not both.") + + parsed = parse_output_name(file_name) + prefix = geometry_prefix(file_name) + block = parsed.block if parsed is not None else None + nodes_file = per_block_path(nodes_file, block) + mapc2p = per_block_path(mapc2p, block) + if nodes_file is not None: + path, kind = nodes_file, "nodes" + elif mapc2p is not None: + path = mapc2p or (f"{prefix}-geo_int_mapc2p.gkyl" if prefix else None) + kind = "mapc2p" + elif prefix is not None: + path, kind = f"{prefix}-geo_int_nodes.gkyl", "nodes" + if not os.path.exists(path): + path, kind = f"{prefix}-geo_int_mapc2p.gkyl", "mapc2p" + else: + path, kind = None, None + + if path is None or not os.path.exists(path): + raise ValueError( + "Could not find a geometry file; pass nodes_file= or mapc2p= explicitly." + ) + + coords, major_r, vert_z, phi = (_read_nodes_geometry(path) if kind == "nodes" + else _read_mapc2p_geometry(path)) + + corner = None + if prefix is not None: + corner_path = f"{prefix}-geo_corn_nodes.gkyl" + if os.path.exists(corner_path): + corner = _read_corner_rz(corner_path) + + geometry = Geometry(coords=coords, + major_r=major_r, + vert_z=vert_z, + phi=phi, + corner=corner) + _validate_geometry(geometry, len(coords)) + return geometry + + +__all__ = [ + "resolve_geometry", "geometry_prefix", "per_block_path", "GKYL_GEOMETRY_ID", + "is_geo_mapc2p" +] diff --git a/src/postgkyl/diagnostics/gk/nodes.py b/src/postgkyl/diagnostics/gk/nodes.py index bed00ff5..3eae5716 100644 --- a/src/postgkyl/diagnostics/gk/nodes.py +++ b/src/postgkyl/diagnostics/gk/nodes.py @@ -14,10 +14,7 @@ import numpy as np from matplotlib.collections import LineCollection -# GKYL_GEOMETRY_ID remains an intentional diagnostics.gk compatibility export. -from postgkyl.operations.gyrokinetics.geometry import ( # noqa: F401 - GKYL_GEOMETRY_ID, is_geo_mapc2p, -) +from .geometry import is_geo_mapc2p from . import utils diff --git a/src/postgkyl/diagnostics/gk/rz.py b/src/postgkyl/diagnostics/gk/rz.py index 169d030f..3002ca6b 100644 --- a/src/postgkyl/diagnostics/gk/rz.py +++ b/src/postgkyl/diagnostics/gk/rz.py @@ -1,20 +1,10 @@ -"""Compatibility aliases for the gyrokinetic R-Z operation. +"""Gkeyll geometry discovery composed with explicit rz mapping.""" +from __future__ import annotations -Canonical imports live in :mod:`postgkyl.operations.gyrokinetics`. This -module remains for the current major version and contains no copied -algorithm or defaults. -""" - -from postgkyl.operations.gyrokinetics.rz import ( - Geometry, - RzProjection, - geometry_prefix, - gk_rz, - map_to_rz, - per_block_path, - resolve_geometry, - resolve_rz_projection, -) +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.operations.map import resolve_rz_projection, map_to_rz +from postgkyl.operations.geometry import RzProjection, validate_mapping_grid +from .geometry import geometry_prefix, per_block_path, resolve_geometry def rz_projections(datasets, @@ -22,12 +12,18 @@ def rz_projections(datasets, mapc2p: str | None = None, nodes_file: str | None = None, z_axis: float = 0.0, - nz_interp: int = 8) -> dict: - """Compatibility batch wrapper using this module's patchable aliases.""" - projections = {} + nz_interp: int = 8) -> dict[str | None, RzProjection]: + """Build one reusable R-Z projection per block geometry. + + Frames from the same block share a projection; distinct blocks resolve + their own geometry. A ``'*'`` in an explicit geometry path is replaced by + the dataset's block index. + """ + projections: dict[str | None, RzProjection] = {} for data in datasets: key = geometry_prefix(data.file_name) if key in projections: + validate_mapping_grid(data, projections[key].computational_grid) continue block = data.ctx.get("block") geometry = resolve_geometry(data.file_name, @@ -40,20 +36,75 @@ def rz_projections(datasets, return projections -def projection_for(projections: dict, data): - """Return the compatibility projection belonging to ``data``'s block.""" +def projection_for(projections: dict[str | None, RzProjection], + data: "GDataState") -> RzProjection: + """Return the projection belonging to ``data``'s block.""" return projections[geometry_prefix(data.file_name)] -__all__ = [ - "Geometry", - "RzProjection", - "geometry_prefix", - "gk_rz", - "map_to_rz", - "per_block_path", - "projection_for", - "resolve_geometry", - "resolve_rz_projection", - "rz_projections", -] +def rz( + data: "GDataState", + *, + mapc2p: str | None = None, + nodes_file: str | None = None, + z_axis: float = 0.0, + phi_tor: float = 0.0, + nz_interp: int = 8, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None, +) -> "GDataState": + """Interpolate one DG component and project it onto a physical R-Z grid. + + ``data`` must be un-interpolated modal DG data with two computational + dimensions, or three for a field-aligned reconstruction. Geometry is + inferred from ``data.file_name``: the pointwise + ``'-geo_int_nodes.gkyl'`` file is preferred, falling back to + ``'-geo_int_mapc2p.gkyl'``. ``nodes_file`` or ``mapc2p`` may + override that choice, but they are mutually exclusive; ``mapc2p=''`` + forces the inferred modal filename. + + Args: + data: Un-interpolated 2-D or 3-D modal DG field. + mapc2p: Optional explicit modal geometry path, or ``''`` for inferred. + nodes_file: Optional explicit nodal geometry path. + z_axis: Magnetic-axis vertical position in meters, added to geometry Z. + phi_tor: Toroidal angle in radians for a 3-D poloidal reconstruction. + nz_interp: Positive integer z-direction up-sampling factor for 3-D data. + comp: Zero-based physical field component to map (default first). + inplace: Replace ``data`` rather than returning a new concrete instance. + tag: Optional result tag; ``None`` preserves the source tag. + label: Optional result label; ``None`` preserves the source label. + + Returns: + The caller's concrete data class, marked ``interpolated=True``. If the + interpolated input counts are ``(Nx, Nz)``, 2-D grid arrays have shape + ``(Nx+1, Nz+1)`` and values ``(Nx, Nz, 1)``. For interpolated 3-D counts + ``(Nx, Ny, Nz)``, grid arrays have shape + ``(Nx+1, nz_interp*Nz+1)`` and values + ``(Nx, nz_interp*Nz, 1)``. + + Raises: + ValueError: For mutually exclusive or missing geometry, input other than + un-interpolated 2-D/3-D modal data, a missing 3-D toroidal angle, + invalid ``comp`` or ``nz_interp``, malformed geometry, or a projection + incompatible with its data grid (when using :func:`map_to_rz`). + """ + geometry = resolve_geometry(data.file_name, + mapc2p=mapc2p, + nodes_file=nodes_file) + projection = resolve_rz_projection(data, + geometry, + z_axis=z_axis, + nz_interp=nz_interp) + return map_to_rz(data, + projection=projection, + phi_tor=phi_tor, + comp=comp, + inplace=inplace, + tag=tag, + label=label) + + +__all__ = ["rz", "rz_projections", "projection_for"] diff --git a/src/postgkyl/gdata/gdata.py b/src/postgkyl/gdata/gdata.py index d5b1ac6c..9547b00b 100644 --- a/src/postgkyl/gdata/gdata.py +++ b/src/postgkyl/gdata/gdata.py @@ -97,8 +97,8 @@ def load(self, # static language servers such as VS Code/Pylance. interpolate = operations.interpolate local_poly = operations.local_poly - gk_rz = operations.gyrokinetics.gk_rz - gk_fluxsurf = operations.gyrokinetics.gk_fluxsurf + map_to_rz = operations.map_to_rz + extract_flux_surface = operations.extract_flux_surface select = operations.select integrate = operations.integrate average = operations.average diff --git a/src/postgkyl/gdatastate/collection.py b/src/postgkyl/gdatastate/collection.py index 7d7625b7..5bd382ac 100644 --- a/src/postgkyl/gdatastate/collection.py +++ b/src/postgkyl/gdatastate/collection.py @@ -43,7 +43,7 @@ def _family_key(data) -> tuple | None: Built from the identity ``GDataState`` stamps at load time (``sim``, ``quantity``, ``frame`` -- see ``io.naming``) plus the dataset's ``tag``, so two differently-tagged results of the same source file (e.g. the raw - load and a ``gk_rz`` projection of it) never merge. ``block`` is + load and a ``map_to_rz`` projection of it) never merge. ``block`` is deliberately absent: it is what family members differ by. Returning ``None`` for single-block data is the property that keeps every diff --git a/src/postgkyl/gdatastate/gdatastate.py b/src/postgkyl/gdatastate/gdatastate.py index f830d2b5..f8a500ae 100644 --- a/src/postgkyl/gdatastate/gdatastate.py +++ b/src/postgkyl/gdatastate/gdatastate.py @@ -95,7 +95,7 @@ def _stamp_output_name(self) -> None: Header metadata wins: ``setdefault`` never overwrites a ``frame`` (or anything else) a reader already read out of the file itself. Because ``clone`` copies ``ctx``, the identity survives every verb, so a - multiblock family is still recognizable after ``interpolate``/``gk_rz`` + multiblock family is still recognizable after ``interpolate``/``map_to_rz`` -- which is what lets terminal verbs draw one field's blocks together (see ``gdatastate.collection.group_blocks``). """ diff --git a/src/postgkyl/numerics/resample.py b/src/postgkyl/numerics/resample.py new file mode 100644 index 00000000..5ac2036a --- /dev/null +++ b/src/postgkyl/numerics/resample.py @@ -0,0 +1,13 @@ +"""Linear tensor-grid resampling with linear extrapolation outside the grid.""" +import numpy as np +from scipy.interpolate import RegularGridInterpolator + + +def resample_grid(values: np.ndarray, src_coords: list[np.ndarray], + dst_coords: list[np.ndarray]) -> np.ndarray: + """Linearly resample ``values`` between tensor-product coordinate grids.""" + mesh = np.meshgrid(*dst_coords, indexing="ij") + return RegularGridInterpolator(tuple(src_coords), + values, + bounds_error=False, + fill_value=None)(tuple(mesh)) diff --git a/src/postgkyl/numerics/rz.py b/src/postgkyl/numerics/rz.py new file mode 100644 index 00000000..4aadc73a --- /dev/null +++ b/src/postgkyl/numerics/rz.py @@ -0,0 +1,39 @@ +"""Fourier reconstruction of a periodic field-aligned poloidal slice. + +The binormal direction is periodic; parallel endpoints obey twist-and-shift. +All coordinate and boundary information is supplied as numerical arrays. +""" +import numpy as np +from scipy.interpolate import PchipInterpolator + + +def fft_poloidal_project(values: np.ndarray, zc: np.ndarray, box: float, + wind: np.ndarray, phi0_zf: np.ndarray, zf: np.ndarray, + phi_tor: float) -> np.ndarray: + """FFT twist-and-shift reconstruction at one physical toroidal angle.""" + nx, ny, nz = values.shape + fk = np.fft.rfft(values, axis=1, norm="forward") + mode_count = fk.shape[1] + + dz = zc[1] - zc[0] + z_extended = np.concatenate(([zc[0] - dz / 2], zc, [zc[-1] + dz / 2])) + fk_extended = np.zeros((nx, mode_count, nz + 2), dtype=complex) + fk_extended[:, :, 1:-1] = fk + phase_shift = (2.0 * np.pi / box) * wind + for mode in range(mode_count): + phase = np.exp(-1j * mode * phase_shift) + fk_extended[:, mode, -1] = 0.5 * (fk[:, mode, -1] + phase * fk[:, mode, 0]) + fk_extended[:, mode, + 0] = 0.5 * (fk[:, mode, 0] + np.conj(phase) * fk[:, mode, -1]) + + fk_zf = (PchipInterpolator(z_extended, fk_extended.real, axis=2)(zf) + + 1j * PchipInterpolator(z_extended, fk_extended.imag, axis=2)(zf)) + + fraction = (phi_tor - phi0_zf) / box + out = np.zeros((nx, len(zf))) + for mode in range(mode_count): + weight = 1.0 if (mode == 0 or + (ny % 2 == 0 and mode == mode_count - 1)) else 2.0 + out += weight * np.real( + fk_zf[:, mode, :] * np.exp(-1j * 2.0 * np.pi * mode * fraction)) + return out diff --git a/src/postgkyl/operations/__init__.py b/src/postgkyl/operations/__init__.py index e07afcf4..36c09abf 100644 --- a/src/postgkyl/operations/__init__.py +++ b/src/postgkyl/operations/__init__.py @@ -10,14 +10,11 @@ ``integrate`` performs full or partial integration inside Gkeyll on modal data (full is terminal; partial stays native and lower-dimensional); ``average`` reduces modal data over a dimension subset via -``gkyl_array_average``, producing a new lower-dimensional modal dataset; -``map`` delegates to the grid-mapping engine in ``dg.map``. Flat modules are -domain-independent core verbs; domain subpackages such as ``gyrokinetics`` -hold transformations that require domain geometry without interpreting field -components as new physical conclusions. Equation-specific physics (the former -``moments``/``agyro``/``current``/``energetics``/``rotate``/ -``transform_frame``/``laguerre`` verbs, folded with the array math they -delegated to) lives one layer up, in ``diagnostics``. +``gkyl_array_average``, producing a new lower-dimensional modal dataset. + +Coordinate transformations live in ``map`` and receive explicit mappings or +geometry. Operations are equation-blind; model-specific auxiliary discovery +and physical compositions belong in ``diagnostics``. The terminal renderers (``plot``, ``animate``, ``plotly``, ``plotly_animate``, and ``pyvista``) are exceptions: @@ -25,7 +22,7 @@ :mod:`postgkyl.render` without wrapping them. """ -from . import arithmetic, gyrokinetics +from . import arithmetic from .interpolate import interpolate from .local_poly import local_poly from .select import select @@ -50,7 +47,9 @@ from .growth import growth from .differentiate import differentiate from .evaluate import available_operators as available_evaluate_operators, evaluate -from .map import map +from .map import (map, map_to_rz, resolve_rz_projection, extract_flux_surface, + resolve_flux_surface_grid) +from .geometry import Geometry, RzProjection, FluxSurfaceGrid # Command metadata is attached at the layer that owns each operation. This # block is deliberately declarative: discovery still walks the public API and @@ -171,11 +170,17 @@ def _resolve_receiver_annotations(*functions) -> None: hidden("registry provider used by evaluate help and validation")( available_evaluate_operators) +for _function in (map_to_rz, resolve_rz_projection, extract_flux_surface, + resolve_flux_surface_grid): + hidden("requires explicit Python geometry or projection objects")(_function) + __all__ = [ "interpolate", "local_poly", "select", "info", "print", "integrate", "average", "eval_at_coord_proj", "plot", "animate", "plotly", "plotly_animate", "pyvista", "arithmetic", "represent", "apply", "fft", "magsq", "relchange", "mask", "collect", "sort", "grid", "val2coord", "extract_input", "fit", "differentiate", "evaluate", - "available_evaluate_operators", "map", "growth", "gyrokinetics" + "available_evaluate_operators", "map", "growth", "map_to_rz", + "resolve_rz_projection", "extract_flux_surface", + "resolve_flux_surface_grid", "Geometry", "RzProjection", "FluxSurfaceGrid" ] diff --git a/src/postgkyl/operations/geometry.py b/src/postgkyl/operations/geometry.py new file mode 100644 index 00000000..9dcc8cf0 --- /dev/null +++ b/src/postgkyl/operations/geometry.py @@ -0,0 +1,204 @@ +"""In-memory coordinate geometry and shared field-mapping contracts. + +These records and helpers do not discover or load auxiliary datasets. +""" +from __future__ import annotations + +from dataclasses import dataclass +import numpy as np + +from postgkyl.dg import num_basis +from postgkyl.gdatastate.gdatastate import GDataState +from postgkyl.numerics import nodal_to_cell_centered_grid +from .interpolate import interpolate + + +@dataclass(frozen=True) +class Geometry: + """Physical ``(R, Z[, phi])`` geometry on its own point grid. + + ``corner`` closes the poloidal domain's ``theta = +/-pi`` ends for a 3-D + R-Z projection. It is optional boundary geometry on its own point grid. + """ + + coords: list[np.ndarray] + major_r: np.ndarray + vert_z: np.ndarray + phi: np.ndarray | None + corner: tuple[list[np.ndarray], np.ndarray, np.ndarray] | None + + +@dataclass(frozen=True) +class RzProjection: + """Precomputed R-Z mapping reusable by fields on one computational grid.""" + + num_dims: int + r: np.ndarray + z: np.ndarray + computational_grid: tuple[np.ndarray, ...] + zc: np.ndarray | None = None + zf: np.ndarray | None = None + box: float | None = None + wind: np.ndarray | None = None + phi0_zf: np.ndarray | None = None + + +@dataclass(frozen=True) +class FluxSurfaceGrid: + """Precomputed toroidal sampling grid for one radial flux surface.""" + + x_idx: int + zc: np.ndarray + zf: np.ndarray + phi_tor_list: np.ndarray + phi_2d: np.ndarray + computational_grid: tuple[np.ndarray, ...] + + +def _validate_geometry(geometry: Geometry, num_dims: int) -> None: + """Validate geometry tensor shapes for a data grid of ``num_dims``.""" + if len(geometry.coords) != num_dims: + raise ValueError( + f"Geometry has {len(geometry.coords)} dimensions but the data is " + f"{num_dims}-D.") + if any( + np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 + for axis in geometry.coords): + raise ValueError( + "Geometry coordinates must be one-dimensional arrays with at least two points." + ) + if any(not (np.all(np.diff(axis) > 0) or np.all(np.diff(axis) < 0)) + for axis in geometry.coords): + raise ValueError("Geometry coordinate arrays must be strictly monotonic.") + shape = tuple(np.asarray(axis).size for axis in geometry.coords) + if geometry.major_r.shape != shape or geometry.vert_z.shape != shape: + raise ValueError( + "Geometry coordinate and R/Z array shapes are incompatible: " + f"expected {shape}, got R{geometry.major_r.shape} and Z{geometry.vert_z.shape}." + ) + if geometry.phi is not None and geometry.phi.shape != shape: + raise ValueError( + f"Geometry toroidal-angle shape {geometry.phi.shape} does not match {shape}." + ) + if geometry.corner is not None: + corner_coords, corner_r, corner_z = geometry.corner + if len(corner_coords) != num_dims: + raise ValueError( + f"Corner geometry has {len(corner_coords)} dimensions; expected {num_dims}." + ) + corner_shape = tuple(np.asarray(axis).size for axis in corner_coords) + if (any( + np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 + for axis in corner_coords) or corner_r.shape != corner_shape + or corner_z.shape != corner_shape): + raise ValueError( + "Corner geometry coordinate and R/Z array shapes are incompatible.") + + +def _validate_positive_int(value: int, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, + (int, np.integer)) or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + return int(value) + + +def _validate_modal_data(data: GDataState, operation: str, + dimensions: tuple[int, ...]) -> None: + """Enforce the shared raw-DG input contract for coordinate projections.""" + if data.num_dims not in dimensions: + expected = " or ".join(f"{dim}-D" for dim in dimensions) + raise ValueError( + f"{operation} requires {expected} data; got {data.num_dims}-D.") + if data.values is None: + raise ValueError(f"{operation} requires a loaded dataset.") + if data.ctx.get("interpolated") or data.ctx.get("value_form", + "modal") != "modal": + raise ValueError(f"{operation} expects un-interpolated modal DG data.") + if not data.ctx.get("basis_type"): + raise ValueError( + f"{operation} requires 'basis_type' metadata on the input data.") + poly_order = data.ctx.get("poly_order") + if isinstance(poly_order, bool) or not isinstance(poly_order, (int, np.integer)) \ + or poly_order < 0: + raise ValueError( + f"{operation} requires a nonnegative integer 'poly_order'.") + if len(data.grid) != data.num_dims or any( + np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 + for axis in data.grid): + raise ValueError( + f"{operation} requires one one-dimensional edge grid per data dimension." + ) + if any(not (np.all(np.diff(axis) > 0) or np.all(np.diff(axis) < 0)) + for axis in data.grid): + raise ValueError( + f"{operation} requires strictly monotonic data edge grids.") + + +def _num_fields(data: GDataState) -> int: + """Return the number of physical fields stored in raw modal data.""" + basis_count = num_basis(data.num_dims, int(data.ctx["poly_order"]), + data.ctx["basis_type"]) + stored = data.values.shape[-1] + if stored % basis_count: + raise ValueError( + f"Data stores {stored} coefficients per cell, which is incompatible " + f"with a {basis_count}-coefficient basis.") + return stored // basis_count + + +def _validate_component(data: GDataState, comp: int) -> int: + if isinstance(comp, bool) or not isinstance(comp, (int, np.integer)): + raise ValueError("comp must be an integer component index.") + comp = int(comp) + num_fields = _num_fields(data) + if not 0 <= comp < num_fields: + raise ValueError( + f"comp {comp} is out of bounds for data with {num_fields} component(s)." + ) + return comp + + +def _interpolation_grid( + data: GDataState) -> tuple[list[np.ndarray], list[np.ndarray]]: + """Return interpolation edges/centers without evaluating field values.""" + num_interp = int(data.ctx["poly_order"]) + 1 + edges = [ + np.linspace(axis[0], axis[-1], + num_interp * (axis.size - 1) + 1) for axis in data.grid + ] + centers = nodal_to_cell_centered_grid( + edges, np.array([axis.size - 1 for axis in edges])) + return edges, centers + + +def _interpolate_component( + data: GDataState, + comp: int) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray]: + """Interpolate and return a component already checked by the public API.""" + field = interpolate(data) + cells = field.values.shape[:-1] + centers = nodal_to_cell_centered_grid(field.grid, cells) + return field.grid, centers, field.values[..., comp] + + +def _same_grid(left: tuple[np.ndarray, ...] | list[np.ndarray], + right: tuple[np.ndarray, ...] | list[np.ndarray]) -> bool: + return len(left) == len(right) and all( + a.shape == b.shape and np.allclose(a, b, rtol=1e-12, atol=1e-14) + for a, b in zip(left, right)) + + +def validate_mapping_grid(data: GDataState, + computational_grid: tuple[np.ndarray, ...]) -> None: + """Require the modal field grid used to construct a reusable mapping.""" + _validate_modal_data(data, "coordinate mapping", (len(computational_grid), )) + edges, _ = _interpolation_grid(data) + if not _same_grid(computational_grid, edges): + raise ValueError( + "Incompatible mapping: data computational grid does not match " + "the grid used to build the projection.") + + +__all__ = [ + "Geometry", "RzProjection", "FluxSurfaceGrid", "validate_mapping_grid" +] diff --git a/src/postgkyl/operations/gyrokinetics/__init__.py b/src/postgkyl/operations/gyrokinetics/__init__.py deleted file mode 100644 index de479aba..00000000 --- a/src/postgkyl/operations/gyrokinetics/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Gyrokinetic data transformations. - -Placement answers two independent questions: ``operations`` says these -functions re-express data rather than derive physical conclusions, while -``gyrokinetics`` identifies the domain knowledge their geometry requires. -""" - -from .geometry import GKYL_GEOMETRY_ID, Geometry, is_geo_mapc2p, resolve_geometry -from .rz import RzProjection, gk_rz, map_to_rz, resolve_rz_projection -from .fluxsurf import ( - FluxSurfaceGrid, - extract_flux_surface, - gk_fluxsurf, - resolve_flux_surface_grid, -) - -from postgkyl.cli_spec import CommandSpec, Execution, Section, command, hidden -from postgkyl.gdatastate.gdatastate import GDataState - -gk_rz.__globals__.setdefault("GDataState", GDataState) -gk_fluxsurf.__globals__.setdefault("GDataState", GDataState) -command(CommandSpec(Section.VERBS, Execution.MAP_REPLACE))(gk_rz) -command(CommandSpec(Section.VERBS, Execution.MAP_REPLACE))(gk_fluxsurf) -for _function in ( - is_geo_mapc2p, - resolve_geometry, - map_to_rz, - resolve_rz_projection, - extract_flux_surface, - resolve_flux_surface_grid, -): - hidden("lower-level geometry API requires Python geometry objects")(_function) - -__all__ = [ - "GKYL_GEOMETRY_ID", - "Geometry", - "is_geo_mapc2p", - "resolve_geometry", - "RzProjection", - "gk_rz", - "map_to_rz", - "resolve_rz_projection", - "FluxSurfaceGrid", - "extract_flux_surface", - "gk_fluxsurf", - "resolve_flux_surface_grid", -] diff --git a/src/postgkyl/operations/gyrokinetics/fluxsurf.py b/src/postgkyl/operations/gyrokinetics/fluxsurf.py deleted file mode 100644 index 0bb203c2..00000000 --- a/src/postgkyl/operations/gyrokinetics/fluxsurf.py +++ /dev/null @@ -1,218 +0,0 @@ -"""Extract theta-phi flux surfaces from gyrokinetic field-aligned data.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -import numpy as np -from scipy.interpolate import PchipInterpolator - -from .geometry import ( - Geometry, - geometry_prefix, - _interpolate_component, - _interpolation_grid, - _resample_grid, - _same_grid, - _validate_component, - _validate_geometry, - _validate_modal_data, - _validate_positive_int, - per_block_path, - resolve_geometry, -) - -if TYPE_CHECKING: - from postgkyl.gdatastate.gdatastate import GDataState - - -@dataclass(frozen=True) -class FluxSurfaceGrid: - """Precomputed toroidal sampling grid for one radial flux surface.""" - - x_idx: int - zc: np.ndarray - zf: np.ndarray - phi_tor_list: np.ndarray - phi_2d: np.ndarray - computational_grid: tuple[np.ndarray, ...] | None = None - - -def resolve_flux_surface_grid(first: "GDataState", - geo: Geometry, - *, - x_idx: int = 0, - nphi: int = 128, - nz_interp: int = 8) -> FluxSurfaceGrid: - """Precompute a flux-surface sampling grid for compatible 3-D fields.""" - _validate_modal_data(first, "gk_fluxsurf", (3, )) - nphi = _validate_positive_int(nphi, "nphi") - nz_interp = _validate_positive_int(nz_interp, "nz_interp") - _validate_geometry(geo, 3) - if geo.phi is None: - raise ValueError( - "The geometry file has no toroidal-angle component; cannot extract a flux surface." - ) - if isinstance(x_idx, bool) or not isinstance(x_idx, (int, np.integer)): - raise ValueError("x_idx must be an integer radial index.") - - edges, centers = _interpolation_grid(first) - xc, yc, zc = centers - x_idx = int(x_idx) - if not 0 <= x_idx < xc.size: - raise ValueError( - f"x_idx {x_idx} is out of bounds for data with Nx={xc.size}.") - if zc.size < 2 or yc.size < 2: - raise ValueError( - "gk_fluxsurf requires at least two interpolated y and z points.") - - zf_edges = np.linspace(edges[2][0], edges[2][-1], nz_interp * zc.size + 1) - zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) - phi = np.unwrap(np.unwrap(np.unwrap(geo.phi, axis=2), axis=1), axis=0) - phi_grid = _resample_grid(phi, geo.coords, [xc, yc, zf]) - phi_2d = phi_grid[x_idx, :, :] - phi_tor_list = np.linspace(0.0, 2.0 * np.pi, nphi, endpoint=False) - return FluxSurfaceGrid(x_idx=x_idx, - zc=zc, - zf=zf, - phi_tor_list=phi_tor_list, - phi_2d=phi_2d, - computational_grid=tuple( - np.array(axis, copy=True) for axis in edges)) - - -def extract_flux_surface(data: "GDataState", - fs_grid: FluxSurfaceGrid, - *, - comp: int = 0, - inplace: bool = False, - tag: str | None = None, - label: str | None = None) -> "GDataState": - """Extract component ``comp`` using a reusable ``fs_grid``.""" - _validate_modal_data(data, "gk_fluxsurf", (3, )) - _validate_component(data, comp) - edges, _ = _interpolation_grid(data) - if fs_grid.computational_grid is not None \ - and not _same_grid(fs_grid.computational_grid, edges): - raise ValueError( - "Incompatible flux-surface grid: data computational grid does not " - "match the grid used to build the projection.") - nx, ny, nz = (axis.size - 1 for axis in edges) - if not 0 <= fs_grid.x_idx < nx: - raise ValueError( - f"x_idx {fs_grid.x_idx} is out of bounds for data with Nx={nx}.") - if (fs_grid.zc.shape != (nz, ) - or fs_grid.phi_2d.shape != (ny, fs_grid.zf.size) - or fs_grid.phi_tor_list.ndim != 1): - raise ValueError( - "Incompatible flux-surface grid: projection and data grid shapes differ." - ) - - _, _, values = _interpolate_component(data, comp) - vals_zf = PchipInterpolator(fs_grid.zc, values, axis=-1, - extrapolate=True)(fs_grid.zf) - vals_2d = vals_zf[fs_grid.x_idx, :, :] - - flux_surf_data = np.empty((fs_grid.phi_tor_list.size, fs_grid.zf.size)) - for iz in range(fs_grid.zf.size): - phi_y = fs_grid.phi_2d[:, iz] - val_y = vals_2d[:, iz] - box = np.mean(np.diff(phi_y)) * ny - if not np.isfinite(box) or np.isclose(box, 0.0): - raise ValueError( - "Toroidal geometry has a zero or non-finite binormal angular span.") - phi_ext = np.concatenate([phi_y - box, phi_y, phi_y + box]) - val_ext = np.concatenate([val_y, val_y, val_y]) - order = np.argsort(phi_ext) - folded = phi_y[0] + np.mod(fs_grid.phi_tor_list - phi_y[0], box) - flux_surf_data[:, iz] = np.interp(folded, phi_ext[order], val_ext[order]) - - return data._result([fs_grid.phi_tor_list, fs_grid.zf], - flux_surf_data[..., np.newaxis], - inplace=inplace, - tag=tag, - label=label, - interpolated=True) - - -def flux_surface_grids(datasets, - *, - mapc2p: str | None = None, - nodes_file: str | None = None, - x_idx: int = 0, - nphi: int = 128, - nz_interp: int = 8) -> dict[str | None, FluxSurfaceGrid]: - """Build one reusable flux-surface grid per block geometry.""" - grids: dict[str | None, FluxSurfaceGrid] = {} - for data in datasets: - key = geometry_prefix(data.file_name) - if key in grids: - continue - block = data.ctx.get("block") - geometry = resolve_geometry(data.file_name, - mapc2p=per_block_path(mapc2p, block), - nodes_file=per_block_path(nodes_file, block)) - grids[key] = resolve_flux_surface_grid(data, - geometry, - x_idx=x_idx, - nphi=nphi, - nz_interp=nz_interp) - return grids - - -def grid_for(grids: dict[str | None, FluxSurfaceGrid], - data: "GDataState") -> FluxSurfaceGrid: - """Return the sampling grid belonging to ``data``'s block.""" - return grids[geometry_prefix(data.file_name)] - - -def gk_fluxsurf(data: "GDataState", - *, - mapc2p: str | None = None, - nodes_file: str | None = None, - x_idx: int = 0, - nphi: int = 128, - nz_interp: int = 8, - comp: int = 0, - inplace: bool = False, - tag: str | None = None, - label: str | None = None) -> "GDataState": - """Extract one field component on a toroidal flux surface. - - Args: - data: Three-dimensional field-aligned modal dataset. - mapc2p: Explicit modal geometry path. - nodes_file: Explicit nodal geometry path. - x_idx: Radial cell index identifying the surface. - nphi: Number of toroidal-angle slices. - nz_interp: Parallel-direction interpolation factor. - comp: Physical field component to extract. - inplace: Mutate and return ``data`` instead of creating a dataset. - tag: Optional tag for the returned dataset. - label: Optional label for the returned dataset. - """ - geometry = resolve_geometry(data.file_name, - mapc2p=mapc2p, - nodes_file=nodes_file) - fs_grid = resolve_flux_surface_grid(data, - geometry, - x_idx=x_idx, - nphi=nphi, - nz_interp=nz_interp) - return extract_flux_surface(data, - fs_grid, - comp=comp, - inplace=inplace, - tag=tag, - label=label) - - -__all__ = [ - "FluxSurfaceGrid", - "extract_flux_surface", - "flux_surface_grids", - "gk_fluxsurf", - "grid_for", - "resolve_flux_surface_grid", -] diff --git a/src/postgkyl/operations/gyrokinetics/geometry.py b/src/postgkyl/operations/gyrokinetics/geometry.py deleted file mode 100644 index c1c99fc0..00000000 --- a/src/postgkyl/operations/gyrokinetics/geometry.py +++ /dev/null @@ -1,355 +0,0 @@ -"""Shared geometry machinery for gyrokinetic data transformations. - -This module owns geometry-file discovery and loading plus the grid helpers -used by both the R-Z and flux-surface operations. It deliberately constructs -the verb-less :class:`~postgkyl.gdatastate.gdatastate.GDataState` and calls the -lower interpolation operation directly; the operation layer never reaches up -through the fluent :class:`postgkyl.gdata.GData` surface. -""" - -from __future__ import annotations - -import os -from dataclasses import dataclass - -import numpy as np -from scipy.interpolate import RegularGridInterpolator - -from postgkyl.dg import num_basis -from postgkyl.gdatastate.gdatastate import GDataState -from postgkyl.io import parse_output_name -from postgkyl.numerics import nodal_to_cell_centered_grid - -from ..interpolate import interpolate - -# Mirrors ``enum gkyl_geometry_id`` in gkeyll/core/zero/gkyl_eqn_type.h. -# This foreign-format fact is shared with the grid-node diagnostic, which -# imports it from here instead of maintaining a second copy. -GKYL_GEOMETRY_ID = [ - "GKYL_GEOMETRY_NONE", - "GKYL_GEOMETRY_TOKAMAK", - "GKYL_GEOMETRY_MIRROR", - "GKYL_GEOMETRY_MAPC2P", - "GKYL_GEOMETRY_FROMFILE", -] -_MAPC2P_IDX = GKYL_GEOMETRY_ID.index("GKYL_GEOMETRY_MAPC2P") - - -@dataclass(frozen=True) -class Geometry: - """Physical ``(R, Z[, phi])`` geometry on its own point grid. - - ``corner`` closes the poloidal domain's ``theta = +/-pi`` ends for a 3-D - R-Z projection. It is ``None`` when no - ``'-geo_corn_nodes.gkyl'`` file exists alongside the field. - """ - - coords: list[np.ndarray] - major_r: np.ndarray - vert_z: np.ndarray - phi: np.ndarray | None - corner: tuple[list[np.ndarray], np.ndarray, np.ndarray] | None - - -def is_geo_mapc2p(ctx: dict) -> bool: - """Whether ``ctx`` identifies user-supplied Cartesian MAPC2P geometry. - - Files without ``geometry_type`` retain the historical MAPC2P default. - """ - return ctx.get("geometry_type", _MAPC2P_IDX) == _MAPC2P_IDX - - -def geometry_prefix(file_name: str | None) -> str | None: - """Return the per-block simulation prefix for ``file_name``. - - Parsing is delegated to :mod:`postgkyl.io.naming`, the authoritative home - of Gkeyll's output-name convention. - """ - name = parse_output_name(file_name) - return name.prefix if name is not None else None - - -def per_block_path(path: str | None, block: int | None) -> str | None: - """Substitute a multiblock index for ``'*'`` in a geometry override.""" - if path is None or block is None or "*" not in path: - return path - return path.replace("*", str(block)) - - -def _gauss_nodes(edges: np.ndarray) -> np.ndarray: - """Physical p1 Gauss-node coordinates for a one-dimensional edge grid.""" - centers = 0.5 * (edges[:-1] + edges[1:]) - offsets = np.diff(edges) / (2.0 * np.sqrt(3.0)) - return np.ravel(np.column_stack([centers - offsets, centers + offsets])) - - -def _pointwise_file( - path: str) -> tuple[list[np.ndarray], np.ndarray, GDataState]: - """Read a point-value geometry file and squeeze singleton dimensions.""" - data = GDataState(path) - grid = [np.squeeze(axis) for axis in data.grid] - return grid, np.squeeze(data.values), data - - -def _geometry_components( - values: np.ndarray, data: GDataState, - path: str) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: - """Interpret one geometry value array as ``R``, ``Z``, and optional phi.""" - required = 3 if is_geo_mapc2p(data.ctx) else 2 - if values.ndim < 2 or values.shape[-1] < required: - kind = "Cartesian X/Y/Z" if required == 3 else "R/Z" - raise ValueError( - f"Geometry file '{path}' must contain at least {required} {kind} components." - ) - - if is_geo_mapc2p(data.ctx): - x, y, z = values[..., 0], values[..., 1], values[..., 2] - return np.sqrt(x**2 + y**2), z, np.arctan2(y, x) - r, z = values[..., 0], values[..., 1] - phi = values[..., 2] if r.ndim == 3 and values.shape[-1] >= 3 else None - return r, z, phi - - -def _read_mapc2p_geometry(path: str): - """Interpolate a modal geometry file to physical ``R``, ``Z``, and phi.""" - source = GDataState(path) - field = interpolate(source) - cells = field.values.shape[:-1] - coords = nodal_to_cell_centered_grid(field.grid, cells) - major_r, vert_z, phi = _geometry_components(field.values, field, path) - return coords, major_r, vert_z, phi - - -def _read_nodes_geometry(path: str): - """Read a p1 pointwise nodal geometry file.""" - grid, values, data = _pointwise_file(path) - coords = [] - for dim, axis in enumerate(grid): - if axis.ndim != 1 or axis.shape[0] != values.shape[dim] + 1 \ - or values.shape[dim] % 2: - raise ValueError(f"Unrecognized nodal geometry layout in '{path}'.") - coords.append(_gauss_nodes(axis[::2])) - major_r, vert_z, phi = _geometry_components(values, data, path) - return coords, major_r, vert_z, phi - - -def _read_corner_rz(path: str): - """Read R and Z from a pointwise ``'-geo_corn_nodes.gkyl'`` file.""" - grid, values, data = _pointwise_file(path) - coords = [ - np.linspace(axis[0], axis[-1], n) - for axis, n in zip(grid, values.shape[:-1]) - ] - major_r, vert_z, _ = _geometry_components(values, data, path) - return coords, major_r, vert_z - - -def _validate_geometry(geometry: Geometry, num_dims: int) -> None: - """Validate geometry tensor shapes for a data grid of ``num_dims``.""" - if len(geometry.coords) != num_dims: - raise ValueError( - f"Geometry has {len(geometry.coords)} dimensions but the data is " - f"{num_dims}-D.") - if any( - np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 - for axis in geometry.coords): - raise ValueError( - "Geometry coordinates must be one-dimensional arrays with at least two points." - ) - if any(not (np.all(np.diff(axis) > 0) or np.all(np.diff(axis) < 0)) - for axis in geometry.coords): - raise ValueError("Geometry coordinate arrays must be strictly monotonic.") - shape = tuple(np.asarray(axis).size for axis in geometry.coords) - if geometry.major_r.shape != shape or geometry.vert_z.shape != shape: - raise ValueError( - "Geometry coordinate and R/Z array shapes are incompatible: " - f"expected {shape}, got R{geometry.major_r.shape} and Z{geometry.vert_z.shape}." - ) - if geometry.phi is not None and geometry.phi.shape != shape: - raise ValueError( - f"Geometry toroidal-angle shape {geometry.phi.shape} does not match {shape}." - ) - if geometry.corner is not None: - corner_coords, corner_r, corner_z = geometry.corner - if len(corner_coords) != num_dims: - raise ValueError( - f"Corner geometry has {len(corner_coords)} dimensions; expected {num_dims}." - ) - corner_shape = tuple(np.asarray(axis).size for axis in corner_coords) - if (any( - np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 - for axis in corner_coords) or corner_r.shape != corner_shape - or corner_z.shape != corner_shape): - raise ValueError( - "Corner geometry coordinate and R/Z array shapes are incompatible.") - - -def resolve_geometry(file_name: str | None, - *, - mapc2p: str | None = None, - nodes_file: str | None = None) -> Geometry: - """Resolve and load the geometry belonging to ``file_name``. - - The exact pointwise ``'-geo_int_nodes.gkyl'`` representation is - preferred, with ``'-geo_int_mapc2p.gkyl'`` as the modal fallback. - ``nodes_file`` and ``mapc2p`` override that lookup and are mutually - exclusive. Passing ``mapc2p=''`` explicitly requests the inferred modal - filename. - - Raises: - ValueError: If both overrides are supplied or no geometry can be found. - """ - if mapc2p is not None and nodes_file is not None: - raise ValueError("Pass either mapc2p= or nodes_file=, not both.") - - parsed = parse_output_name(file_name) - prefix = geometry_prefix(file_name) - block = parsed.block if parsed is not None else None - nodes_file = per_block_path(nodes_file, block) - mapc2p = per_block_path(mapc2p, block) - if nodes_file is not None: - path, kind = nodes_file, "nodes" - elif mapc2p is not None: - path = mapc2p or (f"{prefix}-geo_int_mapc2p.gkyl" if prefix else None) - kind = "mapc2p" - elif prefix is not None: - path, kind = f"{prefix}-geo_int_nodes.gkyl", "nodes" - if not os.path.exists(path): - path, kind = f"{prefix}-geo_int_mapc2p.gkyl", "mapc2p" - else: - path, kind = None, None - - if path is None or not os.path.exists(path): - raise ValueError( - "Could not find a geometry file; pass nodes_file= or mapc2p= explicitly." - ) - - coords, major_r, vert_z, phi = (_read_nodes_geometry(path) if kind == "nodes" - else _read_mapc2p_geometry(path)) - - corner = None - if prefix is not None: - corner_path = f"{prefix}-geo_corn_nodes.gkyl" - if os.path.exists(corner_path): - corner = _read_corner_rz(corner_path) - - geometry = Geometry(coords=coords, - major_r=major_r, - vert_z=vert_z, - phi=phi, - corner=corner) - _validate_geometry(geometry, len(coords)) - return geometry - - -def _validate_positive_int(value: int, name: str) -> int: - if isinstance(value, bool) or not isinstance(value, - (int, np.integer)) or value <= 0: - raise ValueError(f"{name} must be a positive integer.") - return int(value) - - -def _validate_modal_data(data: GDataState, operation: str, - dimensions: tuple[int, ...]) -> None: - """Enforce the shared raw-DG input contract for GK projections.""" - if data.num_dims not in dimensions: - expected = " or ".join(f"{dim}-D" for dim in dimensions) - raise ValueError( - f"{operation} requires {expected} data; got {data.num_dims}-D.") - if data.values is None: - raise ValueError(f"{operation} requires a loaded dataset.") - if data.ctx.get("interpolated") or data.ctx.get("value_form", - "modal") != "modal": - raise ValueError(f"{operation} expects un-interpolated modal DG data.") - if not data.ctx.get("basis_type"): - raise ValueError( - f"{operation} requires 'basis_type' metadata on the input data.") - poly_order = data.ctx.get("poly_order") - if isinstance(poly_order, bool) or not isinstance(poly_order, (int, np.integer)) \ - or poly_order < 0: - raise ValueError( - f"{operation} requires a nonnegative integer 'poly_order'.") - if len(data.grid) != data.num_dims or any( - np.asarray(axis).ndim != 1 or np.asarray(axis).size < 2 - for axis in data.grid): - raise ValueError( - f"{operation} requires one one-dimensional edge grid per data dimension." - ) - if any(not (np.all(np.diff(axis) > 0) or np.all(np.diff(axis) < 0)) - for axis in data.grid): - raise ValueError( - f"{operation} requires strictly monotonic data edge grids.") - - -def _num_fields(data: GDataState) -> int: - """Return the number of physical fields stored in raw modal data.""" - basis_count = num_basis(data.num_dims, int(data.ctx["poly_order"]), - data.ctx["basis_type"]) - stored = data.values.shape[-1] - if stored % basis_count: - raise ValueError( - f"Data stores {stored} coefficients per cell, which is incompatible " - f"with a {basis_count}-coefficient basis.") - return stored // basis_count - - -def _validate_component(data: GDataState, comp: int) -> int: - if isinstance(comp, bool) or not isinstance(comp, (int, np.integer)): - raise ValueError("comp must be an integer component index.") - comp = int(comp) - num_fields = _num_fields(data) - if not 0 <= comp < num_fields: - raise ValueError( - f"comp {comp} is out of bounds for data with {num_fields} component(s)." - ) - return comp - - -def _interpolation_grid( - data: GDataState) -> tuple[list[np.ndarray], list[np.ndarray]]: - """Return interpolation edges/centers without evaluating field values.""" - num_interp = int(data.ctx["poly_order"]) + 1 - edges = [ - np.linspace(axis[0], axis[-1], - num_interp * (axis.size - 1) + 1) for axis in data.grid - ] - centers = nodal_to_cell_centered_grid( - edges, np.array([axis.size - 1 for axis in edges])) - return edges, centers - - -def _interpolate_component( - data: GDataState, - comp: int) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray]: - """Interpolate and return a component already checked by the public API.""" - field = interpolate(data) - cells = field.values.shape[:-1] - centers = nodal_to_cell_centered_grid(field.grid, cells) - return field.grid, centers, field.values[..., comp] - - -def _resample_grid(values: np.ndarray, src_coords: list[np.ndarray], - dst_coords: list[np.ndarray]) -> np.ndarray: - """Linearly resample ``values`` between tensor-product coordinate grids.""" - mesh = np.meshgrid(*dst_coords, indexing="ij") - return RegularGridInterpolator(tuple(src_coords), - values, - bounds_error=False, - fill_value=None)(tuple(mesh)) - - -def _same_grid(left: tuple[np.ndarray, ...] | list[np.ndarray], - right: tuple[np.ndarray, ...] | list[np.ndarray]) -> bool: - return len(left) == len(right) and all( - a.shape == b.shape and np.allclose(a, b, rtol=1e-12, atol=1e-14) - for a, b in zip(left, right)) - - -__all__ = [ - "GKYL_GEOMETRY_ID", - "Geometry", - "geometry_prefix", - "is_geo_mapc2p", - "per_block_path", - "resolve_geometry", -] diff --git a/src/postgkyl/operations/gyrokinetics/rz.py b/src/postgkyl/operations/gyrokinetics/rz.py deleted file mode 100644 index b4b1e051..00000000 --- a/src/postgkyl/operations/gyrokinetics/rz.py +++ /dev/null @@ -1,359 +0,0 @@ -"""Project gyrokinetic DG fields onto a physical poloidal R-Z plane.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -import numpy as np -from scipy.interpolate import PchipInterpolator - -from .geometry import ( - Geometry, - geometry_prefix, - _interpolate_component, - _interpolation_grid, - _resample_grid, - _same_grid, - _validate_component, - _validate_geometry, - _validate_modal_data, - _validate_positive_int, - per_block_path, - resolve_geometry, -) - -if TYPE_CHECKING: - from postgkyl.gdatastate.gdatastate import GDataState - - -@dataclass(frozen=True) -class RzProjection: - """Precomputed R-Z mapping reusable by fields on one computational grid.""" - - num_dims: int - r: np.ndarray - z: np.ndarray - zc: np.ndarray | None = None - zf: np.ndarray | None = None - box: float | None = None - wind: np.ndarray | None = None - phi0_zf: np.ndarray | None = None - computational_grid: tuple[np.ndarray, ...] | None = None - - -def _fft_poloidal_project(values: np.ndarray, zc: np.ndarray, box: float, - wind: np.ndarray, phi0_zf: np.ndarray, zf: np.ndarray, - phi_tor: float) -> np.ndarray: - """FFT twist-and-shift reconstruction at one physical toroidal angle.""" - nx, ny, nz = values.shape - fk = np.fft.rfft(values, axis=1, norm="forward") - mode_count = fk.shape[1] - - dz = zc[1] - zc[0] - z_extended = np.concatenate(([zc[0] - dz / 2], zc, [zc[-1] + dz / 2])) - fk_extended = np.zeros((nx, mode_count, nz + 2), dtype=complex) - fk_extended[:, :, 1:-1] = fk - phase_shift = (2.0 * np.pi / box) * wind - for mode in range(mode_count): - phase = np.exp(-1j * mode * phase_shift) - fk_extended[:, mode, -1] = 0.5 * (fk[:, mode, -1] + phase * fk[:, mode, 0]) - fk_extended[:, mode, - 0] = 0.5 * (fk[:, mode, 0] + np.conj(phase) * fk[:, mode, -1]) - - fk_zf = (PchipInterpolator(z_extended, fk_extended.real, axis=2)(zf) + - 1j * PchipInterpolator(z_extended, fk_extended.imag, axis=2)(zf)) - - fraction = (phi_tor - phi0_zf) / box - out = np.zeros((nx, len(zf))) - for mode in range(mode_count): - weight = 1.0 if (mode == 0 or - (ny % 2 == 0 and mode == mode_count - 1)) else 2.0 - out += weight * np.real( - fk_zf[:, mode, :] * np.exp(-1j * 2.0 * np.pi * mode * fraction)) - return out - - -def resolve_rz_projection(first: "GDataState", - geo: Geometry, - *, - z_axis: float = 0.0, - nz_interp: int = 8) -> RzProjection: - """Build an R-Z projection for ``first``'s grid and ``geo``. - - Only the computational grid and DG metadata are read from ``first``; - projection construction never evaluates or selects its field values. - ``z_axis`` is the magnetic-axis vertical position in meters. - """ - _validate_modal_data(first, "gk_rz", (2, 3)) - nz_interp = _validate_positive_int(nz_interp, "nz_interp") - _validate_geometry(geo, first.num_dims) - edges, centers = _interpolation_grid(first) - vert_z = geo.vert_z + float(z_axis) - - if first.num_dims == 2: - r = _resample_grid(geo.major_r, geo.coords, edges) - z = _resample_grid(vert_z, geo.coords, edges) - return RzProjection(num_dims=2, - r=r, - z=z, - computational_grid=tuple( - np.array(axis, copy=True) for axis in edges)) - - if geo.phi is None: - raise ValueError( - "The geometry file has no toroidal-angle component; 3-D gk_rz requires one." - ) - - xc, yc, zc = centers - nx, ny, nz = xc.size, yc.size, zc.size - if ny < 2 or nz < 2: - raise ValueError( - "3-D gk_rz requires at least two interpolated y and z points.") - - phi = np.unwrap(np.unwrap(np.unwrap(geo.phi, axis=2), axis=1), axis=0) - phi_field = _resample_grid(phi, geo.coords, centers) - box_estimate = np.mean(np.diff(phi_field[nx // 2, :, nz // 2])) * ny - if not np.isfinite(box_estimate) or np.isclose(box_estimate, 0.0): - raise ValueError( - "Toroidal geometry has a zero or non-finite binormal angular span.") - n0 = max(1, int(round(abs(2.0 * np.pi / box_estimate)))) - box = np.sign(box_estimate) * 2.0 * np.pi / n0 - wind = phi_field[:, 0, -1] - phi_field[:, 0, 0] - - xn, _, zn = edges - zf_edges = np.linspace(zn[0], zn[-1], nz_interp * nz + 1) - zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) - phi0_zf = np.array( - [np.interp(zf, zc, phi_field[ix, 0, :]) for ix in range(nx)]) - - gx, _, gz = geo.coords - r2d = geo.major_r[:, 0, :] - z2d = vert_z[:, 0, :] - gz_rz = gz - if geo.corner is not None: - corner_coords, corner_r, corner_z = geo.corner - if len(corner_coords) != 3: - raise ValueError( - "Corner geometry must be three-dimensional for 3-D gk_rz.") - cx, cz = corner_coords[0], corner_coords[2] - corner_r = corner_r[:, 0, :] - corner_z = corner_z[:, 0, :] + float(z_axis) - r2d = np.concatenate([ - np.interp(gx, cx, corner_r[:, 0])[:, None], r2d, - np.interp(gx, cx, corner_r[:, -1])[:, None] - ], - axis=1) - z2d = np.concatenate([ - np.interp(gx, cx, corner_z[:, 0])[:, None], z2d, - np.interp(gx, cx, corner_z[:, -1])[:, None] - ], - axis=1) - gz_rz = np.concatenate([[cz[0]], gz, [cz[-1]]]) - - r = _resample_grid(r2d, [gx, gz_rz], [xn, zf_edges]) - z = _resample_grid(z2d, [gx, gz_rz], [xn, zf_edges]) - return RzProjection(num_dims=3, - r=r, - z=z, - zc=zc, - zf=zf, - box=box, - wind=wind, - phi0_zf=phi0_zf, - computational_grid=tuple( - np.array(axis, copy=True) for axis in edges)) - - -def _validate_projection(data: "GDataState", projection: RzProjection) -> None: - if projection.num_dims not in (2, 3): - raise ValueError( - f"R-Z projection has invalid dimensionality {projection.num_dims}; expected 2 or 3." - ) - if data.num_dims != projection.num_dims: - raise ValueError( - "Incompatible R-Z projection: projection dimensionality does not match the data." - ) - edges, _ = _interpolation_grid(data) - if projection.computational_grid is not None \ - and not _same_grid(projection.computational_grid, edges): - raise ValueError( - "Incompatible R-Z projection: data computational grid does not match " - "the grid used to build the projection.") - if projection.r.shape != projection.z.shape or projection.r.ndim != 2: - raise ValueError( - "Incompatible R-Z projection: R and Z grids must be matching 2-D arrays." - ) - - if projection.num_dims == 2: - expected = (edges[0].size, edges[1].size) - if projection.r.shape != expected: - raise ValueError( - f"Incompatible R-Z projection: expected grid shape {expected}, " - f"got {projection.r.shape}.") - return - - required = (projection.zc, projection.zf, projection.box, projection.wind, - projection.phi0_zf) - if any(value is None for value in required): - raise ValueError( - "Incompatible R-Z projection: 3-D projection metadata is incomplete.") - if not np.isfinite(projection.box) or np.isclose(projection.box, 0.0): - raise ValueError( - "Incompatible R-Z projection: toroidal angular span must be finite and nonzero." - ) - nx, _, nz = (axis.size - 1 for axis in edges) - if (projection.zc.shape != (nz, ) or projection.wind.shape != (nx, ) - or projection.phi0_zf.shape != (nx, projection.zf.size) - or projection.r.shape != (nx + 1, projection.zf.size + 1)): - raise ValueError( - "Incompatible R-Z projection: projection and data grid shapes differ.") - - -def map_to_rz(data: "GDataState", - projection: RzProjection, - *, - phi_tor: float = 0.0, - comp: int = 0, - inplace: bool = False, - tag: str | None = None, - label: str | None = None) -> "GDataState": - """Map one component of ``data`` with a reusable ``projection``. - - ``phi_tor`` is in radians and is used only for 3-D field-aligned input. - The source and projection computational grids must be identical. - """ - _validate_modal_data(data, "gk_rz", (2, 3)) - _validate_component(data, comp) - _validate_projection(data, projection) - _, _, values = _interpolate_component(data, comp) - - if projection.num_dims == 2: - out = values[..., np.newaxis] - else: - out = _fft_poloidal_project(values, projection.zc, projection.box, - projection.wind, projection.phi0_zf, - projection.zf, float(phi_tor))[..., np.newaxis] - - return data._result([projection.r, projection.z], - out, - inplace=inplace, - tag=tag, - label=label, - interpolated=True) - - -def rz_projections(datasets, - *, - mapc2p: str | None = None, - nodes_file: str | None = None, - z_axis: float = 0.0, - nz_interp: int = 8) -> dict[str | None, RzProjection]: - """Build one reusable R-Z projection per block geometry. - - Frames from the same block share a projection; distinct blocks resolve - their own geometry. A ``'*'`` in an explicit geometry path is replaced by - the dataset's block index. - """ - projections: dict[str | None, RzProjection] = {} - for data in datasets: - key = geometry_prefix(data.file_name) - if key in projections: - continue - block = data.ctx.get("block") - geometry = resolve_geometry(data.file_name, - mapc2p=per_block_path(mapc2p, block), - nodes_file=per_block_path(nodes_file, block)) - projections[key] = resolve_rz_projection(data, - geometry, - z_axis=z_axis, - nz_interp=nz_interp) - return projections - - -def projection_for(projections: dict[str | None, RzProjection], - data: "GDataState") -> RzProjection: - """Return the projection belonging to ``data``'s block.""" - return projections[geometry_prefix(data.file_name)] - - -def gk_rz( - data: "GDataState", - *, - mapc2p: str | None = None, - nodes_file: str | None = None, - z_axis: float = 0.0, - phi_tor: float = 0.0, - nz_interp: int = 8, - comp: int = 0, - inplace: bool = False, - tag: str | None = None, - label: str | None = None, -) -> "GDataState": - """Interpolate one DG component and project it onto a physical R-Z grid. - - ``data`` must be un-interpolated modal DG data with two computational - dimensions, or three for a field-aligned reconstruction. Geometry is - inferred from ``data.file_name``: the pointwise - ``'-geo_int_nodes.gkyl'`` file is preferred, falling back to - ``'-geo_int_mapc2p.gkyl'``. ``nodes_file`` or ``mapc2p`` may - override that choice, but they are mutually exclusive; ``mapc2p=''`` - forces the inferred modal filename. - - Args: - data: Un-interpolated 2-D or 3-D modal DG field. - mapc2p: Optional explicit modal geometry path, or ``''`` for inferred. - nodes_file: Optional explicit nodal geometry path. - z_axis: Magnetic-axis vertical position in meters, added to geometry Z. - phi_tor: Toroidal angle in radians for a 3-D poloidal reconstruction. - nz_interp: Positive integer z-direction up-sampling factor for 3-D data. - comp: Zero-based physical field component to map (default first). - inplace: Replace ``data`` rather than returning a new concrete instance. - tag: Optional result tag; ``None`` preserves the source tag. - label: Optional result label; ``None`` preserves the source label. - - Returns: - The caller's concrete data class, marked ``interpolated=True``. If the - interpolated input counts are ``(Nx, Nz)``, 2-D grid arrays have shape - ``(Nx+1, Nz+1)`` and values ``(Nx, Nz, 1)``. For interpolated 3-D counts - ``(Nx, Ny, Nz)``, grid arrays have shape - ``(Nx+1, nz_interp*Nz+1)`` and values - ``(Nx, nz_interp*Nz, 1)``. - - Raises: - ValueError: For mutually exclusive or missing geometry, input other than - un-interpolated 2-D/3-D modal data, a missing 3-D toroidal angle, - invalid ``comp`` or ``nz_interp``, malformed geometry, or a projection - incompatible with its data grid (when using :func:`map_to_rz`). - """ - _validate_modal_data(data, "gk_rz", (2, 3)) - _validate_positive_int(nz_interp, "nz_interp") - _validate_component(data, comp) - geometry = resolve_geometry(data.file_name, - mapc2p=mapc2p, - nodes_file=nodes_file) - projection = resolve_rz_projection(data, - geometry, - z_axis=z_axis, - nz_interp=nz_interp) - return map_to_rz(data, - projection, - phi_tor=phi_tor, - comp=comp, - inplace=inplace, - tag=tag, - label=label) - - -__all__ = [ - "Geometry", - "RzProjection", - "geometry_prefix", - "gk_rz", - "map_to_rz", - "per_block_path", - "projection_for", - "resolve_geometry", - "resolve_rz_projection", - "rz_projections", -] diff --git a/src/postgkyl/operations/map.py b/src/postgkyl/operations/map.py index 17244948..c8c62e25 100644 --- a/src/postgkyl/operations/map.py +++ b/src/postgkyl/operations/map.py @@ -1,4 +1,11 @@ -"""The ``map`` verb -- deform a dataset's grid by evaluating a coordinate map. +"""Coordinate transformations with explicit mapping inputs. + +``map_to_rz`` reconstructs a poloidal slice; ``extract_flux_surface`` samples +a toroidal surface. Both receive precomputed geometry and never discover +auxiliary files. Three-dimensional R-Z reconstruction assumes periodic +field-aligned coordinates with twist-and-shift boundary conditions. + +The ``map`` verb deforms a dataset's grid by evaluating a coordinate map. See ``MAPPING.md`` for the full design. A mapping file is a DG field whose components hold the coefficients of the physical coordinates of each mapped @@ -14,6 +21,24 @@ from typing import TYPE_CHECKING +import numpy as np +from scipy.interpolate import PchipInterpolator + +from postgkyl.numerics.resample import resample_grid +from postgkyl.numerics.rz import fft_poloidal_project +from .geometry import ( + Geometry, + RzProjection, + FluxSurfaceGrid, + _interpolate_component, + _interpolation_grid, + validate_mapping_grid, + _validate_component, + _validate_geometry, + _validate_modal_data, + _validate_positive_int, +) + from postgkyl import dg from postgkyl.gdatastate.gdatastate import GDataState @@ -161,3 +186,262 @@ def map(data: "_GDataState", label=label, grid_type="mapped", mapped_axes=mapped_axes) + + +def resolve_rz_projection(first: "GDataState", + geo: Geometry, + *, + z_axis: float = 0.0, + nz_interp: int = 8) -> RzProjection: + """Build an R-Z projection for ``first``'s grid and ``geo``. + + Only the computational grid and DG metadata are read from ``first``; + projection construction never evaluates or selects its field values. + ``z_axis`` is the magnetic-axis vertical position in meters. + """ + _validate_modal_data(first, "map_to_rz", (2, 3)) + nz_interp = _validate_positive_int(nz_interp, "nz_interp") + _validate_geometry(geo, first.num_dims) + edges, centers = _interpolation_grid(first) + vert_z = geo.vert_z + float(z_axis) + + if first.num_dims == 2: + r = resample_grid(geo.major_r, geo.coords, edges) + z = resample_grid(vert_z, geo.coords, edges) + return RzProjection(num_dims=2, + r=r, + z=z, + computational_grid=tuple( + np.array(axis, copy=True) for axis in edges)) + + if geo.phi is None: + raise ValueError( + "The geometry has no toroidal-angle component; 3-D map_to_rz requires one." + ) + + xc, yc, zc = centers + nx, ny, nz = xc.size, yc.size, zc.size + if ny < 2 or nz < 2: + raise ValueError( + "3-D map_to_rz requires at least two interpolated y and z points.") + + phi = np.unwrap(np.unwrap(np.unwrap(geo.phi, axis=2), axis=1), axis=0) + phi_field = resample_grid(phi, geo.coords, centers) + box_estimate = np.mean(np.diff(phi_field[nx // 2, :, nz // 2])) * ny + if not np.isfinite(box_estimate) or np.isclose(box_estimate, 0.0): + raise ValueError( + "Toroidal geometry has a zero or non-finite binormal angular span.") + n0 = max(1, int(round(abs(2.0 * np.pi / box_estimate)))) + box = np.sign(box_estimate) * 2.0 * np.pi / n0 + wind = phi_field[:, 0, -1] - phi_field[:, 0, 0] + + xn, _, zn = edges + zf_edges = np.linspace(zn[0], zn[-1], nz_interp * nz + 1) + zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) + phi0_zf = np.array( + [np.interp(zf, zc, phi_field[ix, 0, :]) for ix in range(nx)]) + + gx, _, gz = geo.coords + r2d = geo.major_r[:, 0, :] + z2d = vert_z[:, 0, :] + gz_rz = gz + if geo.corner is not None: + corner_coords, corner_r, corner_z = geo.corner + if len(corner_coords) != 3: + raise ValueError( + "Corner geometry must be three-dimensional for 3-D map_to_rz.") + cx, cz = corner_coords[0], corner_coords[2] + corner_r = corner_r[:, 0, :] + corner_z = corner_z[:, 0, :] + float(z_axis) + r2d = np.concatenate([ + np.interp(gx, cx, corner_r[:, 0])[:, None], r2d, + np.interp(gx, cx, corner_r[:, -1])[:, None] + ], + axis=1) + z2d = np.concatenate([ + np.interp(gx, cx, corner_z[:, 0])[:, None], z2d, + np.interp(gx, cx, corner_z[:, -1])[:, None] + ], + axis=1) + gz_rz = np.concatenate([[cz[0]], gz, [cz[-1]]]) + + r = resample_grid(r2d, [gx, gz_rz], [xn, zf_edges]) + z = resample_grid(z2d, [gx, gz_rz], [xn, zf_edges]) + return RzProjection(num_dims=3, + r=r, + z=z, + zc=zc, + zf=zf, + box=box, + wind=wind, + phi0_zf=phi0_zf, + computational_grid=tuple( + np.array(axis, copy=True) for axis in edges)) + + +def _validate_projection(data: "GDataState", projection: RzProjection) -> None: + if projection.num_dims not in (2, 3): + raise ValueError( + f"R-Z projection has invalid dimensionality {projection.num_dims}; expected 2 or 3." + ) + if data.num_dims != projection.num_dims: + raise ValueError( + "Incompatible R-Z projection: projection dimensionality does not match the data." + ) + edges, _ = _interpolation_grid(data) + validate_mapping_grid(data, projection.computational_grid) + if projection.r.shape != projection.z.shape or projection.r.ndim != 2: + raise ValueError( + "Incompatible R-Z projection: R and Z grids must be matching 2-D arrays." + ) + + if projection.num_dims == 2: + expected = (edges[0].size, edges[1].size) + if projection.r.shape != expected: + raise ValueError( + f"Incompatible R-Z projection: expected grid shape {expected}, " + f"got {projection.r.shape}.") + return + + required = (projection.zc, projection.zf, projection.box, projection.wind, + projection.phi0_zf) + if any(value is None for value in required): + raise ValueError( + "Incompatible R-Z projection: 3-D projection metadata is incomplete.") + if not np.isfinite(projection.box) or np.isclose(projection.box, 0.0): + raise ValueError( + "Incompatible R-Z projection: toroidal angular span must be finite and nonzero." + ) + nx, _, nz = (axis.size - 1 for axis in edges) + if (projection.zc.shape != (nz, ) or projection.wind.shape != (nx, ) + or projection.phi0_zf.shape != (nx, projection.zf.size) + or projection.r.shape != (nx + 1, projection.zf.size + 1)): + raise ValueError( + "Incompatible R-Z projection: projection and data grid shapes differ.") + + +def map_to_rz(data: "GDataState", + *, + projection: RzProjection, + phi_tor: float = 0.0, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Map one component of ``data`` with a reusable ``projection``. + + ``phi_tor`` is in radians and is used only for 3-D field-aligned input. + The source and projection computational grids must be identical. + """ + _validate_modal_data(data, "map_to_rz", (2, 3)) + _validate_component(data, comp) + _validate_projection(data, projection) + _, _, values = _interpolate_component(data, comp) + + if projection.num_dims == 2: + out = values[..., np.newaxis] + else: + out = fft_poloidal_project(values, projection.zc, projection.box, + projection.wind, projection.phi0_zf, + projection.zf, float(phi_tor))[..., np.newaxis] + + return data._result([projection.r, projection.z], + out, + inplace=inplace, + tag=tag, + label=label, + interpolated=True) + + +def resolve_flux_surface_grid(first: "GDataState", + geo: Geometry, + *, + x_idx: int = 0, + nphi: int = 128, + nz_interp: int = 8) -> FluxSurfaceGrid: + """Precompute a flux-surface sampling grid for compatible 3-D fields.""" + _validate_modal_data(first, "extract_flux_surface", (3, )) + nphi = _validate_positive_int(nphi, "nphi") + nz_interp = _validate_positive_int(nz_interp, "nz_interp") + _validate_geometry(geo, 3) + if geo.phi is None: + raise ValueError( + "The geometry has no toroidal-angle component; cannot extract a flux surface." + ) + if isinstance(x_idx, bool) or not isinstance(x_idx, (int, np.integer)): + raise ValueError("x_idx must be an integer radial index.") + + edges, centers = _interpolation_grid(first) + xc, yc, zc = centers + x_idx = int(x_idx) + if not 0 <= x_idx < xc.size: + raise ValueError( + f"x_idx {x_idx} is out of bounds for data with Nx={xc.size}.") + if zc.size < 2 or yc.size < 2: + raise ValueError( + "extract_flux_surface requires at least two interpolated y and z points." + ) + + zf_edges = np.linspace(edges[2][0], edges[2][-1], nz_interp * zc.size + 1) + zf = 0.5 * (zf_edges[:-1] + zf_edges[1:]) + phi = np.unwrap(np.unwrap(np.unwrap(geo.phi, axis=2), axis=1), axis=0) + phi_grid = resample_grid(phi, geo.coords, [xc, yc, zf]) + phi_2d = phi_grid[x_idx, :, :] + phi_tor_list = np.linspace(0.0, 2.0 * np.pi, nphi, endpoint=False) + return FluxSurfaceGrid(x_idx=x_idx, + zc=zc, + zf=zf, + phi_tor_list=phi_tor_list, + phi_2d=phi_2d, + computational_grid=tuple( + np.array(axis, copy=True) for axis in edges)) + + +def extract_flux_surface(data: "GDataState", + *, + fs_grid: FluxSurfaceGrid, + comp: int = 0, + inplace: bool = False, + tag: str | None = None, + label: str | None = None) -> "GDataState": + """Extract component ``comp`` using a reusable ``fs_grid``.""" + _validate_modal_data(data, "extract_flux_surface", (3, )) + _validate_component(data, comp) + edges, _ = _interpolation_grid(data) + validate_mapping_grid(data, fs_grid.computational_grid) + nx, ny, nz = (axis.size - 1 for axis in edges) + if not 0 <= fs_grid.x_idx < nx: + raise ValueError( + f"x_idx {fs_grid.x_idx} is out of bounds for data with Nx={nx}.") + if (fs_grid.zc.shape != (nz, ) + or fs_grid.phi_2d.shape != (ny, fs_grid.zf.size) + or fs_grid.phi_tor_list.ndim != 1): + raise ValueError( + "Incompatible flux-surface grid: projection and data grid shapes differ." + ) + + _, _, values = _interpolate_component(data, comp) + vals_zf = PchipInterpolator(fs_grid.zc, values, axis=-1, + extrapolate=True)(fs_grid.zf) + vals_2d = vals_zf[fs_grid.x_idx, :, :] + + flux_surf_data = np.empty((fs_grid.phi_tor_list.size, fs_grid.zf.size)) + for iz in range(fs_grid.zf.size): + phi_y = fs_grid.phi_2d[:, iz] + val_y = vals_2d[:, iz] + box = np.mean(np.diff(phi_y)) * ny + if not np.isfinite(box) or np.isclose(box, 0.0): + raise ValueError( + "Toroidal geometry has a zero or non-finite binormal angular span.") + phi_ext = np.concatenate([phi_y - box, phi_y, phi_y + box]) + val_ext = np.concatenate([val_y, val_y, val_y]) + order = np.argsort(phi_ext) + folded = phi_y[0] + np.mod(fs_grid.phi_tor_list - phi_y[0], box) + flux_surf_data[:, iz] = np.interp(folded, phi_ext[order], val_ext[order]) + + return data._result([fs_grid.phi_tor_list, fs_grid.zf], + flux_surf_data[..., np.newaxis], + inplace=inplace, + tag=tag, + label=label, + interpolated=True) diff --git a/tests/test_diagnostics_gk_fluxsurf.py b/tests/test_diagnostics_gk_fluxsurf.py new file mode 100644 index 00000000..60d2cf29 --- /dev/null +++ b/tests/test_diagnostics_gk_fluxsurf.py @@ -0,0 +1,80 @@ +"""Gkeyll flux-surface composition and per-block discovery.""" + +from __future__ import annotations + +from importlib import import_module +from types import SimpleNamespace + + +def test_flux_surface_grid_collection_caches_by_geometry_prefix(monkeypatch): + fluxsurf = import_module("postgkyl.diagnostics.gk.fluxsurf") + monkeypatch.setattr(fluxsurf, "validate_mapping_grid", lambda *_args: None) + first = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + repeated = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + second = SimpleNamespace(file_name="block-two", ctx={"block": 2}) + calls = [] + monkeypatch.setattr(fluxsurf, "geometry_prefix", lambda path: path) + monkeypatch.setattr( + fluxsurf, "resolve_geometry", lambda path, **kwargs: calls.append( + (path, kwargs)) or path) + monkeypatch.setattr( + fluxsurf, "resolve_flux_surface_grid", + lambda data, geo, **_kwargs: SimpleNamespace(computational_grid=(), + geometry=geo)) + + grids = fluxsurf.flux_surface_grids([first, repeated, second], + mapc2p="map-*.gkyl", + nodes_file="nodes-*.gkyl") + assert grids == { + "block-one": SimpleNamespace(computational_grid=(), geometry="block-one"), + "block-two": SimpleNamespace(computational_grid=(), geometry="block-two"), + } + assert calls == [ + ("block-one", { + "mapc2p": "map-1.gkyl", + "nodes_file": "nodes-1.gkyl" + }), + ("block-two", { + "mapc2p": "map-2.gkyl", + "nodes_file": "nodes-2.gkyl" + }), + ] + assert fluxsurf.grid_for(grids, + first) == SimpleNamespace(computational_grid=(), + geometry="block-one") + + +def test_gk_fluxsurf_composes_geometry_grid_and_extraction(monkeypatch): + fluxsurf = import_module("postgkyl.diagnostics.gk.fluxsurf") + monkeypatch.setattr(fluxsurf, "validate_mapping_grid", lambda *_args: None) + data = SimpleNamespace(file_name="field.gkyl") + calls = [] + monkeypatch.setattr( + fluxsurf, "resolve_geometry", lambda path, **kwargs: calls.append( + ("geometry", path, kwargs)) or "geo") + monkeypatch.setattr( + fluxsurf, "resolve_flux_surface_grid", + lambda source, geo, **kwargs: calls.append( + ("grid", source, geo, kwargs)) or "grid") + monkeypatch.setattr( + fluxsurf, "extract_flux_surface", + lambda source, fs_grid, **kwargs: calls.append( + ("extract", source, fs_grid, kwargs)) or "result") + + result = fluxsurf.fluxsurf(data, + mapc2p="map.gkyl", + x_idx=2, + nphi=16, + nz_interp=3, + comp=4, + inplace=True, + tag="surface", + label="flux") + assert result == "result" + assert [call[0] for call in calls] == ["geometry", "grid", "extract"] + assert calls[-1][-1] == { + "comp": 4, + "inplace": True, + "tag": "surface", + "label": "flux" + } diff --git a/tests/test_operations_gk_geometry.py b/tests/test_diagnostics_gk_geometry.py similarity index 53% rename from tests/test_operations_gk_geometry.py rename to tests/test_diagnostics_gk_geometry.py index e12324cd..ee6cc8a0 100644 --- a/tests/test_operations_gk_geometry.py +++ b/tests/test_diagnostics_gk_geometry.py @@ -1,29 +1,14 @@ -"""Unit contracts for shared gyrokinetic geometry machinery.""" +"""Gkeyll auxiliary geometry discovery and decoding.""" from __future__ import annotations -from dataclasses import replace from importlib import import_module from types import SimpleNamespace import numpy as np import pytest -import postgkyl as pg -from postgkyl import gpython - -geometry = import_module("postgkyl.operations.gyrokinetics.geometry") - -needs_gkeyll = pytest.mark.skipif( - not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") - - -def _valid_geometry(num_dims=2, *, phi=False, corner=None): - coords = [np.array([0.0, 1.0]) for _ in range(num_dims)] - shape = (2, ) * num_dims - values = np.ones(shape) - return geometry.Geometry(coords, values, 2.0 * values, - values if phi else None, corner) +geometry = import_module("postgkyl.diagnostics.gk.geometry") def test_gauss_nodes_are_ordered_inside_each_cell(): @@ -103,32 +88,6 @@ def test_read_corner_geometry_builds_point_coordinates(monkeypatch): assert major_r.shape == vert_z.shape == (3, 4) -@pytest.mark.parametrize(("candidate", "num_dims", "message"), [ - (_valid_geometry(1), 2, "Geometry has 1 dimensions"), - (geometry.Geometry([np.array([[0.0, 1.0]])], np.ones((2, )), np.ones( - (2, )), None, None), 1, "one-dimensional arrays"), - (geometry.Geometry([np.array([0.0, 1.0, 0.5])], np.ones( - (3, )), np.ones((3, )), None, None), 1, "strictly monotonic"), - (geometry.Geometry([np.array([0.0, 1.0])], np.ones((3, )), np.ones( - (2, )), None, None), 1, "R/Z array shapes are incompatible"), - (replace(_valid_geometry(2), phi=np.ones( - (2, ))), 2, "toroidal-angle shape"), - (_valid_geometry(2, - corner=([np.array([0.0, 1.0])], np.ones( - (2, )), np.ones( - (2, )))), 2, "Corner geometry has 1 dimensions"), - (_valid_geometry( - 2, - corner=([np.array([0.0]), np.array([0.0, 1.0])], np.ones( - (1, 2)), np.ones( - (1, 2)))), 2, "Corner geometry coordinate and R/Z"), -]) -def test_validate_geometry_rejects_each_shape_invariant(candidate, num_dims, - message): - with pytest.raises(ValueError, match=message): - geometry._validate_geometry(candidate, num_dims) - - def test_resolve_geometry_honors_explicit_nodes_and_loads_corner( monkeypatch, tmp_path): source = tmp_path / "sim-field_0.gkyl" @@ -153,58 +112,3 @@ def test_resolve_geometry_honors_explicit_nodes_and_loads_corner( def test_resolve_geometry_without_a_name_requires_an_override(): with pytest.raises(ValueError, match="Could not find a geometry file"): geometry.resolve_geometry(None) - - -def test_validate_modal_data_reports_missing_data_and_metadata(): - empty = pg.GData() - with pytest.raises(ValueError, match="loaded dataset"): - geometry._validate_modal_data(empty, "projection", (0, )) - - no_basis = pg.GData() - no_basis.push([np.array([0.0, 1.0])], np.ones((1, 1))) - with pytest.raises(ValueError, match="basis_type"): - geometry._validate_modal_data(no_basis, "projection", (1, )) - - no_basis.ctx["basis_type"] = "serendipity" - no_basis.ctx["poly_order"] = True - with pytest.raises(ValueError, match="nonnegative integer"): - geometry._validate_modal_data(no_basis, "projection", (1, )) - - -def test_validate_modal_data_reports_grid_shape_and_monotonicity(): - data = pg.GData(ctx={ - "basis_type": "serendipity", - "poly_order": 0, - "value_form": "modal", - }) - data.push([np.array([0.0, 1.0])], np.ones((1, 1))) - data._grid = [np.array([0.0])] - with pytest.raises(ValueError, match="one-dimensional edge grid"): - geometry._validate_modal_data(data, "projection", (1, )) - - data._grid = [np.array([0.0, 1.0, 0.5])] - with pytest.raises(ValueError, match="strictly monotonic"): - geometry._validate_modal_data(data, "projection", (1, )) - - -@needs_gkeyll -def test_num_fields_rejects_incompatible_coefficient_count(): - data = pg.GData(ctx={ - "basis_type": "serendipity", - "poly_order": 1, - "value_form": "modal", - }) - data.push([np.array([0.0, 1.0])], np.ones((1, 3))) - with pytest.raises(ValueError, match="incompatible"): - geometry._num_fields(data) - - -def test_validate_component_rejects_boolean_before_basis_lookup(): - with pytest.raises(ValueError, match="integer component"): - geometry._validate_component(pg.GData(), True) - - -def test_same_grid_rejects_dimension_and_shape_mismatches(): - axis = np.array([0.0, 1.0]) - assert not geometry._same_grid([axis], [axis, axis]) - assert not geometry._same_grid([axis], [np.array([0.0, 0.5, 1.0])]) diff --git a/tests/test_diagnostics_gk_rz.py b/tests/test_diagnostics_gk_rz.py new file mode 100644 index 00000000..b8f698ed --- /dev/null +++ b/tests/test_diagnostics_gk_rz.py @@ -0,0 +1,184 @@ +"""Explicit R-Z mapping and Gkeyll geometry compositions.""" + +from __future__ import annotations + +from importlib import import_module +import os +from types import SimpleNamespace + +import click +import numpy as np +import pytest +from click.testing import CliRunner + +import postgkyl as pg +from postgkyl import gpython +from postgkyl.cli.app import COMMANDS + +gk_rz_command = next(command for command in COMMANDS if command.name == "gk_rz") +from postgkyl.cli.state import DataSpace +from postgkyl.diagnostics.gk.geometry import resolve_geometry + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1D = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") +F2D = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-elc_M2par_10.gkyl") +F2D_GEO = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-geo_int_mapc2p.gkyl") +F3D = os.path.join(DATA, "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +def test_geometry_prefers_nodes_and_honors_explicit_modal_override( + tmp_path, monkeypatch): + from postgkyl.diagnostics.gk import geometry as geometry_module + + source = tmp_path / "sim-field_0.gkyl" + nodes = tmp_path / "sim-geo_int_nodes.gkyl" + modal = tmp_path / "sim-geo_int_mapc2p.gkyl" + nodes.touch() + modal.touch() + coords = [np.array([0.0, 1.0]), np.array([-1.0, 1.0])] + arrays = np.ones((2, 2)) + calls = [] + monkeypatch.setattr( + geometry_module, "_read_nodes_geometry", lambda path: (calls.append( + ("nodes", path)) or (coords, arrays, arrays, None))) + monkeypatch.setattr( + geometry_module, "_read_mapc2p_geometry", lambda path: (calls.append( + ("mapc2p", path)) or (coords, arrays, arrays, None))) + + resolve_geometry(str(source)) + assert calls[-1] == ("nodes", str(nodes)) + resolve_geometry(str(source), mapc2p="") + assert calls[-1] == ("mapc2p", str(modal)) + + +@needs_gkeyll +def test_geometry_overrides_and_validation_errors(tmp_path): + data = pg.load(F2D) + with pytest.raises(ValueError, match="either mapc2p=.*nodes_file"): + pg.gk.rz(data, mapc2p=F2D_GEO, nodes_file=F2D_GEO) + explicit = pg.gk.rz(data, mapc2p=F2D_GEO, nz_interp=2) + inferred = pg.gk.rz(data, nz_interp=2) + np.testing.assert_allclose(explicit.values, inferred.values) + + missing = data.clone() + missing._file_name = str(tmp_path / "absent-field_0.gkyl") + with pytest.raises(ValueError, match="Could not find a geometry file"): + pg.gk.rz(missing) + with pytest.raises(ValueError, match="positive integer"): + pg.gk.rz(data, nz_interp=0) + with pytest.raises(ValueError, match="out of bounds"): + pg.gk.rz(data, comp=1) + with pytest.raises(ValueError, match="requires 2-D or 3-D"): + pg.gk.rz(pg.load(F1D), mapc2p=F2D_GEO) + with pytest.raises(ValueError, match="un-interpolated modal DG"): + pg.gk.rz(data.interpolate()) + + +def test_rz_projection_collection_caches_by_geometry_prefix(monkeypatch): + rz = import_module("postgkyl.diagnostics.gk.rz") + monkeypatch.setattr(rz, "validate_mapping_grid", lambda *_args: None) + first = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + repeated = SimpleNamespace(file_name="block-one", ctx={"block": 1}) + second = SimpleNamespace(file_name="block-two", ctx={"block": 2}) + calls = [] + monkeypatch.setattr(rz, "geometry_prefix", lambda path: path) + monkeypatch.setattr( + rz, "resolve_geometry", lambda path, **kwargs: calls.append( + (path, kwargs)) or path) + monkeypatch.setattr( + rz, "resolve_rz_projection", + lambda data, geo, **_kwargs: SimpleNamespace(computational_grid=(), + geometry=geo)) + + projections = rz.rz_projections([first, repeated, second], + mapc2p="map-*.gkyl", + nodes_file="nodes-*.gkyl") + assert projections == { + "block-one": SimpleNamespace(computational_grid=(), geometry="block-one"), + "block-two": SimpleNamespace(computational_grid=(), geometry="block-two"), + } + assert calls == [ + ("block-one", { + "mapc2p": "map-1.gkyl", + "nodes_file": "nodes-1.gkyl" + }), + ("block-two", { + "mapc2p": "map-2.gkyl", + "nodes_file": "nodes-2.gkyl" + }), + ] + assert rz.projection_for(projections, + first) == SimpleNamespace(computational_grid=(), + geometry="block-one") + + +@needs_gkeyll +def test_group_mapping_cli_and_help_section(): + from postgkyl.cli.app import cli + + group = pg.GDataGroup([pg.load(F2D), pg.load(F2D)]) + projection = pg.resolve_rz_projection(group[0], + resolve_geometry(F2D), + nz_interp=2) + mapped = group.map_to_rz(projection=projection) + assert isinstance(mapped, pg.GDataGroup) and len(mapped) == 2 + + space = DataSpace(datasets=[pg.load(F2D)]) + with click.Context(gk_rz_command, obj=space) as ctx: + ctx.invoke(gk_rz_command, + mapc2p=F2D_GEO, + nodes_file=None, + z_axis=0.0, + phi_tor=0.0, + nz_interp=2, + use=None, + tag="rz", + label=None) + expected = pg.gk.rz(pg.load(F2D), mapc2p=F2D_GEO, nz_interp=2, tag="rz") + np.testing.assert_allclose(space.datasets[0].values, expected.values) + + help_text = CliRunner().invoke(cli, ["--help"]).output + verbs = help_text.split("Diagnostics:", 1)[0] + diagnostics = help_text.split("Diagnostics:", 1)[1].split("Render:", 1)[0] + assert "gk_rz" not in verbs and "gk_rz" in diagnostics + + +@needs_gkeyll +@pytest.mark.parametrize("module_name, builder_name", [ + ("rz", "rz_projections"), + ("fluxsurf", "flux_surface_grids"), +]) +def test_cached_mapping_rejects_changed_grid_with_same_prefix( + module_name, builder_name): + module = import_module(f"postgkyl.diagnostics.gk.{module_name}") + first = pg.load(F3D) + changed = first.clone() + changed.grid[0] = changed.grid[0] + 0.1 + with pytest.raises(ValueError, match="computational grid does not match"): + getattr(module, builder_name)([first, changed], nz_interp=2) + + +@needs_gkeyll +def test_flux_surface_diagnostic_cli_matches_explicit_core(): + from postgkyl.cli.app import cli + data = pg.load(F3D) + geometry = pg.gk.resolve_geometry(data.file_name) + grid = pg.resolve_flux_surface_grid(data, geometry, nphi=4, nz_interp=2) + expected = data.extract_flux_surface(fs_grid=grid) + actual = pg.gk.fluxsurf(data, nphi=4, nz_interp=2) + np.testing.assert_allclose(actual.values, expected.values) + command = next(c for c in COMMANDS if c.name == "gk_fluxsurf") + space = DataSpace(datasets=[data]) + with click.Context(command, obj=space) as ctx: + ctx.invoke(command, nphi=4, nz_interp=2, use=None, tag=None, label=None) + np.testing.assert_allclose(space.datasets[0].values, expected.values) + result = CliRunner().invoke(cli, ["--help"]) + assert result.exit_code == 0 + assert "gk_fluxsurf" in result.output.split("Diagnostics:", 1)[1] + assert not hasattr(pg, "gk_fluxsurf") + assert not hasattr(pg.GData, "gk_fluxsurf") diff --git a/tests/test_diagnostics_programs_nodes.py b/tests/test_diagnostics_programs_nodes.py index 741563a4..7886698a 100644 --- a/tests/test_diagnostics_programs_nodes.py +++ b/tests/test_diagnostics_programs_nodes.py @@ -27,6 +27,7 @@ import pytest from postgkyl.diagnostics.gk import utils as gk_utils +from postgkyl.diagnostics.gk import geometry nodes = importlib.import_module("postgkyl.diagnostics.gk.nodes") @@ -39,19 +40,19 @@ class TestGeometryEnum: def test_mapc2p_index_matches_gkeyll_header(self): # gkeyll/core/zero/gkyl_eqn_type.h: GKYL_GEOMETRY_MAPC2P = 3. - assert nodes.GKYL_GEOMETRY_ID.index("GKYL_GEOMETRY_MAPC2P") == 3 + assert geometry.GKYL_GEOMETRY_ID.index("GKYL_GEOMETRY_MAPC2P") == 3 class TestIsGeoMapc2p: def test_defaults_true_when_absent(self): - assert nodes.is_geo_mapc2p({}) is True + assert geometry.is_geo_mapc2p({}) is True def test_true_for_mapc2p(self): - assert nodes.is_geo_mapc2p({"geometry_type": 3}) is True + assert geometry.is_geo_mapc2p({"geometry_type": 3}) is True def test_false_for_tokamak(self): - assert nodes.is_geo_mapc2p({"geometry_type": 1}) is False + assert geometry.is_geo_mapc2p({"geometry_type": 1}) is False class TestNodesToRZ: diff --git a/tests/test_gdata_fluent.py b/tests/test_gdata_fluent.py index 75046d28..cb50f2aa 100644 --- a/tests/test_gdata_fluent.py +++ b/tests/test_gdata_fluent.py @@ -5,7 +5,7 @@ facade re-exports. Physics diagnostics are deliberately not fluent methods. Domain-specific data -transformations such as ``gk_rz`` are operations and do belong on the fluent +transformations with explicit geometry such as ``map_to_rz`` belong on the fluent surface alongside domain-independent core verbs. """ @@ -68,8 +68,8 @@ def _line(cls=MyData, tag: str = "default", value: float = 1.0, n: int = 5): # operations that act on the group as a whole. Every multi-dataset operation # also has a functional spelling on the top-level ``pg`` facade. INSTANCE_VERBS = [ - "load", "interpolate", "local_poly", "gk_rz", "select", "plot", "plotly", - "pyvista", "save", "mul", "div", "integrate", "average", + "load", "interpolate", "local_poly", "map_to_rz", "select", "plot", + "plotly", "pyvista", "save", "mul", "div", "integrate", "average", "eval_at_coord_proj", "to_modal", "to_nodal", "to_quad", "apply", "fft", "magsq", "mask", "val2coord", "extract_input", "fit", "differentiate", "map" ] diff --git a/tests/test_multiblock.py b/tests/test_multiblock.py index 48ed75f7..40bacd01 100644 --- a/tests/test_multiblock.py +++ b/tests/test_multiblock.py @@ -219,7 +219,8 @@ def test_explicit_zlim_still_wins(self): class TestPerBlockGeometry: def test_geometry_prefix_is_per_block(self): - from postgkyl.diagnostics.gk import rz + from importlib import import_module + rz = import_module("postgkyl.diagnostics.gk.rz") assert rz.geometry_prefix("d/sim_b2-elc_M0_3.gkyl") == "d/sim_b2" assert rz.geometry_prefix("d/sim-elc_M0_3.gkyl") == "d/sim" @@ -237,14 +238,19 @@ def test_each_block_resolves_its_own_geometry(self, monkeypatch): # The bug this replaces: geometry was resolved once, from the first # dataset, and that one projection was applied to every block -- drawing # every block at block 0's position. - from postgkyl.diagnostics.gk import rz + from importlib import import_module + rz = import_module("postgkyl.diagnostics.gk.rz") + from types import SimpleNamespace + monkeypatch.setattr(rz, "validate_mapping_grid", lambda *_args: None) seen = [] monkeypatch.setattr( rz, "resolve_geometry", lambda file_name, **kw: seen.append(file_name) or file_name) - monkeypatch.setattr(rz, "resolve_rz_projection", lambda first, geo, **kw: - ("projection", geo)) + monkeypatch.setattr( + rz, "resolve_rz_projection", + lambda first, geo, **kw: SimpleNamespace(computational_grid=(), + geometry=geo)) blocks = _blocks(0) + _blocks(1) # 3 blocks x 2 frames projections = rz.rz_projections(blocks) diff --git a/tests/test_numerics_rz.py b/tests/test_numerics_rz.py new file mode 100644 index 00000000..26fb9e24 --- /dev/null +++ b/tests/test_numerics_rz.py @@ -0,0 +1,17 @@ +"""Analytic checks of periodic poloidal reconstruction.""" +import numpy as np +import pytest + +from postgkyl.numerics.rz import fft_poloidal_project + + +@pytest.mark.parametrize("ny", [7, 8]) +@pytest.mark.parametrize("angle", [0.0, 0.37, np.pi / 2]) +def test_fourier_mode_reconstructs_at_physical_angle(ny, angle): + zc = np.linspace(-0.75, 0.75, 4) + samples = np.cos(2 * np.pi * np.arange(ny) / ny) + values = np.broadcast_to(samples[None, :, None], (2, ny, 4)) + zf = np.linspace(-1.0, 1.0, 9) + reconstructed = fft_poloidal_project(values, zc, 2 * np.pi, np.zeros(2), + np.zeros((2, 9)), zf, angle) + np.testing.assert_allclose(reconstructed, np.cos(angle), atol=1e-14) diff --git a/tests/test_operations_geometry.py b/tests/test_operations_geometry.py new file mode 100644 index 00000000..d3e33f77 --- /dev/null +++ b/tests/test_operations_geometry.py @@ -0,0 +1,106 @@ +"""Contracts for explicit coordinate geometry and modal mapping inputs.""" + +from __future__ import annotations + +from dataclasses import replace + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython + +from postgkyl.operations import geometry as core_geometry + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +def _valid_geometry(num_dims=2, *, phi=False, corner=None): + coords = [np.array([0.0, 1.0]) for _ in range(num_dims)] + shape = (2, ) * num_dims + values = np.ones(shape) + return core_geometry.Geometry(coords, values, 2.0 * values, + values if phi else None, corner) + + +@pytest.mark.parametrize(("candidate", "num_dims", "message"), [ + (_valid_geometry(1), 2, "Geometry has 1 dimensions"), + (core_geometry.Geometry([np.array([[0.0, 1.0]])], np.ones( + (2, )), np.ones((2, )), None, None), 1, "one-dimensional arrays"), + (core_geometry.Geometry([np.array([0.0, 1.0, 0.5])], np.ones( + (3, )), np.ones((3, )), None, None), 1, "strictly monotonic"), + (core_geometry.Geometry([np.array([0.0, 1.0])], np.ones( + (3, )), np.ones( + (2, )), None, None), 1, "R/Z array shapes are incompatible"), + (replace(_valid_geometry(2), phi=np.ones( + (2, ))), 2, "toroidal-angle shape"), + (_valid_geometry(2, + corner=([np.array([0.0, 1.0])], np.ones( + (2, )), np.ones( + (2, )))), 2, "Corner geometry has 1 dimensions"), + (_valid_geometry( + 2, + corner=([np.array([0.0]), np.array([0.0, 1.0])], np.ones( + (1, 2)), np.ones( + (1, 2)))), 2, "Corner geometry coordinate and R/Z"), +]) +def test_validate_geometry_rejects_each_shape_invariant(candidate, num_dims, + message): + with pytest.raises(ValueError, match=message): + core_geometry._validate_geometry(candidate, num_dims) + + +def test_validate_modal_data_reports_missing_data_and_metadata(): + empty = pg.GData() + with pytest.raises(ValueError, match="loaded dataset"): + core_geometry._validate_modal_data(empty, "projection", (0, )) + + no_basis = pg.GData() + no_basis.push([np.array([0.0, 1.0])], np.ones((1, 1))) + with pytest.raises(ValueError, match="basis_type"): + core_geometry._validate_modal_data(no_basis, "projection", (1, )) + + no_basis.ctx["basis_type"] = "serendipity" + no_basis.ctx["poly_order"] = True + with pytest.raises(ValueError, match="nonnegative integer"): + core_geometry._validate_modal_data(no_basis, "projection", (1, )) + + +def test_validate_modal_data_reports_grid_shape_and_monotonicity(): + data = pg.GData(ctx={ + "basis_type": "serendipity", + "poly_order": 0, + "value_form": "modal", + }) + data.push([np.array([0.0, 1.0])], np.ones((1, 1))) + data._grid = [np.array([0.0])] + with pytest.raises(ValueError, match="one-dimensional edge grid"): + core_geometry._validate_modal_data(data, "projection", (1, )) + + data._grid = [np.array([0.0, 1.0, 0.5])] + with pytest.raises(ValueError, match="strictly monotonic"): + core_geometry._validate_modal_data(data, "projection", (1, )) + + +@needs_gkeyll +def test_num_fields_rejects_incompatible_coefficient_count(): + data = pg.GData(ctx={ + "basis_type": "serendipity", + "poly_order": 1, + "value_form": "modal", + }) + data.push([np.array([0.0, 1.0])], np.ones((1, 3))) + with pytest.raises(ValueError, match="incompatible"): + core_geometry._num_fields(data) + + +def test_validate_component_rejects_boolean_before_basis_lookup(): + with pytest.raises(ValueError, match="integer component"): + core_geometry._validate_component(pg.GData(), True) + + +def test_same_grid_rejects_dimension_and_shape_mismatches(): + axis = np.array([0.0, 1.0]) + assert not core_geometry._same_grid([axis], [axis, axis]) + assert not core_geometry._same_grid([axis], [np.array([0.0, 0.5, 1.0])]) diff --git a/tests/test_operations_gk_fluxsurf.py b/tests/test_operations_gk_fluxsurf.py deleted file mode 100644 index 243572d9..00000000 --- a/tests/test_operations_gk_fluxsurf.py +++ /dev/null @@ -1,173 +0,0 @@ -"""Characterization tests for the moved gyrokinetic flux-surface operation.""" - -from __future__ import annotations - -from dataclasses import replace -from importlib import import_module -import os -from types import SimpleNamespace - -import numpy as np -import pytest - -import postgkyl as pg -from postgkyl import gpython -from postgkyl.operations import gyrokinetics as gk_ops - -ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -FIELD = os.path.join(ROOT, "tests", "test_data", - "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") - -needs_gkeyll = pytest.mark.skipif( - not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") - - -@needs_gkeyll -def test_flux_surface_move_preserves_output_and_projection_reuse(): - data = pg.load(FIELD) - geometry = gk_ops.resolve_geometry(data.file_name) - grid = gk_ops.resolve_flux_surface_grid(data, - geometry, - x_idx=0, - nphi=4, - nz_interp=2) - first = gk_ops.extract_flux_surface(data, grid) - second = gk_ops.extract_flux_surface(data.clone(), grid) - assert first.values.shape == (4, 64, 1) - assert first.ctx["interpolated"] is True - np.testing.assert_allclose(first.values, second.values) - - -@needs_gkeyll -@pytest.mark.parametrize(("kwargs", "message"), [ - ({ - "nphi": 0 - }, "nphi must be a positive integer"), - ({ - "nz_interp": 0 - }, "nz_interp must be a positive integer"), - ({ - "x_idx": -1 - }, "out of bounds"), -]) -def test_flux_surface_public_validation(kwargs, message): - data = pg.load(FIELD) - geometry = gk_ops.resolve_geometry(data.file_name) - with pytest.raises(ValueError, match=message): - gk_ops.resolve_flux_surface_grid(data, geometry, **kwargs) - - -@needs_gkeyll -def test_flux_surface_grid_requires_toroidal_geometry_and_integer_index(): - data = pg.load(FIELD) - geometry = gk_ops.resolve_geometry(data.file_name) - with pytest.raises(ValueError, match="no toroidal-angle component"): - gk_ops.resolve_flux_surface_grid(data, replace(geometry, phi=None)) - with pytest.raises(ValueError, match="x_idx must be an integer"): - gk_ops.resolve_flux_surface_grid(data, geometry, x_idx=True) - - -@needs_gkeyll -def test_flux_surface_grid_requires_two_binormal_and_parallel_points(): - data = pg.load(FIELD).clone() - data.ctx["poly_order"] = 0 - data._grid[1] = np.array([0.0, 1.0]) - geometry = gk_ops.resolve_geometry(data.file_name) - with pytest.raises(ValueError, match="at least two interpolated y and z"): - gk_ops.resolve_flux_surface_grid(data, geometry) - - -@needs_gkeyll -def test_extract_flux_surface_validates_reusable_grid_metadata(): - data = pg.load(FIELD) - geometry = gk_ops.resolve_geometry(data.file_name) - grid = gk_ops.resolve_flux_surface_grid(data, geometry, nphi=4, nz_interp=2) - - shifted = data.clone() - shifted.grid[0] = shifted.grid[0] + 0.1 - with pytest.raises(ValueError, match="computational grid does not match"): - gk_ops.extract_flux_surface(shifted, grid) - - with pytest.raises(ValueError, match="out of bounds"): - gk_ops.extract_flux_surface(data, replace(grid, x_idx=10_000)) - - with pytest.raises(ValueError, match="projection and data grid shapes"): - gk_ops.extract_flux_surface(data, replace(grid, phi_2d=np.ones((1, 1)))) - - -@needs_gkeyll -def test_extract_flux_surface_rejects_zero_toroidal_span(): - data = pg.load(FIELD) - geometry = gk_ops.resolve_geometry(data.file_name) - grid = gk_ops.resolve_flux_surface_grid(data, geometry, nphi=4, nz_interp=2) - zero_span = replace(grid, phi_2d=np.zeros_like(grid.phi_2d)) - with pytest.raises(ValueError, match="zero or non-finite"): - gk_ops.extract_flux_surface(data, zero_span) - - -def test_flux_surface_grid_collection_caches_by_geometry_prefix(monkeypatch): - fluxsurf = import_module("postgkyl.operations.gyrokinetics.fluxsurf") - first = SimpleNamespace(file_name="block-one", ctx={"block": 1}) - repeated = SimpleNamespace(file_name="block-one", ctx={"block": 1}) - second = SimpleNamespace(file_name="block-two", ctx={"block": 2}) - calls = [] - monkeypatch.setattr(fluxsurf, "geometry_prefix", lambda path: path) - monkeypatch.setattr( - fluxsurf, "resolve_geometry", lambda path, **kwargs: calls.append( - (path, kwargs)) or path) - monkeypatch.setattr(fluxsurf, "resolve_flux_surface_grid", - lambda data, geo, **_kwargs: f"grid:{geo}") - - grids = fluxsurf.flux_surface_grids([first, repeated, second], - mapc2p="map-*.gkyl", - nodes_file="nodes-*.gkyl") - assert grids == { - "block-one": "grid:block-one", - "block-two": "grid:block-two", - } - assert calls == [ - ("block-one", { - "mapc2p": "map-1.gkyl", - "nodes_file": "nodes-1.gkyl" - }), - ("block-two", { - "mapc2p": "map-2.gkyl", - "nodes_file": "nodes-2.gkyl" - }), - ] - assert fluxsurf.grid_for(grids, first) == "grid:block-one" - - -def test_gk_fluxsurf_composes_geometry_grid_and_extraction(monkeypatch): - fluxsurf = import_module("postgkyl.operations.gyrokinetics.fluxsurf") - data = SimpleNamespace(file_name="field.gkyl") - calls = [] - monkeypatch.setattr( - fluxsurf, "resolve_geometry", lambda path, **kwargs: calls.append( - ("geometry", path, kwargs)) or "geo") - monkeypatch.setattr( - fluxsurf, "resolve_flux_surface_grid", - lambda source, geo, **kwargs: calls.append( - ("grid", source, geo, kwargs)) or "grid") - monkeypatch.setattr( - fluxsurf, "extract_flux_surface", - lambda source, grid, **kwargs: calls.append( - ("extract", source, grid, kwargs)) or "result") - - result = fluxsurf.gk_fluxsurf(data, - mapc2p="map.gkyl", - x_idx=2, - nphi=16, - nz_interp=3, - comp=4, - inplace=True, - tag="surface", - label="flux") - assert result == "result" - assert [call[0] for call in calls] == ["geometry", "grid", "extract"] - assert calls[-1][-1] == { - "comp": 4, - "inplace": True, - "tag": "surface", - "label": "flux" - } diff --git a/tests/test_operations_gk_rz.py b/tests/test_operations_gk_rz.py deleted file mode 100644 index ba50553f..00000000 --- a/tests/test_operations_gk_rz.py +++ /dev/null @@ -1,308 +0,0 @@ -"""Gyrokinetic R-Z operation, public surfaces, and compatibility paths.""" - -from __future__ import annotations - -from dataclasses import replace -from importlib import import_module -import os -from types import SimpleNamespace - -import click -import numpy as np -import pytest -from click.testing import CliRunner - -import postgkyl as pg -from postgkyl import gpython -from postgkyl.cli.app import COMMANDS - -gk_rz_command = next(command for command in COMMANDS if command.name == "gk_rz") -from postgkyl.cli.state import DataSpace -from postgkyl.operations import gyrokinetics as gk_ops - -ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -DATA = os.path.join(ROOT, "tests", "test_data") -F1D = os.path.join( - DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") -F2D = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-elc_M2par_10.gkyl") -F2D_GEO = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-geo_int_mapc2p.gkyl") -F3D = os.path.join(DATA, "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") - -needs_gkeyll = pytest.mark.skipif( - not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") - - -@needs_gkeyll -def test_2d_mapping_reference_grid_and_values(): - mapped = pg.gk_rz(pg.load(F2D), nz_interp=2) - assert [axis.shape for axis in mapped.grid] == [(33, 33), (33, 33)] - assert mapped.values.shape == (32, 32, 1) - np.testing.assert_allclose(mapped.values.flat[:5], [ - 3.12774618e30, 3.15719364e30, 3.23307916e30, 3.18034806e30, 3.15422838e30 - ], - rtol=2e-9) - - -@needs_gkeyll -def test_3d_mapping_reference_and_fft_phase(): - data = pg.load(F3D) - geometry = gk_ops.resolve_geometry(data.file_name) - projection = gk_ops.resolve_rz_projection(data, geometry, nz_interp=2) - at_zero = gk_ops.map_to_rz(data, projection, phi_tor=0.0) - at_quarter = gk_ops.map_to_rz(data, projection, phi_tor=np.pi / 2) - assert [axis.shape for axis in at_zero.grid] == [(97, 65), (97, 65)] - assert at_zero.values.shape == (96, 64, 1) - np.testing.assert_allclose(at_zero.values.flat[:5], [ - 9.97245134e18, 9.71438453e18, 9.57553863e18, 9.50610482e18, 9.81316530e18 - ], - rtol=2e-9) - assert not np.allclose(at_zero.values, at_quarter.values) - - -def test_geometry_prefers_nodes_and_honors_explicit_modal_override( - tmp_path, monkeypatch): - from postgkyl.operations.gyrokinetics import geometry as geometry_module - - source = tmp_path / "sim-field_0.gkyl" - nodes = tmp_path / "sim-geo_int_nodes.gkyl" - modal = tmp_path / "sim-geo_int_mapc2p.gkyl" - nodes.touch() - modal.touch() - coords = [np.array([0.0, 1.0]), np.array([-1.0, 1.0])] - arrays = np.ones((2, 2)) - calls = [] - monkeypatch.setattr( - geometry_module, "_read_nodes_geometry", lambda path: (calls.append( - ("nodes", path)) or (coords, arrays, arrays, None))) - monkeypatch.setattr( - geometry_module, "_read_mapc2p_geometry", lambda path: (calls.append( - ("mapc2p", path)) or (coords, arrays, arrays, None))) - - gk_ops.resolve_geometry(str(source)) - assert calls[-1] == ("nodes", str(nodes)) - gk_ops.resolve_geometry(str(source), mapc2p="") - assert calls[-1] == ("mapc2p", str(modal)) - - -@needs_gkeyll -def test_geometry_overrides_and_validation_errors(tmp_path): - data = pg.load(F2D) - with pytest.raises(ValueError, match="either mapc2p=.*nodes_file"): - pg.gk_rz(data, mapc2p=F2D_GEO, nodes_file=F2D_GEO) - explicit = pg.gk_rz(data, mapc2p=F2D_GEO, nz_interp=2) - inferred = pg.gk_rz(data, nz_interp=2) - np.testing.assert_allclose(explicit.values, inferred.values) - - missing = data.clone() - missing._file_name = str(tmp_path / "absent-field_0.gkyl") - with pytest.raises(ValueError, match="Could not find a geometry file"): - pg.gk_rz(missing) - with pytest.raises(ValueError, match="positive integer"): - pg.gk_rz(data, nz_interp=0) - with pytest.raises(ValueError, match="out of bounds"): - pg.gk_rz(data, comp=1) - with pytest.raises(ValueError, match="requires 2-D or 3-D"): - pg.gk_rz(pg.load(F1D)) - with pytest.raises(ValueError, match="un-interpolated modal DG"): - pg.gk_rz(data.interpolate()) - - -@needs_gkeyll -def test_missing_toroidal_geometry_and_incompatible_projection_fail_clearly(): - data = pg.load(F3D) - coords = [np.array([0.0, 1.0])] * 3 - values = np.ones((2, 2, 2)) - no_phi = gk_ops.Geometry(coords=coords, - major_r=values, - vert_z=values, - phi=None, - corner=None) - with pytest.raises(ValueError, match="no toroidal-angle component"): - gk_ops.resolve_rz_projection(data, no_phi) - - geometry = gk_ops.resolve_geometry(data.file_name) - projection = gk_ops.resolve_rz_projection(data, geometry, nz_interp=2) - shifted = data.clone() - shifted.grid[0] = shifted.grid[0] + 0.01 - with pytest.raises(ValueError, match="computational grid does not match"): - gk_ops.map_to_rz(shifted, projection) - - -@needs_gkeyll -def test_3d_projection_rejects_thin_and_zero_span_geometry(): - data = pg.load(F3D) - geometry = gk_ops.resolve_geometry(data.file_name) - - thin = data.clone() - thin.ctx["poly_order"] = 0 - thin._grid[1] = np.array([0.0, 1.0]) - with pytest.raises(ValueError, match="at least two interpolated y and z"): - gk_ops.resolve_rz_projection(thin, geometry) - - zero_span = replace(geometry, phi=np.zeros_like(geometry.phi)) - with pytest.raises(ValueError, match="zero or non-finite"): - gk_ops.resolve_rz_projection(data, zero_span) - - -@needs_gkeyll -def test_3d_projection_uses_corner_geometry_when_available(): - data = pg.load(F3D) - geometry = gk_ops.resolve_geometry(data.file_name) - corner_coords = [np.array(axis, copy=True) for axis in geometry.coords] - dz = geometry.coords[2][-1] - geometry.coords[2][0] - corner_coords[2] = np.array( - [geometry.coords[2][0] - dz, geometry.coords[2][-1] + dz]) - corner_r = np.stack([geometry.major_r[..., 0], geometry.major_r[..., -1]], - axis=-1) - corner_z = np.stack([geometry.vert_z[..., 0], geometry.vert_z[..., -1]], - axis=-1) - with_corner = replace(geometry, corner=(corner_coords, corner_r, corner_z)) - projection = gk_ops.resolve_rz_projection(data, with_corner, nz_interp=2) - assert projection.r.shape == projection.z.shape == (97, 65) - - -@needs_gkeyll -def test_reusable_rz_projection_validates_every_shape_contract(): - data_2d = pg.load(F2D) - geometry_2d = gk_ops.resolve_geometry(data_2d.file_name) - projection_2d = gk_ops.resolve_rz_projection(data_2d, geometry_2d) - - invalid_2d = [ - (replace(projection_2d, num_dims=4), "invalid dimensionality"), - (replace(projection_2d, num_dims=3), "dimensionality does not match"), - (replace(projection_2d, - z=projection_2d.z[:, :-1]), "matching 2-D arrays"), - (replace(projection_2d, r=projection_2d.r[:-1], - z=projection_2d.z[:-1]), "expected grid shape"), - ] - for projection, message in invalid_2d: - with pytest.raises(ValueError, match=message): - gk_ops.map_to_rz(data_2d, projection) - - data_3d = pg.load(F3D) - geometry_3d = gk_ops.resolve_geometry(data_3d.file_name) - projection_3d = gk_ops.resolve_rz_projection(data_3d, - geometry_3d, - nz_interp=2) - invalid_3d = [ - (replace(projection_3d, zc=None), "metadata is incomplete"), - (replace(projection_3d, box=0.0), "span must be finite and nonzero"), - (replace(projection_3d, wind=projection_3d.wind[:-1]), - "projection and data grid shapes differ"), - ] - for projection, message in invalid_3d: - with pytest.raises(ValueError, match=message): - gk_ops.map_to_rz(data_3d, projection) - - -def test_rz_projection_collection_caches_by_geometry_prefix(monkeypatch): - rz = import_module("postgkyl.operations.gyrokinetics.rz") - first = SimpleNamespace(file_name="block-one", ctx={"block": 1}) - repeated = SimpleNamespace(file_name="block-one", ctx={"block": 1}) - second = SimpleNamespace(file_name="block-two", ctx={"block": 2}) - calls = [] - monkeypatch.setattr(rz, "geometry_prefix", lambda path: path) - monkeypatch.setattr( - rz, "resolve_geometry", lambda path, **kwargs: calls.append( - (path, kwargs)) or path) - monkeypatch.setattr(rz, "resolve_rz_projection", - lambda data, geo, **_kwargs: f"projection:{geo}") - - projections = rz.rz_projections([first, repeated, second], - mapc2p="map-*.gkyl", - nodes_file="nodes-*.gkyl") - assert projections == { - "block-one": "projection:block-one", - "block-two": "projection:block-two", - } - assert calls == [ - ("block-one", { - "mapc2p": "map-1.gkyl", - "nodes_file": "nodes-1.gkyl" - }), - ("block-two", { - "mapc2p": "map-2.gkyl", - "nodes_file": "nodes-2.gkyl" - }), - ] - assert rz.projection_for(projections, first) == "projection:block-one" - - -@needs_gkeyll -def test_state_propagation_projection_reuse_and_public_surfaces(): - - class DerivedData(pg.GData): - pass - - source = DerivedData(F2D, tag="source", label="original") - original = source.values.copy() - geometry = gk_ops.resolve_geometry(source.file_name) - projection = gk_ops.resolve_rz_projection(source, geometry, nz_interp=2) - first = gk_ops.map_to_rz(source, projection, tag="rz", label="mapped") - second = gk_ops.map_to_rz(source.clone(), projection) - assert isinstance(first, DerivedData) - assert first is not source - assert first.file_name == source.file_name - assert (first.tag, first.label, first.ctx["interpolated"]) == ("rz", "mapped", - True) - np.testing.assert_array_equal(source.values, original) - np.testing.assert_allclose(first.values, second.values) - - fluent = source.gk_rz(mapc2p=F2D_GEO, nz_interp=2) - functional = pg.gk_rz(source, mapc2p=F2D_GEO, nz_interp=2) - np.testing.assert_allclose(fluent.values, functional.values) - assert pg.gk_rz is gk_ops.gk_rz - - inplace = source.clone() - result = pg.gk_rz(inplace, mapc2p=F2D_GEO, nz_interp=2, inplace=True) - assert result is inplace and result.ctx["interpolated"] is True - - -@needs_gkeyll -def test_comp_selects_an_explicit_physical_field(): - source = pg.load(F2D) - multi = pg.GData(ctx={ - key: value - for key, value in source.ctx.items() if key != "num_comps" - }) - multi.push([axis.copy() for axis in source.grid], - np.concatenate([source.values, 2.0 * source.values], axis=-1)) - multi._file_name = source.file_name - first = pg.gk_rz(multi, comp=0, nz_interp=2) - second = pg.gk_rz(multi, comp=1, nz_interp=2) - np.testing.assert_allclose(second.values, 2.0 * first.values) - - -@needs_gkeyll -def test_group_compatibility_cli_and_help_section(): - from postgkyl.cli.app import cli - from postgkyl.diagnostics.gk import fluxsurf as old_fluxsurf - from postgkyl.diagnostics.gk import rz as old_rz - - group = pg.GDataGroup([pg.load(F2D), pg.load(F2D)]) - mapped = group.gk_rz(mapc2p=F2D_GEO, nz_interp=2) - assert isinstance(mapped, pg.GDataGroup) and len(mapped) == 2 - - assert old_rz.gk_rz is gk_ops.gk_rz - assert old_rz.RzProjection is gk_ops.RzProjection - assert old_fluxsurf.FluxSurfaceGrid is gk_ops.FluxSurfaceGrid - assert old_fluxsurf.extract_flux_surface is gk_ops.extract_flux_surface - - space = DataSpace(datasets=[pg.load(F2D)]) - with click.Context(gk_rz_command, obj=space) as ctx: - ctx.invoke(gk_rz_command, - mapc2p=F2D_GEO, - nodes_file=None, - z_axis=0.0, - phi_tor=0.0, - nz_interp=2, - use=None, - tag="rz", - label=None) - expected = pg.gk_rz(pg.load(F2D), mapc2p=F2D_GEO, nz_interp=2, tag="rz") - np.testing.assert_allclose(space.datasets[0].values, expected.values) - - help_text = CliRunner().invoke(cli, ["--help"]).output - verbs = help_text.split("Diagnostics:", 1)[0] - diagnostics = help_text.split("Diagnostics:", 1)[1].split("Render:", 1)[0] - assert "gk_rz" in verbs and "gk_rz" not in diagnostics diff --git a/tests/test_operations_map_fluxsurf.py b/tests/test_operations_map_fluxsurf.py new file mode 100644 index 00000000..d24acb4a --- /dev/null +++ b/tests/test_operations_map_fluxsurf.py @@ -0,0 +1,105 @@ +"""Explicit flux-surface mapping and reusable grid contracts.""" + +from __future__ import annotations + +from dataclasses import replace +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl import operations as mapping +from postgkyl.diagnostics.gk.geometry import resolve_geometry + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +FIELD = os.path.join(ROOT, "tests", "test_data", + "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +@needs_gkeyll +def test_flux_surface_move_preserves_output_and_projection_reuse(): + data = pg.load(FIELD) + geometry = resolve_geometry(data.file_name) + grid = mapping.resolve_flux_surface_grid(data, + geometry, + x_idx=0, + nphi=4, + nz_interp=2) + first = mapping.extract_flux_surface(data, fs_grid=grid) + second = mapping.extract_flux_surface(data.clone(), fs_grid=grid) + assert first.values.shape == (4, 64, 1) + assert first.ctx["interpolated"] is True + np.testing.assert_allclose(first.values, second.values) + + +@needs_gkeyll +@pytest.mark.parametrize(("kwargs", "message"), [ + ({ + "nphi": 0 + }, "nphi must be a positive integer"), + ({ + "nz_interp": 0 + }, "nz_interp must be a positive integer"), + ({ + "x_idx": -1 + }, "out of bounds"), +]) +def test_flux_surface_public_validation(kwargs, message): + data = pg.load(FIELD) + geometry = resolve_geometry(data.file_name) + with pytest.raises(ValueError, match=message): + mapping.resolve_flux_surface_grid(data, geometry, **kwargs) + + +@needs_gkeyll +def test_flux_surface_grid_requires_toroidal_geometry_and_integer_index(): + data = pg.load(FIELD) + geometry = resolve_geometry(data.file_name) + with pytest.raises(ValueError, match="no toroidal-angle component"): + mapping.resolve_flux_surface_grid(data, replace(geometry, phi=None)) + with pytest.raises(ValueError, match="x_idx must be an integer"): + mapping.resolve_flux_surface_grid(data, geometry, x_idx=True) + + +@needs_gkeyll +def test_flux_surface_grid_requires_two_binormal_and_parallel_points(): + data = pg.load(FIELD).clone() + data.ctx["poly_order"] = 0 + data._grid[1] = np.array([0.0, 1.0]) + geometry = resolve_geometry(data.file_name) + with pytest.raises(ValueError, match="at least two interpolated y and z"): + mapping.resolve_flux_surface_grid(data, geometry) + + +@needs_gkeyll +def test_extract_flux_surface_validates_reusable_grid_metadata(): + data = pg.load(FIELD) + geometry = resolve_geometry(data.file_name) + grid = mapping.resolve_flux_surface_grid(data, geometry, nphi=4, nz_interp=2) + + shifted = data.clone() + shifted.grid[0] = shifted.grid[0] + 0.1 + with pytest.raises(ValueError, match="computational grid does not match"): + mapping.extract_flux_surface(shifted, fs_grid=grid) + + with pytest.raises(ValueError, match="out of bounds"): + mapping.extract_flux_surface(data, fs_grid=replace(grid, x_idx=10_000)) + + with pytest.raises(ValueError, match="projection and data grid shapes"): + mapping.extract_flux_surface(data, + fs_grid=replace(grid, phi_2d=np.ones((1, 1)))) + + +@needs_gkeyll +def test_extract_flux_surface_rejects_zero_toroidal_span(): + data = pg.load(FIELD) + geometry = resolve_geometry(data.file_name) + grid = mapping.resolve_flux_surface_grid(data, geometry, nphi=4, nz_interp=2) + zero_span = replace(grid, phi_2d=np.zeros_like(grid.phi_2d)) + with pytest.raises(ValueError, match="zero or non-finite"): + mapping.extract_flux_surface(data, fs_grid=zero_span) diff --git a/tests/test_operations_map_rz.py b/tests/test_operations_map_rz.py new file mode 100644 index 00000000..b80520c6 --- /dev/null +++ b/tests/test_operations_map_rz.py @@ -0,0 +1,223 @@ +"""Explicit R-Z mapping and Gkeyll geometry compositions.""" + +from __future__ import annotations + +from dataclasses import replace +import os + +import numpy as np +import pytest + +import postgkyl as pg +from postgkyl import gpython +from postgkyl import operations as mapping +from postgkyl.diagnostics.gk.geometry import resolve_geometry + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +DATA = os.path.join(ROOT, "tests", "test_data") +F1D = os.path.join( + DATA, "rt_gk_tcv_iwl_adapt_source_1x2v_p1-ion_HamiltonianMoments_250.gkyl") +F2D = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-elc_M2par_10.gkyl") +F2D_GEO = os.path.join(DATA, "gk_ltx_iwl_2x2v_p1-geo_int_mapc2p.gkyl") +F3D = os.path.join(DATA, "rt_gk_tcv_nt_iwl_3x2v_p1-elc_M0_5.gkyl") + +needs_gkeyll = pytest.mark.skipif( + not gpython.available(), reason="no compiled Gkeyll (libg0core.so) found") + + +@needs_gkeyll +def test_2d_mapping_reference_grid_and_values(): + mapped = pg.gk.rz(pg.load(F2D), nz_interp=2) + assert [axis.shape for axis in mapped.grid] == [(33, 33), (33, 33)] + assert mapped.values.shape == (32, 32, 1) + np.testing.assert_allclose(mapped.values.flat[:5], [ + 3.12774618e30, 3.15719364e30, 3.23307916e30, 3.18034806e30, 3.15422838e30 + ], + rtol=2e-9) + + +@needs_gkeyll +def test_3d_mapping_reference_and_fft_phase(): + data = pg.load(F3D) + geometry = resolve_geometry(data.file_name) + projection = mapping.resolve_rz_projection(data, geometry, nz_interp=2) + at_zero = data.map_to_rz(projection=projection, phi_tor=0.0) + at_quarter = mapping.map_to_rz(data, projection=projection, phi_tor=np.pi / 2) + assert [axis.shape for axis in at_zero.grid] == [(97, 65), (97, 65)] + assert at_zero.values.shape == (96, 64, 1) + np.testing.assert_allclose(at_zero.values.flat[:5], [ + 9.97245134e18, 9.71438453e18, 9.57553863e18, 9.50610482e18, 9.81316530e18 + ], + rtol=2e-9) + assert not np.allclose(at_zero.values, at_quarter.values) + + +@needs_gkeyll +def test_missing_toroidal_geometry_and_incompatible_projection_fail_clearly(): + data = pg.load(F3D) + coords = [np.array([0.0, 1.0])] * 3 + values = np.ones((2, 2, 2)) + no_phi = mapping.Geometry(coords=coords, + major_r=values, + vert_z=values, + phi=None, + corner=None) + with pytest.raises(ValueError, match="no toroidal-angle component"): + mapping.resolve_rz_projection(data, no_phi) + + geometry = resolve_geometry(data.file_name) + projection = mapping.resolve_rz_projection(data, geometry, nz_interp=2) + shifted = data.clone() + shifted.grid[0] = shifted.grid[0] + 0.01 + with pytest.raises(ValueError, match="computational grid does not match"): + mapping.map_to_rz(shifted, projection=projection) + + +@needs_gkeyll +def test_3d_projection_rejects_thin_and_zero_span_geometry(): + data = pg.load(F3D) + geometry = resolve_geometry(data.file_name) + + thin = data.clone() + thin.ctx["poly_order"] = 0 + thin._grid[1] = np.array([0.0, 1.0]) + with pytest.raises(ValueError, match="at least two interpolated y and z"): + mapping.resolve_rz_projection(thin, geometry) + + zero_span = replace(geometry, phi=np.zeros_like(geometry.phi)) + with pytest.raises(ValueError, match="zero or non-finite"): + mapping.resolve_rz_projection(data, zero_span) + + +@needs_gkeyll +def test_3d_projection_uses_corner_geometry_when_available(): + data = pg.load(F3D) + geometry = resolve_geometry(data.file_name) + corner_coords = [np.array(axis, copy=True) for axis in geometry.coords] + dz = geometry.coords[2][-1] - geometry.coords[2][0] + corner_coords[2] = np.array( + [geometry.coords[2][0] - dz, geometry.coords[2][-1] + dz]) + corner_r = np.stack([geometry.major_r[..., 0], geometry.major_r[..., -1]], + axis=-1) + corner_z = np.stack([geometry.vert_z[..., 0], geometry.vert_z[..., -1]], + axis=-1) + with_corner = replace(geometry, corner=(corner_coords, corner_r, corner_z)) + projection = mapping.resolve_rz_projection(data, with_corner, nz_interp=2) + assert projection.r.shape == projection.z.shape == (97, 65) + + +@needs_gkeyll +def test_reusable_rz_projection_validates_every_shape_contract(): + data_2d = pg.load(F2D) + geometry_2d = resolve_geometry(data_2d.file_name) + projection_2d = mapping.resolve_rz_projection(data_2d, geometry_2d) + + invalid_2d = [ + (replace(projection_2d, num_dims=4), "invalid dimensionality"), + (replace(projection_2d, num_dims=3), "dimensionality does not match"), + (replace(projection_2d, + z=projection_2d.z[:, :-1]), "matching 2-D arrays"), + (replace(projection_2d, r=projection_2d.r[:-1], + z=projection_2d.z[:-1]), "expected grid shape"), + ] + for projection, message in invalid_2d: + with pytest.raises(ValueError, match=message): + mapping.map_to_rz(data_2d, projection=projection) + + data_3d = pg.load(F3D) + geometry_3d = resolve_geometry(data_3d.file_name) + projection_3d = mapping.resolve_rz_projection(data_3d, + geometry_3d, + nz_interp=2) + invalid_3d = [ + (replace(projection_3d, zc=None), "metadata is incomplete"), + (replace(projection_3d, box=0.0), "span must be finite and nonzero"), + (replace(projection_3d, wind=projection_3d.wind[:-1]), + "projection and data grid shapes differ"), + ] + for projection, message in invalid_3d: + with pytest.raises(ValueError, match=message): + mapping.map_to_rz(data_3d, projection=projection) + + +@needs_gkeyll +def test_state_propagation_projection_reuse_and_public_surfaces(): + + class DerivedData(pg.GData): + pass + + source = DerivedData(F2D, tag="source", label="original") + original = source.values.copy() + geometry = resolve_geometry(source.file_name) + projection = mapping.resolve_rz_projection(source, geometry, nz_interp=2) + first = mapping.map_to_rz(source, + projection=projection, + tag="rz", + label="mapped") + second = mapping.map_to_rz(source.clone(), projection=projection) + assert isinstance(first, DerivedData) + assert first is not source + assert first.file_name == source.file_name + assert (first.tag, first.label, first.ctx["interpolated"]) == ("rz", "mapped", + True) + np.testing.assert_array_equal(source.values, original) + np.testing.assert_allclose(first.values, second.values) + + fluent = source.map_to_rz(projection=projection) + functional = pg.gk.rz(source, mapc2p=F2D_GEO, nz_interp=2) + np.testing.assert_allclose(fluent.values, functional.values) + assert pg.map_to_rz is mapping.map_to_rz is pg.GData.map_to_rz + assert not hasattr(pg, "gk_rz") + assert not hasattr(pg.GData, "gk_rz") + + inplace = source.clone() + result = pg.gk.rz(inplace, mapc2p=F2D_GEO, nz_interp=2, inplace=True) + assert result is inplace and result.ctx["interpolated"] is True + + +@needs_gkeyll +def test_comp_selects_an_explicit_physical_field(): + source = pg.load(F2D) + multi = pg.GData(ctx={ + key: value + for key, value in source.ctx.items() if key != "num_comps" + }) + multi.push([axis.copy() for axis in source.grid], + np.concatenate([source.values, 2.0 * source.values], axis=-1)) + multi._file_name = source.file_name + first = pg.gk.rz(multi, comp=0, nz_interp=2) + second = pg.gk.rz(multi, comp=1, nz_interp=2) + np.testing.assert_allclose(second.values, 2.0 * first.values) + + +@needs_gkeyll +@pytest.mark.parametrize("num_dims", [2, 3]) +def test_mapping_accepts_filename_free_fields_and_explicit_geometry(num_dims): + axes = [np.linspace(0.0, 1.0, 5) for _ in range(num_dims)] + data = pg.GData(ctx={ + "basis_type": "serendipity", + "poly_order": 0, + "value_form": "modal" + }) + # The normalized p0 basis is 2**(-ndim/2); these coefficients represent 7. + data.push(axes, np.full((4, ) * num_dims + (1, ), 7.0 * 2**(num_dims / 2))) + mesh = np.meshgrid(*axes, indexing="ij") + geometry = pg.Geometry(coords=axes, + major_r=2.0 + mesh[0], + vert_z=mesh[-1], + phi=2 * np.pi * mesh[1] if num_dims == 3 else None, + corner=None) + projection = pg.resolve_rz_projection(data, geometry, nz_interp=2) + mapped = data.map_to_rz(projection=projection, phi_tor=0.37) + assert not data.file_name and not mapped.file_name + assert mapped.backend == "numpy" and mapped.is_interpolated + assert mapped.values.shape == (4, 8 if num_dims == 3 else 4, 1) + np.testing.assert_allclose(mapped.values, 7.0) + np.testing.assert_allclose(mapped.grid[0][:, 0], 2.0 + axes[0]) + np.testing.assert_allclose(mapped.grid[1][0], + np.linspace(0, 1, mapped.values.shape[1] + 1)) + if num_dims == 3: + surface = pg.resolve_flux_surface_grid(data, geometry, nphi=5, nz_interp=2) + result = data.extract_flux_surface(fs_grid=surface) + assert result.values.shape == (5, 8, 1) + np.testing.assert_allclose(result.values, 7.0) diff --git a/tests/test_postgkyl.py b/tests/test_postgkyl.py index 03b62702..46567238 100644 --- a/tests/test_postgkyl.py +++ b/tests/test_postgkyl.py @@ -449,11 +449,8 @@ def test_cli_abbreviation_and_info(): # transform_frame/laguerre) moved up # into diagnostics, folded with the # models/ array math they delegated to; - # flat modules are equation-blind core - # verbs; domain subpackages (currently - # gyrokinetics) own transformations that - # need domain geometry without deriving - # a physical conclusion + # flat modules are equation-blind core verbs; model-specific auxiliary + # discovery and interpretation belong to diagnostics. "diagnostics": { "gdatastate", "operations", "numerics", "gdata", "render", "io", "cli_spec" @@ -685,3 +682,9 @@ def test_foreign_floor_offenders_flags_ctypes_and_gpython_outside_gpython( offenders = _foreign_floor_offenders(pkg_root) assert any(o.endswith(": ctypes") for o in offenders) assert any(o.endswith(": _gpython") for o in offenders) + + +def test_operations_has_no_domain_subpackages(): + """Core transformations stay flat; model-family compositions live above.""" + operations = Path(SRC) / "postgkyl" / "operations" + assert list(operations.glob("*/__init__.py")) == []