From 69fbb2b41b96318f12a9d0d180d66900ee4f49fa Mon Sep 17 00:00:00 2001 From: mrquell Date: Fri, 13 Mar 2026 17:46:41 -0400 Subject: [PATCH 01/31] Add gk_distf command for building real distribution function from Jf data. This gets complex in gyrokinetics because of the velocity jacobian. Also, we are currently not capable of using a mapc2p or non-uniform mesh map for plotting distribution function. I will add non-uniform mesh grids in a future commit --- src/postgkyl/commands/__init__.py | 1 + src/postgkyl/commands/gk_distf.py | 137 ++++++++++++++++++++++++++++++ src/postgkyl/pgkyl.py | 1 + 3 files changed, 139 insertions(+) create mode 100644 src/postgkyl/commands/gk_distf.py diff --git a/src/postgkyl/commands/__init__.py b/src/postgkyl/commands/__init__.py index 5793303..ea279ce 100644 --- a/src/postgkyl/commands/__init__.py +++ b/src/postgkyl/commands/__init__.py @@ -30,6 +30,7 @@ from postgkyl.commands.mhd import mhd from postgkyl.commands.parrotate import parrotate from postgkyl.commands.gk_energy_balance import gk_energy_balance +from postgkyl.commands.gk_distf import gk_distf from postgkyl.commands.gk_particle_balance import gk_particle_balance from postgkyl.commands.perprotate import perprotate from postgkyl.commands.plot import plot diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py new file mode 100644 index 0000000..f7b126f --- /dev/null +++ b/src/postgkyl/commands/gk_distf.py @@ -0,0 +1,137 @@ +import os + +import click +import numpy as np + +from postgkyl.data import GData +from postgkyl.data import GInterpModal +from postgkyl.utils import verb_print + + +def _resolve_path(path: str, file_name: str) -> str: + if os.path.isabs(file_name): + return file_name + return os.path.join(path, file_name) + + +def _resolve_files(name: str, species: str, frame: int, path: str): + frame_file = _resolve_path(path, f"{name}-{species}_{frame}.gkyl") + mapc2p_vel_file = _resolve_path(path, f"{name}-{species}_mapc2p_vel.gkyl") + jacobvel_file = _resolve_path(path, f"{name}-{species}_jacobvel.gkyl") + jacobtot_inv_file = _resolve_path(path, f"{name}-jacobtot_inv.gkyl") + return frame_file, mapc2p_vel_file, jacobvel_file, jacobtot_inv_file + + +def _assert_files_exist(files: dict): + for kind, file_name in files.items(): + if not os.path.exists(file_name): + raise FileNotFoundError(f"Missing {kind} file: {file_name}") + + +def _compute_fjx(jf_values: np.ndarray, jac_values: np.ndarray) -> np.ndarray: + if jf_values.shape == jac_values.shape: + return jf_values / jac_values + + if jf_values.shape[:-1] != jac_values.shape[:-1]: + raise ValueError( + "Jf and jacobvel spatial shapes do not match: " + f"{jf_values.shape} vs {jac_values.shape}" + ) + + if jac_values.shape[-1] != 1: + raise ValueError( + "jacobvel component axis must be 1 or match Jf: " + f"{jf_values.shape} vs {jac_values.shape}" + ) + + return jf_values / jac_values + + +def _broadcast_multiply(distf: np.ndarray, jacobtot_inv: np.ndarray) -> np.ndarray: + if jacobtot_inv.ndim > distf.ndim: + raise ValueError( + "jacobtot_inv has more dimensions than distribution function: " + f"{jacobtot_inv.ndim} > {distf.ndim}" + ) + + expected_prefix = distf.shape[:jacobtot_inv.ndim] + if expected_prefix != jacobtot_inv.shape: + raise ValueError( + "Cannot broadcast jacobtot_inv over distribution function. " + f"Expected prefix {expected_prefix}, got {jacobtot_inv.shape}" + ) + + jacob_shape = jacobtot_inv.shape + (1,) * (distf.ndim - jacobtot_inv.ndim) + return distf * jacobtot_inv.reshape(jacob_shape) + + +def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): + fjx_interp = GInterpModal(fjx_data, 1, "gkhyb") + jacob_interp = GInterpModal(jacobtot_inv_data, 1, "ms") + + grid, fjx_values = fjx_interp.interpolate() + _, jacob_values = jacob_interp.interpolate() + + fjx_values = np.squeeze(fjx_values) + jacob_values = np.squeeze(jacob_values) + return grid, fjx_values, jacob_values + + +@click.command() +@click.option("--name", "-n", required=True, type=click.STRING, + help="Simulation name prefix (e.g. gk_lorentzian_mirror).") +@click.option("--species", "-s", required=True, type=click.STRING, + help="Species name (e.g. ion or elc).") +@click.option("--frame", "-f", required=True, type=click.INT, + help="Frame number.") +@click.option("--path", "-p", default="./", type=click.STRING, + help="Path to simulation data.") +@click.option("--tag", "-t", default="df", type=click.STRING, + help="Tag for output dataset.") +@click.option("--c2p-vel", default=False, + help="Use -_mapc2p_vel.gkyl when loading Jf.") +@click.option("--debug", is_flag=True, + help="Print resolved file names and shape diagnostics.") +@click.pass_context +def gk_distf(ctx, **kwargs): + """Gyrokinetics: build real distribution function from saved Jf data.""" + verb_print(ctx, "Starting gk_distf") + data = ctx.obj["data"] + + files = {} + jf_file, mapc2p_vel_file, jacobvel_file, jacobtot_inv_file = _resolve_files( + kwargs["name"], kwargs["species"], kwargs["frame"], kwargs["path"]) + files["Jf"] = jf_file + files["jacobvel"] = jacobvel_file + files["jacobtot_inv"] = jacobtot_inv_file + if kwargs["c2p_vel"]: + files["mapc2p_vel"] = mapc2p_vel_file + + _assert_files_exist(files) + + if kwargs["debug"]: + click.echo(f"gk_distf: Jf={jf_file}") + click.echo(f"gk_distf: jacobvel={jacobvel_file}") + click.echo(f"gk_distf: jacobtot_inv={jacobtot_inv_file}") + if kwargs["c2p_vel"]: + click.echo(f"gk_distf: mapc2p_vel={mapc2p_vel_file}") + + jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file if kwargs["c2p_vel"] else "") + jacobvel_data = GData(jacobvel_file) + jacobtot_inv_data = GData(jacobtot_inv_file) + + fjx_data = GData(tag=kwargs["tag"], ctx=jf_data.ctx) + fjx_values = _compute_fjx(jf_data.get_values(), jacobvel_data.get_values()) + fjx_data.push(jf_data.get_grid(), fjx_values) + + out_grid, fjx_interp, jacob_interp = _interpolate_fjx_and_jacob(fjx_data, jacobtot_inv_data) + distf_values = _broadcast_multiply(fjx_interp, jacob_interp) + + if kwargs["debug"]: + click.echo(f"gk_distf: output shape={distf_values.shape}") + + out = GData(tag=kwargs["tag"], ctx=jf_data.ctx) + out.push(out_grid, np.asarray(distf_values)[..., np.newaxis]) + data.add(out) + + verb_print(ctx, "Finishing gk_distf") \ No newline at end of file diff --git a/src/postgkyl/pgkyl.py b/src/postgkyl/pgkyl.py index 29533c3..3f32d12 100755 --- a/src/postgkyl/pgkyl.py +++ b/src/postgkyl/pgkyl.py @@ -149,6 +149,7 @@ def cli(ctx, **kwargs): cli.add_command(cmd.extractinput) cli.add_command(cmd.fft) cli.add_command(cmd.gk_nodes) +cli.add_command(cmd.gk_distf) cli.add_command(cmd.grid) cli.add_command(cmd.growth) cli.add_command(cmd.info) From 11acdc3d86daf9402779fee908ee0da15f1a2fc6 Mon Sep 17 00:00:00 2001 From: mrquell Date: Fri, 13 Mar 2026 18:05:18 -0400 Subject: [PATCH 02/31] Add support for mc2nu. This has been something I've wanted for a long time. Finally, we can plot both non-uniform velocity and position space grids in command line pgkyl --- src/postgkyl/commands/gk_distf.py | 89 +++++++++++++++++++++++++++++-- 1 file changed, 85 insertions(+), 4 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index f7b126f..67e0d80 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -17,9 +17,77 @@ def _resolve_path(path: str, file_name: str) -> str: def _resolve_files(name: str, species: str, frame: int, path: str): frame_file = _resolve_path(path, f"{name}-{species}_{frame}.gkyl") mapc2p_vel_file = _resolve_path(path, f"{name}-{species}_mapc2p_vel.gkyl") + mc2nu_file = _resolve_path(path, f"{name}-mc2nu_pos_deflated.gkyl") jacobvel_file = _resolve_path(path, f"{name}-{species}_jacobvel.gkyl") jacobtot_inv_file = _resolve_path(path, f"{name}-jacobtot_inv.gkyl") - return frame_file, mapc2p_vel_file, jacobvel_file, jacobtot_inv_file + return frame_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file + + +def _resolve_mapping_kwargs(use_c2p_vel: bool, mapc2p_vel_file: str) -> dict: + if use_c2p_vel: + return {"mapc2p_vel_name": mapc2p_vel_file} + return {"mapc2p_vel_name": ""} + + +def _cell_centers_to_nodes(cell_centers: np.ndarray) -> np.ndarray: + nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) + nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) + nodes[0] = cell_centers[0] - (nodes[1] - cell_centers[0]) + nodes[-1] = cell_centers[-1] + (cell_centers[-1] - nodes[-2]) + return nodes + + +def _get_dim_split(out_grid: list, mc2nu_data: GData) -> tuple[int, int]: + total_dims = len(out_grid) + cdim = mc2nu_data.get_num_dims() + if cdim < 1 or cdim > total_dims: + raise ValueError(f"Invalid cdim={cdim} for total_dims={total_dims}") + vdim = total_dims - cdim + return cdim, vdim + + +def _extract_mapped_axis(mapped_values: np.ndarray, axis: int, cdim: int) -> np.ndarray: + if cdim == 1: + return np.asarray(mapped_values[..., axis]).reshape(-1) + + # Extract one mapped coordinate direction at a reference point in the + # remaining configuration directions to construct a 1D axis grid. + idx = [0] * (cdim + 1) + idx[axis] = slice(None) + idx[-1] = axis + return np.asarray(mapped_values[tuple(idx)]).reshape(-1) + + +def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, debug: bool) -> tuple[list, int, int]: + mc2nu_data = GData(mc2nu_file) + cdim, vdim = _get_dim_split(out_grid, mc2nu_data) + + mc2nu_interp = GInterpModal(mc2nu_data, 1, "ms") + _, mc2nu_values = mc2nu_interp.interpolate(tuple(range(cdim))) + mapped_values = np.asarray(mc2nu_values) + + deformed_grid = list(out_grid) + for d in range(cdim): + mapped_cfg = _extract_mapped_axis(mapped_values, d, cdim) + + old_cfg = np.asarray(out_grid[d]) + if mapped_cfg.size == old_cfg.size: + new_cfg = mapped_cfg + elif mapped_cfg.size + 1 == old_cfg.size: + new_cfg = _cell_centers_to_nodes(mapped_cfg) + else: + raise ValueError( + "mc2nu mapping size is incompatible with configuration grid on axis " + f"{d}: {mapped_cfg.size} vs {old_cfg.size}" + ) + + deformed_grid[d] = new_cfg + + if debug: + click.echo(f"gk_distf: cdim={cdim}, vdim={vdim}") + click.echo(f"gk_distf: mc2nu mapped {cdim} configuration axis/axes") + + return deformed_grid, cdim, vdim def _assert_files_exist(files: dict): @@ -88,8 +156,10 @@ def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): help="Path to simulation data.") @click.option("--tag", "-t", default="df", type=click.STRING, help="Tag for output dataset.") -@click.option("--c2p-vel", default=False, +@click.option("--c2p-vel/--no-c2p-vel", default=False, help="Use -_mapc2p_vel.gkyl when loading Jf.") +@click.option("--mc2nu", is_flag=True, + help="Use -mc2nu_pos_deflated.gkyl to deform configuration-space grid.") @click.option("--debug", is_flag=True, help="Print resolved file names and shape diagnostics.") @click.pass_context @@ -99,24 +169,31 @@ def gk_distf(ctx, **kwargs): data = ctx.obj["data"] files = {} - jf_file, mapc2p_vel_file, jacobvel_file, jacobtot_inv_file = _resolve_files( + jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _resolve_files( kwargs["name"], kwargs["species"], kwargs["frame"], kwargs["path"]) files["Jf"] = jf_file files["jacobvel"] = jacobvel_file files["jacobtot_inv"] = jacobtot_inv_file + if kwargs["c2p_vel"]: files["mapc2p_vel"] = mapc2p_vel_file + if kwargs["mc2nu"]: + files["mc2nu"] = mc2nu_file _assert_files_exist(files) + map_kwargs = _resolve_mapping_kwargs(kwargs["c2p_vel"], mapc2p_vel_file) + if kwargs["debug"]: click.echo(f"gk_distf: Jf={jf_file}") click.echo(f"gk_distf: jacobvel={jacobvel_file}") click.echo(f"gk_distf: jacobtot_inv={jacobtot_inv_file}") if kwargs["c2p_vel"]: click.echo(f"gk_distf: mapc2p_vel={mapc2p_vel_file}") + if kwargs["mc2nu"]: + click.echo(f"gk_distf: mc2nu={mc2nu_file}") - jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file if kwargs["c2p_vel"] else "") + jf_data = GData(jf_file, mapc2p_vel_name=map_kwargs["mapc2p_vel_name"]) jacobvel_data = GData(jacobvel_file) jacobtot_inv_data = GData(jacobtot_inv_file) @@ -127,6 +204,10 @@ def gk_distf(ctx, **kwargs): out_grid, fjx_interp, jacob_interp = _interpolate_fjx_and_jacob(fjx_data, jacobtot_inv_data) distf_values = _broadcast_multiply(fjx_interp, jacob_interp) + if kwargs["mc2nu"]: + out_grid, cdim, vdim = _apply_mc2nu_grid(out_grid, mc2nu_file, kwargs["debug"]) + jf_data.ctx["grid_type"] = "mc2nu" + if kwargs["debug"]: click.echo(f"gk_distf: output shape={distf_values.shape}") From d61df541115db8a4026bd0d94a4bb0e0e5323bab Mon Sep 17 00:00:00 2001 From: mrquell Date: Fri, 13 Mar 2026 18:25:39 -0400 Subject: [PATCH 03/31] Add a source boolean to plot sources --- src/postgkyl/commands/gk_distf.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 67e0d80..e6d944d 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -14,8 +14,11 @@ def _resolve_path(path: str, file_name: str) -> str: return os.path.join(path, file_name) -def _resolve_files(name: str, species: str, frame: int, path: str): - frame_file = _resolve_path(path, f"{name}-{species}_{frame}.gkyl") +def _resolve_files(name: str, species: str, frame: int, path: str, source: bool = False): + if source: + frame_file = _resolve_path(path, f"{name}-{species}_source_{frame}.gkyl") + else: + frame_file = _resolve_path(path, f"{name}-{species}_{frame}.gkyl") mapc2p_vel_file = _resolve_path(path, f"{name}-{species}_mapc2p_vel.gkyl") mc2nu_file = _resolve_path(path, f"{name}-mc2nu_pos_deflated.gkyl") jacobvel_file = _resolve_path(path, f"{name}-{species}_jacobvel.gkyl") @@ -152,6 +155,8 @@ def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): help="Species name (e.g. ion or elc).") @click.option("--frame", "-f", required=True, type=click.INT, help="Frame number.") +@click.option("--source", is_flag=True, + help="Use -_source_.gkyl as the input distribution.") @click.option("--path", "-p", default="./", type=click.STRING, help="Path to simulation data.") @click.option("--tag", "-t", default="df", type=click.STRING, @@ -170,7 +175,7 @@ def gk_distf(ctx, **kwargs): files = {} jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _resolve_files( - kwargs["name"], kwargs["species"], kwargs["frame"], kwargs["path"]) + kwargs["name"], kwargs["species"], kwargs["frame"], kwargs["path"], kwargs["source"]) files["Jf"] = jf_file files["jacobvel"] = jacobvel_file files["jacobtot_inv"] = jacobtot_inv_file From 3a2caeb9160b1d02452da84bed53b768028e5a6f Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 11:31:44 -0400 Subject: [PATCH 04/31] Update tool to have a version which is callable from python. Here is an example script import matplotlib.pyplot as plt import postgkyl as pg distf = pg.commands.build_gk_distf( name="gk_mirror_boltz_elc_poa_2x2v_p1", species="ion", frame=0) pg.data.select(distf, z0=0, z1=0.0, overwrite=True) pg.output.plot(distf) plt.show() --- src/postgkyl/__init__.py | 1 + src/postgkyl/commands/__init__.py | 1 + src/postgkyl/commands/gk_distf.py | 92 +++++++++++++++++++------------ 3 files changed, 58 insertions(+), 36 deletions(-) diff --git a/src/postgkyl/__init__.py b/src/postgkyl/__init__.py index 4151d08..17c1475 100644 --- a/src/postgkyl/__init__.py +++ b/src/postgkyl/__init__.py @@ -8,6 +8,7 @@ __version__ = "1.7.5" # import submodules +from postgkyl import commands from postgkyl import data from postgkyl import utils from postgkyl import tools diff --git a/src/postgkyl/commands/__init__.py b/src/postgkyl/commands/__init__.py index ea279ce..16a7641 100644 --- a/src/postgkyl/commands/__init__.py +++ b/src/postgkyl/commands/__init__.py @@ -30,6 +30,7 @@ from postgkyl.commands.mhd import mhd from postgkyl.commands.parrotate import parrotate from postgkyl.commands.gk_energy_balance import gk_energy_balance +from postgkyl.commands.gk_distf import build_gk_distf from postgkyl.commands.gk_distf import gk_distf from postgkyl.commands.gk_particle_balance import gk_particle_balance from postgkyl.commands.perprotate import perprotate diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index e6d944d..9b587d4 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -148,76 +148,96 @@ def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): return grid, fjx_values, jacob_values -@click.command() -@click.option("--name", "-n", required=True, type=click.STRING, - help="Simulation name prefix (e.g. gk_lorentzian_mirror).") -@click.option("--species", "-s", required=True, type=click.STRING, - help="Species name (e.g. ion or elc).") -@click.option("--frame", "-f", required=True, type=click.INT, - help="Frame number.") -@click.option("--source", is_flag=True, - help="Use -_source_.gkyl as the input distribution.") -@click.option("--path", "-p", default="./", type=click.STRING, - help="Path to simulation data.") -@click.option("--tag", "-t", default="df", type=click.STRING, - help="Tag for output dataset.") -@click.option("--c2p-vel/--no-c2p-vel", default=False, - help="Use -_mapc2p_vel.gkyl when loading Jf.") -@click.option("--mc2nu", is_flag=True, - help="Use -mc2nu_pos_deflated.gkyl to deform configuration-space grid.") -@click.option("--debug", is_flag=True, - help="Print resolved file names and shape diagnostics.") -@click.pass_context -def gk_distf(ctx, **kwargs): - """Gyrokinetics: build real distribution function from saved Jf data.""" - verb_print(ctx, "Starting gk_distf") - data = ctx.obj["data"] +def build_gk_distf(name: str, species: str, frame: int, path: str = "./", + tag: str = "df", source: bool = False, use_c2p_vel: bool = False, + use_mc2nu: bool = False, debug: bool = False) -> GData: + """Build a real distribution function from saved Jf data. + This is the Python API used by the ``gk-distf`` CLI command. + """ files = {} jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _resolve_files( - kwargs["name"], kwargs["species"], kwargs["frame"], kwargs["path"], kwargs["source"]) + name, species, frame, path, source) files["Jf"] = jf_file files["jacobvel"] = jacobvel_file files["jacobtot_inv"] = jacobtot_inv_file - if kwargs["c2p_vel"]: + if use_c2p_vel: files["mapc2p_vel"] = mapc2p_vel_file - if kwargs["mc2nu"]: + if use_mc2nu: files["mc2nu"] = mc2nu_file _assert_files_exist(files) - map_kwargs = _resolve_mapping_kwargs(kwargs["c2p_vel"], mapc2p_vel_file) + map_kwargs = _resolve_mapping_kwargs(use_c2p_vel, mapc2p_vel_file) - if kwargs["debug"]: + if debug: click.echo(f"gk_distf: Jf={jf_file}") click.echo(f"gk_distf: jacobvel={jacobvel_file}") click.echo(f"gk_distf: jacobtot_inv={jacobtot_inv_file}") - if kwargs["c2p_vel"]: + if use_c2p_vel: click.echo(f"gk_distf: mapc2p_vel={mapc2p_vel_file}") - if kwargs["mc2nu"]: + if use_mc2nu: click.echo(f"gk_distf: mc2nu={mc2nu_file}") jf_data = GData(jf_file, mapc2p_vel_name=map_kwargs["mapc2p_vel_name"]) jacobvel_data = GData(jacobvel_file) jacobtot_inv_data = GData(jacobtot_inv_file) - fjx_data = GData(tag=kwargs["tag"], ctx=jf_data.ctx) + fjx_data = GData(tag=tag, ctx=jf_data.ctx) fjx_values = _compute_fjx(jf_data.get_values(), jacobvel_data.get_values()) fjx_data.push(jf_data.get_grid(), fjx_values) out_grid, fjx_interp, jacob_interp = _interpolate_fjx_and_jacob(fjx_data, jacobtot_inv_data) distf_values = _broadcast_multiply(fjx_interp, jacob_interp) - if kwargs["mc2nu"]: - out_grid, cdim, vdim = _apply_mc2nu_grid(out_grid, mc2nu_file, kwargs["debug"]) + if use_mc2nu: + out_grid, _, _ = _apply_mc2nu_grid(out_grid, mc2nu_file, debug) jf_data.ctx["grid_type"] = "mc2nu" - if kwargs["debug"]: + if debug: click.echo(f"gk_distf: output shape={distf_values.shape}") - out = GData(tag=kwargs["tag"], ctx=jf_data.ctx) + out = GData(tag=tag, ctx=jf_data.ctx) out.push(out_grid, np.asarray(distf_values)[..., np.newaxis]) + return out + + +@click.command() +@click.option("--name", "-n", required=True, type=click.STRING, + help="Simulation name prefix (e.g. gk_lorentzian_mirror).") +@click.option("--species", "-s", required=True, type=click.STRING, + help="Species name (e.g. ion or elc).") +@click.option("--frame", "-f", required=True, type=click.INT, + help="Frame number.") +@click.option("--source", is_flag=True, + help="Use -_source_.gkyl as the input distribution.") +@click.option("--path", "-p", default="./", type=click.STRING, + help="Path to simulation data.") +@click.option("--tag", "-t", default="df", type=click.STRING, + help="Tag for output dataset.") +@click.option("--c2p-vel/--no-c2p-vel", default=False, + help="Use -_mapc2p_vel.gkyl when loading Jf.") +@click.option("--mc2nu", is_flag=True, + help="Use -mc2nu_pos_deflated.gkyl to deform configuration-space grid.") +@click.option("--debug", is_flag=True, + help="Print resolved file names and shape diagnostics.") +@click.pass_context +def gk_distf(ctx, **kwargs): + """Gyrokinetics: build real distribution function from saved Jf data.""" + verb_print(ctx, "Starting gk_distf") + data = ctx.obj["data"] + + out = build_gk_distf( + name=kwargs["name"], + species=kwargs["species"], + frame=kwargs["frame"], + path=kwargs["path"], + tag=kwargs["tag"], + source=kwargs["source"], + use_c2p_vel=kwargs["c2p_vel"], + use_mc2nu=kwargs["mc2nu"], + debug=kwargs["debug"]) data.add(out) verb_print(ctx, "Finishing gk_distf") \ No newline at end of file From c35a7e2efd730b7e8d22e3df9d75539f1698cfa8 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 12:00:11 -0400 Subject: [PATCH 05/31] Add grid type detection to GData context based on grid shape. This was giving an error when collecting in time, then the grid gets messed up --- src/postgkyl/data/gdata.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/postgkyl/data/gdata.py b/src/postgkyl/data/gdata.py index ac96eb7..f2fe60d 100644 --- a/src/postgkyl/data/gdata.py +++ b/src/postgkyl/data/gdata.py @@ -268,6 +268,11 @@ def set_grid(self, grid: list) -> None: up[d] = self._grid[d].max() self.ctx["lower"] = lo self.ctx["upper"] = up + if "grid_type" not in self.ctx: + if all(len(np.asarray(coord).shape) == 1 for coord in self._grid): + self.ctx["grid_type"] = "uniform" + else: + self.ctx["grid_type"] = "mapped" grid = property(get_grid, set_grid) From 74b15708e81fc7da9afd58c9cde85087a1978cf5 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 12:00:54 -0400 Subject: [PATCH 06/31] Enhance gk_distf to support frame ranges and improve frame resolution handling --- src/postgkyl/commands/gk_distf.py | 123 ++++++++++++++++++++++++++---- 1 file changed, 110 insertions(+), 13 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 9b587d4..60a7c0a 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -1,3 +1,4 @@ +import glob import os import click @@ -26,6 +27,85 @@ def _resolve_files(name: str, species: str, frame: int, path: str, source: bool return frame_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file +def _resolve_available_frames(name: str, species: str, path: str, + source: bool = False) -> list[int]: + if source: + prefix = _resolve_path(path, f"{name}-{species}_source_") + else: + prefix = _resolve_path(path, f"{name}-{species}_") + + files = glob.glob(f"{glob.escape(prefix)}*.gkyl") + frames = [] + for file_name in files: + suffix = file_name.removeprefix(prefix) + if suffix.endswith(".gkyl"): + frame_str = suffix[:-5] + try: + frames.append(int(frame_str)) + except ValueError: + pass + + return sorted(set(frames)) + + +def _parse_frame_bounds(frame: str) -> tuple[int | None, int | None, int]: + parts = frame.split(":") + if len(parts) < 2 or len(parts) > 3: + raise ValueError( + "Frame range must use the form 'start:stop[:step]' or ':' for all frames." + ) + + try: + start = int(parts[0]) if parts[0] else None + stop = int(parts[1]) if parts[1] else None + step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 + except ValueError as err: + raise ValueError( + "Frame range must use integer bounds, e.g. ':', '40:50', or '40:50:2'." + ) from err + + if step <= 0: + raise ValueError("Frame range step must be a positive integer.") + + return start, stop, step + + +def _resolve_frames(name: str, species: str, frame: str, path: str, + source: bool = False) -> list[int]: + frame = frame.strip() + if not frame: + raise ValueError("Frame cannot be empty.") + + if ":" not in frame: + try: + return [int(frame)] + except ValueError as err: + raise ValueError( + "Frame must be an integer or a slice like ':', '40:50', or '40:50:2'." + ) from err + + available_frames = _resolve_available_frames(name, species, path, source) + if not available_frames: + raise FileNotFoundError( + f"No frame files found for {name}-{species} in {path}." + ) + + start, stop, step = _parse_frame_bounds(frame) + lower = available_frames[0] if start is None else start + upper = available_frames[-1] + 1 if stop is None else stop + + frames = [ + frame_num for frame_num in available_frames + if lower <= frame_num < upper and (frame_num - lower) % step == 0 + ] + if not frames: + raise FileNotFoundError( + f"No frames matched '{frame}' for {name}-{species} in {path}." + ) + + return frames + + def _resolve_mapping_kwargs(use_c2p_vel: bool, mapc2p_vel_file: str) -> dict: if use_c2p_vel: return {"mapc2p_vel_name": mapc2p_vel_file} @@ -208,8 +288,8 @@ def build_gk_distf(name: str, species: str, frame: int, path: str = "./", help="Simulation name prefix (e.g. gk_lorentzian_mirror).") @click.option("--species", "-s", required=True, type=click.STRING, help="Species name (e.g. ion or elc).") -@click.option("--frame", "-f", required=True, type=click.INT, - help="Frame number.") +@click.option("--frame", "-f", required=True, type=click.STRING, + help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") @click.option("--source", is_flag=True, help="Use -_source_.gkyl as the input distribution.") @click.option("--path", "-p", default="./", type=click.STRING, @@ -228,16 +308,33 @@ def gk_distf(ctx, **kwargs): verb_print(ctx, "Starting gk_distf") data = ctx.obj["data"] - out = build_gk_distf( - name=kwargs["name"], - species=kwargs["species"], - frame=kwargs["frame"], - path=kwargs["path"], - tag=kwargs["tag"], - source=kwargs["source"], - use_c2p_vel=kwargs["c2p_vel"], - use_mc2nu=kwargs["mc2nu"], - debug=kwargs["debug"]) - data.add(out) + try: + frames = _resolve_frames( + name=kwargs["name"], + species=kwargs["species"], + frame=kwargs["frame"], + path=kwargs["path"], + source=kwargs["source"]) + except (FileNotFoundError, ValueError) as err: + ctx.fail(click.style(str(err), fg="red")) + + if kwargs["debug"] and len(frames) > 1: + click.echo(f"gk_distf: resolved frames={frames}") + + for frame in frames: + out = build_gk_distf( + name=kwargs["name"], + species=kwargs["species"], + frame=frame, + path=kwargs["path"], + tag=kwargs["tag"], + source=kwargs["source"], + use_c2p_vel=kwargs["c2p_vel"], + use_mc2nu=kwargs["mc2nu"], + debug=kwargs["debug"]) + data.add(out) + + if len(frames) > 1: + data.set_unique_labels() verb_print(ctx, "Finishing gk_distf") \ No newline at end of file From 1bf42c6b570cdb2c297f31a6ebb19ab64f9b71b1 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 12:04:50 -0400 Subject: [PATCH 07/31] rename build_gk_distf to load_gk_distf, which is more in line with how we do data loading. This command is basically a loading command, but has some computation underneath, but is just loading some data. --- src/postgkyl/commands/__init__.py | 2 +- src/postgkyl/commands/gk_distf.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/postgkyl/commands/__init__.py b/src/postgkyl/commands/__init__.py index 16a7641..83b960c 100644 --- a/src/postgkyl/commands/__init__.py +++ b/src/postgkyl/commands/__init__.py @@ -30,7 +30,7 @@ from postgkyl.commands.mhd import mhd from postgkyl.commands.parrotate import parrotate from postgkyl.commands.gk_energy_balance import gk_energy_balance -from postgkyl.commands.gk_distf import build_gk_distf +from postgkyl.commands.gk_distf import load_gk_distf from postgkyl.commands.gk_distf import gk_distf from postgkyl.commands.gk_particle_balance import gk_particle_balance from postgkyl.commands.perprotate import perprotate diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 60a7c0a..e20ffbe 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -228,7 +228,7 @@ def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): return grid, fjx_values, jacob_values -def build_gk_distf(name: str, species: str, frame: int, path: str = "./", +def load_gk_distf(name: str, species: str, frame: int, path: str = "./", tag: str = "df", source: bool = False, use_c2p_vel: bool = False, use_mc2nu: bool = False, debug: bool = False) -> GData: """Build a real distribution function from saved Jf data. @@ -322,7 +322,7 @@ def gk_distf(ctx, **kwargs): click.echo(f"gk_distf: resolved frames={frames}") for frame in frames: - out = build_gk_distf( + out = load_gk_distf( name=kwargs["name"], species=kwargs["species"], frame=frame, From 3144a145046865460875a015a759a54c8f021d06 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 12:37:46 -0400 Subject: [PATCH 08/31] Add multiblock implementation --- src/postgkyl/commands/gk_distf.py | 208 ++++++++++++++++++++++++++---- 1 file changed, 183 insertions(+), 25 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index e20ffbe..f67608b 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -7,6 +7,7 @@ from postgkyl.data import GData from postgkyl.data import GInterpModal from postgkyl.utils import verb_print +import postgkyl.utils.gk_utils as gku def _resolve_path(path: str, file_name: str) -> str: @@ -15,24 +16,51 @@ def _resolve_path(path: str, file_name: str) -> str: return os.path.join(path, file_name) -def _resolve_files(name: str, species: str, frame: int, path: str, source: bool = False): +def _discover_block_indices(name: str, species: str, path: str, + source: bool = False) -> list[int]: + """Return sorted unique block indices by parsing filenames on disk. + + Matches ``{name}_b{N}-{species}_*.gkyl`` and extracts N from each filename, + so the result is independent of how many frame files exist per block. + """ + src_infix = "source_" if source else "" + template = _resolve_path(path, f"{name}_b*-{species}_{src_infix}*.gkyl") + files = glob.glob(template) + prefix = f"{name}_b" + sep = f"-{species}_" + indices = set() + for f in files: + base = os.path.basename(f) + if base.startswith(prefix) and sep in base: + block_str = base[len(prefix):base.index(sep)] + try: + indices.add(int(block_str)) + except ValueError: + pass + return sorted(indices) + + +def _resolve_files(name: str, species: str, frame: int, path: str, + source: bool = False, block_idx: int | None = None): + prefix = f"{name}_b{block_idx}" if block_idx is not None else name if source: - frame_file = _resolve_path(path, f"{name}-{species}_source_{frame}.gkyl") + frame_file = _resolve_path(path, f"{prefix}-{species}_source_{frame}.gkyl") else: - frame_file = _resolve_path(path, f"{name}-{species}_{frame}.gkyl") - mapc2p_vel_file = _resolve_path(path, f"{name}-{species}_mapc2p_vel.gkyl") - mc2nu_file = _resolve_path(path, f"{name}-mc2nu_pos_deflated.gkyl") - jacobvel_file = _resolve_path(path, f"{name}-{species}_jacobvel.gkyl") - jacobtot_inv_file = _resolve_path(path, f"{name}-jacobtot_inv.gkyl") + frame_file = _resolve_path(path, f"{prefix}-{species}_{frame}.gkyl") + mapc2p_vel_file = _resolve_path(path, f"{prefix}-{species}_mapc2p_vel.gkyl") + mc2nu_file = _resolve_path(path, f"{prefix}-mc2nu_pos_deflated.gkyl") + jacobvel_file = _resolve_path(path, f"{prefix}-{species}_jacobvel.gkyl") + jacobtot_inv_file = _resolve_path(path, f"{prefix}-jacobtot_inv.gkyl") return frame_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file def _resolve_available_frames(name: str, species: str, path: str, - source: bool = False) -> list[int]: + source: bool = False, block_idx: int | None = None) -> list[int]: + prefix_name = f"{name}_b{block_idx}" if block_idx is not None else name if source: - prefix = _resolve_path(path, f"{name}-{species}_source_") + prefix = _resolve_path(path, f"{prefix_name}-{species}_source_") else: - prefix = _resolve_path(path, f"{name}-{species}_") + prefix = _resolve_path(path, f"{prefix_name}-{species}_") files = glob.glob(f"{glob.escape(prefix)}*.gkyl") frames = [] @@ -71,7 +99,7 @@ def _parse_frame_bounds(frame: str) -> tuple[int | None, int | None, int]: def _resolve_frames(name: str, species: str, frame: str, path: str, - source: bool = False) -> list[int]: + source: bool = False, block_idx: int | None = None) -> list[int]: frame = frame.strip() if not frame: raise ValueError("Frame cannot be empty.") @@ -84,7 +112,8 @@ def _resolve_frames(name: str, species: str, frame: str, path: str, "Frame must be an integer or a slice like ':', '40:50', or '40:50:2'." ) from err - available_frames = _resolve_available_frames(name, species, path, source) + available_frames = _resolve_available_frames(name, species, path, source, + block_idx=block_idx) if not available_frames: raise FileNotFoundError( f"No frame files found for {name}-{species} in {path}." @@ -216,6 +245,65 @@ def _broadcast_multiply(distf: np.ndarray, jacobtot_inv: np.ndarray) -> np.ndarr return distf * jacobtot_inv.reshape(jacob_shape) +def _find_stitch_axis(grids: list) -> int: + """Return the axis along which consecutive blocks are non-overlapping. + + The stitch axis is the dimension where every consecutive block pair starts + no earlier than the previous block ends. Velocity-space axes always share + the same range across blocks so they never qualify. + """ + ndim = len(grids[0]) + for axis in range(ndim): + is_stitch = True + for b in range(len(grids) - 1): + max_curr = float(np.asarray(grids[b][axis]).max()) + min_next = float(np.asarray(grids[b + 1][axis]).min()) + tol = 1e-10 * max(abs(max_curr), 1.0) + if min_next < max_curr - tol: + is_stitch = False + break + if is_stitch: + return axis + raise ValueError( + "Could not determine stitch axis: no single axis is non-overlapping " + "across all consecutive block pairs." + ) + + +def _stitch_block_data(block_outputs: list, tag: str) -> GData: + """Concatenate per-block GData objects into one dataset along the stitch axis.""" + grids = [out.get_grid() for out in block_outputs] + stitch_axis = _find_stitch_axis(grids) + + # Sort blocks by their minimum coordinate on the stitch axis. + order = sorted(range(len(block_outputs)), + key=lambda i: float(np.asarray(grids[i][stitch_axis]).min())) + block_outputs = [block_outputs[i] for i in order] + grids = [block_outputs[i].get_grid() for i in range(len(block_outputs))] + values = [block_outputs[i].get_values() for i in range(len(block_outputs))] + + # Build merged grid. Node coordinates share a boundary, so drop the + # first node of each subsequent block along the stitch axis. + ndim = len(grids[0]) + merged_grid = [] + for d in range(ndim): + if d == stitch_axis: + merged_axis = np.concatenate( + [np.asarray(grids[0][d])] + + [np.asarray(grids[b][d][1:]) for b in range(1, len(grids))] + ) + else: + merged_axis = np.asarray(grids[0][d]) + merged_grid.append(merged_axis) + + # Concatenate values along the spatial stitch axis. + merged_values = np.concatenate(values, axis=stitch_axis) + + result = GData(tag=tag, ctx=block_outputs[0].ctx) + result.push(merged_grid, merged_values) + return result + + def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): fjx_interp = GInterpModal(fjx_data, 1, "gkhyb") jacob_interp = GInterpModal(jacobtot_inv_data, 1, "ms") @@ -230,14 +318,15 @@ def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): def load_gk_distf(name: str, species: str, frame: int, path: str = "./", tag: str = "df", source: bool = False, use_c2p_vel: bool = False, - use_mc2nu: bool = False, debug: bool = False) -> GData: + use_mc2nu: bool = False, debug: bool = False, + block_idx: int | None = None) -> GData: """Build a real distribution function from saved Jf data. This is the Python API used by the ``gk-distf`` CLI command. """ files = {} jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _resolve_files( - name, species, frame, path, source) + name, species, frame, path, source, block_idx=block_idx) files["Jf"] = jf_file files["jacobvel"] = jacobvel_file files["jacobtot_inv"] = jacobtot_inv_file @@ -283,6 +372,40 @@ def load_gk_distf(name: str, species: str, frame: int, path: str = "./", return out +def load_gk_distf_mb(name: str, species: str, frame: int, path: str = "./", + tag: str = "df", source: bool = False, use_c2p_vel: bool = False, + use_mc2nu: bool = False, debug: bool = False, + mb_blocks: str = "-1") -> GData: + """Load and stitch a multiblock distribution function. + + Files are expected as ``{name}_b{N}-{species}_{frame}.gkyl`` with the + same ``_b{N}`` prefix on all auxiliary files (jacobvel, jacobtot_inv, + mapc2p_vel, mc2nu). + """ + if mb_blocks == "-1": + blocks = _discover_block_indices(name, species, path, source) + else: + blocks = gku.get_block_indices(mb_blocks, "") + + if not blocks: + raise FileNotFoundError( + f"No block files found for {name}_b*-{species} in {path}") + + if debug: + click.echo(f"gk_distf: multiblock mode, blocks={blocks}") + + block_outputs = [] + for block_idx in blocks: + out = load_gk_distf( + name=name, species=species, frame=frame, path=path, + tag=f"{tag}_b{block_idx}", source=source, + use_c2p_vel=use_c2p_vel, use_mc2nu=use_mc2nu, + debug=debug, block_idx=block_idx) + block_outputs.append(out) + + return _stitch_block_data(block_outputs, tag) + + @click.command() @click.option("--name", "-n", required=True, type=click.STRING, help="Simulation name prefix (e.g. gk_lorentzian_mirror).") @@ -300,6 +423,11 @@ def load_gk_distf(name: str, species: str, frame: int, path: str = "./", help="Use -_mapc2p_vel.gkyl when loading Jf.") @click.option("--mc2nu", is_flag=True, help="Use -mc2nu_pos_deflated.gkyl to deform configuration-space grid.") +@click.option("--mb", "-b", is_flag=True, default=False, + help="Multiblock mode. Files are expected as {name}_b{N}-{species}_{frame}.gkyl.") +@click.option("--mb-blocks", type=click.STRING, default=None, + help=("Restrict multiblock to specific blocks: comma-separated list or slice " + "(start:stop:step). Requires --mb. Defaults to all discovered blocks.")) @click.option("--debug", is_flag=True, help="Print resolved file names and shape diagnostics.") @click.pass_context @@ -308,13 +436,30 @@ def gk_distf(ctx, **kwargs): verb_print(ctx, "Starting gk_distf") data = ctx.obj["data"] + is_multib = kwargs["mb"] + mb_blocks_str = kwargs["mb_blocks"] if kwargs["mb_blocks"] is not None else "-1" + + # For frame-range resolution, use block 0 as the reference when multiblock. + ref_block = None + if is_multib: + if mb_blocks_str == "-1": + blocks = _discover_block_indices( + kwargs["name"], kwargs["species"], kwargs["path"], kwargs["source"]) + else: + blocks = gku.get_block_indices(mb_blocks_str, "") + if not blocks: + ctx.fail(click.style( + f"No block files found for {kwargs['name']} in {kwargs['path']}", fg="red")) + ref_block = blocks[0] + try: frames = _resolve_frames( name=kwargs["name"], species=kwargs["species"], frame=kwargs["frame"], path=kwargs["path"], - source=kwargs["source"]) + source=kwargs["source"], + block_idx=ref_block) except (FileNotFoundError, ValueError) as err: ctx.fail(click.style(str(err), fg="red")) @@ -322,16 +467,29 @@ def gk_distf(ctx, **kwargs): click.echo(f"gk_distf: resolved frames={frames}") for frame in frames: - out = load_gk_distf( - name=kwargs["name"], - species=kwargs["species"], - frame=frame, - path=kwargs["path"], - tag=kwargs["tag"], - source=kwargs["source"], - use_c2p_vel=kwargs["c2p_vel"], - use_mc2nu=kwargs["mc2nu"], - debug=kwargs["debug"]) + if is_multib: + out = load_gk_distf_mb( + name=kwargs["name"], + species=kwargs["species"], + frame=frame, + path=kwargs["path"], + tag=kwargs["tag"], + source=kwargs["source"], + use_c2p_vel=kwargs["c2p_vel"], + use_mc2nu=kwargs["mc2nu"], + debug=kwargs["debug"], + mb_blocks=mb_blocks_str) + else: + out = load_gk_distf( + name=kwargs["name"], + species=kwargs["species"], + frame=frame, + path=kwargs["path"], + tag=kwargs["tag"], + source=kwargs["source"], + use_c2p_vel=kwargs["c2p_vel"], + use_mc2nu=kwargs["mc2nu"], + debug=kwargs["debug"]) data.add(out) if len(frames) > 1: From 81610a7b347f903d26c5b8750c42d644c9151b12 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 12:42:41 -0400 Subject: [PATCH 09/31] Update kwargs for some shortcuts --- src/postgkyl/commands/gk_distf.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index f67608b..cadf436 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -411,23 +411,23 @@ def load_gk_distf_mb(name: str, species: str, frame: int, path: str = "./", help="Simulation name prefix (e.g. gk_lorentzian_mirror).") @click.option("--species", "-s", required=True, type=click.STRING, help="Species name (e.g. ion or elc).") -@click.option("--frame", "-f", required=True, type=click.STRING, - help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") @click.option("--source", is_flag=True, help="Use -_source_.gkyl as the input distribution.") +@click.option("--frame", "-f", required=True, type=click.STRING, + help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") @click.option("--path", "-p", default="./", type=click.STRING, help="Path to simulation data.") @click.option("--tag", "-t", default="df", type=click.STRING, help="Tag for output dataset.") -@click.option("--c2p-vel/--no-c2p-vel", default=False, +@click.option("--c2p-vel", "-v", is_flag=True, default=False, help="Use -_mapc2p_vel.gkyl when loading Jf.") -@click.option("--mc2nu", is_flag=True, +@click.option("--mc2nu", "-m", is_flag=True, help="Use -mc2nu_pos_deflated.gkyl to deform configuration-space grid.") -@click.option("--mb", "-b", is_flag=True, default=False, +@click.option("--multi-block", "-b", is_flag=True, default=False, help="Multiblock mode. Files are expected as {name}_b{N}-{species}_{frame}.gkyl.") -@click.option("--mb-blocks", type=click.STRING, default=None, +@click.option("--block", type=click.STRING, default=None, help=("Restrict multiblock to specific blocks: comma-separated list or slice " - "(start:stop:step). Requires --mb. Defaults to all discovered blocks.")) + "(start:stop:step). Requires --multi-block. Defaults to all discovered blocks.")) @click.option("--debug", is_flag=True, help="Print resolved file names and shape diagnostics.") @click.pass_context @@ -436,8 +436,8 @@ def gk_distf(ctx, **kwargs): verb_print(ctx, "Starting gk_distf") data = ctx.obj["data"] - is_multib = kwargs["mb"] - mb_blocks_str = kwargs["mb_blocks"] if kwargs["mb_blocks"] is not None else "-1" + is_multib = kwargs["multi_block"] + mb_blocks_str = kwargs["block"] if kwargs["block"] is not None else "-1" # For frame-range resolution, use block 0 as the reference when multiblock. ref_block = None From 5c6e1107b9c2154002207308dad6d75bdb02e255 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 20:55:03 -0400 Subject: [PATCH 10/31] Refactor the file for organization. Match a few more conventions of the code base. Reduce the mutliblock infrastructure to not stitch together distribution functions, but just to specify 1 block to plot at a time --- src/postgkyl/commands/gk_distf.py | 669 +++++++++++++----------------- 1 file changed, 295 insertions(+), 374 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index cadf436..799b9a8 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -4,175 +4,286 @@ import click import numpy as np -from postgkyl.data import GData -from postgkyl.data import GInterpModal +from postgkyl.data import GData, GInterpModal from postgkyl.utils import verb_print -import postgkyl.utils.gk_utils as gku +# --------------------------------------------------------------------------- +# File/path helpers +# --------------------------------------------------------------------------- + def _resolve_path(path: str, file_name: str) -> str: if os.path.isabs(file_name): return file_name + # end return os.path.join(path, file_name) +# end -def _discover_block_indices(name: str, species: str, path: str, - source: bool = False) -> list[int]: - """Return sorted unique block indices by parsing filenames on disk. - - Matches ``{name}_b{N}-{species}_*.gkyl`` and extracts N from each filename, - so the result is independent of how many frame files exist per block. - """ - src_infix = "source_" if source else "" - template = _resolve_path(path, f"{name}_b*-{species}_{src_infix}*.gkyl") - files = glob.glob(template) - prefix = f"{name}_b" - sep = f"-{species}_" - indices = set() - for f in files: - base = os.path.basename(f) - if base.startswith(prefix) and sep in base: - block_str = base[len(prefix):base.index(sep)] - try: - indices.add(int(block_str)) - except ValueError: - pass - return sorted(indices) - - -def _resolve_files(name: str, species: str, frame: int, path: str, - source: bool = False, block_idx: int | None = None): +def _resolve_files( + name: str, species: str, frame: int, path: str, + source: bool = False, block_idx: int | None = None, +) -> tuple[str, str, str, str, str]: prefix = f"{name}_b{block_idx}" if block_idx is not None else name - if source: - frame_file = _resolve_path(path, f"{prefix}-{species}_source_{frame}.gkyl") - else: - frame_file = _resolve_path(path, f"{prefix}-{species}_{frame}.gkyl") + frame_infix = "source_" if source else "" + frame_file = _resolve_path(path, f"{prefix}-{species}_{frame_infix}{frame}.gkyl") mapc2p_vel_file = _resolve_path(path, f"{prefix}-{species}_mapc2p_vel.gkyl") - mc2nu_file = _resolve_path(path, f"{prefix}-mc2nu_pos_deflated.gkyl") - jacobvel_file = _resolve_path(path, f"{prefix}-{species}_jacobvel.gkyl") + mc2nu_file = _resolve_path(path, f"{prefix}-mc2nu_pos_deflated.gkyl") + jacobvel_file = _resolve_path(path, f"{prefix}-{species}_jacobvel.gkyl") jacobtot_inv_file = _resolve_path(path, f"{prefix}-jacobtot_inv.gkyl") return frame_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file +# end + + +def _error_check(check: str, **kwargs): + """Centralized validation for gk_distf helpers. + + This keeps error checks in one place while helper functions focus on data flow. + """ + if check == "frame_nonempty": + if not kwargs["frame"]: + raise ValueError("Frame cannot be empty.") + # end + # end + + if check == "single_frame": + try: + int(kwargs["frame"]) + except ValueError as err: + raise ValueError( + "Frame must be an integer or a slice like ':', '40:50', or '40:50:2'." + ) from err + # end + # end + + if check == "frame_bounds": + parts = kwargs["frame"].split(":") + if len(parts) < 2 or len(parts) > 3: + raise ValueError( + "Frame range must use the form 'start:stop[:step]' or ':' for all frames." + ) + # end + + try: + int(parts[0]) if parts[0] else None + int(parts[1]) if parts[1] else None + step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 + except ValueError as err: + raise ValueError( + "Frame range must use integer bounds, e.g. ':', '40:50', or '40:50:2'." + ) from err + # end + + if step <= 0: + raise ValueError("Frame range step must be a positive integer.") + # end + # end + + if check == "available_frames": + if not kwargs["available"]: + raise FileNotFoundError( + f"No frame files found for {kwargs['name']}-{kwargs['species']} in {kwargs['path']}." + ) + # end + # end + + if check == "matched_frames": + if not kwargs["frames"]: + raise FileNotFoundError( + f"No frames matched '{kwargs['frame']}' for {kwargs['name']}-{kwargs['species']} in {kwargs['path']}." + ) + # end + # end + + if check == "required_files": + missing = [ + f"{label}: {file_name}" + for label, file_name in kwargs["files"].items() + if not os.path.exists(file_name) + ] + if missing: + raise FileNotFoundError("Missing required files:\n - " + "\n - ".join(missing)) + # end + # end + + if check == "fjx_shape": + jf_shape = kwargs["jf_shape"] + jac_shape = kwargs["jac_shape"] + if jf_shape != jac_shape: + if jf_shape[:-1] != jac_shape[:-1]: + raise ValueError( + "Jf and jacobvel spatial shapes do not match: " + f"{jf_shape} vs {jac_shape}" + ) + # end + if jac_shape[-1] != 1: + raise ValueError( + "jacobvel component axis must be 1 or match Jf: " + f"{jf_shape} vs {jac_shape}" + ) + # end + # end + # end + + if check == "broadcast_shape": + distf_shape = kwargs["distf_shape"] + jac_shape = kwargs["jac_shape"] + if len(jac_shape) > len(distf_shape): + raise ValueError( + "jacobtot_inv has more dimensions than distribution function: " + f"{len(jac_shape)} > {len(distf_shape)}" + ) + # end + + if distf_shape[:len(jac_shape)] != jac_shape: + raise ValueError( + "Cannot broadcast jacobtot_inv over distribution function. " + f"Expected prefix {distf_shape[:len(jac_shape)]}, got {jac_shape}" + ) + # end + # end + + if check == "mc2nu_cdim": + cdim = kwargs["cdim"] + total_dims = kwargs["total_dims"] + if cdim < 1 or cdim > total_dims: + raise ValueError(f"Invalid cdim={cdim} for total_dims={total_dims}") + # end + # end + + if check == "mc2nu_axis": + mapped_size = kwargs["mapped_size"] + old_size = kwargs["old_size"] + axis = kwargs["axis"] + if mapped_size != old_size and mapped_size + 1 != old_size: + raise ValueError( + f"mc2nu mapping size is incompatible with configuration grid on axis {axis}: " + f"{mapped_size} vs {old_size}" + ) + # end + # end + + raise ValueError(f"Unknown precheck type '{check}'") +# end -def _resolve_available_frames(name: str, species: str, path: str, - source: bool = False, block_idx: int | None = None) -> list[int]: +# --------------------------------------------------------------------------- +# Frame resolution helpers +# --------------------------------------------------------------------------- + +def _resolve_available_frames( + name: str, species: str, path: str, + source: bool = False, block_idx: int | None = None,) -> list[int]: prefix_name = f"{name}_b{block_idx}" if block_idx is not None else name - if source: - prefix = _resolve_path(path, f"{prefix_name}-{species}_source_") - else: - prefix = _resolve_path(path, f"{prefix_name}-{species}_") + frame_infix = "source_" if source else "" + prefix = _resolve_path(path, f"{prefix_name}-{species}_{frame_infix}") - files = glob.glob(f"{glob.escape(prefix)}*.gkyl") frames = [] - for file_name in files: + for file_name in glob.glob(f"{glob.escape(prefix)}*.gkyl"): suffix = file_name.removeprefix(prefix) if suffix.endswith(".gkyl"): - frame_str = suffix[:-5] try: - frames.append(int(frame_str)) + frames.append(int(suffix[:-5])) except ValueError: pass - + # end + # end + # end return sorted(set(frames)) +# end -def _parse_frame_bounds(frame: str) -> tuple[int | None, int | None, int]: - parts = frame.split(":") - if len(parts) < 2 or len(parts) > 3: - raise ValueError( - "Frame range must use the form 'start:stop[:step]' or ':' for all frames." - ) - - try: - start = int(parts[0]) if parts[0] else None - stop = int(parts[1]) if parts[1] else None - step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 - except ValueError as err: - raise ValueError( - "Frame range must use integer bounds, e.g. ':', '40:50', or '40:50:2'." - ) from err - - if step <= 0: - raise ValueError("Frame range step must be a positive integer.") - - return start, stop, step - - -def _resolve_frames(name: str, species: str, frame: str, path: str, - source: bool = False, block_idx: int | None = None) -> list[int]: +def _resolve_frames( + name: str, species: str, frame: str, path: str, + source: bool = False, block_idx: int | None = None, +) -> list[int]: frame = frame.strip() - if not frame: - raise ValueError("Frame cannot be empty.") + _error_check("frame_nonempty", frame=frame) if ":" not in frame: - try: - return [int(frame)] - except ValueError as err: - raise ValueError( - "Frame must be an integer or a slice like ':', '40:50', or '40:50:2'." - ) from err + _error_check("single_frame", frame=frame) + return [int(frame)] + # end - available_frames = _resolve_available_frames(name, species, path, source, - block_idx=block_idx) - if not available_frames: - raise FileNotFoundError( - f"No frame files found for {name}-{species} in {path}." - ) + available = _resolve_available_frames(name, species, path, source, block_idx=block_idx) + _error_check("available_frames", available=available, name=name, species=species, path=path) - start, stop, step = _parse_frame_bounds(frame) - lower = available_frames[0] if start is None else start - upper = available_frames[-1] + 1 if stop is None else stop + _error_check("frame_bounds", frame=frame) + parts = frame.split(":") + start = int(parts[0]) if parts[0] else None + stop = int(parts[1]) if parts[1] else None + step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 + lower = available[0] if start is None else start + upper = available[-1] + 1 if stop is None else stop frames = [ - frame_num for frame_num in available_frames - if lower <= frame_num < upper and (frame_num - lower) % step == 0 + f for f in available + if lower <= f < upper and (f - lower) % step == 0 ] - if not frames: - raise FileNotFoundError( - f"No frames matched '{frame}' for {name}-{species} in {path}." - ) - + _error_check("matched_frames", frames=frames, frame=frame, name=name, species=species, path=path) return frames +# end +# --------------------------------------------------------------------------- +# Physics / DG helpers +# --------------------------------------------------------------------------- + +def _compute_fjx(jf_values: np.ndarray, jac_values: np.ndarray) -> np.ndarray: + """Divide Jf by jacobvel to get f * J_x.""" + _error_check("fjx_shape", jf_shape=jf_values.shape, jac_shape=jac_values.shape) + return jf_values / jac_values +# end -def _resolve_mapping_kwargs(use_c2p_vel: bool, mapc2p_vel_file: str) -> dict: - if use_c2p_vel: - return {"mapc2p_vel_name": mapc2p_vel_file} - return {"mapc2p_vel_name": ""} +def _broadcast_multiply(distf: np.ndarray, jacobtot_inv: np.ndarray) -> np.ndarray: + """Multiply distf by jacobtot_inv, broadcasting over the component axis.""" + _error_check("broadcast_shape", distf_shape=distf.shape, jac_shape=jacobtot_inv.shape) + jacob_shape = jacobtot_inv.shape + (1,) * (distf.ndim - jacobtot_inv.ndim) + return distf * jacobtot_inv.reshape(jacob_shape) +# end + + +def _interpolate_distf_components( + fjx_data: GData, jacobtot_inv_data: GData, +) -> tuple[list, np.ndarray, np.ndarray]: + """Interpolate f/J_x (gkhyb basis) and jacobtot_inv (ms basis).""" + grid, fjx_values = GInterpModal(fjx_data, 1, "gkhyb").interpolate() + _, jacob_values = GInterpModal(jacobtot_inv_data, 1, "ms").interpolate() + return grid, np.squeeze(fjx_values), np.squeeze(jacob_values) +# end + +# --------------------------------------------------------------------------- +# mc2nu grid deformation helpers +# --------------------------------------------------------------------------- def _cell_centers_to_nodes(cell_centers: np.ndarray) -> np.ndarray: nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) - nodes[0] = cell_centers[0] - (nodes[1] - cell_centers[0]) + nodes[0] = cell_centers[0] - (nodes[1] - cell_centers[0]) nodes[-1] = cell_centers[-1] + (cell_centers[-1] - nodes[-2]) return nodes - - -def _get_dim_split(out_grid: list, mc2nu_data: GData) -> tuple[int, int]: - total_dims = len(out_grid) - cdim = mc2nu_data.get_num_dims() - if cdim < 1 or cdim > total_dims: - raise ValueError(f"Invalid cdim={cdim} for total_dims={total_dims}") - vdim = total_dims - cdim - return cdim, vdim +# end def _extract_mapped_axis(mapped_values: np.ndarray, axis: int, cdim: int) -> np.ndarray: + """Extract a 1D mapped coordinate array for the given configuration axis.""" if cdim == 1: return np.asarray(mapped_values[..., axis]).reshape(-1) + # end - # Extract one mapped coordinate direction at a reference point in the - # remaining configuration directions to construct a 1D axis grid. + # Pick a reference point in all other config directions. idx = [0] * (cdim + 1) idx[axis] = slice(None) idx[-1] = axis return np.asarray(mapped_values[tuple(idx)]).reshape(-1) +# end -def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, debug: bool) -> tuple[list, int, int]: +def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, ctx=None) -> list: + """Replace configuration-space axes in out_grid with mc2nu-mapped coordinates.""" mc2nu_data = GData(mc2nu_file) - cdim, vdim = _get_dim_split(out_grid, mc2nu_data) + cdim = mc2nu_data.get_num_dims() + _error_check("mc2nu_cdim", cdim=cdim, total_dims=len(out_grid)) mc2nu_interp = GInterpModal(mc2nu_data, 1, "ms") _, mc2nu_values = mc2nu_interp.interpolate(tuple(range(cdim))) @@ -181,318 +292,128 @@ def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, debug: bool) -> tuple[lis deformed_grid = list(out_grid) for d in range(cdim): mapped_cfg = _extract_mapped_axis(mapped_values, d, cdim) - old_cfg = np.asarray(out_grid[d]) + _error_check("mc2nu_axis", mapped_size=mapped_cfg.size, old_size=old_cfg.size, axis=d) if mapped_cfg.size == old_cfg.size: - new_cfg = mapped_cfg - elif mapped_cfg.size + 1 == old_cfg.size: - new_cfg = _cell_centers_to_nodes(mapped_cfg) + deformed_grid[d] = mapped_cfg else: - raise ValueError( - "mc2nu mapping size is incompatible with configuration grid on axis " - f"{d}: {mapped_cfg.size} vs {old_cfg.size}" - ) - - deformed_grid[d] = new_cfg - - if debug: - click.echo(f"gk_distf: cdim={cdim}, vdim={vdim}") - click.echo(f"gk_distf: mc2nu mapped {cdim} configuration axis/axes") + deformed_grid[d] = _cell_centers_to_nodes(mapped_cfg) + # end + # end - return deformed_grid, cdim, vdim + verb_print(ctx, f"gk_distf: mc2nu mapped {cdim} configuration axis/axes") + return deformed_grid +# end -def _assert_files_exist(files: dict): - for kind, file_name in files.items(): - if not os.path.exists(file_name): - raise FileNotFoundError(f"Missing {kind} file: {file_name}") - - -def _compute_fjx(jf_values: np.ndarray, jac_values: np.ndarray) -> np.ndarray: - if jf_values.shape == jac_values.shape: - return jf_values / jac_values - - if jf_values.shape[:-1] != jac_values.shape[:-1]: - raise ValueError( - "Jf and jacobvel spatial shapes do not match: " - f"{jf_values.shape} vs {jac_values.shape}" - ) - - if jac_values.shape[-1] != 1: - raise ValueError( - "jacobvel component axis must be 1 or match Jf: " - f"{jf_values.shape} vs {jac_values.shape}" - ) - - return jf_values / jac_values - - -def _broadcast_multiply(distf: np.ndarray, jacobtot_inv: np.ndarray) -> np.ndarray: - if jacobtot_inv.ndim > distf.ndim: - raise ValueError( - "jacobtot_inv has more dimensions than distribution function: " - f"{jacobtot_inv.ndim} > {distf.ndim}" - ) - - expected_prefix = distf.shape[:jacobtot_inv.ndim] - if expected_prefix != jacobtot_inv.shape: - raise ValueError( - "Cannot broadcast jacobtot_inv over distribution function. " - f"Expected prefix {expected_prefix}, got {jacobtot_inv.shape}" - ) - - jacob_shape = jacobtot_inv.shape + (1,) * (distf.ndim - jacobtot_inv.ndim) - return distf * jacobtot_inv.reshape(jacob_shape) - - -def _find_stitch_axis(grids: list) -> int: - """Return the axis along which consecutive blocks are non-overlapping. - - The stitch axis is the dimension where every consecutive block pair starts - no earlier than the previous block ends. Velocity-space axes always share - the same range across blocks so they never qualify. - """ - ndim = len(grids[0]) - for axis in range(ndim): - is_stitch = True - for b in range(len(grids) - 1): - max_curr = float(np.asarray(grids[b][axis]).max()) - min_next = float(np.asarray(grids[b + 1][axis]).min()) - tol = 1e-10 * max(abs(max_curr), 1.0) - if min_next < max_curr - tol: - is_stitch = False - break - if is_stitch: - return axis - raise ValueError( - "Could not determine stitch axis: no single axis is non-overlapping " - "across all consecutive block pairs." - ) - +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- -def _stitch_block_data(block_outputs: list, tag: str) -> GData: - """Concatenate per-block GData objects into one dataset along the stitch axis.""" - grids = [out.get_grid() for out in block_outputs] - stitch_axis = _find_stitch_axis(grids) - - # Sort blocks by their minimum coordinate on the stitch axis. - order = sorted(range(len(block_outputs)), - key=lambda i: float(np.asarray(grids[i][stitch_axis]).min())) - block_outputs = [block_outputs[i] for i in order] - grids = [block_outputs[i].get_grid() for i in range(len(block_outputs))] - values = [block_outputs[i].get_values() for i in range(len(block_outputs))] - - # Build merged grid. Node coordinates share a boundary, so drop the - # first node of each subsequent block along the stitch axis. - ndim = len(grids[0]) - merged_grid = [] - for d in range(ndim): - if d == stitch_axis: - merged_axis = np.concatenate( - [np.asarray(grids[0][d])] + - [np.asarray(grids[b][d][1:]) for b in range(1, len(grids))] - ) - else: - merged_axis = np.asarray(grids[0][d]) - merged_grid.append(merged_axis) - - # Concatenate values along the spatial stitch axis. - merged_values = np.concatenate(values, axis=stitch_axis) - - result = GData(tag=tag, ctx=block_outputs[0].ctx) - result.push(merged_grid, merged_values) - return result - - -def _interpolate_fjx_and_jacob(fjx_data: GData, jacobtot_inv_data: GData): - fjx_interp = GInterpModal(fjx_data, 1, "gkhyb") - jacob_interp = GInterpModal(jacobtot_inv_data, 1, "ms") - - grid, fjx_values = fjx_interp.interpolate() - _, jacob_values = jacob_interp.interpolate() - - fjx_values = np.squeeze(fjx_values) - jacob_values = np.squeeze(jacob_values) - return grid, fjx_values, jacob_values - - -def load_gk_distf(name: str, species: str, frame: int, path: str = "./", +def load_gk_distf( + name: str, species: str, frame: int, path: str = "./", tag: str = "df", source: bool = False, use_c2p_vel: bool = False, - use_mc2nu: bool = False, debug: bool = False, - block_idx: int | None = None) -> GData: - """Build a real distribution function from saved Jf data. - - This is the Python API used by the ``gk-distf`` CLI command. - """ - files = {} + use_mc2nu: bool = False, + block_idx: int | None = None, + ctx=None, +) -> GData: + """Build a real distribution function from saved Jf data.""" jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _resolve_files( name, species, frame, path, source, block_idx=block_idx) - files["Jf"] = jf_file - files["jacobvel"] = jacobvel_file - files["jacobtot_inv"] = jacobtot_inv_file + required_files = {"Jf": jf_file, "jacobvel": jacobvel_file, "jacobtot_inv": jacobtot_inv_file} if use_c2p_vel: - files["mapc2p_vel"] = mapc2p_vel_file + required_files["mapc2p_vel"] = mapc2p_vel_file + # end if use_mc2nu: - files["mc2nu"] = mc2nu_file + required_files["mc2nu"] = mc2nu_file + # end + _error_check("required_files", files=required_files) - _assert_files_exist(files) + for kind, fname in required_files.items(): + verb_print(ctx, f"gk_distf: {kind}={fname}") + # end - map_kwargs = _resolve_mapping_kwargs(use_c2p_vel, mapc2p_vel_file) - - if debug: - click.echo(f"gk_distf: Jf={jf_file}") - click.echo(f"gk_distf: jacobvel={jacobvel_file}") - click.echo(f"gk_distf: jacobtot_inv={jacobtot_inv_file}") - if use_c2p_vel: - click.echo(f"gk_distf: mapc2p_vel={mapc2p_vel_file}") - if use_mc2nu: - click.echo(f"gk_distf: mc2nu={mc2nu_file}") - - jf_data = GData(jf_file, mapc2p_vel_name=map_kwargs["mapc2p_vel_name"]) + jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file if use_c2p_vel else "") jacobvel_data = GData(jacobvel_file) jacobtot_inv_data = GData(jacobtot_inv_file) fjx_data = GData(tag=tag, ctx=jf_data.ctx) - fjx_values = _compute_fjx(jf_data.get_values(), jacobvel_data.get_values()) - fjx_data.push(jf_data.get_grid(), fjx_values) + fjx_data.push(jf_data.get_grid(), _compute_fjx(jf_data.get_values(), jacobvel_data.get_values())) - out_grid, fjx_interp, jacob_interp = _interpolate_fjx_and_jacob(fjx_data, jacobtot_inv_data) - distf_values = _broadcast_multiply(fjx_interp, jacob_interp) + out_grid, fjx_values, jacob_values = _interpolate_distf_components(fjx_data, jacobtot_inv_data) + distf_values = _broadcast_multiply(fjx_values, jacob_values) if use_mc2nu: - out_grid, _, _ = _apply_mc2nu_grid(out_grid, mc2nu_file, debug) + out_grid = _apply_mc2nu_grid(out_grid, mc2nu_file, ctx) jf_data.ctx["grid_type"] = "mc2nu" + # end - if debug: - click.echo(f"gk_distf: output shape={distf_values.shape}") + verb_print(ctx, f"gk_distf: output shape={distf_values.shape}") out = GData(tag=tag, ctx=jf_data.ctx) out.push(out_grid, np.asarray(distf_values)[..., np.newaxis]) return out +# end - -def load_gk_distf_mb(name: str, species: str, frame: int, path: str = "./", - tag: str = "df", source: bool = False, use_c2p_vel: bool = False, - use_mc2nu: bool = False, debug: bool = False, - mb_blocks: str = "-1") -> GData: - """Load and stitch a multiblock distribution function. - - Files are expected as ``{name}_b{N}-{species}_{frame}.gkyl`` with the - same ``_b{N}`` prefix on all auxiliary files (jacobvel, jacobtot_inv, - mapc2p_vel, mc2nu). - """ - if mb_blocks == "-1": - blocks = _discover_block_indices(name, species, path, source) - else: - blocks = gku.get_block_indices(mb_blocks, "") - - if not blocks: - raise FileNotFoundError( - f"No block files found for {name}_b*-{species} in {path}") - - if debug: - click.echo(f"gk_distf: multiblock mode, blocks={blocks}") - - block_outputs = [] - for block_idx in blocks: - out = load_gk_distf( - name=name, species=species, frame=frame, path=path, - tag=f"{tag}_b{block_idx}", source=source, - use_c2p_vel=use_c2p_vel, use_mc2nu=use_mc2nu, - debug=debug, block_idx=block_idx) - block_outputs.append(out) - - return _stitch_block_data(block_outputs, tag) - +# --------------------------------------------------------------------------- +# CLI command +# --------------------------------------------------------------------------- @click.command() -@click.option("--name", "-n", required=True, type=click.STRING, +@click.option("--name", "-n", required=True, type=click.STRING, help="Simulation name prefix (e.g. gk_lorentzian_mirror).") -@click.option("--species", "-s", required=True, type=click.STRING, +@click.option("--species", "-s", required=True, type=click.STRING, help="Species name (e.g. ion or elc).") -@click.option("--source", is_flag=True, - help="Use -_source_.gkyl as the input distribution.") -@click.option("--frame", "-f", required=True, type=click.STRING, - help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") -@click.option("--path", "-p", default="./", type=click.STRING, +@click.option("--frame", "-f", required=True, type=click.STRING, + help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") +@click.option("--path", "-p", default="./", type=click.STRING, help="Path to simulation data.") -@click.option("--tag", "-t", default="df", type=click.STRING, +@click.option("--tag", "-t", default="df", type=click.STRING, help="Tag for output dataset.") -@click.option("--c2p-vel", "-v", is_flag=True, default=False, +@click.option("--source", is_flag=True, + help="Use -_source_.gkyl as the input distribution.") +@click.option("--c2p-vel", "-v", is_flag=True, default=False, help="Use -_mapc2p_vel.gkyl when loading Jf.") -@click.option("--mc2nu", "-m", is_flag=True, - help="Use -mc2nu_pos_deflated.gkyl to deform configuration-space grid.") -@click.option("--multi-block", "-b", is_flag=True, default=False, - help="Multiblock mode. Files are expected as {name}_b{N}-{species}_{frame}.gkyl.") -@click.option("--block", type=click.STRING, default=None, - help=("Restrict multiblock to specific blocks: comma-separated list or slice " - "(start:stop:step). Requires --multi-block. Defaults to all discovered blocks.")) -@click.option("--debug", is_flag=True, - help="Print resolved file names and shape diagnostics.") +@click.option("--mc2nu", "-m", is_flag=True, + help="Use -mc2nu_pos_deflated.gkyl to deform the configuration-space grid.") +@click.option("--block", "-b", default=None, type=click.INT, + help="Use block-specific files with _b prefix, e.g. -b 1 loads _b1-*.gkyl.") @click.pass_context def gk_distf(ctx, **kwargs): """Gyrokinetics: build real distribution function from saved Jf data.""" verb_print(ctx, "Starting gk_distf") data = ctx.obj["data"] - is_multib = kwargs["multi_block"] - mb_blocks_str = kwargs["block"] if kwargs["block"] is not None else "-1" - - # For frame-range resolution, use block 0 as the reference when multiblock. - ref_block = None - if is_multib: - if mb_blocks_str == "-1": - blocks = _discover_block_indices( - kwargs["name"], kwargs["species"], kwargs["path"], kwargs["source"]) - else: - blocks = gku.get_block_indices(mb_blocks_str, "") - if not blocks: - ctx.fail(click.style( - f"No block files found for {kwargs['name']} in {kwargs['path']}", fg="red")) - ref_block = blocks[0] - try: frames = _resolve_frames( - name=kwargs["name"], - species=kwargs["species"], - frame=kwargs["frame"], - path=kwargs["path"], - source=kwargs["source"], - block_idx=ref_block) + name=kwargs["name"], species=kwargs["species"], + frame=kwargs["frame"], path=kwargs["path"], + source=kwargs["source"], block_idx=kwargs["block"]) except (FileNotFoundError, ValueError) as err: ctx.fail(click.style(str(err), fg="red")) + # end - if kwargs["debug"] and len(frames) > 1: - click.echo(f"gk_distf: resolved frames={frames}") - + if len(frames) > 1: + verb_print(ctx, f"gk_distf: resolved frames={frames}") + # end + + common = dict( + name=kwargs["name"], species=kwargs["species"], + path=kwargs["path"], tag=kwargs["tag"], + source=kwargs["source"], use_c2p_vel=kwargs["c2p_vel"], + use_mc2nu=kwargs["mc2nu"], + block_idx=kwargs["block"], + ctx=ctx, + ) for frame in frames: - if is_multib: - out = load_gk_distf_mb( - name=kwargs["name"], - species=kwargs["species"], - frame=frame, - path=kwargs["path"], - tag=kwargs["tag"], - source=kwargs["source"], - use_c2p_vel=kwargs["c2p_vel"], - use_mc2nu=kwargs["mc2nu"], - debug=kwargs["debug"], - mb_blocks=mb_blocks_str) - else: - out = load_gk_distf( - name=kwargs["name"], - species=kwargs["species"], - frame=frame, - path=kwargs["path"], - tag=kwargs["tag"], - source=kwargs["source"], - use_c2p_vel=kwargs["c2p_vel"], - use_mc2nu=kwargs["mc2nu"], - debug=kwargs["debug"]) + out = load_gk_distf(frame=frame, **common) data.add(out) + # end if len(frames) > 1: data.set_unique_labels() + # end - verb_print(ctx, "Finishing gk_distf") \ No newline at end of file + verb_print(ctx, "Finishing gk_distf") +# end \ No newline at end of file From 0311199ec2a2251d255fe8b7461ff59bf0542a51 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sat, 14 Mar 2026 21:08:20 -0400 Subject: [PATCH 11/31] Enhance documentation and variable naming in gk_distf.py for clarity and consistency --- src/postgkyl/commands/gk_distf.py | 34 ++++++++++++++++--------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 799b9a8..6200baf 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -174,6 +174,7 @@ def _error_check(check: str, **kwargs): def _resolve_available_frames( name: str, species: str, path: str, source: bool = False, block_idx: int | None = None,) -> list[int]: + """Scan the path for available frame files matching the naming pattern.""" prefix_name = f"{name}_b{block_idx}" if block_idx is not None else name frame_infix = "source_" if source else "" prefix = _resolve_path(path, f"{prefix_name}-{species}_{frame_infix}") @@ -197,6 +198,7 @@ def _resolve_frames( name: str, species: str, frame: str, path: str, source: bool = False, block_idx: int | None = None, ) -> list[int]: + """Parse frame input and resolve to a list of available frame numbers.""" frame = frame.strip() _error_check("frame_nonempty", frame=frame) @@ -207,8 +209,8 @@ def _resolve_frames( available = _resolve_available_frames(name, species, path, source, block_idx=block_idx) _error_check("available_frames", available=available, name=name, species=species, path=path) - _error_check("frame_bounds", frame=frame) + parts = frame.split(":") start = int(parts[0]) if parts[0] else None stop = int(parts[1]) if parts[1] else None @@ -229,14 +231,14 @@ def _resolve_frames( # --------------------------------------------------------------------------- def _compute_fjx(jf_values: np.ndarray, jac_values: np.ndarray) -> np.ndarray: - """Divide Jf by jacobvel to get f * J_x.""" + """Divide Jf by jacobvel to get f * J_x * B. Could be improved with weak division""" _error_check("fjx_shape", jf_shape=jf_values.shape, jac_shape=jac_values.shape) return jf_values / jac_values # end def _broadcast_multiply(distf: np.ndarray, jacobtot_inv: np.ndarray) -> np.ndarray: - """Multiply distf by jacobtot_inv, broadcasting over the component axis.""" + """Multiply f * J_x * B by jacobtot_inv, broadcasting over the component axis.""" _error_check("broadcast_shape", distf_shape=distf.shape, jac_shape=jacobtot_inv.shape) jacob_shape = jacobtot_inv.shape + (1,) * (distf.ndim - jacobtot_inv.ndim) return distf * jacobtot_inv.reshape(jacob_shape) @@ -246,7 +248,7 @@ def _broadcast_multiply(distf: np.ndarray, jacobtot_inv: np.ndarray) -> np.ndarr def _interpolate_distf_components( fjx_data: GData, jacobtot_inv_data: GData, ) -> tuple[list, np.ndarray, np.ndarray]: - """Interpolate f/J_x (gkhyb basis) and jacobtot_inv (ms basis).""" + """Interpolate f * J_x * B (gkhyb basis) and jacobtot_inv (ms basis).""" grid, fjx_values = GInterpModal(fjx_data, 1, "gkhyb").interpolate() _, jacob_values = GInterpModal(jacobtot_inv_data, 1, "ms").interpolate() return grid, np.squeeze(fjx_values), np.squeeze(jacob_values) @@ -317,7 +319,7 @@ def load_gk_distf( block_idx: int | None = None, ctx=None, ) -> GData: - """Build a real distribution function from saved Jf data.""" + """Build a real distribution function from saved JBf data.""" jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _resolve_files( name, species, frame, path, source, block_idx=block_idx) @@ -338,10 +340,10 @@ def load_gk_distf( jacobvel_data = GData(jacobvel_file) jacobtot_inv_data = GData(jacobtot_inv_file) - fjx_data = GData(tag=tag, ctx=jf_data.ctx) - fjx_data.push(jf_data.get_grid(), _compute_fjx(jf_data.get_values(), jacobvel_data.get_values())) + fjxB_data = GData(tag=tag, ctx=jf_data.ctx) + fjxB_data.push(jf_data.get_grid(), _compute_fjx(jf_data.get_values(), jacobvel_data.get_values())) - out_grid, fjx_values, jacob_values = _interpolate_distf_components(fjx_data, jacobtot_inv_data) + out_grid, fjx_values, jacob_values = _interpolate_distf_components(fjxB_data, jacobtot_inv_data) distf_values = _broadcast_multiply(fjx_values, jacob_values) if use_mc2nu: @@ -361,21 +363,21 @@ def load_gk_distf( # --------------------------------------------------------------------------- @click.command() -@click.option("--name", "-n", required=True, type=click.STRING, +@click.option("--name", "-n", required=True, type=click.STRING, help="Simulation name prefix (e.g. gk_lorentzian_mirror).") -@click.option("--species", "-s", required=True, type=click.STRING, +@click.option("--species", "-s", required=True, type=click.STRING, help="Species name (e.g. ion or elc).") -@click.option("--frame", "-f", required=True, type=click.STRING, +@click.option("--frame", "-f", required=True, type=click.STRING, help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") -@click.option("--path", "-p", default="./", type=click.STRING, +@click.option("--path", "-p", default="./", type=click.STRING, help="Path to simulation data.") -@click.option("--tag", "-t", default="df", type=click.STRING, +@click.option("--tag", "-t", default="df", type=click.STRING, help="Tag for output dataset.") -@click.option("--source", is_flag=True, +@click.option("--source", is_flag=True, help="Use -_source_.gkyl as the input distribution.") -@click.option("--c2p-vel", "-v", is_flag=True, default=False, +@click.option("--c2p-vel", "-v", is_flag=True, default=False, help="Use -_mapc2p_vel.gkyl when loading Jf.") -@click.option("--mc2nu", "-m", is_flag=True, +@click.option("--mc2nu", "-m", is_flag=True, help="Use -mc2nu_pos_deflated.gkyl to deform the configuration-space grid.") @click.option("--block", "-b", default=None, type=click.INT, help="Use block-specific files with _b prefix, e.g. -b 1 loads _b1-*.gkyl.") From b52f28559f33eb8af98bf1cb0de02b8c41a8f78d Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 13:38:40 -0400 Subject: [PATCH 12/31] Refactor file handling in gk_distf.py: streamline file resolution and error checking, enhance frame parsing, and improve code organization for clarity. --- src/postgkyl/commands/gk_distf.py | 331 ++++++------------------------ 1 file changed, 66 insertions(+), 265 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 6200baf..1de7d68 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -1,215 +1,69 @@ import glob -import os import click import numpy as np from postgkyl.data import GData, GInterpModal -from postgkyl.utils import verb_print -# --------------------------------------------------------------------------- -# File/path helpers -# --------------------------------------------------------------------------- - -def _resolve_path(path: str, file_name: str) -> str: - if os.path.isabs(file_name): - return file_name - # end - return os.path.join(path, file_name) -# end - - -def _resolve_files( - name: str, species: str, frame: int, path: str, +def _parse_file_names( + name: str, species: str, frame: int, source: bool = False, block_idx: int | None = None, + use_c2p_vel: bool = False, use_mc2nu: bool = False ) -> tuple[str, str, str, str, str]: + "Find all relevant files for the given frame, checking for existence and returning paths." prefix = f"{name}_b{block_idx}" if block_idx is not None else name frame_infix = "source_" if source else "" - frame_file = _resolve_path(path, f"{prefix}-{species}_{frame_infix}{frame}.gkyl") - mapc2p_vel_file = _resolve_path(path, f"{prefix}-{species}_mapc2p_vel.gkyl") - mc2nu_file = _resolve_path(path, f"{prefix}-mc2nu_pos_deflated.gkyl") - jacobvel_file = _resolve_path(path, f"{prefix}-{species}_jacobvel.gkyl") - jacobtot_inv_file = _resolve_path(path, f"{prefix}-jacobtot_inv.gkyl") - return frame_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file -# end - - -def _error_check(check: str, **kwargs): - """Centralized validation for gk_distf helpers. - - This keeps error checks in one place while helper functions focus on data flow. - """ - if check == "frame_nonempty": - if not kwargs["frame"]: - raise ValueError("Frame cannot be empty.") + frame_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" + mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" + mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" + jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" + jacobtot_inv_file = f"{prefix}-jacobtot_inv.gkyl" + + # Check all files exist. + for file in [frame_file, jacobvel_file, jacobtot_inv_file]: + if not glob.glob(file): + raise FileNotFoundError(f"Required file not found: {file}") # end # end - - if check == "single_frame": - try: - int(kwargs["frame"]) - except ValueError as err: - raise ValueError( - "Frame must be an integer or a slice like ':', '40:50', or '40:50:2'." - ) from err - # end - # end - - if check == "frame_bounds": - parts = kwargs["frame"].split(":") - if len(parts) < 2 or len(parts) > 3: - raise ValueError( - "Frame range must use the form 'start:stop[:step]' or ':' for all frames." - ) - # end - - try: - int(parts[0]) if parts[0] else None - int(parts[1]) if parts[1] else None - step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 - except ValueError as err: - raise ValueError( - "Frame range must use integer bounds, e.g. ':', '40:50', or '40:50:2'." - ) from err - # end - - if step <= 0: - raise ValueError("Frame range step must be a positive integer.") - # end - # end - - if check == "available_frames": - if not kwargs["available"]: - raise FileNotFoundError( - f"No frame files found for {kwargs['name']}-{kwargs['species']} in {kwargs['path']}." - ) - # end - # end - - if check == "matched_frames": - if not kwargs["frames"]: - raise FileNotFoundError( - f"No frames matched '{kwargs['frame']}' for {kwargs['name']}-{kwargs['species']} in {kwargs['path']}." - ) - # end - # end - - if check == "required_files": - missing = [ - f"{label}: {file_name}" - for label, file_name in kwargs["files"].items() - if not os.path.exists(file_name) - ] - if missing: - raise FileNotFoundError("Missing required files:\n - " + "\n - ".join(missing)) - # end - # end - - if check == "fjx_shape": - jf_shape = kwargs["jf_shape"] - jac_shape = kwargs["jac_shape"] - if jf_shape != jac_shape: - if jf_shape[:-1] != jac_shape[:-1]: - raise ValueError( - "Jf and jacobvel spatial shapes do not match: " - f"{jf_shape} vs {jac_shape}" - ) - # end - if jac_shape[-1] != 1: - raise ValueError( - "jacobvel component axis must be 1 or match Jf: " - f"{jf_shape} vs {jac_shape}" - ) - # end - # end - # end - - if check == "broadcast_shape": - distf_shape = kwargs["distf_shape"] - jac_shape = kwargs["jac_shape"] - if len(jac_shape) > len(distf_shape): - raise ValueError( - "jacobtot_inv has more dimensions than distribution function: " - f"{len(jac_shape)} > {len(distf_shape)}" - ) - # end - - if distf_shape[:len(jac_shape)] != jac_shape: - raise ValueError( - "Cannot broadcast jacobtot_inv over distribution function. " - f"Expected prefix {distf_shape[:len(jac_shape)]}, got {jac_shape}" - ) - # end - # end - - if check == "mc2nu_cdim": - cdim = kwargs["cdim"] - total_dims = kwargs["total_dims"] - if cdim < 1 or cdim > total_dims: - raise ValueError(f"Invalid cdim={cdim} for total_dims={total_dims}") - # end - # end - - if check == "mc2nu_axis": - mapped_size = kwargs["mapped_size"] - old_size = kwargs["old_size"] - axis = kwargs["axis"] - if mapped_size != old_size and mapped_size + 1 != old_size: - raise ValueError( - f"mc2nu mapping size is incompatible with configuration grid on axis {axis}: " - f"{mapped_size} vs {old_size}" - ) + if use_c2p_vel: + if not glob.glob(mapc2p_vel_file): + raise FileNotFoundError(f"Required file not found: {mapc2p_vel_file}") # end # end - raise ValueError(f"Unknown precheck type '{check}'") -# end - - -# --------------------------------------------------------------------------- -# Frame resolution helpers -# --------------------------------------------------------------------------- - -def _resolve_available_frames( - name: str, species: str, path: str, - source: bool = False, block_idx: int | None = None,) -> list[int]: - """Scan the path for available frame files matching the naming pattern.""" - prefix_name = f"{name}_b{block_idx}" if block_idx is not None else name - frame_infix = "source_" if source else "" - prefix = _resolve_path(path, f"{prefix_name}-{species}_{frame_infix}") - - frames = [] - for file_name in glob.glob(f"{glob.escape(prefix)}*.gkyl"): - suffix = file_name.removeprefix(prefix) - if suffix.endswith(".gkyl"): - try: - frames.append(int(suffix[:-5])) - except ValueError: - pass - # end + if use_mc2nu: + if not glob.glob(mc2nu_file): + raise FileNotFoundError(f"Required file not found: {mc2nu_file}") # end # end - return sorted(set(frames)) + return frame_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file # end - -def _resolve_frames( - name: str, species: str, frame: str, path: str, +def _parse_frames_in_directory( + name: str, species: str, frame: str, source: bool = False, block_idx: int | None = None, ) -> list[int]: """Parse frame input and resolve to a list of available frame numbers.""" frame = frame.strip() - _error_check("frame_nonempty", frame=frame) + # If frame is a single number, return it as a list. if ":" not in frame: - _error_check("single_frame", frame=frame) return [int(frame)] # end - available = _resolve_available_frames(name, species, path, source, block_idx=block_idx) - _error_check("available_frames", available=available, name=name, species=species, path=path) - _error_check("frame_bounds", frame=frame) + prefix_name = f"{name}_b{block_idx}" if block_idx is not None else name + frame_infix = "source_" if source else "" + prefix = f"{prefix_name}-{species}_{frame_infix}" + + # Keep only files whose suffix is an integer frame number. + frames = [ + int(suffix[:-5]) + for file_name in glob.glob(f"{glob.escape(prefix)}*.gkyl") + if (suffix := file_name.removeprefix(prefix)).endswith(".gkyl") + and suffix[:-5].isdigit() + ] + available = sorted(set(frames)) parts = frame.split(":") start = int(parts[0]) if parts[0] else None @@ -222,43 +76,15 @@ def _resolve_frames( f for f in available if lower <= f < upper and (f - lower) % step == 0 ] - _error_check("matched_frames", frames=frames, frame=frame, name=name, species=species, path=path) return frames # end -# --------------------------------------------------------------------------- -# Physics / DG helpers -# --------------------------------------------------------------------------- - -def _compute_fjx(jf_values: np.ndarray, jac_values: np.ndarray) -> np.ndarray: - """Divide Jf by jacobvel to get f * J_x * B. Could be improved with weak division""" - _error_check("fjx_shape", jf_shape=jf_values.shape, jac_shape=jac_values.shape) - return jf_values / jac_values -# end - - -def _broadcast_multiply(distf: np.ndarray, jacobtot_inv: np.ndarray) -> np.ndarray: - """Multiply f * J_x * B by jacobtot_inv, broadcasting over the component axis.""" - _error_check("broadcast_shape", distf_shape=distf.shape, jac_shape=jacobtot_inv.shape) - jacob_shape = jacobtot_inv.shape + (1,) * (distf.ndim - jacobtot_inv.ndim) - return distf * jacobtot_inv.reshape(jacob_shape) -# end - - -def _interpolate_distf_components( - fjx_data: GData, jacobtot_inv_data: GData, -) -> tuple[list, np.ndarray, np.ndarray]: - """Interpolate f * J_x * B (gkhyb basis) and jacobtot_inv (ms basis).""" - grid, fjx_values = GInterpModal(fjx_data, 1, "gkhyb").interpolate() - _, jacob_values = GInterpModal(jacobtot_inv_data, 1, "ms").interpolate() - return grid, np.squeeze(fjx_values), np.squeeze(jacob_values) -# end - # --------------------------------------------------------------------------- # mc2nu grid deformation helpers # --------------------------------------------------------------------------- def _cell_centers_to_nodes(cell_centers: np.ndarray) -> np.ndarray: + """ Given an array of cell centers, return the corresponding node coordinates by extrapolating half a cell width at the boundaries.""" nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) nodes[0] = cell_centers[0] - (nodes[1] - cell_centers[0]) @@ -282,10 +108,9 @@ def _extract_mapped_axis(mapped_values: np.ndarray, axis: int, cdim: int) -> np. def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, ctx=None) -> list: - """Replace configuration-space axes in out_grid with mc2nu-mapped coordinates.""" + """Replace computational configuration-space grid with non-uniform spatial coordinates.""" mc2nu_data = GData(mc2nu_file) cdim = mc2nu_data.get_num_dims() - _error_check("mc2nu_cdim", cdim=cdim, total_dims=len(out_grid)) mc2nu_interp = GInterpModal(mc2nu_data, 1, "ms") _, mc2nu_values = mc2nu_interp.interpolate(tuple(range(cdim))) @@ -294,16 +119,9 @@ def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, ctx=None) -> list: deformed_grid = list(out_grid) for d in range(cdim): mapped_cfg = _extract_mapped_axis(mapped_values, d, cdim) - old_cfg = np.asarray(out_grid[d]) - _error_check("mc2nu_axis", mapped_size=mapped_cfg.size, old_size=old_cfg.size, axis=d) - if mapped_cfg.size == old_cfg.size: - deformed_grid[d] = mapped_cfg - else: - deformed_grid[d] = _cell_centers_to_nodes(mapped_cfg) + deformed_grid[d] = _cell_centers_to_nodes(mapped_cfg) # end # end - - verb_print(ctx, f"gk_distf: mc2nu mapped {cdim} configuration axis/axes") return deformed_grid # end @@ -313,46 +131,42 @@ def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, ctx=None) -> list: # --------------------------------------------------------------------------- def load_gk_distf( - name: str, species: str, frame: int, path: str = "./", + name: str, species: str, frame: int, tag: str = "df", source: bool = False, use_c2p_vel: bool = False, use_mc2nu: bool = False, block_idx: int | None = None, ctx=None, ) -> GData: """Build a real distribution function from saved JBf data.""" - jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _resolve_files( - name, species, frame, path, source, block_idx=block_idx) - - required_files = {"Jf": jf_file, "jacobvel": jacobvel_file, "jacobtot_inv": jacobtot_inv_file} - if use_c2p_vel: - required_files["mapc2p_vel"] = mapc2p_vel_file - # end - if use_mc2nu: - required_files["mc2nu"] = mc2nu_file - # end - _error_check("required_files", files=required_files) - - for kind, fname in required_files.items(): - verb_print(ctx, f"gk_distf: {kind}={fname}") - # end + jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _parse_file_names( + name, species, frame, source, block_idx=block_idx, use_c2p_vel=use_c2p_vel, use_mc2nu=use_mc2nu) - jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file if use_c2p_vel else "") + jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file) jacobvel_data = GData(jacobvel_file) jacobtot_inv_data = GData(jacobtot_inv_file) + # Divide Jf by jacobvel to get f * J_x * B + # Could be improved using weak division fjxB_data = GData(tag=tag, ctx=jf_data.ctx) - fjxB_data.push(jf_data.get_grid(), _compute_fjx(jf_data.get_values(), jacobvel_data.get_values())) + fjxB_values = jf_data.get_values() / jacobvel_data.get_values() + fjxB_data.push(jf_data.get_grid(), fjxB_values) - out_grid, fjx_values, jacob_values = _interpolate_distf_components(fjxB_data, jacobtot_inv_data) - distf_values = _broadcast_multiply(fjx_values, jacob_values) + # Interpolate f * J_x * B and jacobtot_inv to the same grid. + out_grid, fjxB_values = GInterpModal(fjxB_data, 1, "gkhyb").interpolate() + _, jacobtot_inv_values = GInterpModal(jacobtot_inv_data, 1, "ms").interpolate() + fjxB_values = np.squeeze(fjxB_values) + jacobtot_inv_values = np.squeeze(jacobtot_inv_values) + + # Reshape jacobtot_inv so that we can multiply f * J_x * B by 1/(J_x B) + vdim = fjxB_values.ndim - jacobtot_inv_values.ndim + jacob_shape = jacobtot_inv_values.shape + (1,) * vdim # Single array index in velocity directions + distf_values = fjxB_values * jacobtot_inv_values.reshape(jacob_shape) if use_mc2nu: out_grid = _apply_mc2nu_grid(out_grid, mc2nu_file, ctx) jf_data.ctx["grid_type"] = "mc2nu" # end - verb_print(ctx, f"gk_distf: output shape={distf_values.shape}") - out = GData(tag=tag, ctx=jf_data.ctx) out.push(out_grid, np.asarray(distf_values)[..., np.newaxis]) return out @@ -367,42 +181,31 @@ def load_gk_distf( help="Simulation name prefix (e.g. gk_lorentzian_mirror).") @click.option("--species", "-s", required=True, type=click.STRING, help="Species name (e.g. ion or elc).") -@click.option("--frame", "-f", required=True, type=click.STRING, - help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") -@click.option("--path", "-p", default="./", type=click.STRING, - help="Path to simulation data.") -@click.option("--tag", "-t", default="df", type=click.STRING, - help="Tag for output dataset.") @click.option("--source", is_flag=True, help="Use -_source_.gkyl as the input distribution.") +@click.option("--frame", "-f", required=True, type=click.STRING, + help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") @click.option("--c2p-vel", "-v", is_flag=True, default=False, help="Use -_mapc2p_vel.gkyl when loading Jf.") @click.option("--mc2nu", "-m", is_flag=True, help="Use -mc2nu_pos_deflated.gkyl to deform the configuration-space grid.") @click.option("--block", "-b", default=None, type=click.INT, help="Use block-specific files with _b prefix, e.g. -b 1 loads _b1-*.gkyl.") +@click.option("--tag", "-t", default="f", type=click.STRING, + help="Tag for output dataset.") @click.pass_context def gk_distf(ctx, **kwargs): """Gyrokinetics: build real distribution function from saved Jf data.""" - verb_print(ctx, "Starting gk_distf") data = ctx.obj["data"] - try: - frames = _resolve_frames( - name=kwargs["name"], species=kwargs["species"], - frame=kwargs["frame"], path=kwargs["path"], - source=kwargs["source"], block_idx=kwargs["block"]) - except (FileNotFoundError, ValueError) as err: - ctx.fail(click.style(str(err), fg="red")) - # end - - if len(frames) > 1: - verb_print(ctx, f"gk_distf: resolved frames={frames}") - # end + frames = _parse_frames_in_directory( + name=kwargs["name"], species=kwargs["species"], + frame=kwargs["frame"], + source=kwargs["source"], block_idx=kwargs["block"]) common = dict( name=kwargs["name"], species=kwargs["species"], - path=kwargs["path"], tag=kwargs["tag"], + tag=kwargs["tag"], source=kwargs["source"], use_c2p_vel=kwargs["c2p_vel"], use_mc2nu=kwargs["mc2nu"], block_idx=kwargs["block"], @@ -416,6 +219,4 @@ def gk_distf(ctx, **kwargs): if len(frames) > 1: data.set_unique_labels() # end - - verb_print(ctx, "Finishing gk_distf") # end \ No newline at end of file From e1d1ef346faf4b0a19941497058a103344018736 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 13:45:30 -0400 Subject: [PATCH 13/31] Refactor gk_distf.py: remove redundant file parsing functions, streamline frame resolution logic, and enhance grid application for clarity and efficiency. --- src/postgkyl/commands/gk_distf.py | 165 ++++++++++-------------------- 1 file changed, 54 insertions(+), 111 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 1de7d68..ee5bd7d 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -4,80 +4,7 @@ import numpy as np from postgkyl.data import GData, GInterpModal - - -def _parse_file_names( - name: str, species: str, frame: int, - source: bool = False, block_idx: int | None = None, - use_c2p_vel: bool = False, use_mc2nu: bool = False -) -> tuple[str, str, str, str, str]: - "Find all relevant files for the given frame, checking for existence and returning paths." - prefix = f"{name}_b{block_idx}" if block_idx is not None else name - frame_infix = "source_" if source else "" - frame_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" - mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" - mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" - jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" - jacobtot_inv_file = f"{prefix}-jacobtot_inv.gkyl" - - # Check all files exist. - for file in [frame_file, jacobvel_file, jacobtot_inv_file]: - if not glob.glob(file): - raise FileNotFoundError(f"Required file not found: {file}") - # end - # end - if use_c2p_vel: - if not glob.glob(mapc2p_vel_file): - raise FileNotFoundError(f"Required file not found: {mapc2p_vel_file}") - # end - # end - - if use_mc2nu: - if not glob.glob(mc2nu_file): - raise FileNotFoundError(f"Required file not found: {mc2nu_file}") - # end - # end - return frame_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file -# end - -def _parse_frames_in_directory( - name: str, species: str, frame: str, - source: bool = False, block_idx: int | None = None, -) -> list[int]: - """Parse frame input and resolve to a list of available frame numbers.""" - frame = frame.strip() - - # If frame is a single number, return it as a list. - if ":" not in frame: - return [int(frame)] - # end - - prefix_name = f"{name}_b{block_idx}" if block_idx is not None else name - frame_infix = "source_" if source else "" - prefix = f"{prefix_name}-{species}_{frame_infix}" - - # Keep only files whose suffix is an integer frame number. - frames = [ - int(suffix[:-5]) - for file_name in glob.glob(f"{glob.escape(prefix)}*.gkyl") - if (suffix := file_name.removeprefix(prefix)).endswith(".gkyl") - and suffix[:-5].isdigit() - ] - available = sorted(set(frames)) - - parts = frame.split(":") - start = int(parts[0]) if parts[0] else None - stop = int(parts[1]) if parts[1] else None - step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 - lower = available[0] if start is None else start - upper = available[-1] + 1 if stop is None else stop - - frames = [ - f for f in available - if lower <= f < upper and (f - lower) % step == 0 - ] - return frames -# end +from postgkyl.utils import verb_print # --------------------------------------------------------------------------- # mc2nu grid deformation helpers @@ -107,20 +34,19 @@ def _extract_mapped_axis(mapped_values: np.ndarray, axis: int, cdim: int) -> np. # end -def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, ctx=None) -> list: +def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str) -> list: """Replace computational configuration-space grid with non-uniform spatial coordinates.""" mc2nu_data = GData(mc2nu_file) cdim = mc2nu_data.get_num_dims() - mc2nu_interp = GInterpModal(mc2nu_data, 1, "ms") - _, mc2nu_values = mc2nu_interp.interpolate(tuple(range(cdim))) + _, mc2nu_values = GInterpModal(mc2nu_data, 1, "ms").interpolate(tuple(range(cdim))) mapped_values = np.asarray(mc2nu_values) deformed_grid = list(out_grid) for d in range(cdim): mapped_cfg = _extract_mapped_axis(mapped_values, d, cdim) - deformed_grid[d] = _cell_centers_to_nodes(mapped_cfg) - # end + old_cfg = np.asarray(out_grid[d]) + deformed_grid[d] = mapped_cfg if mapped_cfg.size == old_cfg.size else _cell_centers_to_nodes(mapped_cfg) # end return deformed_grid # end @@ -132,38 +58,41 @@ def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str, ctx=None) -> list: def load_gk_distf( name: str, species: str, frame: int, - tag: str = "df", source: bool = False, use_c2p_vel: bool = False, - use_mc2nu: bool = False, - block_idx: int | None = None, - ctx=None, + tag: str = "f", source: bool = False, use_c2p_vel: bool = False, + use_mc2nu: bool = False, block_idx: int | None = None, ) -> GData: """Build a real distribution function from saved JBf data.""" - jf_file, mapc2p_vel_file, mc2nu_file, jacobvel_file, jacobtot_inv_file = _parse_file_names( - name, species, frame, source, block_idx=block_idx, use_c2p_vel=use_c2p_vel, use_mc2nu=use_mc2nu) + prefix = f"{name}_b{block_idx}" if block_idx is not None else name + frame_infix = "source_" if source else "" + + jf_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" + mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" + mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" + jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" + jacobtot_inv_file = f"{prefix}-jacobtot_inv.gkyl" - jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file) - jacobvel_data = GData(jacobvel_file) + jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file if use_c2p_vel else None) + jacobvel_data = GData(jacobvel_file) jacobtot_inv_data = GData(jacobtot_inv_file) - # Divide Jf by jacobvel to get f * J_x * B - # Could be improved using weak division + # Divide Jf by jacobvel to get f * J_x * B. fjxB_data = GData(tag=tag, ctx=jf_data.ctx) fjxB_values = jf_data.get_values() / jacobvel_data.get_values() fjxB_data.push(jf_data.get_grid(), fjxB_values) # Interpolate f * J_x * B and jacobtot_inv to the same grid. - out_grid, fjxB_values = GInterpModal(fjxB_data, 1, "gkhyb").interpolate() - _, jacobtot_inv_values = GInterpModal(jacobtot_inv_data, 1, "ms").interpolate() - fjxB_values = np.squeeze(fjxB_values) - jacobtot_inv_values = np.squeeze(jacobtot_inv_values) + out_grid, fjxB_values = GInterpModal(fjxB_data, 1, "gkhyb").interpolate() + _, jacobtot_inv_values = GInterpModal(jacobtot_inv_data, 1, "ms").interpolate() + fjxB_values = np.squeeze(fjxB_values) + jacobtot_inv_values = np.squeeze(jacobtot_inv_values) - # Reshape jacobtot_inv so that we can multiply f * J_x * B by 1/(J_x B) + # Reshape jacobtot_inv to broadcast over velocity dimensions, then multiply. vdim = fjxB_values.ndim - jacobtot_inv_values.ndim - jacob_shape = jacobtot_inv_values.shape + (1,) * vdim # Single array index in velocity directions - distf_values = fjxB_values * jacobtot_inv_values.reshape(jacob_shape) + jacobtot_inv_reshaped = jacobtot_inv_values.reshape(jacobtot_inv_values.shape + (1,) * vdim) + distf_values = fjxB_values * jacobtot_inv_reshaped if use_mc2nu: - out_grid = _apply_mc2nu_grid(out_grid, mc2nu_file, ctx) + out_grid = _apply_mc2nu_grid(out_grid, mc2nu_file) jf_data.ctx["grid_type"] = "mc2nu" # end @@ -198,21 +127,35 @@ def gk_distf(ctx, **kwargs): """Gyrokinetics: build real distribution function from saved Jf data.""" data = ctx.obj["data"] - frames = _parse_frames_in_directory( - name=kwargs["name"], species=kwargs["species"], - frame=kwargs["frame"], - source=kwargs["source"], block_idx=kwargs["block"]) - - common = dict( - name=kwargs["name"], species=kwargs["species"], - tag=kwargs["tag"], - source=kwargs["source"], use_c2p_vel=kwargs["c2p_vel"], - use_mc2nu=kwargs["mc2nu"], - block_idx=kwargs["block"], - ctx=ctx, - ) + verb_print(ctx, "Building distribution function for " + kwargs["name"]) + + # Resolve list of frame numbers. + frame_spec = kwargs["frame"].strip() + if ":" not in frame_spec: + frames = [int(frame_spec)] + else: + prefix = f"{kwargs['name']}_b{kwargs['block']}" if kwargs["block"] is not None else kwargs["name"] + frame_infix = "source_" if kwargs["source"] else "" + stem = f"{prefix}-{kwargs['species']}_{frame_infix}" + available = sorted({ + int(f.removeprefix(stem)[:-5]) + for f in glob.glob(f"{glob.escape(stem)}*.gkyl") + if f.removeprefix(stem)[:-5].isdigit() + }) + parts = frame_spec.split(":") + lower = int(parts[0]) if parts[0] else available[0] + upper = int(parts[1]) if parts[1] else available[-1] + 1 + step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 + frames = [f for f in available if lower <= f < upper and (f - lower) % step == 0] + # end + for frame in frames: - out = load_gk_distf(frame=frame, **common) + out = load_gk_distf( + name=kwargs["name"], species=kwargs["species"], frame=frame, + tag=kwargs["tag"], source=kwargs["source"], + use_c2p_vel=kwargs["c2p_vel"], use_mc2nu=kwargs["mc2nu"], + block_idx=kwargs["block"], + ) data.add(out) # end From 6f825af332d0a3015650b24aec750095d0d4d96d Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 14:12:12 -0400 Subject: [PATCH 14/31] Refactor gk_distf.py: improve documentation by clarifying grid deformation helper functions and removing unnecessary comment separators for better readability. --- src/postgkyl/commands/gk_distf.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index ee5bd7d..b891b86 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -6,10 +6,9 @@ from postgkyl.data import GData, GInterpModal from postgkyl.utils import verb_print -# --------------------------------------------------------------------------- # mc2nu grid deformation helpers -# --------------------------------------------------------------------------- - +# This is a result of the gkyl_reader not having support for both mapc2p and mapc2p-vel grids. +# Particularly, the gkyl_reader, does not support mapping phase space arrays with mapc2p def _cell_centers_to_nodes(cell_centers: np.ndarray) -> np.ndarray: """ Given an array of cell centers, return the corresponding node coordinates by extrapolating half a cell width at the boundaries.""" nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) @@ -51,11 +50,7 @@ def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str) -> list: return deformed_grid # end - -# --------------------------------------------------------------------------- # Public API -# --------------------------------------------------------------------------- - def load_gk_distf( name: str, species: str, frame: int, tag: str = "f", source: bool = False, use_c2p_vel: bool = False, From 24200f5bc836b6ae1b2c58d38aa17dbabb05dbe0 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 14:19:58 -0400 Subject: [PATCH 15/31] Refactor gk_distf.py: remove unnecessary CLI command comment separators and add frame loading log for improved clarity --- src/postgkyl/commands/gk_distf.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index b891b86..661364d 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -96,10 +96,6 @@ def load_gk_distf( return out # end -# --------------------------------------------------------------------------- -# CLI command -# --------------------------------------------------------------------------- - @click.command() @click.option("--name", "-n", required=True, type=click.STRING, help="Simulation name prefix (e.g. gk_lorentzian_mirror).") @@ -143,6 +139,7 @@ def gk_distf(ctx, **kwargs): step = int(parts[2]) if len(parts) == 3 and parts[2] else 1 frames = [f for f in available if lower <= f < upper and (f - lower) % step == 0] # end + verb_print(ctx, f"Loading frames: {frames}") for frame in frames: out = load_gk_distf( From ecb53b0db0400b6b4c610f9afb7746f44964ce37 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 14:24:16 -0400 Subject: [PATCH 16/31] Fix load_gk_distf function: reorder mc2nu_file assignment for consistency in file naming --- src/postgkyl/commands/gk_distf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 661364d..6d2d36b 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -62,8 +62,8 @@ def load_gk_distf( jf_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" - mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" + mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" jacobtot_inv_file = f"{prefix}-jacobtot_inv.gkyl" jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file if use_c2p_vel else None) From d1520e61cd87069b4fe232ba7dc289a7d0005f1c Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 19:57:13 -0400 Subject: [PATCH 17/31] Simplify non-uniform map --- src/postgkyl/commands/gk_distf.py | 23 +++-------------------- 1 file changed, 3 insertions(+), 20 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 6d2d36b..ab283cf 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -9,7 +9,7 @@ # mc2nu grid deformation helpers # This is a result of the gkyl_reader not having support for both mapc2p and mapc2p-vel grids. # Particularly, the gkyl_reader, does not support mapping phase space arrays with mapc2p -def _cell_centers_to_nodes(cell_centers: np.ndarray) -> np.ndarray: +def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: """ Given an array of cell centers, return the corresponding node coordinates by extrapolating half a cell width at the boundaries.""" nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) @@ -18,34 +18,17 @@ def _cell_centers_to_nodes(cell_centers: np.ndarray) -> np.ndarray: return nodes # end - -def _extract_mapped_axis(mapped_values: np.ndarray, axis: int, cdim: int) -> np.ndarray: - """Extract a 1D mapped coordinate array for the given configuration axis.""" - if cdim == 1: - return np.asarray(mapped_values[..., axis]).reshape(-1) - # end - - # Pick a reference point in all other config directions. - idx = [0] * (cdim + 1) - idx[axis] = slice(None) - idx[-1] = axis - return np.asarray(mapped_values[tuple(idx)]).reshape(-1) -# end - - def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str) -> list: """Replace computational configuration-space grid with non-uniform spatial coordinates.""" mc2nu_data = GData(mc2nu_file) cdim = mc2nu_data.get_num_dims() _, mc2nu_values = GInterpModal(mc2nu_data, 1, "ms").interpolate(tuple(range(cdim))) - mapped_values = np.asarray(mc2nu_values) + mc2nu_values = np.squeeze(mc2nu_values) deformed_grid = list(out_grid) for d in range(cdim): - mapped_cfg = _extract_mapped_axis(mapped_values, d, cdim) - old_cfg = np.asarray(out_grid[d]) - deformed_grid[d] = mapped_cfg if mapped_cfg.size == old_cfg.size else _cell_centers_to_nodes(mapped_cfg) + deformed_grid[d] = _convert_cell_centered_to_nodal(mc2nu_values) # end return deformed_grid # end From 4a54724699daf859c521af4ec69cea6a1fc403f4 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 20:23:24 -0400 Subject: [PATCH 18/31] Simplify the non-uniform mapping part of this workflow --- src/postgkyl/commands/gk_distf.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index ab283cf..b409f4e 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -10,7 +10,8 @@ # This is a result of the gkyl_reader not having support for both mapc2p and mapc2p-vel grids. # Particularly, the gkyl_reader, does not support mapping phase space arrays with mapc2p def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: - """ Given an array of cell centers, return the corresponding node coordinates by extrapolating half a cell width at the boundaries.""" + """ Given an array defined at cell centers, return the corresponding nodal values + by interpolating half a cell width at the boundaries.""" nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) nodes[0] = cell_centers[0] - (nodes[1] - cell_centers[0]) @@ -18,19 +19,27 @@ def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: return nodes # end -def _apply_mc2nu_grid(out_grid: list, mc2nu_file: str) -> list: +def _extract_values_along_dimension(mapped_values: np.ndarray, axis: int, cdim: int) -> np.ndarray: + """Decompose mapped_values into a 1D array along the specified axis""" + idx = [0] * (cdim + 1) # Initilize index array. mc2nu has cdim+1 dimensions. + idx[axis] = slice(None) # Define a slice along the desired axis. + idx[-1] = axis # Select the appropriate component of mc2nu + return mapped_values[tuple(idx)].reshape(-1) # Apply indices and flatten to 1D. +# end + +def _apply_mc2nu_grid(uniform_grid: list, mc2nu_file: str) -> list: """Replace computational configuration-space grid with non-uniform spatial coordinates.""" mc2nu_data = GData(mc2nu_file) cdim = mc2nu_data.get_num_dims() _, mc2nu_values = GInterpModal(mc2nu_data, 1, "ms").interpolate(tuple(range(cdim))) - mc2nu_values = np.squeeze(mc2nu_values) - deformed_grid = list(out_grid) + nonuniform_grid = list(uniform_grid) for d in range(cdim): - deformed_grid[d] = _convert_cell_centered_to_nodal(mc2nu_values) + mc2nu_single_axis = _extract_values_along_dimension(mc2nu_values, d, cdim) + nonuniform_grid[d] = _convert_cell_centered_to_nodal(mc2nu_single_axis) # end - return deformed_grid + return nonuniform_grid # end # Public API From 42c0eb882de083982f452b10165690f5e0320577 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 20:24:38 -0400 Subject: [PATCH 19/31] re-order the subtraction --- src/postgkyl/commands/gk_distf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index b409f4e..28b46e1 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -15,7 +15,7 @@ def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) nodes[0] = cell_centers[0] - (nodes[1] - cell_centers[0]) - nodes[-1] = cell_centers[-1] + (cell_centers[-1] - nodes[-2]) + nodes[-1] = cell_centers[-1] - (nodes[-2] - cell_centers[-1]) return nodes # end From e9d9630d6ae315a9f06da934738c303b1c0c9c35 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 20:40:46 -0400 Subject: [PATCH 20/31] Improve clarity in the cell centered to nodal operation. Remove tag in the loading of the data that didn't need to be there. Clarify why the jacobian needed to be resnaped. Also add a reshaping to the distribution function because it has 1 component after being interpolated --- src/postgkyl/commands/gk_distf.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 28b46e1..b9dcd82 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -8,14 +8,14 @@ # mc2nu grid deformation helpers # This is a result of the gkyl_reader not having support for both mapc2p and mapc2p-vel grids. -# Particularly, the gkyl_reader, does not support mapping phase space arrays with mapc2p +# Particularly, the gkyl_reader does not support mapping phase space arrays with mapc2p def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: """ Given an array defined at cell centers, return the corresponding nodal values by interpolating half a cell width at the boundaries.""" nodes = np.zeros(cell_centers.size + 1, dtype=cell_centers.dtype) nodes[1:-1] = 0.5 * (cell_centers[:-1] + cell_centers[1:]) - nodes[0] = cell_centers[0] - (nodes[1] - cell_centers[0]) - nodes[-1] = cell_centers[-1] - (nodes[-2] - cell_centers[-1]) + nodes[0] = cell_centers[0] + (cell_centers[0] - nodes[1]) # Cell center plus half a cell width + nodes[-1] = cell_centers[-1] + (cell_centers[-1] - nodes[-2]) # Cell center plus half a cell width return nodes # end @@ -63,7 +63,7 @@ def load_gk_distf( jacobtot_inv_data = GData(jacobtot_inv_file) # Divide Jf by jacobvel to get f * J_x * B. - fjxB_data = GData(tag=tag, ctx=jf_data.ctx) + fjxB_data = GData(ctx=jf_data.ctx) fjxB_values = jf_data.get_values() / jacobvel_data.get_values() fjxB_data.push(jf_data.get_grid(), fjxB_values) @@ -73,18 +73,24 @@ def load_gk_distf( fjxB_values = np.squeeze(fjxB_values) jacobtot_inv_values = np.squeeze(jacobtot_inv_values) - # Reshape jacobtot_inv to broadcast over velocity dimensions, then multiply. + # Reshape jacobtot_inv to have 1 component over velocity dimensions, then multiply. vdim = fjxB_values.ndim - jacobtot_inv_values.ndim jacobtot_inv_reshaped = jacobtot_inv_values.reshape(jacobtot_inv_values.shape + (1,) * vdim) - distf_values = fjxB_values * jacobtot_inv_reshaped + f_values = fjxB_values * jacobtot_inv_reshaped + # Add 1 dimension to represent 1 component + f_values = f_values.reshape(f_values.shape + (1,)) if use_mc2nu: out_grid = _apply_mc2nu_grid(out_grid, mc2nu_file) - jf_data.ctx["grid_type"] = "mc2nu" + if use_c2p_vel: + jf_data.ctx["grid_type"] = "c2p_vel + mc2nu" + else: + jf_data.ctx["grid_type"] = "mc2nu" + # end # end out = GData(tag=tag, ctx=jf_data.ctx) - out.push(out_grid, np.asarray(distf_values)[..., np.newaxis]) + out.push(out_grid, f_values) return out # end From 0763f7156502fd0eb3cd1bd15203af91a6d4f095 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 20:42:57 -0400 Subject: [PATCH 21/31] Fix typo in comment for index initialization in _extract_values_along_dimension function --- src/postgkyl/commands/gk_distf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index b9dcd82..afcfd02 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -21,7 +21,7 @@ def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: def _extract_values_along_dimension(mapped_values: np.ndarray, axis: int, cdim: int) -> np.ndarray: """Decompose mapped_values into a 1D array along the specified axis""" - idx = [0] * (cdim + 1) # Initilize index array. mc2nu has cdim+1 dimensions. + idx = [0] * (cdim + 1) # Initialize indexing array. mc2nu has cdim+1 dimensions. idx[axis] = slice(None) # Define a slice along the desired axis. idx[-1] = axis # Select the appropriate component of mc2nu return mapped_values[tuple(idx)].reshape(-1) # Apply indices and flatten to 1D. From eda9e402db93d312aea5a3fbe1f2c47e70f3b90b Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 20:45:09 -0400 Subject: [PATCH 22/31] Add clarification comment for GData object in load_gk_distf function --- src/postgkyl/commands/gk_distf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index afcfd02..1a3a1a0 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -63,7 +63,7 @@ def load_gk_distf( jacobtot_inv_data = GData(jacobtot_inv_file) # Divide Jf by jacobvel to get f * J_x * B. - fjxB_data = GData(ctx=jf_data.ctx) + fjxB_data = GData(ctx=jf_data.ctx) # Inside a GData object so we can interpolate fjxB_values = jf_data.get_values() / jacobvel_data.get_values() fjxB_data.push(jf_data.get_grid(), fjxB_values) From efc49bc859b6aa91dbbdd3cfedaf07678194e2b6 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 20:50:46 -0400 Subject: [PATCH 23/31] Add two clarifying comments --- src/postgkyl/commands/gk_distf.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 1a3a1a0..11386c5 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -118,11 +118,11 @@ def gk_distf(ctx, **kwargs): verb_print(ctx, "Building distribution function for " + kwargs["name"]) - # Resolve list of frame numbers. frame_spec = kwargs["frame"].strip() if ":" not in frame_spec: - frames = [int(frame_spec)] + frames = [int(frame_spec)] # Stick to the frame specified on input else: + # Figure out how many frames are possible to read based on what files are available prefix = f"{kwargs['name']}_b{kwargs['block']}" if kwargs["block"] is not None else kwargs["name"] frame_infix = "source_" if kwargs["source"] else "" stem = f"{prefix}-{kwargs['species']}_{frame_infix}" @@ -131,6 +131,7 @@ def gk_distf(ctx, **kwargs): for f in glob.glob(f"{glob.escape(stem)}*.gkyl") if f.removeprefix(stem)[:-5].isdigit() }) + # Slice the data accordingly parts = frame_spec.split(":") lower = int(parts[0]) if parts[0] else available[0] upper = int(parts[1]) if parts[1] else available[-1] + 1 From 96b97a83d781b643845fd02bee07b3879781559e Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Sun, 15 Mar 2026 20:53:43 -0400 Subject: [PATCH 24/31] Set a default for the non-uniform map. I didn't realize that this didn't have a default --- src/postgkyl/commands/gk_distf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 11386c5..77759f6 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -105,7 +105,7 @@ def load_gk_distf( help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") @click.option("--c2p-vel", "-v", is_flag=True, default=False, help="Use -_mapc2p_vel.gkyl when loading Jf.") -@click.option("--mc2nu", "-m", is_flag=True, +@click.option("--mc2nu", "-m", is_flag=True, default=False, help="Use -mc2nu_pos_deflated.gkyl to deform the configuration-space grid.") @click.option("--block", "-b", default=None, type=click.INT, help="Use block-specific files with _b prefix, e.g. -b 1 loads _b1-*.gkyl.") From ca5bc189b480e46e14a8f27c84ebec051829e203 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Mon, 16 Mar 2026 13:36:02 -0400 Subject: [PATCH 25/31] Add a comma separated list of frames to load --- src/postgkyl/commands/gk_distf.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 77759f6..c8ef266 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -121,6 +121,8 @@ def gk_distf(ctx, **kwargs): frame_spec = kwargs["frame"].strip() if ":" not in frame_spec: frames = [int(frame_spec)] # Stick to the frame specified on input + elif "," in frame_spec: + frames = [int(f.strip()) for f in frame_spec.split(",")] # List of frames specified on input else: # Figure out how many frames are possible to read based on what files are available prefix = f"{kwargs['name']}_b{kwargs['block']}" if kwargs["block"] is not None else kwargs["name"] From dbf77d83a417168498bddca30f6d56eea7f2a2d1 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Mon, 16 Mar 2026 13:43:23 -0400 Subject: [PATCH 26/31] Fix typo --- src/postgkyl/commands/gk_distf.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index c8ef266..99179da 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -119,10 +119,10 @@ def gk_distf(ctx, **kwargs): verb_print(ctx, "Building distribution function for " + kwargs["name"]) frame_spec = kwargs["frame"].strip() - if ":" not in frame_spec: - frames = [int(frame_spec)] # Stick to the frame specified on input - elif "," in frame_spec: + if "," in frame_spec: frames = [int(f.strip()) for f in frame_spec.split(",")] # List of frames specified on input + elif ":" not in frame_spec: + frames = [int(frame_spec)] # Stick to the frame specified on input else: # Figure out how many frames are possible to read based on what files are available prefix = f"{kwargs['name']}_b{kwargs['block']}" if kwargs["block"] is not None else kwargs["name"] From 7e950c3645714df615508c41bd918048bb8571b0 Mon Sep 17 00:00:00 2001 From: Maxwell-Rosen Date: Mon, 16 Mar 2026 16:47:29 -0400 Subject: [PATCH 27/31] Add comments to describe the LLM involvement in this PR --- src/postgkyl/commands/gk_distf.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 99179da..0c32a58 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -1,3 +1,5 @@ +# MR was heavily inspired by LLMs (copilot) in writing this file. Highly specific and detailed prompts were used, with several iterations of refactors. Copilot inserted several unnecessary error checks because it didn't understand the assumptions we can make about the data. It was incredibly helpful in checking. I had to remove a lot of functions it used to make load_gk_distf more concise. + import glob import click @@ -9,6 +11,7 @@ # mc2nu grid deformation helpers # This is a result of the gkyl_reader not having support for both mapc2p and mapc2p-vel grids. # Particularly, the gkyl_reader does not support mapping phase space arrays with mapc2p +# Nearly 100% by LLMs, commented and verified by MR 3/16/26 def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: """ Given an array defined at cell centers, return the corresponding nodal values by interpolating half a cell width at the boundaries.""" @@ -19,6 +22,7 @@ def _convert_cell_centered_to_nodal(cell_centers: np.ndarray) -> np.ndarray: return nodes # end +# Nearly 100% by LLMs, commented and verified by MR 3/16/26 def _extract_values_along_dimension(mapped_values: np.ndarray, axis: int, cdim: int) -> np.ndarray: """Decompose mapped_values into a 1D array along the specified axis""" idx = [0] * (cdim + 1) # Initialize indexing array. mc2nu has cdim+1 dimensions. @@ -27,6 +31,7 @@ def _extract_values_along_dimension(mapped_values: np.ndarray, axis: int, cdim: return mapped_values[tuple(idx)].reshape(-1) # Apply indices and flatten to 1D. # end +# Nearly 100% by LLMs, commented and verified by MR 3/16/26, removing extra code. def _apply_mc2nu_grid(uniform_grid: list, mc2nu_file: str) -> list: """Replace computational configuration-space grid with non-uniform spatial coordinates.""" mc2nu_data = GData(mc2nu_file) @@ -49,6 +54,7 @@ def load_gk_distf( use_mc2nu: bool = False, block_idx: int | None = None, ) -> GData: """Build a real distribution function from saved JBf data.""" + # Mostly by LLMs, but heavily refactored and verified by MR 3/16/26 prefix = f"{name}_b{block_idx}" if block_idx is not None else name frame_infix = "source_" if source else "" @@ -94,6 +100,7 @@ def load_gk_distf( return out # end +# Generated by LLMs, commented and verified by MR 3/16/26 @click.command() @click.option("--name", "-n", required=True, type=click.STRING, help="Simulation name prefix (e.g. gk_lorentzian_mirror).") @@ -102,7 +109,8 @@ def load_gk_distf( @click.option("--source", is_flag=True, help="Use -_source_.gkyl as the input distribution.") @click.option("--frame", "-f", required=True, type=click.STRING, - help="Frame number or range. Use ':' for all frames and 'start:stop[:step]' for ranges.") + help="Frame number, comma separated values, or range. Use ':' for all frames\n" + " and 'start:stop[:step]' for ranges.") @click.option("--c2p-vel", "-v", is_flag=True, default=False, help="Use -_mapc2p_vel.gkyl when loading Jf.") @click.option("--mc2nu", "-m", is_flag=True, default=False, @@ -125,6 +133,7 @@ def gk_distf(ctx, **kwargs): frames = [int(frame_spec)] # Stick to the frame specified on input else: # Figure out how many frames are possible to read based on what files are available + # Generated by LLMs prefix = f"{kwargs['name']}_b{kwargs['block']}" if kwargs["block"] is not None else kwargs["name"] frame_infix = "source_" if kwargs["source"] else "" stem = f"{prefix}-{kwargs['species']}_{frame_infix}" From 329a75d4bd6c0791221b04a406dae27b26dc9e8c Mon Sep 17 00:00:00 2001 From: mrquell Date: Thu, 26 Mar 2026 13:34:32 -0400 Subject: [PATCH 28/31] Remove changes to gdata. Add a generalized suffix to the species so you can do --suffix source instead of --source, which is more general --- src/postgkyl/commands/gk_distf.py | 12 ++++++------ src/postgkyl/data/gdata.py | 5 ----- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 0c32a58..f1adb33 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -50,13 +50,13 @@ def _apply_mc2nu_grid(uniform_grid: list, mc2nu_file: str) -> list: # Public API def load_gk_distf( name: str, species: str, frame: int, - tag: str = "f", source: bool = False, use_c2p_vel: bool = False, + tag: str = "f", suffix: str = "", use_c2p_vel: bool = False, use_mc2nu: bool = False, block_idx: int | None = None, ) -> GData: """Build a real distribution function from saved JBf data.""" # Mostly by LLMs, but heavily refactored and verified by MR 3/16/26 prefix = f"{name}_b{block_idx}" if block_idx is not None else name - frame_infix = "source_" if source else "" + frame_infix = f"{suffix}_" if suffix else "" jf_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" @@ -106,8 +106,8 @@ def load_gk_distf( help="Simulation name prefix (e.g. gk_lorentzian_mirror).") @click.option("--species", "-s", required=True, type=click.STRING, help="Species name (e.g. ion or elc).") -@click.option("--source", is_flag=True, - help="Use -_source_.gkyl as the input distribution.") +@click.option("--suffix", default="", type=click.STRING, + help="Use -__.gkyl as the input distribution.") @click.option("--frame", "-f", required=True, type=click.STRING, help="Frame number, comma separated values, or range. Use ':' for all frames\n" " and 'start:stop[:step]' for ranges.") @@ -135,7 +135,7 @@ def gk_distf(ctx, **kwargs): # Figure out how many frames are possible to read based on what files are available # Generated by LLMs prefix = f"{kwargs['name']}_b{kwargs['block']}" if kwargs["block"] is not None else kwargs["name"] - frame_infix = "source_" if kwargs["source"] else "" + frame_infix = f"{kwargs['suffix']}_" if kwargs["suffix"] else "" stem = f"{prefix}-{kwargs['species']}_{frame_infix}" available = sorted({ int(f.removeprefix(stem)[:-5]) @@ -154,7 +154,7 @@ def gk_distf(ctx, **kwargs): for frame in frames: out = load_gk_distf( name=kwargs["name"], species=kwargs["species"], frame=frame, - tag=kwargs["tag"], source=kwargs["source"], + tag=kwargs["tag"], suffix=kwargs["suffix"], use_c2p_vel=kwargs["c2p_vel"], use_mc2nu=kwargs["mc2nu"], block_idx=kwargs["block"], ) diff --git a/src/postgkyl/data/gdata.py b/src/postgkyl/data/gdata.py index f2fe60d..ac96eb7 100644 --- a/src/postgkyl/data/gdata.py +++ b/src/postgkyl/data/gdata.py @@ -268,11 +268,6 @@ def set_grid(self, grid: list) -> None: up[d] = self._grid[d].max() self.ctx["lower"] = lo self.ctx["upper"] = up - if "grid_type" not in self.ctx: - if all(len(np.asarray(coord).shape) == 1 for coord in self._grid): - self.ctx["grid_type"] = "uniform" - else: - self.ctx["grid_type"] = "mapped" grid = property(get_grid, set_grid) From 57254794b84677b3188d860577f393ff2c460f4e Mon Sep 17 00:00:00 2001 From: mrquell Date: Thu, 26 Mar 2026 13:43:11 -0400 Subject: [PATCH 29/31] Enhance load_gk_distf function with optional filename overrides and update usage example --- src/postgkyl/commands/gk_distf.py | 53 +++++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 6 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index f1adb33..aa34eb7 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -1,5 +1,16 @@ # MR was heavily inspired by LLMs (copilot) in writing this file. Highly specific and detailed prompts were used, with several iterations of refactors. Copilot inserted several unnecessary error checks because it didn't understand the assumptions we can make about the data. It was incredibly helpful in checking. I had to remove a lot of functions it used to make load_gk_distf more concise. +""" +Script example of usage in python + +distf = pg.commands.load_gk_distf( + name="gk_lorentzian_mirror", + species="ion", + frame=0) +pg.data.select(distf, z0=0.0, overwrite=True) +pg.output.plot(distf) +plt.show() +""" import glob import click @@ -52,17 +63,32 @@ def load_gk_distf( name: str, species: str, frame: int, tag: str = "f", suffix: str = "", use_c2p_vel: bool = False, use_mc2nu: bool = False, block_idx: int | None = None, + jf_file: str | None = None, + mapc2p_vel_file: str | None = None, + jacobvel_file: str | None = None, + mc2nu_file: str | None = None, + jacobtot_inv_file: str | None = None, ) -> GData: """Build a real distribution function from saved JBf data.""" # Mostly by LLMs, but heavily refactored and verified by MR 3/16/26 prefix = f"{name}_b{block_idx}" if block_idx is not None else name frame_infix = f"{suffix}_" if suffix else "" - jf_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" - mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" - jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" - mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" - jacobtot_inv_file = f"{prefix}-jacobtot_inv.gkyl" + if jf_file is None: + jf_file = f"{prefix}-{species}_{frame_infix}{frame}.gkyl" + # end + if mapc2p_vel_file is None: + mapc2p_vel_file = f"{prefix}-{species}_mapc2p_vel.gkyl" + # end + if jacobvel_file is None: + jacobvel_file = f"{prefix}-{species}_jacobvel.gkyl" + # end + if mc2nu_file is None: + mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" + # end + if jacobtot_inv_file is None: + jacobtot_inv_file = f"{prefix}-jacobtot_inv.gkyl" + # end jf_data = GData(jf_file, mapc2p_vel_name=mapc2p_vel_file if use_c2p_vel else None) jacobvel_data = GData(jacobvel_file) @@ -107,7 +133,17 @@ def load_gk_distf( @click.option("--species", "-s", required=True, type=click.STRING, help="Species name (e.g. ion or elc).") @click.option("--suffix", default="", type=click.STRING, - help="Use -__.gkyl as the input distribution.") + help="Use -__.gkyl as the input distribution.") +@click.option("--jf-file", default=None, type=click.STRING, + help="Literal Jf filename override. If omitted, the default naming convention is used.") +@click.option("--mapc2p-vel-file", default=None, type=click.STRING, + help="Literal mapc2p_vel filename override. If omitted, the default naming convention is used.") +@click.option("--jacobvel-file", default=None, type=click.STRING, + help="Literal jacobvel filename override. If omitted, the default naming convention is used.") +@click.option("--mc2nu-file", default=None, type=click.STRING, + help="Literal mc2nu filename override. If omitted, the default naming convention is used.") +@click.option("--jacobtot-inv-file", default=None, type=click.STRING, + help="Literal jacobtot_inv filename override. If omitted, the default naming convention is used.") @click.option("--frame", "-f", required=True, type=click.STRING, help="Frame number, comma separated values, or range. Use ':' for all frames\n" " and 'start:stop[:step]' for ranges.") @@ -157,6 +193,11 @@ def gk_distf(ctx, **kwargs): tag=kwargs["tag"], suffix=kwargs["suffix"], use_c2p_vel=kwargs["c2p_vel"], use_mc2nu=kwargs["mc2nu"], block_idx=kwargs["block"], + jf_file=kwargs["jf_file"], + mapc2p_vel_file=kwargs["mapc2p_vel_file"], + jacobvel_file=kwargs["jacobvel_file"], + mc2nu_file=kwargs["mc2nu_file"], + jacobtot_inv_file=kwargs["jacobtot_inv_file"], ) data.add(out) # end From e29c3f3bc5d27e793168baf5dbd37fb60e9deaf4 Mon Sep 17 00:00:00 2001 From: mrquell Date: Thu, 26 Mar 2026 13:47:23 -0400 Subject: [PATCH 30/31] Add support for optional mapc2p file in load_gk_distf function --- src/postgkyl/commands/gk_distf.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index aa34eb7..22e5a80 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -62,11 +62,12 @@ def _apply_mc2nu_grid(uniform_grid: list, mc2nu_file: str) -> list: def load_gk_distf( name: str, species: str, frame: int, tag: str = "f", suffix: str = "", use_c2p_vel: bool = False, - use_mc2nu: bool = False, block_idx: int | None = None, + use_mc2nu: bool = False, use_mapc2p: bool = False, block_idx: int | None = None, jf_file: str | None = None, mapc2p_vel_file: str | None = None, jacobvel_file: str | None = None, mc2nu_file: str | None = None, + mapc2p_file: str | None = None, jacobtot_inv_file: str | None = None, ) -> GData: """Build a real distribution function from saved JBf data.""" @@ -86,6 +87,9 @@ def load_gk_distf( if mc2nu_file is None: mc2nu_file = f"{prefix}-mc2nu_pos_deflated.gkyl" # end + if mapc2p_file is None: + mapc2p_file = f"{prefix}-mapc2p_deflated.gkyl" + # end if jacobtot_inv_file is None: jacobtot_inv_file = f"{prefix}-jacobtot_inv.gkyl" # end @@ -119,6 +123,13 @@ def load_gk_distf( else: jf_data.ctx["grid_type"] = "mc2nu" # end + elif use_mapc2p: + out_grid = _apply_mc2nu_grid(out_grid, mapc2p_file) + if use_c2p_vel: + jf_data.ctx["grid_type"] = "c2p_vel + mapc2p" + else: + jf_data.ctx["grid_type"] = "mapc2p" + # end # end out = GData(tag=tag, ctx=jf_data.ctx) @@ -142,6 +153,8 @@ def load_gk_distf( help="Literal jacobvel filename override. If omitted, the default naming convention is used.") @click.option("--mc2nu-file", default=None, type=click.STRING, help="Literal mc2nu filename override. If omitted, the default naming convention is used.") +@click.option("--mapc2p-file", default=None, type=click.STRING, + help="Literal mapc2p filename override. If omitted, the default naming convention is used.") @click.option("--jacobtot-inv-file", default=None, type=click.STRING, help="Literal jacobtot_inv filename override. If omitted, the default naming convention is used.") @click.option("--frame", "-f", required=True, type=click.STRING, @@ -151,6 +164,8 @@ def load_gk_distf( help="Use -_mapc2p_vel.gkyl when loading Jf.") @click.option("--mc2nu", "-m", is_flag=True, default=False, help="Use -mc2nu_pos_deflated.gkyl to deform the configuration-space grid.") +@click.option("--mapc2p", "-p", is_flag=True, default=False, + help="Use -mapc2p_deflated.gkyl to deform the configuration-space grid.") @click.option("--block", "-b", default=None, type=click.INT, help="Use block-specific files with _b prefix, e.g. -b 1 loads _b1-*.gkyl.") @click.option("--tag", "-t", default="f", type=click.STRING, @@ -191,12 +206,14 @@ def gk_distf(ctx, **kwargs): out = load_gk_distf( name=kwargs["name"], species=kwargs["species"], frame=frame, tag=kwargs["tag"], suffix=kwargs["suffix"], - use_c2p_vel=kwargs["c2p_vel"], use_mc2nu=kwargs["mc2nu"], + use_c2p_vel=kwargs["c2p_vel"], + use_mc2nu=kwargs["mc2nu"], use_mapc2p=kwargs["mapc2p"], block_idx=kwargs["block"], jf_file=kwargs["jf_file"], mapc2p_vel_file=kwargs["mapc2p_vel_file"], jacobvel_file=kwargs["jacobvel_file"], mc2nu_file=kwargs["mc2nu_file"], + mapc2p_file=kwargs["mapc2p_file"], jacobtot_inv_file=kwargs["jacobtot_inv_file"], ) data.add(out) From 951ee6c6f1c9b65f563ae6516cea39350ce24de3 Mon Sep 17 00:00:00 2001 From: mrquell Date: Thu, 26 Mar 2026 13:51:43 -0400 Subject: [PATCH 31/31] Remove unused import and clean up load_gk_distf function parameters --- src/postgkyl/__init__.py | 1 - src/postgkyl/commands/gk_distf.py | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/postgkyl/__init__.py b/src/postgkyl/__init__.py index 17c1475..4151d08 100644 --- a/src/postgkyl/__init__.py +++ b/src/postgkyl/__init__.py @@ -8,7 +8,6 @@ __version__ = "1.7.5" # import submodules -from postgkyl import commands from postgkyl import data from postgkyl import utils from postgkyl import tools diff --git a/src/postgkyl/commands/gk_distf.py b/src/postgkyl/commands/gk_distf.py index 22e5a80..325fce6 100644 --- a/src/postgkyl/commands/gk_distf.py +++ b/src/postgkyl/commands/gk_distf.py @@ -1,7 +1,10 @@ # MR was heavily inspired by LLMs (copilot) in writing this file. Highly specific and detailed prompts were used, with several iterations of refactors. Copilot inserted several unnecessary error checks because it didn't understand the assumptions we can make about the data. It was incredibly helpful in checking. I had to remove a lot of functions it used to make load_gk_distf more concise. """ -Script example of usage in python +# Script example of usage in python +import postgkyl as pg +from postgkyl.commands import load_gk_distf +import matplotlib.pyplot as plt distf = pg.commands.load_gk_distf( name="gk_lorentzian_mirror", @@ -153,8 +156,6 @@ def load_gk_distf( help="Literal jacobvel filename override. If omitted, the default naming convention is used.") @click.option("--mc2nu-file", default=None, type=click.STRING, help="Literal mc2nu filename override. If omitted, the default naming convention is used.") -@click.option("--mapc2p-file", default=None, type=click.STRING, - help="Literal mapc2p filename override. If omitted, the default naming convention is used.") @click.option("--jacobtot-inv-file", default=None, type=click.STRING, help="Literal jacobtot_inv filename override. If omitted, the default naming convention is used.") @click.option("--frame", "-f", required=True, type=click.STRING, @@ -213,7 +214,6 @@ def gk_distf(ctx, **kwargs): mapc2p_vel_file=kwargs["mapc2p_vel_file"], jacobvel_file=kwargs["jacobvel_file"], mc2nu_file=kwargs["mc2nu_file"], - mapc2p_file=kwargs["mapc2p_file"], jacobtot_inv_file=kwargs["jacobtot_inv_file"], ) data.add(out)