"""Read and write runtime summaries, diagnostics, restart state, and tables."""
from __future__ import annotations
from dataclasses import dataclass, replace
import json
from pathlib import Path
from typing import Any
from jax.typing import ArrayLike
import numpy as np
from gkx.artifacts.spectral_layout import (
_condense_kx,
_condense_ky,
_condense_kykx,
_dealiased_kx_count,
_dealiased_kx_indices,
_dealiased_ky_count,
_require_netcdf4,
_validate_netcdf_schema_version,
)
from gkx.core_ky_layout import (
FULL,
HALF,
KY_AXIS,
KyLayout,
is_half,
nyc_from_ny,
to_full,
to_half,
)
from gkx.diagnostics import (
ResolvedDiagnostics,
SimulationDiagnostics,
total_energy,
)
from gkx.workflows.runtime.diagnostic_arrays import (
timestep_cost_payload,
validate_finite_runtime_diagnostics,
)
from gkx.solvers_linear_implicit import implicit_solve_payload
from gkx.solvers_linear_krylov import eigen_status_payload
_RUNTIME_FIELD_NAMES = ("phi", "apar", "bpar")
def _artifact_base(path: Path) -> Path:
if path.suffix.lower() in {".json", ".csv", ".npy", ".npz"}:
return path.with_suffix("")
return path
def _is_netcdf_output_target(path: Path) -> bool:
suffixes = [suffix.lower() for suffix in path.suffixes]
return bool(suffixes and suffixes[-1] == ".nc")
def _netcdf_bundle_base(path: Path) -> Path:
name = path.name
for suffix in (".out.nc", ".big.nc", ".restart.nc"):
if name.lower().endswith(suffix):
return path.with_name(name[: -len(suffix)])
if path.suffix.lower() == ".nc":
return path.with_suffix("")
return path
def _ensure_parent(path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
def _flatten_series(series: np.ndarray) -> np.ndarray:
arr = np.asarray(series)
if arr.ndim == 1:
return arr
arr = arr.reshape(arr.shape[0], -1)
if arr.shape[1] == 1:
return arr[:, 0]
return np.mean(arr, axis=1)
def _write_json(path: Path, payload: dict[str, Any]) -> None:
_ensure_parent(path)
path.write_text(
json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
)
def _write_csv(path: Path, headers: list[str], cols: list[np.ndarray]) -> None:
_ensure_parent(path)
data_out = np.column_stack(cols)
np.savetxt(path, data_out, delimiter=",", header=",".join(headers), comments="")
def _write_state(base: Path, state: np.ndarray | None) -> Path | None:
if state is None:
return None
path = Path(f"{base}.state.npy")
_ensure_parent(path)
np.save(path, np.asarray(state))
return path
[docs]
def validate_finite_array(value: Any, *, label: str) -> None:
"""Raise if an optional artifact array contains NaN or infinite values."""
if value is None:
return
arr = np.asarray(value)
if arr.size == 0 or np.isfinite(arr).all():
return
raise RuntimeError(f"{label} contains non-finite values")
[docs]
def validate_finite_runtime_result(result: Any, *, label: str) -> None:
"""Validate nonlinear runtime result payloads before artifact writes."""
if result.diagnostics is not None:
validate_finite_runtime_diagnostics(result.diagnostics, label=label)
validate_finite_array(result.state, label=f"{label} state")
fields = result.fields
if fields is None:
return
for name in _RUNTIME_FIELD_NAMES:
validate_finite_array(getattr(fields, name, None), label=f"{label} {name}")
@dataclass(frozen=True)
class _LinearArtifactTargets:
base: Path
summary_path: Path
timeseries_path: Path
@dataclass(frozen=True)
class _LinearScanTargets:
base: Path
summary_path: Path
scan_path: Path
@dataclass(frozen=True)
class _LinearScanArrays:
ky: np.ndarray
gamma: np.ndarray
omega: np.ndarray
@dataclass(frozen=True)
class _QuasilinearSpectrumColumns:
ky: np.ndarray
mode_ky: np.ndarray
gamma: np.ndarray
omega: np.ndarray
kperp_eff2: np.ndarray
heat: np.ndarray
particle: np.ndarray
amp2: np.ndarray
saturated_heat: np.ndarray
saturated_particle: np.ndarray
[docs]
def write_quasilinear_artifacts(
out: str | Path, quasilinear: dict[str, Any]
) -> dict[str, str]:
"""Write quasilinear summary and species tables."""
out_path = Path(out)
base = _artifact_base(out_path)
summary_path = (
out_path
if out_path.suffix.lower() == ".json"
else Path(f"{base}.quasilinear.summary.json")
)
_write_json(summary_path, quasilinear)
paths = {"quasilinear_summary": str(summary_path)}
heat = np.asarray(quasilinear.get("heat_flux_weight_species", []), dtype=float)
particle = np.asarray(
quasilinear.get("particle_flux_weight_species", []), dtype=float
)
if heat.size or particle.size:
n = max(int(heat.size), int(particle.size))
heat_col = np.full(n, np.nan, dtype=float)
particle_col = np.full(n, np.nan, dtype=float)
if heat.size:
heat_col[: heat.size] = heat
if particle.size:
particle_col[: particle.size] = particle
sat_heat = np.full(n, np.nan, dtype=float)
sat_particle = np.full(n, np.nan, dtype=float)
sat_heat_raw = quasilinear.get("saturated_heat_flux_species")
sat_particle_raw = quasilinear.get("saturated_particle_flux_species")
if sat_heat_raw is not None:
sat = np.asarray(sat_heat_raw, dtype=float)
sat_heat[: sat.size] = sat
if sat_particle_raw is not None:
sat = np.asarray(sat_particle_raw, dtype=float)
sat_particle[: sat.size] = sat
species_path = Path(f"{base}.quasilinear_species.csv")
_write_csv(
species_path,
[
"species_index",
"heat_flux_weight",
"particle_flux_weight",
"saturated_heat_flux",
"saturated_particle_flux",
],
[
np.arange(n, dtype=float),
heat_col,
particle_col,
sat_heat,
sat_particle,
],
)
paths["quasilinear_species"] = str(species_path)
return paths
def _runtime_linear_scan_targets(out: str | Path) -> _LinearScanTargets:
out_path = Path(out)
base = _artifact_base(out_path)
summary_path = (
out_path if out_path.suffix.lower() == ".json" else Path(f"{base}.summary.json")
)
scan_path = (
out_path if out_path.suffix.lower() == ".csv" else Path(f"{base}.scan.csv")
)
return _LinearScanTargets(base=base, summary_path=summary_path, scan_path=scan_path)
def _runtime_linear_scan_arrays(result: Any) -> _LinearScanArrays:
return _LinearScanArrays(
ky=np.asarray(result.ky, dtype=float),
gamma=np.asarray(result.gamma, dtype=float),
omega=np.asarray(result.omega, dtype=float),
)
def _runtime_linear_scan_summary(
*,
arrays: _LinearScanArrays,
result: Any,
ql_payloads: tuple[Any, ...],
) -> dict[str, Any]:
summary = {
"kind": "linear_scan",
"n_ky": int(arrays.ky.size),
"ky_min": None if arrays.ky.size == 0 else float(np.min(arrays.ky)),
"ky_max": None if arrays.ky.size == 0 else float(np.max(arrays.ky)),
"has_quasilinear": bool(ql_payloads),
}
parallel = getattr(result, "parallel", None)
if isinstance(parallel, dict):
summary["parallel"] = parallel
warm_start = getattr(result, "warm_start", None)
if isinstance(warm_start, dict):
summary["warm_start"] = warm_start
return summary
def _payload_float(payload: Any, key: str) -> float:
return float(payload.get(key, np.nan))
def _payload_optional_float(payload: Any, key: str) -> float:
value = payload.get(key)
return np.nan if value is None else float(value)
def _quasilinear_scan_columns(
*,
ky: np.ndarray,
ql_payloads: tuple[Any, ...],
) -> _QuasilinearSpectrumColumns:
# The scan coordinate is the user-requested target ky. Individual linear
# payloads also carry the selected signed grid-mode ky, which can differ
# for linked-boundary layouts.
ql_ky = (
np.asarray(ky, dtype=float)
if len(ql_payloads) == int(ky.size)
else np.asarray([_payload_float(p, "ky") for p in ql_payloads], dtype=float)
)
return _QuasilinearSpectrumColumns(
ky=ql_ky,
mode_ky=np.asarray([_payload_float(p, "ky") for p in ql_payloads], dtype=float),
gamma=np.asarray(
[_payload_float(p, "gamma") for p in ql_payloads], dtype=float
),
omega=np.asarray(
[_payload_float(p, "omega") for p in ql_payloads], dtype=float
),
kperp_eff2=np.asarray(
[_payload_float(p, "kperp_eff2") for p in ql_payloads], dtype=float
),
heat=np.asarray(
[_payload_float(p, "heat_flux_weight_total") for p in ql_payloads],
dtype=float,
),
particle=np.asarray(
[_payload_float(p, "particle_flux_weight_total") for p in ql_payloads],
dtype=float,
),
amp2=np.asarray(
[_payload_optional_float(p, "amplitude2") for p in ql_payloads],
dtype=float,
),
saturated_heat=np.asarray(
[
_payload_optional_float(p, "saturated_heat_flux_total")
for p in ql_payloads
],
dtype=float,
),
saturated_particle=np.asarray(
[
_payload_optional_float(p, "saturated_particle_flux_total")
for p in ql_payloads
],
dtype=float,
),
)
def _write_quasilinear_scan_spectrum(
*,
base: Path,
ky: np.ndarray,
ql_payloads: tuple[Any, ...],
) -> str | None:
if not ql_payloads:
return None
ql_path = Path(f"{base}.quasilinear_spectrum.csv")
columns = _quasilinear_scan_columns(ky=ky, ql_payloads=ql_payloads)
_write_csv(
ql_path,
[
"ky",
"mode_ky",
"gamma",
"omega",
"kperp_eff2",
"heat_flux_weight_total",
"particle_flux_weight_total",
"amplitude2",
"saturated_heat_flux_total",
"saturated_particle_flux_total",
],
[
columns.ky,
columns.mode_ky,
columns.gamma,
columns.omega,
columns.kperp_eff2,
columns.heat,
columns.particle,
columns.amp2,
columns.saturated_heat,
columns.saturated_particle,
],
)
return str(ql_path)
def _runtime_linear_artifact_targets(out: str | Path) -> _LinearArtifactTargets:
out_path = Path(out)
base = _artifact_base(out_path)
summary_path = (
out_path if out_path.suffix.lower() == ".json" else Path(f"{base}.summary.json")
)
timeseries_path = (
out_path
if out_path.suffix.lower() == ".csv"
else Path(f"{base}.timeseries.csv")
)
return _LinearArtifactTargets(
base=base,
summary_path=summary_path,
timeseries_path=timeseries_path,
)
def _optional_float(value: Any) -> float | None:
"""JSON-safe float: ``None`` for missing or non-finite fit diagnostics."""
if value is None:
return None
number = float(value)
return number if np.isfinite(number) else None
def _runtime_linear_summary(result: Any) -> dict[str, Any]:
summary = {
"kind": "linear",
"ky": float(result.ky),
"gamma": float(result.gamma),
"omega": float(result.omega),
"fit_window_tmin": None
if result.fit_window_tmin is None
else float(result.fit_window_tmin),
"fit_window_tmax": None
if result.fit_window_tmax is None
else float(result.fit_window_tmax),
"fit_signal_used": result.fit_signal_used,
# Fit-quality diagnostics; absent (None) for eigensolver results.
"gamma_stderr": _optional_float(getattr(result, "gamma_stderr", None)),
"omega_stderr": _optional_float(getattr(result, "omega_stderr", None)),
"fit_r2": _optional_float(getattr(result, "fit_r2", None)),
"fit_settled": getattr(result, "fit_settled", None),
"selection": {
"ky_index": int(result.selection.ky_index),
"kx_index": int(result.selection.kx_index),
"z_index": int(result.selection.z_index),
},
"n_samples": 0 if result.t is None else int(np.asarray(result.t).size),
"n_state_shape": None
if result.state is None
else list(np.asarray(result.state).shape),
"has_eigenfunction": bool(
result.z is not None and result.eigenfunction is not None
),
"has_quasilinear": bool(getattr(result, "quasilinear", None) is not None),
# Solver status, with the same key names the in-memory result reports.
# A saved run used to carry no convergence channel at all, so a reader
# could not tell a certified eigenpair or a converged implicit scan from
# an unchecked one (queue row Q29).
**eigen_status_payload(getattr(result, "eigen_status", None)),
**implicit_solve_payload(getattr(result, "implicit_solve", None)),
}
if getattr(result, "quasilinear", None) is not None:
summary["quasilinear"] = result.quasilinear
return summary
def _write_runtime_linear_timeseries(path: Path, result: Any) -> str | None:
if result.t is None or result.signal is None:
return None
signal = _flatten_series(np.asarray(result.signal))
_write_csv(
path,
headers=["t", "signal_real", "signal_imag", "signal_abs"],
cols=[
_flatten_series(np.asarray(result.t)),
np.real(signal),
np.imag(signal),
np.abs(signal),
],
)
return str(path)
def _write_runtime_linear_eigenfunction(base: Path, result: Any) -> str | None:
if result.z is None or result.eigenfunction is None:
return None
eig_path = Path(f"{base}.eigenfunction.csv")
eig = np.asarray(result.eigenfunction)
_write_csv(
eig_path,
headers=["z", "eigen_real", "eigen_imag", "eigen_abs"],
cols=[
np.asarray(result.z, dtype=float),
np.real(eig),
np.imag(eig),
np.abs(eig),
],
)
return str(eig_path)
def _write_runtime_linear_state(base: Path, result: Any) -> str | None:
state_path = _write_state(
base, None if result.state is None else np.asarray(result.state)
)
return None if state_path is None else str(state_path)
def _write_runtime_linear_optional_artifacts(
*,
targets: _LinearArtifactTargets,
result: Any,
) -> dict[str, str]:
paths: dict[str, str] = {}
timeseries = _write_runtime_linear_timeseries(targets.timeseries_path, result)
if timeseries is not None:
paths["timeseries"] = timeseries
eigenfunction = _write_runtime_linear_eigenfunction(targets.base, result)
if eigenfunction is not None:
paths["eigenfunction"] = eigenfunction
state = _write_runtime_linear_state(targets.base, result)
if state is not None:
paths["state"] = state
if getattr(result, "quasilinear", None) is not None:
paths.update(write_quasilinear_artifacts(targets.base, result.quasilinear))
return paths
[docs]
def write_runtime_linear_scan_artifacts(out: str | Path, result: Any) -> dict[str, str]:
"""Write ky-scan growth/frequency and optional quasilinear spectra."""
targets = _runtime_linear_scan_targets(out)
arrays = _runtime_linear_scan_arrays(result)
ql_payloads = tuple(getattr(result, "quasilinear", None) or ())
_write_json(
targets.summary_path,
_runtime_linear_scan_summary(
arrays=arrays,
result=result,
ql_payloads=ql_payloads,
),
)
_write_csv(
targets.scan_path,
["ky", "gamma", "omega"],
[arrays.ky, arrays.gamma, arrays.omega],
)
paths = {"summary": str(targets.summary_path), "scan": str(targets.scan_path)}
ql_spectrum = _write_quasilinear_scan_spectrum(
base=targets.base,
ky=arrays.ky,
ql_payloads=ql_payloads,
)
if ql_spectrum is not None:
paths["quasilinear_spectrum"] = ql_spectrum
return paths
[docs]
def write_runtime_linear_artifacts(out: str | Path, result: Any) -> dict[str, str]:
"""Write summary/timeseries/state artifacts for a linear runtime run."""
targets = _runtime_linear_artifact_targets(out)
_write_json(targets.summary_path, _runtime_linear_summary(result))
paths = {"summary": str(targets.summary_path)}
paths.update(
_write_runtime_linear_optional_artifacts(targets=targets, result=result)
)
return paths
def _last_or_zero(series: Any) -> float:
arr = np.asarray(series)
return float(arr[-1]) if arr.size else 0.0
def _nonlinear_summary(result: Any) -> dict[str, Any]:
diag = result.diagnostics
payload: dict[str, Any] = {
"kind": "nonlinear",
"ky_selected": None
if result.ky_selected is None
else float(result.ky_selected),
"kx_selected": None
if result.kx_selected is None
else float(result.kx_selected),
"n_state_shape": None
if result.state is None
else list(np.asarray(result.state).shape),
# Every implicit GMRES solve of an IMEX run, as the in-memory result
# reports it; None for explicit methods (queue row Q29).
**implicit_solve_payload(getattr(result, "implicit_solve", None)),
}
if diag is not None:
payload.update(
{
"n_samples": int(np.asarray(diag.t).size),
"t_last": _last_or_zero(diag.t),
"dt_mean": float(np.asarray(diag.dt_mean)),
"gamma_last": _last_or_zero(diag.gamma_t),
"omega_last": _last_or_zero(diag.omega_t),
"Wg_last": _last_or_zero(diag.Wg_t),
"Wphi_last": _last_or_zero(diag.Wphi_t),
"Wapar_last": _last_or_zero(diag.Wapar_t),
"heat_flux_last": _last_or_zero(diag.heat_flux_t),
"particle_flux_last": _last_or_zero(diag.particle_flux_t),
# Cost per unit of simulated time plus the CFL attribution
# behind it, so a slow surface explains itself from the summary
# a user already reads.
"timestep_cost": timestep_cost_payload(
diag, wall_seconds=getattr(result, "wall_seconds", None)
),
}
)
elif result.phi2 is not None:
payload.update(
{
"n_samples": 0,
"t_last": 0.0,
"phi2_last": float(np.asarray(result.phi2)),
}
)
saturation = getattr(result, "saturation", None)
if saturation is not None:
payload["saturation"] = saturation
return payload
[docs]
def write_runtime_nonlinear_table_artifacts(
out: str | Path, result: Any
) -> dict[str, str]:
"""Write non-NetCDF nonlinear summary/diagnostics/state artifacts."""
out_path = Path(out)
base = _artifact_base(out_path)
summary_path = (
out_path if out_path.suffix.lower() == ".json" else Path(f"{base}.summary.json")
)
csv_path = (
out_path
if out_path.suffix.lower() == ".csv"
else Path(f"{base}.diagnostics.csv")
)
_write_json(summary_path, _nonlinear_summary(result))
paths = {"summary": str(summary_path)}
diag: SimulationDiagnostics | None = result.diagnostics
if diag is not None:
cols = [
_flatten_series(np.asarray(diag.t)),
_flatten_series(np.asarray(diag.dt_t)),
_flatten_series(np.asarray(diag.gamma_t)),
_flatten_series(np.asarray(diag.omega_t)),
_flatten_series(np.asarray(diag.Wg_t)),
_flatten_series(np.asarray(diag.Wphi_t)),
_flatten_series(np.asarray(diag.Wapar_t)),
_flatten_series(np.asarray(diag.energy_t)),
_flatten_series(np.asarray(diag.heat_flux_t)),
_flatten_series(np.asarray(diag.particle_flux_t)),
]
headers = [
"t",
"dt",
"gamma",
"omega",
"Wg",
"Wphi",
"Wapar",
"energy",
"heat_flux",
"particle_flux",
]
if diag.turbulent_heating_t is not None:
cols.append(_flatten_series(np.asarray(diag.turbulent_heating_t)))
headers.append("turbulent_heating")
if diag.heat_flux_species_t is not None:
heat_s = np.asarray(diag.heat_flux_species_t)
if heat_s.ndim == 1:
heat_s = heat_s[:, None]
for i in range(heat_s.shape[1]):
cols.append(heat_s[:, i])
headers.append(f"heat_flux_s{i}")
if diag.particle_flux_species_t is not None:
pflux_s = np.asarray(diag.particle_flux_species_t)
if pflux_s.ndim == 1:
pflux_s = pflux_s[:, None]
for i in range(pflux_s.shape[1]):
cols.append(pflux_s[:, i])
headers.append(f"particle_flux_s{i}")
if diag.turbulent_heating_species_t is not None:
turb_heat_s = np.asarray(diag.turbulent_heating_species_t)
if turb_heat_s.ndim == 1:
turb_heat_s = turb_heat_s[:, None]
for i in range(turb_heat_s.shape[1]):
cols.append(turb_heat_s[:, i])
headers.append(f"turbulent_heating_s{i}")
_write_csv(csv_path, headers=headers, cols=cols)
paths["diagnostics"] = str(csv_path)
state_path = _write_state(
base, None if result.state is None else np.asarray(result.state)
)
if state_path is not None:
paths["state"] = str(state_path)
return paths
def _expand_positive_ky_to_full(
state_positive_ky: np.ndarray, *, ny_full: int
) -> np.ndarray:
state = np.asarray(state_positive_ky)
if state.ndim != 6:
raise ValueError("state_positive_ky must have shape (Ns, Nl, Nm, Nyc, Nx, Nz)")
nyc = state.shape[3]
expected_nyc = nyc_from_ny(int(ny_full))
if nyc != expected_nyc:
raise ValueError(
f"positive-ky state Nyc={nyc} does not match ny_full={ny_full}"
)
return to_full(state, ny_full=int(ny_full))
def _expand_netcdf_restart_state_to_full_positive_ky(
state_active: np.ndarray,
*,
ny_full: int,
nx_full: int,
) -> np.ndarray:
state = np.asarray(state_active)
if state.ndim != 6:
raise ValueError("state_active must have shape (Ns, Nl, Nm, Naky, Nakx, Nz)")
nspec, nl, nm, naky, nakx, nz = state.shape
nyc_full = int(ny_full) // 2 + 1
expected_naky = _dealiased_ky_count(int(ny_full))
expected_nakx = _dealiased_kx_count(int(nx_full))
if naky != expected_naky:
raise ValueError(
f"restart Nky={naky} does not match ny_full={ny_full} (expected {expected_naky})"
)
if nakx != expected_nakx:
raise ValueError(
f"restart Nkx={nakx} does not match nx_full={nx_full} (expected {expected_nakx})"
)
out = np.zeros((nspec, nl, nm, nyc_full, int(nx_full), nz), dtype=np.complex64)
out[..., :naky, _dealiased_kx_indices(int(nx_full)), :] = state
return out
def _expand_netcdf_restart_state_full_ky(
state_active: np.ndarray,
*,
nx_full: int,
) -> np.ndarray:
"""Expand a NetCDF restart that already stores the full ``ky`` axis."""
state = np.asarray(state_active)
if state.ndim != 6:
raise ValueError("state_active must have shape (Ns, Nl, Nm, Ny, Nakx, Nz)")
nspec, nl, nm, ny_full, nakx, nz = state.shape
expected_nakx = _dealiased_kx_count(int(nx_full))
if nakx != expected_nakx:
raise ValueError(
f"restart Nkx={nakx} does not match nx_full={nx_full} (expected {expected_nakx})"
)
out = np.zeros((nspec, nl, nm, int(ny_full), int(nx_full), nz), dtype=np.complex64)
out[..., _dealiased_kx_indices(int(nx_full)), :] = state
return out
[docs]
def write_netcdf_restart_state(
path: str | Path, state: ArrayLike, *, ny_full: int | None = None
) -> Path:
"""Write a restart state in flat complex64 restart layout.
The file carries no header, so the reader tells its two accepted forms
apart by size alone: ``Ny`` stored ``ky`` rows is GKX's own state order,
and ``Nyc = 1 + Ny // 2`` rows is the packed positive-``ky`` interchange
order, whose axes are permuted (:doc:`inputs`). A half-spectrum state
would land on the second size while holding the first order, and the
reader has nothing to tell the two apart with -- it would return a
scrambled state of the right shape and dtype.
So a half-spectrum state is widened before it is written, which needs
``ny_full``: the length of the two-sided axis. The widening is exact and
costs only file size, because every row it adds is the conjugate of a row
already there. Nothing about an existing file changes: a two-sided state
is written exactly as before, byte for byte.
``ny_full=None`` writes whatever it is handed, which is what a caller that
already holds a two-sided state does.
"""
arr = np.asarray(state, dtype=np.complex64)
if ny_full is not None and is_half(int(arr.shape[KY_AXIS]), int(ny_full)):
arr = np.asarray(to_full(arr, ny_full=int(ny_full)), dtype=np.complex64)
out = Path(path)
out.parent.mkdir(parents=True, exist_ok=True)
arr.tofile(out)
return out
[docs]
def load_netcdf_restart_state(
path: str | Path,
*,
nspecies: int,
Nl: int,
Nm: int,
ny: int,
nx: int,
nz: int,
ky_layout: KyLayout = FULL,
) -> np.ndarray:
"""Load a NetCDF restart file onto ``ky_layout``.
``ny`` is the length of the two-sided axis in both layouts; ``ky_layout``
says how many rows the caller wants back, ``Ny`` or
``Nyc = 1 + Ny // 2``. The file itself is unchanged by the layout: it has
always stored the dealiased ``ky >= 0`` block, which is why a run started
on either axis can continue from a file written by the other, and why a
2.1.0 restart still loads (plan 5.3 N3, queue row Q10).
"""
try:
from netCDF4 import Dataset
except ImportError as exc: # pragma: no cover
raise ImportError("netCDF4 is required to load NetCDF restart files") from exc
with Dataset(Path(path), "r") as root:
_validate_netcdf_schema_version(root)
if "G" not in root.variables:
raise ValueError(f"restart file {path} does not contain variable 'G'")
raw = np.asarray(root.variables["G"][:], dtype=float)
if raw.ndim != 7 or raw.shape[-1] != 2:
raise ValueError(f"unexpected NetCDF restart G shape {raw.shape}")
state_active = raw[..., 0] + 1j * raw[..., 1]
state_active = np.asarray(
np.transpose(state_active, (0, 2, 1, 5, 4, 3)), dtype=np.complex64
)
if state_active.shape[:3] != (int(nspecies), int(Nl), int(Nm)):
raise ValueError(
f"restart state shape {state_active.shape[:3]} does not match requested {(int(nspecies), int(Nl), int(Nm))}"
)
if state_active.shape[-1] != int(nz):
raise ValueError(
f"restart Nz={state_active.shape[-1]} does not match requested {int(nz)}"
)
if state_active.shape[3] == int(ny):
full = _expand_netcdf_restart_state_full_ky(state_active, nx_full=nx)
return to_half(full, ny_full=int(ny)) if ky_layout == HALF else full
positive_ky = _expand_netcdf_restart_state_to_full_positive_ky(
state_active, ny_full=ny, nx_full=nx
)
if ky_layout == HALF:
return positive_ky
return _expand_positive_ky_to_full(positive_ky, ny_full=ny)
def _resolved_species_time(arr: Any | None, *, fallback: np.ndarray) -> np.ndarray:
if arr is None:
return np.asarray(fallback, dtype=np.float32)
return np.sum(np.asarray(arr, dtype=np.float32), axis=-1)
def _read_optional_var(group: Any, name: str) -> np.ndarray | None:
if name not in group.variables:
return None
var = group.variables[name]
arr = np.asarray(var[:])
dims = tuple(getattr(var, "dimensions", ()))
if dims and dims[-1] == "ri":
return np.asarray(arr[..., 0] + 1j * arr[..., 1])
return arr
def _resolve_restart_path(out: str | Path, cfg: Any, *, for_write: bool) -> Path:
configured = (
cfg.output.restart_to_file if for_write else cfg.output.restart_from_file
)
if configured is not None:
return Path(configured)
base = _netcdf_bundle_base(Path(out))
return Path(f"{base}.restart.nc")
def _condense_resolved_for_output(
resolved: ResolvedDiagnostics | None, *, ny_full: int | None = None
) -> ResolvedDiagnostics | None:
"""Condense fresh diagnostics onto the axes an existing bundle already uses.
``ny_full`` is the length of the two-sided ``ky`` axis. The number of
dealiased rows is a property of ``Ny``; taking it from the array's own
length keeps ``1 + (Nyc - 1) // 3`` rows of a half-spectrum diagnostic,
which is neither the published length nor anything else, and the append
then fails on a shape mismatch against the loaded history -- or, at a
resolution where the two happen to agree, does not.
"""
if resolved is None:
return None
payload: dict[str, np.ndarray | None] = {}
for field in ResolvedDiagnostics.__dataclass_fields__.values():
value = getattr(resolved, field.name)
if value is None:
payload[field.name] = None
elif field.name.endswith(("_kxt", "_kxst")):
payload[field.name] = _condense_kx(np.asarray(value))
elif field.name.endswith(("_kyt", "_kyst")):
payload[field.name] = _condense_ky(np.asarray(value), ny_full=ny_full)
elif field.name.endswith(("_kxkyt", "_kxkyst")):
payload[field.name] = _condense_kykx(np.asarray(value), ny_full=ny_full)
else:
payload[field.name] = np.asarray(value)
return ResolvedDiagnostics(**payload)
def _condense_diagnostics_for_netcdf_output(
diag: SimulationDiagnostics, *, ny_full: int | None = None
) -> SimulationDiagnostics:
# Nonlinear NetCDF output artifacts do not persist the monitored complex mode trace.
# Drop it when appending from an existing artifact so restart concatenation
# preserves the exact on-disk schema instead of mixing persisted and transient
# diagnostics.
return replace(
diag,
phi_mode_t=None,
resolved=_condense_resolved_for_output(diag.resolved, ny_full=ny_full),
)
def load_nonlinear_netcdf_diagnostics(path: str | Path) -> SimulationDiagnostics:
Dataset = _require_netcdf4()
with Dataset(Path(path), "r") as root:
_validate_netcdf_schema_version(root)
grids = root.groups["Grids"]
diag_group = root.groups["Diagnostics"]
time_vals = np.asarray(grids.variables["time"][:], dtype=np.float64)
wg_st = np.asarray(diag_group.variables["Wg_st"][:], dtype=np.float32)
wphi_st = np.asarray(diag_group.variables["Wphi_st"][:], dtype=np.float32)
wapar_st = np.asarray(diag_group.variables["Wapar_st"][:], dtype=np.float32)
heat_st = np.asarray(diag_group.variables["HeatFlux_st"][:], dtype=np.float32)
pflux_st = np.asarray(
diag_group.variables["ParticleFlux_st"][:], dtype=np.float32
)
turb_heat_st = _read_optional_var(diag_group, "TurbulentHeating_st")
resolved_payload = {
field.name: _read_optional_var(diag_group, field.name)
for field in ResolvedDiagnostics.__dataclass_fields__.values()
}
if turb_heat_st is None:
turb_heat_st = np.zeros_like(heat_st)
dt_t = (
np.diff(np.concatenate(([0.0], time_vals)))
if time_vals.size
else np.asarray([], dtype=np.float64)
)
dt_mean = np.asarray(
np.mean(dt_t[dt_t > 0.0]) if np.any(dt_t > 0.0) else 0.0, dtype=np.float64
)
Wg_t = np.sum(wg_st, axis=1)
Wphi_t = np.sum(wphi_st, axis=1)
Wapar_t = np.sum(wapar_st, axis=1)
heat_t = np.sum(heat_st, axis=1)
pflux_t = np.sum(pflux_st, axis=1)
turb_heat_t = np.sum(np.asarray(turb_heat_st, dtype=np.float32), axis=1)
return SimulationDiagnostics(
t=time_vals,
dt_t=dt_t,
dt_mean=dt_mean,
gamma_t=np.zeros_like(time_vals, dtype=np.float32),
omega_t=np.zeros_like(time_vals, dtype=np.float32),
Wg_t=Wg_t,
Wphi_t=Wphi_t,
Wapar_t=Wapar_t,
heat_flux_t=heat_t,
particle_flux_t=pflux_t,
energy_t=np.asarray(total_energy(Wg_t, Wphi_t, Wapar_t), dtype=np.float32),
heat_flux_species_t=heat_st,
particle_flux_species_t=pflux_st,
turbulent_heating_t=turb_heat_t,
turbulent_heating_species_t=np.asarray(turb_heat_st, dtype=np.float32),
phi_mode_t=None,
resolved=ResolvedDiagnostics(**resolved_payload),
)
__all__ = [
"_artifact_base",
"_condense_diagnostics_for_netcdf_output",
"_ensure_parent",
"_flatten_series",
"_is_netcdf_output_target",
"_netcdf_bundle_base",
"_resolve_restart_path",
"_resolved_species_time",
"_write_csv",
"_write_json",
"_write_state",
"load_netcdf_restart_state",
"load_nonlinear_netcdf_diagnostics",
"validate_finite_array",
"validate_finite_runtime_result",
"write_netcdf_restart_state",
"write_quasilinear_artifacts",
"write_runtime_linear_artifacts",
"write_runtime_linear_scan_artifacts",
"write_runtime_nonlinear_table_artifacts",
]