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")) == []