From be22a3e0305d8fcec2519c1cd0edcd77bc142954 Mon Sep 17 00:00:00 2001 From: Yun Date: Mon, 7 Sep 2026 17:25:19 +0900 Subject: [PATCH 1/2] Read TRANSP output: a transcript that resolves grids by name and converts nothing vaft.code.transp reads a TRANSP run's .CDF and stops there. There is no config, no prepare_* and no run_*: VAFT reads a TRANSP run, it does not launch one, so there is no $TRANSPHOME either. Three things the file forces, each of which the code this replaces got wrong: - A variable's grid comes from its dimension name, never its length. X (zone centres) and XB (zone outer boundaries) are the *same length* in a real run -- 20 and 20 for the MAST case -- so the legacy length comparison silently returned every boundary quantity as a centre one. on_x/on_xb refuse the wrong grid and say that interpolating between them is the caller's explicit choice. - Units come from each variable's own attribute: NE says N/CM**3, TQIN says Nt-M/CM3, PLFLX says Wb/rad. Nothing takes them from an argument, so no wrong flag can cost a factor of a million. - The profiles are densities, not per-zone integrals. Worth stating because the sibling adapter is the other way round -- NUBEAM's Plasma State writes per-zone integrals -- and assuming the family convention carries is an error of order the zone volume. Checked rather than assumed: sum(TQIN) is 5.6e-7, which matches nothing in the file, while sum(TQIN*DVOL) is -0.544 N m, the order of the run's own newton-metre torque scalars (BPHXB = -0.713). TQTOTNB is refused by name, pointing at TQIN: it is a dimensionless zero. The file's own integrity statements are checked on open -- PLFLX2PI/PLFLX is 2*pi, X and XB increase and interleave -- because a run killed mid-write leaves a readable header over truncated data that a classic netCDF reads back as silent zeros. psi_norm is PLFLX/PLFLXA, not (P-P[0])/(P[-1]-P[0]), which would call the innermost zone boundary the axis. enclosed_torque is the one place the grids interact: TQIN and DVOL are both cm-based and zone-centre, so their product is already newton metres, and the cumulative sum lands on the zone boundaries. Read through xarray's scipy backend, explicitly: netCDF4 is not a VAFT dependency and letting xarray auto-detect would exercise whichever backend happens to be installed. Reading is lazy -- a production file holds ~1900 variables. Verified against a real restricted run (45453X01, not committed): TQTOTNB is a dimensionless scalar, PLFLX[-1] == PLFLXA, PLFLX2PI/PLFLX = 6.2831855, the 750 ms enclosed torque totals -0.5441 N m, and psi_norm[0] = 0.0049462188 -- the value the committed reference .kin begins at. Co-Authored-By: Claude Opus 5 --- docs/_guide/API_reference.md | 20 +- test/test_transp_outputs.py | 229 ++++++++++++++++++++ test/transp_cdf_fixtures.py | 125 +++++++++++ vaft/code/__init__.py | 22 ++ vaft/code/transp/__init__.py | 53 +++++ vaft/code/transp/config.py | 59 ++++++ vaft/code/transp/outputs.py | 391 +++++++++++++++++++++++++++++++++++ vaft/code/transp/torque.py | 56 +++++ 8 files changed, 954 insertions(+), 1 deletion(-) create mode 100644 test/test_transp_outputs.py create mode 100644 test/transp_cdf_fixtures.py create mode 100644 vaft/code/transp/__init__.py create mode 100644 vaft/code/transp/config.py create mode 100644 vaft/code/transp/outputs.py create mode 100644 vaft/code/transp/torque.py diff --git a/docs/_guide/API_reference.md b/docs/_guide/API_reference.md index bc6d944d..fca21ff4 100644 --- a/docs/_guide/API_reference.md +++ b/docs/_guide/API_reference.md @@ -59,7 +59,7 @@ flowchart TD | `vaft.formula` | Pure physics functions: equilibrium, stability, Green's functions, constants | [Formula reference]({{ site.baseurl }}/reference/formula/) | | `vaft.machine_mapping` | Raw VEST DAQ to IMAS IDS mapping, plus uncertainty defaults | this page | | `vaft.plot` | Matplotlib figures straight from an ODS/ODC | this page | -| `vaft.code` | Adapters for external codes (EFIT, CHEASE, GPEC, TES, NUBEAM) | this page | +| `vaft.code` | Adapters for external codes (EFIT, CHEASE, GPEC, TES, NUBEAM, TRANSP) | this page | | `vaft.data` | GEQDSK read/write and packaged sample files | this page | | `vaft.imas` | OMAS to IMAS Access Layer bridge | [Data structures]({{ site.baseurl }}/guide/Data_structures/) | @@ -416,8 +416,26 @@ outputs = collect_efit_outputs(workdir, cfg) | GPEC (perturbed equilibrium, 3-D response) | `GPECSuiteConfig`, `GPECCaseInputs`, `GPECModuleRun`, `GPECSuiteResult`, `prepare_gpec_suite_case`, `run_gpec_suite_case`, `run_gpec`, `collect_gpec_suite_outputs`, `format_gfile_header_for_gpec` | | TES (forward equilibrium) | `TESConfig`, `TESInputs`, `TESResult`, `prepare_tes_inputs`, `run_tes`, `collect_tes_outputs`, `scan_tes`, `parse_result_scalars`, `parse_result_coils` | | NUBEAM (neutral-beam Monte Carlo) | `NUBEAMConfig`, `NUBEAMInputs`, `NUBEAMResult`, `find_nubeam_executable`, `prepare_nubeam_inputs`, `run_nubeam`, `run_nubeam_case`, `collect_nubeam_outputs` | +| TRANSP (transport, **read-only**) | `TranspOutput`, `TranspSlice`, `TranspVariable`, `TRANSPResult`, `read_transp_output`, `collect_transp_outputs`, `enclosed_torque`, `input_torque_density`, `zone_volume` | | Base classes | `CodeConfig`, `CodeInputs`, `CodeResult`, `CodeRunner` | +TRANSP is the one adapter with no `*Config`, no `prepare_*` and no `run_*`: VAFT reads a TRANSP run, +it does not launch one, so there is no `TRANSPHOME`. `read_transp_output` opens a `.CDF` +lazily — a production file holds roughly 1900 variables — and `.slice(time_s)` returns one time, +reporting which sample it took. + +This layer keeps TRANSP's own names, grids and units and converts nothing: `NE` is still per cubic +centimetre when it reaches you. Two things it *does* enforce, because the file makes them easy to +get wrong. A variable's grid comes from its dimension name and never its length — `X` (zone centres) +and `XB` (zone outer boundaries) have the *same length* in a real run, so `state.on_x("PLFLX")` +raises rather than quietly handing back a boundary quantity as a centre one. And `TQTOTNB` is +refused by name: it is an empty dimensionless placeholder, and the total input torque is `TQIN`. + +`enclosed_torque(state)` is the one place the two grids interact: `TQIN` is a density in +`N m / cm^3` and `DVOL` a volume in `cm^3`, both zone-centre, so the product is already newton +metres and the cumulative sum lands on the zone *boundaries*. Pairing them once here is deliberate — +converting one to SI and not the other is a factor of a million. + `run_nubeam_case(input_dir, gfile=..., workdir=...)` is the NUBEAM equivalent: it stages a case, builds its Plasma State, and runs INIT then STEP. Results come back as a native container, and `vaft.machine_mapping.core_sources.core_sources_from_nubeam` maps the heating, current-drive and diff --git a/test/test_transp_outputs.py b/test/test_transp_outputs.py new file mode 100644 index 00000000..f3a3542e --- /dev/null +++ b/test/test_transp_outputs.py @@ -0,0 +1,229 @@ +"""The native TRANSP transcript: grids, units, time, and the torque quantities.""" + +from __future__ import annotations + +import numpy as np +import pytest +from transp_cdf_fixtures import UNITS, write_transp_cdf + +from vaft.code import transp +from vaft.code.transp.outputs import TranspFormatError + + +@pytest.fixture() +def run(tmp_path): + expected = write_transp_cdf(tmp_path) + with transp.read_transp_output(expected["path"]) as output: + yield output, expected + + +# --- grids: resolved by dimension name, never by length ----------------------- + + +def test_the_two_radial_grids_are_the_same_length(run): + """Which is exactly why nothing may resolve a grid by counting. + + A real MAST run has len(X) == len(XB) == 20, and the legacy converter + picked a variable's grid by comparing array lengths -- so every + zone-boundary quantity came back as a zone-centre one. + """ + output, expected = run + assert len(expected["x"]) == len(expected["xb"]) + assert output.grid_of("NE") == "X" + assert output.grid_of("PLFLX") == "XB" + assert np.all(expected["x"] < expected["xb"]) # interleaved + + +def test_a_zone_centre_variable_is_refused_on_the_boundary_grid(run): + output, _ = run + state = output.slice(0.20) + with pytest.raises(TranspFormatError, match="NE is on the X grid"): + state.on_xb("NE") + with pytest.raises(TranspFormatError, match="PLFLX is on the XB grid"): + state.on_x("PLFLX") + + +def test_the_refusal_says_that_interpolating_is_the_callers_choice(run): + output, _ = run + with pytest.raises(TranspFormatError, match="explicitly"): + output.slice(0.20).on_xb("NE") + + +def test_a_variable_that_is_not_a_profile_has_no_grid(run): + output, _ = run + assert output.grid_of("PLFLXA") is None + with pytest.raises(TranspFormatError, match="not a radial profile"): + output.slice(0.20).on_x("PLFLXA") + + +# --- units: from the file, never from an argument ----------------------------- + + +def test_units_come_from_the_variables_own_attribute(run): + output, _ = run + assert output.units("NE") == UNITS["NE"] == "N/CM**3" + assert output.units("TQIN") == "Nt-M/CM3" + assert output.units("PLFLX") == "Wb/rad" + + +def test_nothing_is_converted_on_the_way_in(run): + """This layer is a transcript: NE is still per cubic centimetre.""" + output, expected = run + state = output.slice(0.20) + np.testing.assert_allclose(state.on_x("NE"), expected["n_e"][1], rtol=1e-6) + assert state.on_x("NE").max() > 1e12 # cm^-3, not m^-3 + + +def test_a_variable_with_no_units_attribute_reads_as_empty(run): + output, _ = run + assert output.units("X") == "" + + +def test_a_catalogued_variable_is_described_by_this_adapter(run): + output, _ = run + assert output.describe("NE") == "electron density [cm^-3]" + assert "TQIN" in output.describe("TQTOTNB") # the placeholder points at its replacement + + +def test_an_uncatalogued_variable_is_described_from_what_the_file_says(run): + """Q is in the file and not in VARIABLE_DESCRIPTIONS, so the file answers.""" + output, _ = run + described = output.describe("Q") + assert "SAFETY FACTOR" in described + assert "no units declared" in described # Q carries an empty units attribute + assert "TRANSP's own description" in described + + +def test_describing_a_variable_the_file_lacks_does_not_raise(run): + output, _ = run + assert "not in" in output.describe("NOT_A_TRANSP_VARIABLE") + + +# --- time --------------------------------------------------------------------- + + +def test_the_chosen_sample_is_part_of_the_answer(run): + """Samples are irregular, so which one was taken has to travel with it.""" + output, expected = run + state = output.slice(0.24) + assert state.time_index == 1 + assert state.time_s == pytest.approx(expected["time"][1]) + + +def test_both_time_dimensions_are_accepted(run): + output, _ = run + assert set(transp.TIME_DIMENSIONS) == {"TIME", "TIME3"} + assert "TIME" in output.dims and "TIME3" in output.dims + np.testing.assert_allclose(output.time, output.variable("TIME3").values) + + +def test_a_variable_without_a_time_axis_is_returned_whole(run): + output, expected = run + value = output.slice(0.20).variable("PLFLXA") + assert np.asarray(value).shape == () + + +# --- integrity ---------------------------------------------------------------- + + +def test_the_files_own_wb_per_radian_statement_is_checked(tmp_path): + """PLFLX2PI / PLFLX is 2*pi in a sound file; if it is not, the flux is + not what it says and psi_norm would be silently wrong.""" + expected = write_transp_cdf(tmp_path, plflx2pi_factor=3.0) + with pytest.raises(TranspFormatError, match=r"PLFLX2PI / PLFLX is 3\.0"): + transp.read_transp_output(expected["path"]) + + +def test_a_file_that_is_not_transp_output_is_refused(tmp_path): + import xarray as xr + + path = tmp_path / "notatransp.CDF" + xr.Dataset({"foo": (("bar",), np.arange(3.0))}).to_netcdf(path, format="NETCDF3_CLASSIC") + with pytest.raises(TranspFormatError, match="does not look like"): + transp.read_transp_output(path) + + +def test_a_run_directory_is_refused_with_a_useful_message(tmp_path): + with pytest.raises(IsADirectoryError, match="collect_transp_outputs"): + transp.read_transp_output(tmp_path) + + +# --- torque ------------------------------------------------------------------- + + +def test_tqtotnb_is_refused_by_name_and_points_at_tqin(run): + """It is a dimensionless zero in every run checked -- an empty placeholder.""" + output, _ = run + with pytest.raises(TranspFormatError, match="empty placeholder"): + output.variable("TQTOTNB") + with pytest.raises(TranspFormatError, match="Use TQIN"): + output.variable("TQTOTNB") + + +def test_enclosed_torque_pairs_the_two_cm_based_quantities_once(run): + """TQIN is N m / cm^3 and DVOL is cm^3, so the product is already N m. + + Converting one of them to SI and not the other is a factor of a million. + """ + output, expected = run + state = output.slice(0.20) + psi_norm, torque = transp.enclosed_torque(state) + + per_zone = expected["torque_density"][1] * expected["zone_volume"][1] + np.testing.assert_allclose(torque, np.cumsum(per_zone), rtol=1e-6) + assert torque[-1] == pytest.approx(per_zone.sum(), rel=1e-6) + + +def test_enclosed_torque_lands_on_the_zone_boundaries(run): + """Summing zones 1..i gives the torque inside XB[i], not inside X[i].""" + output, expected = run + state = output.slice(0.20) + psi_norm, torque = transp.enclosed_torque(state) + assert len(torque) == len(expected["xb"]) == len(state.xb) + np.testing.assert_allclose(psi_norm, expected["xb"], rtol=1e-6) + assert psi_norm[-1] == pytest.approx(1.0) + + +def test_the_input_torque_density_keeps_its_own_units(run): + output, expected = run + state = output.slice(0.20) + np.testing.assert_allclose(transp.input_torque_density(state), expected["torque_density"][1]) + assert state.units("TQIN") == "Nt-M/CM3" + + +# --- normalized flux ---------------------------------------------------------- + + +def test_psi_norm_is_plflx_over_the_enclosed_flux(run): + """Not (P - P[0]) / (P[-1] - P[0]), which would call the innermost zone + boundary the magnetic axis and move the whole grid inward.""" + output, expected = run + state = output.slice(0.20) + np.testing.assert_allclose(state.psi_norm_xb, expected["plflx"][1] / expected["plflxa"][1], rtol=1e-6) + assert state.psi_norm_xb[0] > 0.0 + edge_axis = (expected["plflx"][1] - expected["plflx"][1][0]) / ( + expected["plflx"][1][-1] - expected["plflx"][1][0] + ) + assert edge_axis[0] == 0.0 + assert state.psi_norm_xb[0] != pytest.approx(edge_axis[0]) + + +# --- collection --------------------------------------------------------------- + + +def test_collecting_a_directory_reports_what_it_holds(tmp_path): + expected = write_transp_cdf(tmp_path) + result = transp.collect_transp_outputs(tmp_path) + assert result.returncode is None # nothing was run + assert result.runid == expected["runid"] + assert expected["path"] in result.outputs["cdf"] + + +def test_collecting_a_missing_directory_is_an_error(tmp_path): + with pytest.raises(FileNotFoundError, match="does not exist"): + transp.collect_transp_outputs(tmp_path / "nope") + + +def test_an_empty_directory_is_not_an_error(tmp_path): + result = transp.collect_transp_outputs(tmp_path) + assert result.outputs["cdf"] == () and result.runid == "" diff --git a/test/transp_cdf_fixtures.py b/test/transp_cdf_fixtures.py new file mode 100644 index 00000000..b42e2898 --- /dev/null +++ b/test/transp_cdf_fixtures.py @@ -0,0 +1,125 @@ +"""Synthetic TRANSP output files, shaped like a real ``.CDF``. + +The layout mirrors a MAST run (45453X01, TRANSP output, 1884 variables): a +``TIME`` axis for the scalars and a ``TIME3`` axis for the profiles carrying +the same values, zone centres on ``X`` and zone outer boundaries on ``XB``, +and CGS units declared on each variable. + +The awkward details are deliberate, and each was found in the real file rather +than imagined: + +* ``X`` and ``XB`` are **the same length**, so a reader that resolves a grid + by counting cannot tell them apart -- and they interleave, ``X`` inside + ``XB`` zone by zone; +* the grids are declared ``(TIME3, X)``, i.e. time-varying, though their + values do not change; +* units are CGS and are declared per variable (``N/CM**3``, ``Nt-M/CM3``, + ``CM**3``), which is the only place they are stated; +* ``PLFLX`` is in ``Wb/rad`` and measured from the axis, with ``PLFLXA`` + carrying the enclosed flux separately, so normalized flux is their ratio; +* ``PLFLX2PI`` is the same flux in webers, exactly ``2*pi`` times ``PLFLX``; +* ``TQTOTNB`` is present, dimensionless and zero -- an empty placeholder. + +Nothing here is copied from a real run: the numbers are synthetic. +""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +#: TRANSP declares CGS on every variable it writes; nothing infers them. +UNITS = { + "TIME": "SECONDS", + "TIME3": "SECONDS", + "NE": "N/CM**3", + "NI": "N/CM**3", + "TE": "EV", + "TI": "EV", + "OMEGA": "RAD/SEC", + "DVOL": "CM**3", + "VRPOT": "VOLTS", + "PLFLX": "Wb/rad", + "PLFLX2PI": "WEBERS", + "PLFLXA": "Wb/rad", + "TQIN": "Nt-M/CM3", +} + + +def write_transp_cdf( + path, + *, + runid="90001A01", + times=(0.10, 0.20, 0.30), + zones=6, + torque_density=-2.0e-8, + zone_volume=1.0e4, + flux_edge=0.05, + plflx2pi_factor=2.0 * np.pi, + include_placeholder=True, +): + """Write a miniature ``.CDF``; returns the ground-truth arrays.""" + time = np.asarray(times, dtype="float32") + # Zone centres inside their own outer boundaries, same length -- as the + # real file has them. + edges = np.linspace(0.0, 1.0, zones + 1)[1:] + centres = edges - 0.5 / zones + x = np.tile(centres, (time.size, 1)).astype("float32") + xb = np.tile(edges, (time.size, 1)).astype("float32") + + flux = np.tile(edges * flux_edge, (time.size, 1)).astype("float32") + edge_flux = np.full(time.size, flux_edge, dtype="float32") + density = np.tile(np.linspace(4.0e13, 4.0e12, zones), (time.size, 1)).astype("float32") + temperature = np.tile(np.linspace(1.2e3, 1.0e2, zones), (time.size, 1)).astype("float32") + omega = np.tile(np.linspace(-1.2e4, -2.0e3, zones), (time.size, 1)).astype("float32") + potential = np.tile(np.linspace(0.0, -80.0, zones), (time.size, 1)).astype("float32") + torque = np.full((time.size, zones), torque_density, dtype="float32") + volume = np.full((time.size, zones), zone_volume, dtype="float32") + + data = { + "X": (("TIME3", "X"), x, {"units": ""}), + "XB": (("TIME3", "XB"), xb, {"units": ""}), + "NE": (("TIME3", "X"), density, {"units": UNITS["NE"], "long_name": "ELECTRON DENSITY"}), + "TE": (("TIME3", "X"), temperature, {"units": UNITS["TE"]}), + "TI": (("TIME3", "X"), temperature * 1.1, {"units": UNITS["TI"]}), + "OMEGA": (("TIME3", "X"), omega, {"units": UNITS["OMEGA"]}), + "DVOL": (("TIME3", "X"), volume, {"units": UNITS["DVOL"]}), + "TQIN": (("TIME3", "X"), torque, {"units": UNITS["TQIN"], "long_name": "TOTAL INPUT TORQUE"}), + # Present in the file, deliberately not in VARIABLE_DESCRIPTIONS: a + # reader must be able to describe it from what the file says. + "Q": (("TIME3", "X"), np.tile(np.linspace(1.0, 5.0, zones), (time.size, 1)).astype("float32"), + {"units": "", "long_name": "SAFETY FACTOR"}), + "VRPOT": (("TIME3", "XB"), potential, {"units": UNITS["VRPOT"]}), + "PLFLX": (("TIME3", "XB"), flux, {"units": UNITS["PLFLX"]}), + "PLFLX2PI": (("TIME3", "XB"), (flux * plflx2pi_factor).astype("float32"), {"units": UNITS["PLFLX2PI"]}), + "PLFLXA": (("TIME3",), edge_flux, {"units": UNITS["PLFLXA"]}), + } + if include_placeholder: + # Dimensionless and zero, exactly as the real file writes it. + data["TQTOTNB"] = ((), np.float32(0.0), {"units": UNITS["TQIN"]}) + + dataset = xr.Dataset( + data, + coords={ + "TIME": ("TIME", time, {"units": UNITS["TIME"]}), + "TIME3": ("TIME3", time, {"units": UNITS["TIME3"]}), + }, + attrs={"title": "synthetic TRANSP output"}, + ) + target = path / f"{runid}.CDF" + dataset.to_netcdf(target, format="NETCDF3_CLASSIC") + return { + "path": target, + "runid": runid, + "time": time, + "x": centres, + "xb": edges, + "n_e": density, + "T_e": temperature, + "omega": omega, + "plflx": flux, + "plflxa": edge_flux, + "torque_density": torque, + "zone_volume": volume, + "zones": zones, + } diff --git a/vaft/code/__init__.py b/vaft/code/__init__.py index 9a2d4f66..2714c144 100644 --- a/vaft/code/__init__.py +++ b/vaft/code/__init__.py @@ -73,6 +73,17 @@ "parse_result_coils", "scan_tes", "tokamaker", + "transp", + "TRANSPResult", + "TranspFormatError", + "TranspOutput", + "TranspSlice", + "TranspVariable", + "collect_transp_outputs", + "enclosed_torque", + "input_torque_density", + "read_transp_output", + "zone_volume", "TokaMakerConfig", "TokaMakerInputs", "TokaMakerResult", @@ -180,6 +191,16 @@ "parse_result_scalars": (".tes", "parse_result_scalars"), "parse_result_coils": (".tes", "parse_result_coils"), "scan_tes": (".tes", "scan_tes"), + "TRANSPResult": (".transp", "TRANSPResult"), + "TranspFormatError": (".transp", "TranspFormatError"), + "TranspOutput": (".transp", "TranspOutput"), + "TranspSlice": (".transp", "TranspSlice"), + "TranspVariable": (".transp", "TranspVariable"), + "collect_transp_outputs": (".transp", "collect_transp_outputs"), + "enclosed_torque": (".transp", "enclosed_torque"), + "input_torque_density": (".transp", "input_torque_density"), + "read_transp_output": (".transp", "read_transp_output"), + "zone_volume": (".transp", "zone_volume"), "TokaMakerConfig": (".tokamaker", "TokaMakerConfig"), "TokaMakerInputs": (".tokamaker", "TokaMakerInputs"), "TokaMakerResult": (".tokamaker", "TokaMakerResult"), @@ -235,6 +256,7 @@ def __getattr__(name: str): "snakemake", "tes", "tokamaker", + "transp", }: module = import_module(f".{name}", __name__) globals()[name] = module diff --git a/vaft/code/transp/__init__.py b/vaft/code/transp/__init__.py new file mode 100644 index 00000000..1226a73b --- /dev/null +++ b/vaft/code/transp/__init__.py @@ -0,0 +1,53 @@ +"""TRANSP output adapter -- read-only. + +TRANSP is run elsewhere; VAFT reads what it produced. There is no +``inputs.py``, no ``runner.py`` and no ``$TRANSPHOME``, because there is +nothing to configure or launch: + + .CDF ── read_transp_output ──▶ TranspOutput (lazy, ~1900 variables) + ── .slice(time_s) ───────▶ TranspSlice (one time, both grids) + ── enclosed_torque ──────▶ (psi_norm, torque [N m]) + +This layer keeps TRANSP's own names, grids and units and converts nothing -- +``NE`` is still per cubic centimetre when it reaches you. The conversion into +VAFT's kinetic-profile container is a separate module, so that what the file +said and what someone made of it stay distinguishable. + +Typical use:: + + from vaft.code import transp + + with transp.read_transp_output("45453X01.CDF") as output: + state = output.slice(0.750) # nearest sample; state.time_s says which + density = state.on_x("NE") # cm^-3, refuses a zone-boundary variable + psi_norm, torque = transp.enclosed_torque(state) +""" + +from .config import TRANSPResult, collect_transp_outputs +from .outputs import ( + PROFILE_GRIDS, + TIME_DIMENSIONS, + VARIABLE_DESCRIPTIONS, + TranspFormatError, + TranspOutput, + TranspSlice, + TranspVariable, + read_transp_output, +) +from .torque import enclosed_torque, input_torque_density, zone_volume + +__all__ = [ + "PROFILE_GRIDS", + "TIME_DIMENSIONS", + "TRANSPResult", + "TranspFormatError", + "TranspOutput", + "TranspSlice", + "TranspVariable", + "VARIABLE_DESCRIPTIONS", + "collect_transp_outputs", + "enclosed_torque", + "input_torque_density", + "read_transp_output", + "zone_volume", +] diff --git a/vaft/code/transp/config.py b/vaft/code/transp/config.py new file mode 100644 index 00000000..3f75277f --- /dev/null +++ b/vaft/code/transp/config.py @@ -0,0 +1,59 @@ +"""Discovering what a TRANSP run directory holds. + +VAFT does not run TRANSP, so there is no configuration to write and no +executable to find; this module exists for the one half of the usual code +adapter that still applies -- saying which output files a directory contains. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Mapping + +from vaft.code.base import CodeResult + +__all__ = ["TRANSPResult", "collect_transp_outputs"] + +#: TRANSP's output file, ``.CDF``. The ``PH`` sibling +#: (``PH.CDF``) holds the particle histories and is a separate product. +OUTPUT_SUFFIX = ".CDF" + + +@dataclass +class TRANSPResult(CodeResult): + """What a TRANSP run directory holds. + + ``returncode`` is ``None``: nothing was run. ``outputs["cdf"]`` lists the + output files found, most recent first. + """ + + runid: str = "" + native: Mapping[str, str] = field(default_factory=dict) + + +def collect_transp_outputs(directory: str | Path) -> TRANSPResult: + """List the TRANSP output files in ``directory``. + + Best-effort about content -- it does not open the files -- but a missing + directory is an error, matching the other ``collect_*`` adapters. + """ + workdir = Path(directory).expanduser() + if not workdir.is_dir(): + raise FileNotFoundError(f"TRANSP output directory does not exist: {workdir}") + + # A run directory is routinely shared with other codes' products, so match + # TRANSP's own naming rather than every file that ends in .CDF: a runid is + # a shot number followed by a letter-and-digits sequence. + candidates = sorted( + (path for path in workdir.glob(f"*{OUTPUT_SUFFIX}") if path.is_file()), + key=lambda path: path.stat().st_mtime, + reverse=True, + ) + runid = candidates[0].stem if candidates else "" + return TRANSPResult( + returncode=None, + workdir=workdir, + outputs={"cdf": tuple(candidates)}, + runid=runid, + ) diff --git a/vaft/code/transp/outputs.py b/vaft/code/transp/outputs.py new file mode 100644 index 00000000..e055f534 --- /dev/null +++ b/vaft/code/transp/outputs.py @@ -0,0 +1,391 @@ +"""Native TRANSP output containers: a transcript of one ``.CDF``, unconverted. + +TRANSP writes its whole run into a single netCDF file -- roughly 1900 +variables for a MAST case -- carrying scalars, time series, and radial +profiles on two interleaved grids. This layer reads it and stops there. +Nothing here writes an IDS and nothing here converts units, following the +split that :mod:`vaft.code.nubeam` and :mod:`vaft.code.gpec` already keep: +the layer above owns the interpretation, and it cannot own it if this one has +already quietly reinterpreted the file. + +Three things the file forces, each of which the code this replaces got wrong: + +- **A variable's grid comes from its dimension name, never its length.** + ``X`` holds zone centres and ``XB`` zone outer boundaries, and in a real run + they are *the same length* -- 20 and 20 for the MAST case -- so a reader + that compares array lengths silently returns every boundary quantity as a + centre quantity. They interleave: ``X = 0.025, 0.075, ...`` against + ``XB = 0.05, 0.10, ...``, so ``XB[i]`` bounds zone ``i`` from outside, and a + cumulative sum over zones is therefore an ``XB`` quantity. +- **Units come from the variable's own attribute.** ``NE`` says + ``N/CM**3``, ``TQIN`` says ``Nt-M/CM3``, ``PLFLX`` says ``Wb/rad``. A + reader that takes them from a flag instead is one wrong argument away from a + silent factor of a million. +- **The profiles are densities, not per-zone integrals.** This is worth + stating because the sibling adapter is the other way round: NUBEAM's Plasma + State writes per-zone integrals, and assuming the family convention carries + is an error of order the zone volume. Checked rather than assumed -- + ``sum(TQIN)`` is 5.6e-7, which matches nothing in the file, while + ``sum(TQIN * DVOL)`` is -0.544 N m, the order of the run's own newton-metre + torque scalars (``BPHXB`` = -0.713). + +The file is netCDF-3 classic, so it is opened through xarray's scipy backend: +``netCDF4`` is not a VAFT dependency and leaving the engine to auto-detection +would exercise an undeclared backend wherever one happens to be installed. +Reading is lazy -- a production ``.CDF`` runs to hundreds of megabytes and a +caller usually wants a handful of variables at one time. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping, Optional + +import numpy as np + +from vaft.code.gpec._netcdf import float_attr, scalar_attr + +__all__ = [ + "PROFILE_GRIDS", + "TIME_DIMENSIONS", + "VARIABLE_DESCRIPTIONS", + "TranspOutput", + "TranspSlice", + "TranspVariable", + "read_transp_output", +] + +#: Dimension names TRANSP uses for its time axis. ``TIME`` carries the +#: scalars and ``TIME3`` the profiles; in the runs checked they hold the same +#: 56 values, but they are separate dimensions and a reader must accept both. +TIME_DIMENSIONS: tuple[str, ...] = ("TIME", "TIME3") + +#: The two radial grids, and what a point on each one means. ``X`` and ``XB`` +#: have the same length in a real file, which is why nothing here resolves a +#: grid by counting. +PROFILE_GRIDS: Mapping[str, str] = { + "X": "zone centres, where a zone-averaged quantity belongs", + "XB": "zone outer boundaries; XB[i] bounds zone i, so a cumulative sum over zones lands here", +} + +#: What the variables this adapter names actually are, in the units TRANSP +#: itself declares on them -- read from the ``units`` attribute of a real run, +#: not from prose documentation. An index for discovery, not a mapping to +#: IMAS. +#: +#: Every profile here is a *density*: a per-zone integral would make +#: ``sum(TQIN)`` a torque, and it is 5.6e-7 against a machine that produces +#: order 0.5 N m. This is the opposite of the Plasma State convention in +#: :mod:`vaft.code.nubeam`, where the same-looking profiles *are* per-zone +#: integrals; do not carry that reading across. +VARIABLE_DESCRIPTIONS: Mapping[str, str] = { + "TIME": "time base of the scalars [s]", + "TIME3": "time base of the profiles [s]", + "X": "zone-centre radial coordinate, sqrt of normalized toroidal flux [-]", + "XB": "zone-boundary radial coordinate [-]", + "DVOL": "volume of each zone [cm^3]", + "DAREA": "cross-sectional area of each zone [cm^2]", + "NE": "electron density [cm^-3]", + "NI": "total ion density [cm^-3]", + "TE": "electron temperature [eV]", + "TI": "ion temperature [eV]", + "OMEGA": "toroidal angular velocity [rad/s]", + "VRPOT": "radial electrostatic potential, on zone boundaries [V]", + "PLFLX": "poloidal flux enclosed, measured from the axis, on zone boundaries [Wb/rad]", + "PLFLX2PI": "the same flux in webers; PLFLX2PI / PLFLX is 2*pi [Wb]", + "PLFLXA": "poloidal flux enclosed by the boundary; equals PLFLX[-1] [Wb/rad]", + "TQIN": "total input torque density [N m / cm^3]", + "TQTOTNB": ( + "an empty placeholder -- a dimensionless scalar of value zero in the runs " + "checked. The total input torque is TQIN" + ), +} + +#: ``TQTOTNB`` is written but carries nothing; asking for it is a mistake +#: worth catching by name rather than returning a scalar zero. +EMPTY_PLACEHOLDERS: Mapping[str, str] = { + "TQTOTNB": "TQIN", +} + + +class TranspFormatError(ValueError): + """The file is not shaped like a TRANSP output.""" + + +@dataclass(frozen=True) +class TranspVariable: + """One variable, with the name, dimensions and units TRANSP gave it.""" + + name: str + dims: tuple[str, ...] + units: str + long_name: str + values: np.ndarray + + @property + def grid(self) -> Optional[str]: + """``"X"``, ``"XB"`` or ``None`` -- from the dimension name.""" + for dim in self.dims: + if dim in PROFILE_GRIDS: + return dim + return None + + @property + def is_profile(self) -> bool: + return self.grid is not None + + +@dataclass +class TranspSlice: + """One time of a TRANSP run: its grids and the variables on them. + + Values are the file's own -- ``NE`` is still per cubic centimetre here. + """ + + time_s: float + time_index: int + x: np.ndarray + xb: np.ndarray + _output: "TranspOutput" = field(repr=False) + + def units(self, name: str) -> str: + return self._output.units(name) + + def describe(self, name: str) -> str: + return self._output.describe(name) + + def grid_of(self, name: str) -> Optional[str]: + return self._output.grid_of(name) + + def variable(self, name: str) -> np.ndarray: + """One variable at this time, in the file's own units. + + A variable with no time dimension is returned whole. + """ + variable = self._output.variable(name) + if any(dim in TIME_DIMENSIONS for dim in variable.dims): + return np.asarray(variable.values[self.time_index]) + return np.asarray(variable.values) + + def _on(self, name: str, grid: str) -> np.ndarray: + actual = self.grid_of(name) + if actual is None: + raise TranspFormatError( + f"{name} is not a radial profile; its dimensions are " + f"{self._output.variable(name).dims}" + ) + if actual != grid: + raise TranspFormatError( + f"{name} is on the {actual} grid ({PROFILE_GRIDS[actual]}), not {grid}. " + "Interpolating between the two is a choice a caller has to make " + "explicitly -- the grids have the same length, so nothing here " + "will do it silently" + ) + return self.variable(name) + + def on_x(self, name: str) -> np.ndarray: + """A zone-centre variable, refusing a zone-boundary one.""" + return self._on(name, "X") + + def on_xb(self, name: str) -> np.ndarray: + """A zone-boundary variable, refusing a zone-centre one.""" + return self._on(name, "XB") + + @property + def psi_norm_xb(self) -> np.ndarray: + """Normalized poloidal flux on the zone boundaries, ``PLFLX / PLFLXA``. + + ``PLFLX`` is measured from the magnetic axis and ``PLFLXA`` is the flux + enclosed by the boundary, so their ratio is the normalized flux + directly. Normalizing by the *first and last samples* instead -- + ``(P - P[0]) / (P[-1] - P[0])`` -- declares the innermost zone + boundary to be the axis, which moves the whole grid inward by + ``PLFLX[0] / PLFLXA`` (0.0049 for the MAST reference run). + """ + flux = self.on_xb("PLFLX") + edge = float(np.asarray(self.variable("PLFLXA")).reshape(-1)[0]) + if edge == 0: + raise TranspFormatError("PLFLXA is zero at this time; psi_norm is undefined") + return flux / edge + + +@dataclass +class TranspOutput: + """A TRANSP run's output file, read lazily. + + Use as a context manager; the underlying dataset stays open so a caller + can pull a few of the file's ~1900 variables without materializing it. + """ + + path: Path + _dataset: Any = field(repr=False, default=None) + _cache: dict = field(repr=False, default_factory=dict) + + # -- lifecycle ------------------------------------------------------- + def __enter__(self) -> "TranspOutput": + return self + + def __exit__(self, *exception) -> None: + self.close() + + def close(self) -> None: + if self._dataset is not None: + self._dataset.close() + self._dataset = None + + # -- structure ------------------------------------------------------- + @property + def variables(self) -> tuple[str, ...]: + return tuple(self._dataset.variables) + + @property + def dims(self) -> Mapping[str, int]: + return {str(name): int(size) for name, size in self._dataset.sizes.items()} + + def units(self, name: str) -> str: + """The file's own ``units`` attribute, stripped; ``""`` when it has none.""" + return str(scalar_attr(self._dataset[name].attrs.get("units", "")) or "").strip() + + def long_name(self, name: str) -> str: + return str(scalar_attr(self._dataset[name].attrs.get("long_name", "")) or "").strip() + + def describe(self, name: str) -> str: + """What a variable is, or a note that this adapter does not catalogue it.""" + described = VARIABLE_DESCRIPTIONS.get(name) + if described is not None: + return described + if name not in self._dataset.variables: + return f"{name}: not in {self.path.name}" + long_name = self.long_name(name) + units = self.units(name) + if long_name or units: + return f"{long_name or name} [{units or 'no units declared'}] (TRANSP's own description)" + return f"{name}: native TRANSP quantity, meaning not catalogued here" + + def grid_of(self, name: str) -> Optional[str]: + """Which radial grid a variable is on, from its dimension name.""" + return self.variable(name).grid + + def variable(self, name: str) -> TranspVariable: + """One variable, materialized and cached.""" + if name in self._cache: + return self._cache[name] + replacement = EMPTY_PLACEHOLDERS.get(name) + if replacement is not None: + raise TranspFormatError( + f"{name} is an empty placeholder in TRANSP output -- {VARIABLE_DESCRIPTIONS[name]}. " + f"Use {replacement}" + ) + if name not in self._dataset.variables: + raise KeyError(f"{self.path.name} has no variable {name!r}") + entry = self._dataset[name] + variable = TranspVariable( + name=name, + dims=tuple(str(dim) for dim in entry.dims), + units=self.units(name), + long_name=self.long_name(name), + values=np.asarray(entry.values), + ) + self._cache[name] = variable + return variable + + # -- time ------------------------------------------------------------ + @property + def time(self) -> np.ndarray: + """The scalar time base [s].""" + return np.asarray(self.variable("TIME").values, dtype=float) + + def time_index(self, time_s: float) -> int: + """Index of the sample nearest ``time_s``.""" + times = self.time + if times.size == 0: + raise TranspFormatError(f"{self.path.name} carries no times") + return int(np.argmin(np.abs(times - float(time_s)))) + + def slice(self, time_s: float) -> TranspSlice: + """The run at the sample nearest ``time_s``. + + The chosen time is on the result as ``time_s``/``time_index``: a run's + samples are irregular, so which one was taken is part of the answer. + """ + index = self.time_index(time_s) + return TranspSlice( + time_s=float(self.time[index]), + time_index=index, + x=np.asarray(self.variable("X").values[index], dtype=float), + xb=np.asarray(self.variable("XB").values[index], dtype=float), + _output=self, + ) + + +def _check_integrity(output: TranspOutput) -> None: + """Assertions the file itself makes possible, checked once on open. + + A TRANSP run killed mid-write leaves a readable header over truncated + data, and a classic netCDF reads the missing tail back as silent zeros + rather than raising, so presence of a variable is not evidence it is + whole. + """ + for name in ("TIME", "X", "XB"): + if name not in output.variables: + raise TranspFormatError( + f"{output.path.name} has no {name!r} variable; it does not look like " + "a TRANSP output file" + ) + + x = np.asarray(output.variable("X").values, dtype=float) + xb = np.asarray(output.variable("XB").values, dtype=float) + if x.shape != xb.shape: + raise TranspFormatError( + f"X is {x.shape} and XB is {xb.shape}; this reader expects one zone " + "centre per zone boundary" + ) + first_x, first_xb = x.reshape(-1, x.shape[-1])[0], xb.reshape(-1, xb.shape[-1])[0] + if not (np.all(np.diff(first_x) > 0) and np.all(np.diff(first_xb) > 0)): + raise TranspFormatError("X and XB must both increase outward") + if not np.all(first_x < first_xb): + raise TranspFormatError( + "every zone centre must lie inside its own outer boundary; X and XB " + "look swapped or misaligned" + ) + + if "PLFLX" in output.variables and "PLFLX2PI" in output.variables: + flux = np.asarray(output.variable("PLFLX").values, dtype=float) + webers = np.asarray(output.variable("PLFLX2PI").values, dtype=float) + finite = np.isfinite(flux) & np.isfinite(webers) & (flux != 0) + if np.any(finite): + ratio = float(np.nanmedian(webers[finite] / flux[finite])) + if not np.isclose(ratio, 2.0 * np.pi, rtol=1e-4): + raise TranspFormatError( + f"PLFLX2PI / PLFLX is {ratio:.6f}, not 2*pi: the file's own " + "Wb versus Wb/rad statement does not hold, so its flux cannot " + "be trusted" + ) + + +def read_transp_output(path: str | Path) -> TranspOutput: + """Open a TRANSP output ``.CDF``. + + Lazy: the dataset stays open and variables materialize on first use. Use + it as a context manager, or call :meth:`TranspOutput.close`. + """ + import xarray as xr + + path = Path(path).expanduser() + if path.is_dir(): + raise IsADirectoryError( + f"{path} is a directory; read_transp_output takes one .CDF file " + "(collect_transp_outputs takes a run directory)" + ) + # engine="scipy" pins the declared backend and the netCDF-3 classic format + # TRANSP writes; netCDF4 is not a VAFT dependency, so letting xarray + # auto-detect would use whichever backend happens to be installed. + dataset = xr.open_dataset(path, engine="scipy", decode_times=False) + output = TranspOutput(path=path, _dataset=dataset) + try: + _check_integrity(output) + except Exception: + output.close() + raise + return output diff --git a/vaft/code/transp/torque.py b/vaft/code/transp/torque.py new file mode 100644 index 00000000..1752d6c6 --- /dev/null +++ b/vaft/code/transp/torque.py @@ -0,0 +1,56 @@ +"""Torque from a TRANSP run: the input density, and what it encloses. + +Two facts decide everything here, and the code this replaces had both wrong. + +``TQTOTNB`` is an empty placeholder -- a dimensionless scalar of value zero in +every run checked -- so the total input torque is ``TQIN``. It is *not* the +beam sum: at 750 ms of the MAST reference run ``sum(TQIN)`` is -5.607e-7 +against ``sum(TQTOT01 + TQTOT02)`` = -5.731e-7, a pointwise difference of the +same order as the signal. ``TQIN`` is what TRANSP calls the total input +torque; the decomposition is a separate question and is not asserted here. + +``TQIN`` is a density in ``N m / cm^3`` and ``DVOL`` a volume in ``cm^3``, +both on the zone-centre grid, so their product is already newton metres. +Converting one of them to SI and not the other is wrong by a factor of a +million, which is why the pairing happens once, here, rather than at each call +site. +""" + +from __future__ import annotations + +import numpy as np + +from .outputs import TranspSlice + +__all__ = ["enclosed_torque", "input_torque_density"] + + +def input_torque_density(slice: TranspSlice) -> np.ndarray: + """``TQIN`` at this time, in TRANSP's own units [N m / cm^3], on ``X``.""" + return slice.on_x("TQIN") + + +def zone_volume(slice: TranspSlice) -> np.ndarray: + """``DVOL`` at this time [cm^3], on ``X`` -- the volume of each zone.""" + return slice.on_x("DVOL") + + +def enclosed_torque(slice: TranspSlice) -> tuple[np.ndarray, np.ndarray]: + """Torque enclosed by each zone boundary. + + Returns ``(psi_norm, torque)``: the normalized poloidal flux of the zone + boundaries and the cumulative torque inside each, in newton metres. + + The result is an ``XB`` quantity even though its ingredients are on ``X``. + Summing zones one to ``i`` gives the torque inside the *outer* boundary of + zone ``i``, which is ``XB[i]`` -- the grids interleave, and this is the one + place that relationship is load-bearing rather than incidental. + """ + density = np.asarray(input_torque_density(slice), dtype=float) + volume = np.asarray(zone_volume(slice), dtype=float) + if density.shape != volume.shape: + raise ValueError( + f"TQIN is {density.shape} and DVOL is {volume.shape}; both are zone-centre " + "quantities and must agree" + ) + return slice.psi_norm_xb, np.cumsum(density * volume) From 57d0779293a32ae3962cba2e45bfde0131d50747 Mon Sep 17 00:00:00 2001 From: Yun Date: Mon, 7 Sep 2026 17:55:39 +0900 Subject: [PATCH 2/2] Review fixes: TIME/TIME3 agreement, truncation, and a fixture that can fail Review of the TRANSP adapter found the reader trusting three things the file does not guarantee, and a fixture too tame to catch any of them. Correctness: - TIME and TIME3 are separate dimensions. A sample index chosen on TIME was applied to profiles on TIME3 with nothing checking they agree: a mismatched pair either raised a bare IndexError from inside the adapter or, worse, returned a profile from one instant while reporting another. They are now required to hold the same instants, which is what the runs looked at do. - X and XB were indexed on axis 0 without checking they carry a time axis. A run writing its grids once as plain (X,) made state.x a single radial point while the profiles stayed full length. Both now go through _at_time, as TranspSlice.variable already did. - psi_norm_xb read PLFLXA through the time-sliced accessor, which is right, but .reshape(-1)[0] read like "take the first sample" and no test could tell the difference: the fixture held PLFLXA constant. The step is now explicit and refuses a whole series, and the fixture's edge flux rises through the run as the reference file's does (0.0272 to 0.0765 Wb/rad). - close() left the object half-working -- cached variables answered, others died on a NoneType deep in xarray. A closed file now says so. Integrity. _check_integrity's docstring claimed a truncated file would be caught and nothing in it did that. Truncation is now detected twice: the backend's own failure to parse a short classic file is translated into a message that names the cause, and a size floor catches what it tolerates (5,566,599 bytes declared against 5,921,444 on disk for the reference run, so no false positive). Grid monotonicity and interleaving are checked at every sample rather than the first -- a zeroed tail is exactly a row that has stopped increasing -- and a PLFLX holding no non-zero finite value is refused instead of skipping the Wb/rad check and reporting the file sound. Cross-file: - collect_transp_outputs described a runid filter it never implemented, so it picked up the PH.CDF sibling and let mtime decide the runid. It now matches TRANSP's naming and orders by name; other .CDF files are listed under "other_cdf" rather than dropped. - The private vaft.code.gpec._netcdf import pulled the whole GPEC package in behind a reader that needs nothing from it (1.43 s against 0.61 s for nubeam), and imported float_attr without using it. Replaced by a local twelve-line helper. - TQTOTNB is not dimensionless: it declares Nt-M/CM3 and is empty. Corrected in the module, torque.py, the docs and the fixture docstring. - TRANSPResult carried a `native` field nothing populated; dropped, and the returncode/ok convention is now stated. zone_volume and TranspFormatError added to the __all__ lists that omitted them. - The TRANSP prose sat between the entry-point table and "the NUBEAM equivalent", stranding that phrase's antecedent; moved after the NUBEAM section. Tests. The fixture now carries what the reference file carries and this one did not: PLFLXA on TIME rather than TIME3 and varying through the run, a flux profile that is not proportional to XB (so psi_norm and the radial coordinate are distinguishable, 0.0049 against 0.05 at the innermost boundary in the real run), a variable on a third radial dimension, a scalar time series, a bare dimensionless scalar, and knobs to break one invariant at a time. Tests that asserted the fixture rather than the code are rewritten; the integrity branches that had no test -- shape, monotonicity, interleaving -- now have one. 24 tests to 37, and the three mutations that survived the review (PLFLXA[0] regardless of time, no shape check, no monotonicity check) are killed. The reference run still reads: -0.5441 N m enclosed at 750 ms, psi_norm[-1] exactly 1.0, psi_norm[0] 0.004946219. Co-Authored-By: Claude Opus 5 --- docs/_guide/API_reference.md | 35 +++--- test/test_transp_outputs.py | 185 +++++++++++++++++++++++++----- test/transp_cdf_fixtures.py | 118 ++++++++++++++----- vaft/code/transp/__init__.py | 2 + vaft/code/transp/config.py | 45 +++++--- vaft/code/transp/outputs.py | 212 ++++++++++++++++++++++++++++------- vaft/code/transp/torque.py | 7 +- 7 files changed, 469 insertions(+), 135 deletions(-) diff --git a/docs/_guide/API_reference.md b/docs/_guide/API_reference.md index fca21ff4..d61e955a 100644 --- a/docs/_guide/API_reference.md +++ b/docs/_guide/API_reference.md @@ -419,23 +419,6 @@ outputs = collect_efit_outputs(workdir, cfg) | TRANSP (transport, **read-only**) | `TranspOutput`, `TranspSlice`, `TranspVariable`, `TRANSPResult`, `read_transp_output`, `collect_transp_outputs`, `enclosed_torque`, `input_torque_density`, `zone_volume` | | Base classes | `CodeConfig`, `CodeInputs`, `CodeResult`, `CodeRunner` | -TRANSP is the one adapter with no `*Config`, no `prepare_*` and no `run_*`: VAFT reads a TRANSP run, -it does not launch one, so there is no `TRANSPHOME`. `read_transp_output` opens a `.CDF` -lazily — a production file holds roughly 1900 variables — and `.slice(time_s)` returns one time, -reporting which sample it took. - -This layer keeps TRANSP's own names, grids and units and converts nothing: `NE` is still per cubic -centimetre when it reaches you. Two things it *does* enforce, because the file makes them easy to -get wrong. A variable's grid comes from its dimension name and never its length — `X` (zone centres) -and `XB` (zone outer boundaries) have the *same length* in a real run, so `state.on_x("PLFLX")` -raises rather than quietly handing back a boundary quantity as a centre one. And `TQTOTNB` is -refused by name: it is an empty dimensionless placeholder, and the total input torque is `TQIN`. - -`enclosed_torque(state)` is the one place the two grids interact: `TQIN` is a density in -`N m / cm^3` and `DVOL` a volume in `cm^3`, both zone-centre, so the product is already newton -metres and the cumulative sum lands on the zone *boundaries*. Pairing them once here is deliberate — -converting one to SI and not the other is a factor of a million. - `run_nubeam_case(input_dir, gfile=..., workdir=...)` is the NUBEAM equivalent: it stages a case, builds its Plasma State, and runs INIT then STEP. Results come back as a native container, and `vaft.machine_mapping.core_sources.core_sources_from_nubeam` maps the heating, current-drive and @@ -469,6 +452,24 @@ deposition markers, lost fast ions, the step log's power budget -- stays in the and is drawn by `vaft.plot.nubeam`, which is why those particular views are not in the plot catalog while `nbi_profile_*` are. +TRANSP is the one adapter with no `*Config`, no `prepare_*` and no `run_*`: VAFT reads a TRANSP run, +it does not launch one, so there is no `TRANSPHOME`. `read_transp_output` opens a `.CDF` +lazily — a production file holds roughly 1900 variables — and `.slice(time_s)` returns one time, +reporting which sample it took. + +This layer keeps TRANSP's own names, grids and units and converts nothing: `NE` is still per cubic +centimetre when it reaches you. Two things it *does* enforce, because the file makes them easy to +get wrong. A variable's grid comes from its dimension name and never its length — `X` (zone centres) +and `XB` (zone outer boundaries) have the *same length* in a real run, so `state.on_x("PLFLX")` +raises rather than quietly handing back a boundary quantity as a centre one. And `TQTOTNB` is +refused by name: it is an empty placeholder — a zero-valued scalar with no radial or time axis, +declared torque-density units notwithstanding — and the total input torque is `TQIN`. + +`enclosed_torque(state)` is the one place the two grids interact: `TQIN` is a density in +`N m / cm^3` and `DVOL` a volume in `cm^3`, both zone-centre, so the product is already newton +metres and the cumulative sum lands on the zone *boundaries*. Pairing them once here is deliberate — +converting one to SI and not the other is a factor of a million. + `refine_equilibrium(source, config=None)` is the one-shot CHEASE convenience: g-file or ODS in, refined equilibrium out. `scan_tes(ods, base_config, values, param="ip0_kA")` sweeps a single TES parameter and collects every result. Snakemake rules should start with diff --git a/test/test_transp_outputs.py b/test/test_transp_outputs.py index f3a3542e..acbaf4ed 100644 --- a/test/test_transp_outputs.py +++ b/test/test_transp_outputs.py @@ -20,18 +20,14 @@ def run(tmp_path): # --- grids: resolved by dimension name, never by length ----------------------- -def test_the_two_radial_grids_are_the_same_length(run): - """Which is exactly why nothing may resolve a grid by counting. - - A real MAST run has len(X) == len(XB) == 20, and the legacy converter +def test_a_grid_is_resolved_by_dimension_name_not_by_length(run): + """A real MAST run has len(X) == len(XB) == 20, and the legacy converter picked a variable's grid by comparing array lengths -- so every - zone-boundary quantity came back as a zone-centre one. - """ - output, expected = run - assert len(expected["x"]) == len(expected["xb"]) + zone-boundary quantity came back as a zone-centre one.""" + output, _ = run + assert output.dims["X"] == output.dims["XB"] # nothing can be told apart by counting assert output.grid_of("NE") == "X" assert output.grid_of("PLFLX") == "XB" - assert np.all(expected["x"] < expected["xb"]) # interleaved def test_a_zone_centre_variable_is_refused_on_the_boundary_grid(run): @@ -56,6 +52,27 @@ def test_a_variable_that_is_not_a_profile_has_no_grid(run): output.slice(0.20).on_x("PLFLXA") +def test_a_variable_on_a_third_radial_dimension_has_no_grid(run): + """79 of the reference file's variables are on RMAJM, THETA and the like.""" + output, expected = run + assert output.grid_of("BDENS") is None + state = output.slice(0.20) + np.testing.assert_allclose(state.variable("BDENS"), expected["beam_density"][1], rtol=1e-6) + with pytest.raises(TranspFormatError, match="not a radial profile"): + state.on_x("BDENS") + + +def test_grids_written_without_a_time_axis_are_returned_whole(tmp_path): + """X and XB are (TIME3, X) in the runs looked at, but need not be, and + indexing axis 0 of a plain (X,) grid would return one radial point.""" + expected = write_transp_cdf(tmp_path, static_grids=True) + with transp.read_transp_output(expected["path"]) as output: + state = output.slice(0.20) + np.testing.assert_allclose(state.x, expected["x"], rtol=1e-6) + np.testing.assert_allclose(state.xb, expected["xb"], rtol=1e-6) + assert len(state.on_x("NE")) == len(state.x) + + # --- units: from the file, never from an argument ----------------------------- @@ -74,9 +91,12 @@ def test_nothing_is_converted_on_the_way_in(run): assert state.on_x("NE").max() > 1e12 # cm^-3, not m^-3 -def test_a_variable_with_no_units_attribute_reads_as_empty(run): - output, _ = run - assert output.units("X") == "" +def test_a_variable_with_no_units_attribute_reads_as_empty(tmp_path): + expected = write_transp_cdf(tmp_path, no_units_variable=True) + with transp.read_transp_output(expected["path"]) as output: + assert "units" not in output._open["NOUNITS"].attrs + assert output.units("NOUNITS") == "" + assert "no units declared" in output.describe("NOUNITS") def test_a_catalogued_variable_is_described_by_this_adapter(run): @@ -90,7 +110,6 @@ def test_an_uncatalogued_variable_is_described_from_what_the_file_says(run): output, _ = run described = output.describe("Q") assert "SAFETY FACTOR" in described - assert "no units declared" in described # Q carries an empty units attribute assert "TRANSP's own description" in described @@ -110,17 +129,33 @@ def test_the_chosen_sample_is_part_of_the_answer(run): assert state.time_s == pytest.approx(expected["time"][1]) -def test_both_time_dimensions_are_accepted(run): +def test_a_scalar_time_series_and_a_profile_are_sampled_at_the_same_instant(run): + """BPHXB is on TIME and NE on TIME3; a slice must reach both at its own time.""" + output, expected = run + state = output.slice(0.30) + assert state.time_index == 2 + assert float(state.variable("BPHXB")) == pytest.approx(-0.8, rel=1e-6) + np.testing.assert_allclose(state.on_x("NE"), expected["n_e"][2], rtol=1e-6) + + +def test_a_dimensionless_scalar_is_returned_whole(run): + """423 of the reference file's variables carry no axis at all.""" output, _ = run - assert set(transp.TIME_DIMENSIONS) == {"TIME", "TIME3"} - assert "TIME" in output.dims and "TIME3" in output.dims - np.testing.assert_allclose(output.time, output.variable("TIME3").values) + assert np.asarray(output.slice(0.20).variable("NLTAUP")).shape == () -def test_a_variable_without_a_time_axis_is_returned_whole(run): - output, expected = run - value = output.slice(0.20).variable("PLFLXA") - assert np.asarray(value).shape == () +def test_the_two_time_axes_must_hold_the_same_instants(tmp_path): + """A sample index is chosen on TIME and applied to profiles on TIME3, so a + file where they disagree cannot say which time a profile is from.""" + expected = write_transp_cdf(tmp_path, times=(0.1, 0.2, 0.3), times3=(0.5, 0.6, 0.7)) + with pytest.raises(TranspFormatError, match="do not hold the same instants"): + transp.read_transp_output(expected["path"]) + + +def test_time_axes_of_different_lengths_are_refused(tmp_path): + expected = write_transp_cdf(tmp_path, times=(0.1, 0.2, 0.3, 0.4), times3=(0.1, 0.2, 0.3)) + with pytest.raises(TranspFormatError, match="do not hold the same instants"): + transp.read_transp_output(expected["path"]) # --- integrity ---------------------------------------------------------------- @@ -134,6 +169,43 @@ def test_the_files_own_wb_per_radian_statement_is_checked(tmp_path): transp.read_transp_output(expected["path"]) +def test_a_flux_that_is_entirely_zero_is_refused_rather_than_skipped(tmp_path): + """A zeroed tail is what a truncated classic netCDF reads back as, so an + unusable PLFLX must not simply skip the check and report the file sound.""" + expected = write_transp_cdf(tmp_path, flux_edge=(0.0, 0.0, 0.0)) + with pytest.raises(TranspFormatError, match="no non-zero finite value"): + transp.read_transp_output(expected["path"]) + + +def test_a_truncated_file_is_refused_rather_than_read_back_as_zeros(tmp_path): + expected = write_transp_cdf(tmp_path) + whole = expected["path"].read_bytes() + expected["path"].write_bytes(whole[: len(whole) - 400]) + with pytest.raises(TranspFormatError, match="killed mid-write"): + transp.read_transp_output(expected["path"]) + + +def test_a_grid_that_stops_increasing_at_a_later_sample_is_refused(tmp_path): + """The grids are declared time-varying, so checking only the first sample + would pass a run whose tail was never written.""" + expected = write_transp_cdf(tmp_path, reverse_xb_at=2) + with pytest.raises(TranspFormatError, match="increase outward; they do not at sample 2"): + transp.read_transp_output(expected["path"]) + + +def test_grids_of_different_lengths_are_refused(tmp_path): + expected = write_transp_cdf(tmp_path, boundary_points=8) + with pytest.raises(TranspFormatError, match="one zone centre per zone boundary"): + transp.read_transp_output(expected["path"]) + + +def test_a_zone_centre_outside_its_own_boundary_is_refused(tmp_path): + """Both grids still increase, so only the interleaving catches this.""" + expected = write_transp_cdf(tmp_path, centres_outside=True) + with pytest.raises(TranspFormatError, match="swapped or misaligned"): + transp.read_transp_output(expected["path"]) + + def test_a_file_that_is_not_transp_output_is_refused(tmp_path): import xarray as xr @@ -148,12 +220,28 @@ def test_a_run_directory_is_refused_with_a_useful_message(tmp_path): transp.read_transp_output(tmp_path) +def test_a_closed_file_says_so_rather_than_half_answering(tmp_path): + """Whatever was read before close() would still answer out of the cache, + and anything else would fail deep in xarray on a NoneType.""" + expected = write_transp_cdf(tmp_path) + output = transp.read_transp_output(expected["path"]) + output.variable("NE") # cached + output.close() + with pytest.raises(TranspFormatError, match="is closed"): + output.variable("TE") + with pytest.raises(TranspFormatError, match="is closed"): + output.units("NE") + output.close() # idempotent + + # --- torque ------------------------------------------------------------------- def test_tqtotnb_is_refused_by_name_and_points_at_tqin(run): - """It is a dimensionless zero in every run checked -- an empty placeholder.""" + """It is a zero-valued scalar in every run checked -- an empty placeholder -- + even though it declares torque-density units.""" output, _ = run + assert output.units("TQTOTNB") == "Nt-M/CM3" with pytest.raises(TranspFormatError, match="empty placeholder"): output.variable("TQTOTNB") with pytest.raises(TranspFormatError, match="Use TQIN"): @@ -167,7 +255,7 @@ def test_enclosed_torque_pairs_the_two_cm_based_quantities_once(run): """ output, expected = run state = output.slice(0.20) - psi_norm, torque = transp.enclosed_torque(state) + _, torque = transp.enclosed_torque(state) per_zone = expected["torque_density"][1] * expected["zone_volume"][1] np.testing.assert_allclose(torque, np.cumsum(per_zone), rtol=1e-6) @@ -179,16 +267,17 @@ def test_enclosed_torque_lands_on_the_zone_boundaries(run): output, expected = run state = output.slice(0.20) psi_norm, torque = transp.enclosed_torque(state) - assert len(torque) == len(expected["xb"]) == len(state.xb) - np.testing.assert_allclose(psi_norm, expected["xb"], rtol=1e-6) + assert len(torque) == len(state.xb) == expected["zones"] assert psi_norm[-1] == pytest.approx(1.0) -def test_the_input_torque_density_keeps_its_own_units(run): +def test_the_input_torque_density_and_zone_volume_keep_their_own_units(run): output, expected = run state = output.slice(0.20) np.testing.assert_allclose(transp.input_torque_density(state), expected["torque_density"][1]) + np.testing.assert_allclose(transp.zone_volume(state), expected["zone_volume"][1]) assert state.units("TQIN") == "Nt-M/CM3" + assert state.units("DVOL") == "CM**3" # --- normalized flux ---------------------------------------------------------- @@ -199,15 +288,39 @@ def test_psi_norm_is_plflx_over_the_enclosed_flux(run): boundary the magnetic axis and move the whole grid inward.""" output, expected = run state = output.slice(0.20) - np.testing.assert_allclose(state.psi_norm_xb, expected["plflx"][1] / expected["plflxa"][1], rtol=1e-6) - assert state.psi_norm_xb[0] > 0.0 + np.testing.assert_allclose( + state.psi_norm_xb, expected["plflx"][1] / expected["plflxa"][1], rtol=1e-6 + ) edge_axis = (expected["plflx"][1] - expected["plflx"][1][0]) / ( expected["plflx"][1][-1] - expected["plflx"][1][0] ) - assert edge_axis[0] == 0.0 + assert state.psi_norm_xb[0] > 0.0 assert state.psi_norm_xb[0] != pytest.approx(edge_axis[0]) +def test_psi_norm_normalizes_by_this_times_edge_flux(run): + """PLFLXA rises through a run -- 0.0272 to 0.0765 Wb/rad in the reference + file -- so taking its first sample regardless of time puts the last closed + surface off the end of the grid.""" + output, expected = run + assert expected["plflxa"][0] != pytest.approx(expected["plflxa"][-1]) + for index, time_s in enumerate(expected["time"]): + state = output.slice(float(time_s)) + assert state.psi_norm_xb[-1] == pytest.approx(1.0, rel=1e-5) + np.testing.assert_allclose( + state.psi_norm_xb, expected["plflx"][index] / expected["plflxa"][index], rtol=1e-5 + ) + + +def test_psi_norm_is_not_the_radial_coordinate(run): + """XB is sqrt of normalized toroidal flux; psi_norm is poloidal. In the + reference run they are 0.0049 against 0.05 at the innermost boundary.""" + output, expected = run + state = output.slice(0.20) + assert state.psi_norm_xb[0] != pytest.approx(state.xb[0], rel=1e-3) + np.testing.assert_allclose(state.psi_norm_xb, expected["psi_norm"], rtol=1e-5) + + # --- collection --------------------------------------------------------------- @@ -219,7 +332,19 @@ def test_collecting_a_directory_reports_what_it_holds(tmp_path): assert expected["path"] in result.outputs["cdf"] -def test_collecting_a_missing_directory_is_an_error(tmp_path): +def test_the_particle_history_sibling_is_not_taken_for_the_run(tmp_path): + """PH.CDF is a separate product, and it is routinely written last -- + so a reader that picked the newest .CDF would report the wrong runid.""" + run = write_transp_cdf(tmp_path, runid="45453X01") + history = tmp_path / "45453X01PH.CDF" + history.write_bytes(run["path"].read_bytes()) + result = transp.collect_transp_outputs(tmp_path) + assert result.runid == "45453X01" + assert result.outputs["cdf"] == (run["path"],) + assert result.outputs["other_cdf"] == (history,) + + +def test_a_missing_directory_is_an_error(tmp_path): with pytest.raises(FileNotFoundError, match="does not exist"): transp.collect_transp_outputs(tmp_path / "nope") diff --git a/test/transp_cdf_fixtures.py b/test/transp_cdf_fixtures.py index b42e2898..a1bcdeb4 100644 --- a/test/transp_cdf_fixtures.py +++ b/test/transp_cdf_fixtures.py @@ -12,13 +12,23 @@ by counting cannot tell them apart -- and they interleave, ``X`` inside ``XB`` zone by zone; * the grids are declared ``(TIME3, X)``, i.e. time-varying, though their - values do not change; + values do not change -- as they are, and do not, in the reference run; * units are CGS and are declared per variable (``N/CM**3``, ``Nt-M/CM3``, ``CM**3``), which is the only place they are stated; * ``PLFLX`` is in ``Wb/rad`` and measured from the axis, with ``PLFLXA`` - carrying the enclosed flux separately, so normalized flux is their ratio; -* ``PLFLX2PI`` is the same flux in webers, exactly ``2*pi`` times ``PLFLX``; -* ``TQTOTNB`` is present, dimensionless and zero -- an empty placeholder. + carrying the enclosed flux separately on the ``TIME`` axis -- so normalized + flux is their ratio, and reading ``PLFLXA`` at the wrong sample is a + mistake a fixture with a constant edge flux could not catch. It rises + through the run, as it does in the reference file (0.0272 to 0.0765); +* the flux profile is **not** proportional to ``XB``, so ``psi_norm`` and the + radial coordinate are distinguishable -- in the reference run they are + 0.0049 against 0.05 at the innermost boundary; +* ``TQTOTNB`` is present and zero, carrying no radial or time axis, but + declaring torque-density units all the same -- an empty placeholder, not a + dimensionless flag; +* variables live on a third radial dimension (``RMAJM``) as well, on the + scalar ``TIME`` axis alone, and as bare dimensionless scalars: 79, 474 and + 423 of the reference file's variables respectively. Nothing here is copied from a real run: the numbers are synthetic. """ @@ -43,6 +53,8 @@ "PLFLX2PI": "WEBERS", "PLFLXA": "Wb/rad", "TQIN": "Nt-M/CM3", + "BPHXB": "NT-M", + "BDENS": "N/CM**3", } @@ -51,34 +63,68 @@ def write_transp_cdf( *, runid="90001A01", times=(0.10, 0.20, 0.30), + times3=None, zones=6, torque_density=-2.0e-8, zone_volume=1.0e4, - flux_edge=0.05, + flux_edge=(0.030, 0.045, 0.050), plflx2pi_factor=2.0 * np.pi, include_placeholder=True, + static_grids=False, + boundary_points=None, + reverse_xb_at=None, + centres_outside=False, + no_units_variable=False, ): - """Write a miniature ``.CDF``; returns the ground-truth arrays.""" + """Write a miniature ``.CDF``; returns the ground-truth arrays. + + The keyword arguments exist to break one thing at a time: ``times3`` + writes a profile axis that disagrees with the scalar one, ``static_grids`` + writes ``X``/``XB`` without a time dimension, ``boundary_points`` gives + ``XB`` a length of its own, ``reverse_xb_at`` makes one sample of the + boundary grid run backwards, and ``centres_outside`` puts every zone + centre outside its own boundary while leaving both grids increasing. + """ time = np.asarray(times, dtype="float32") + time3 = np.asarray(times if times3 is None else times3, dtype="float32") # Zone centres inside their own outer boundaries, same length -- as the # real file has them. edges = np.linspace(0.0, 1.0, zones + 1)[1:] - centres = edges - 0.5 / zones - x = np.tile(centres, (time.size, 1)).astype("float32") - xb = np.tile(edges, (time.size, 1)).astype("float32") - - flux = np.tile(edges * flux_edge, (time.size, 1)).astype("float32") - edge_flux = np.full(time.size, flux_edge, dtype="float32") - density = np.tile(np.linspace(4.0e13, 4.0e12, zones), (time.size, 1)).astype("float32") - temperature = np.tile(np.linspace(1.2e3, 1.0e2, zones), (time.size, 1)).astype("float32") - omega = np.tile(np.linspace(-1.2e4, -2.0e3, zones), (time.size, 1)).astype("float32") - potential = np.tile(np.linspace(0.0, -80.0, zones), (time.size, 1)).astype("float32") - torque = np.full((time.size, zones), torque_density, dtype="float32") - volume = np.full((time.size, zones), zone_volume, dtype="float32") + centres = edges + 0.5 / zones if centres_outside else edges - 0.5 / zones + boundaries = edges if boundary_points is None else np.linspace(0.0, 1.0, boundary_points + 1)[1:] + + def profile(row, axis_size=None, dtype="float32"): + return np.tile(row, ((time3.size if axis_size is None else axis_size), 1)).astype(dtype) + + x = centres.astype("float32") if static_grids else profile(centres) + xb = boundaries.astype("float32") if static_grids else profile(boundaries) + if reverse_xb_at is not None: + assert not static_grids, "reverse_xb_at needs a time-varying XB" + xb = xb.copy() + xb[reverse_xb_at] = xb[reverse_xb_at][::-1] + + # Flux rises through the run and is not proportional to XB, so neither the + # sample nor the coordinate can be mistaken for the other. + scale = np.asarray(flux_edge, dtype="float64").reshape(-1) + edge_flux = np.resize(scale, time.size).astype("float32") + flux_shape = 0.3 * boundaries + 0.7 * boundaries**2 + flux = (np.resize(scale, time3.size)[:, None] * flux_shape[None, :]).astype("float32") + + density = profile(np.linspace(4.0e13, 4.0e12, zones)) + temperature = profile(np.linspace(1.2e3, 1.0e2, zones)) + omega = profile(np.linspace(-1.2e4, -2.0e3, zones)) + potential = profile(np.linspace(0.0, -80.0, boundaries.size)) + torque = np.full((time3.size, zones), torque_density, dtype="float32") + volume = np.full((time3.size, zones), zone_volume, dtype="float32") + # A third radial dimension, as 79 of the reference file's variables have. + major_radius_points = 2 * zones + 1 + beam_density = profile(np.linspace(1.0e11, 0.0, major_radius_points)) + x_dims = ("X",) if static_grids else ("TIME3", "X") + xb_dims = ("XB",) if static_grids else ("TIME3", "XB") data = { - "X": (("TIME3", "X"), x, {"units": ""}), - "XB": (("TIME3", "XB"), xb, {"units": ""}), + "X": (x_dims, x, {"units": ""}), + "XB": (xb_dims, xb, {"units": ""}), "NE": (("TIME3", "X"), density, {"units": UNITS["NE"], "long_name": "ELECTRON DENSITY"}), "TE": (("TIME3", "X"), temperature, {"units": UNITS["TE"]}), "TI": (("TIME3", "X"), temperature * 1.1, {"units": UNITS["TI"]}), @@ -87,22 +133,39 @@ def write_transp_cdf( "TQIN": (("TIME3", "X"), torque, {"units": UNITS["TQIN"], "long_name": "TOTAL INPUT TORQUE"}), # Present in the file, deliberately not in VARIABLE_DESCRIPTIONS: a # reader must be able to describe it from what the file says. - "Q": (("TIME3", "X"), np.tile(np.linspace(1.0, 5.0, zones), (time.size, 1)).astype("float32"), + "Q": (("TIME3", "X"), profile(np.linspace(1.0, 5.0, zones)), {"units": "", "long_name": "SAFETY FACTOR"}), "VRPOT": (("TIME3", "XB"), potential, {"units": UNITS["VRPOT"]}), "PLFLX": (("TIME3", "XB"), flux, {"units": UNITS["PLFLX"]}), - "PLFLX2PI": (("TIME3", "XB"), (flux * plflx2pi_factor).astype("float32"), {"units": UNITS["PLFLX2PI"]}), - "PLFLXA": (("TIME3",), edge_flux, {"units": UNITS["PLFLXA"]}), + "PLFLX2PI": (("TIME3", "XB"), (flux * plflx2pi_factor).astype("float32"), + {"units": UNITS["PLFLX2PI"]}), + # On TIME, as the reference run writes it -- not on the profile axis. + "PLFLXA": (("TIME",), edge_flux, {"units": UNITS["PLFLXA"]}), + # A scalar time series, and a variable on a third radial dimension. + "BPHXB": (("TIME",), np.linspace(-0.5, -0.8, time.size).astype("float32"), + {"units": UNITS["BPHXB"], "long_name": "TOTAL PLASMA TORQUE"}), + "BDENS": (("TIME3", "RMAJM"), beam_density, + {"units": UNITS["BDENS"], "long_name": "BEAM ION DENSITY VS MAJOR RADIUS"}), + # A dimensionless scalar that is not a placeholder: 423 of the + # reference file's variables are shaped like this. + "NLTAUP": ((), np.float32(1.0), {"units": "", "long_name": "PARTICLE CONFINEMENT FLAG"}), } + if no_units_variable: + # Every variable in the reference run declares units, so this one is a + # deliberate divergence: the reader's "no units attribute" default has + # to be reachable from somewhere. + data["NOUNITS"] = (("TIME3", "X"), profile(np.ones(zones)), + {"long_name": "A QUANTITY WITH NO DECLARED UNITS"}) if include_placeholder: - # Dimensionless and zero, exactly as the real file writes it. + # Zero and without an axis, but declaring torque-density units, exactly + # as the real file writes it. data["TQTOTNB"] = ((), np.float32(0.0), {"units": UNITS["TQIN"]}) dataset = xr.Dataset( data, coords={ "TIME": ("TIME", time, {"units": UNITS["TIME"]}), - "TIME3": ("TIME3", time, {"units": UNITS["TIME3"]}), + "TIME3": ("TIME3", time3, {"units": UNITS["TIME3"]}), }, attrs={"title": "synthetic TRANSP output"}, ) @@ -112,13 +175,16 @@ def write_transp_cdf( "path": target, "runid": runid, "time": time, + "time3": time3, "x": centres, - "xb": edges, + "xb": boundaries, "n_e": density, "T_e": temperature, "omega": omega, "plflx": flux, "plflxa": edge_flux, + "psi_norm": flux_shape, + "beam_density": beam_density, "torque_density": torque, "zone_volume": volume, "zones": zones, diff --git a/vaft/code/transp/__init__.py b/vaft/code/transp/__init__.py index 1226a73b..613f817d 100644 --- a/vaft/code/transp/__init__.py +++ b/vaft/code/transp/__init__.py @@ -25,6 +25,7 @@ from .config import TRANSPResult, collect_transp_outputs from .outputs import ( + EMPTY_PLACEHOLDERS, PROFILE_GRIDS, TIME_DIMENSIONS, VARIABLE_DESCRIPTIONS, @@ -37,6 +38,7 @@ from .torque import enclosed_torque, input_torque_density, zone_volume __all__ = [ + "EMPTY_PLACEHOLDERS", "PROFILE_GRIDS", "TIME_DIMENSIONS", "TRANSPResult", diff --git a/vaft/code/transp/config.py b/vaft/code/transp/config.py index 3f75277f..605a801b 100644 --- a/vaft/code/transp/config.py +++ b/vaft/code/transp/config.py @@ -7,29 +7,40 @@ from __future__ import annotations -from dataclasses import dataclass, field +import re +from dataclasses import dataclass from pathlib import Path -from typing import Mapping from vaft.code.base import CodeResult -__all__ = ["TRANSPResult", "collect_transp_outputs"] +__all__ = ["OUTPUT_SUFFIX", "RUNID_PATTERN", "TRANSPResult", "collect_transp_outputs"] #: TRANSP's output file, ``.CDF``. The ``PH`` sibling #: (``PH.CDF``) holds the particle histories and is a separate product. OUTPUT_SUFFIX = ".CDF" +#: A TRANSP runid: a shot number, a run letter, and a run number -- ``45453X01``. +#: Matching this rather than every ``.CDF`` is what keeps the ``PH`` sibling +#: (``45453X01PH``) from being taken for the run itself. +RUNID_PATTERN = re.compile(r"[0-9]+[A-Za-z][0-9]+") + @dataclass class TRANSPResult(CodeResult): """What a TRANSP run directory holds. - ``returncode`` is ``None``: nothing was run. ``outputs["cdf"]`` lists the - output files found, most recent first. + ``returncode`` is ``None`` because nothing was run -- this is a listing of + files that already exist. That also makes the inherited ``ok`` ``False`` + for a perfectly good directory, as it does for the other read-only + ``collect_*`` adapters: ask whether ``outputs["cdf"]`` is empty instead. + + ``outputs["cdf"]`` lists the runid-shaped output files in name order and + ``outputs["other_cdf"]`` the remaining ``.CDF`` files, so a run directory + shared with another code's products loses nothing but does not have its + runid decided by one. """ runid: str = "" - native: Mapping[str, str] = field(default_factory=dict) def collect_transp_outputs(directory: str | Path) -> TRANSPResult: @@ -42,18 +53,18 @@ def collect_transp_outputs(directory: str | Path) -> TRANSPResult: if not workdir.is_dir(): raise FileNotFoundError(f"TRANSP output directory does not exist: {workdir}") - # A run directory is routinely shared with other codes' products, so match - # TRANSP's own naming rather than every file that ends in .CDF: a runid is - # a shot number followed by a letter-and-digits sequence. - candidates = sorted( - (path for path in workdir.glob(f"*{OUTPUT_SUFFIX}") if path.is_file()), - key=lambda path: path.stat().st_mtime, - reverse=True, - ) - runid = candidates[0].stem if candidates else "" + # A run directory holds the PH sibling as well, and is routinely shared + # with other codes' products, so match TRANSP's own naming rather than + # every file that ends in .CDF. Ordering is by name and not by mtime: a + # copy or a checkout normalizes mtimes, and picking the runid by the newest + # timestamp would then hand back "45453X01PH" for a directory whose run is + # 45453X01. + every = sorted(path for path in workdir.glob(f"*{OUTPUT_SUFFIX}") if path.is_file()) + runs = tuple(path for path in every if RUNID_PATTERN.fullmatch(path.stem)) + others = tuple(path for path in every if path not in runs) return TRANSPResult( returncode=None, workdir=workdir, - outputs={"cdf": tuple(candidates)}, - runid=runid, + outputs={"cdf": runs, "other_cdf": others}, + runid=runs[0].stem if runs else "", ) diff --git a/vaft/code/transp/outputs.py b/vaft/code/transp/outputs.py index e055f534..77b82d79 100644 --- a/vaft/code/transp/outputs.py +++ b/vaft/code/transp/outputs.py @@ -44,18 +44,35 @@ import numpy as np -from vaft.code.gpec._netcdf import float_attr, scalar_attr - __all__ = [ + "EMPTY_PLACEHOLDERS", "PROFILE_GRIDS", "TIME_DIMENSIONS", "VARIABLE_DESCRIPTIONS", + "TranspFormatError", "TranspOutput", "TranspSlice", "TranspVariable", "read_transp_output", ] + +def _first(value: Any) -> Any: + """First element of an attribute that may be a scalar or a one-element list. + + netCDF backends differ in whether a one-element attribute surfaces as a + scalar or as a length-1 array, so every attribute read goes through this + rather than calling ``str`` on whatever the file happened to carry. + (:func:`vaft.code.gpec._netcdf.scalar_attr` does the same job for the + GPEC suite; it is not imported here because importing it would pull the + whole ``vaft.code.gpec`` package in behind a reader that needs nothing + from it.) + """ + if isinstance(value, (list, tuple)): + return _first(value[0]) if value else None + array = np.asarray(value) + return array.reshape(-1)[0] if array.ndim else array.item() + #: Dimension names TRANSP uses for its time axis. ``TIME`` carries the #: scalars and ``TIME3`` the profiles; in the runs checked they hold the same #: 56 values, but they are separate dimensions and a reader must accept both. @@ -97,8 +114,9 @@ "PLFLXA": "poloidal flux enclosed by the boundary; equals PLFLX[-1] [Wb/rad]", "TQIN": "total input torque density [N m / cm^3]", "TQTOTNB": ( - "an empty placeholder -- a dimensionless scalar of value zero in the runs " - "checked. The total input torque is TQIN" + "an empty placeholder -- a zero-valued scalar carrying no radial or time " + "axis in the runs checked, though it declares torque-density units " + "(Nt-M/CM3). The total input torque is TQIN" ), } @@ -109,6 +127,19 @@ } +def _at_time(variable: "TranspVariable", index: int) -> np.ndarray: + """``variable`` at time ``index``, or whole when it carries no time axis. + + The dimension is checked rather than assumed: ``X`` and ``XB`` are + ``(TIME3, X)`` in the runs looked at, but TRANSP may write a run's grids + once as plain ``(X,)``, and indexing axis 0 of *that* silently returns a + single radial point where a whole grid was asked for. + """ + if any(dim in TIME_DIMENSIONS for dim in variable.dims): + return np.asarray(variable.values[index]) + return np.asarray(variable.values) + + class TranspFormatError(ValueError): """The file is not shaped like a TRANSP output.""" @@ -163,10 +194,7 @@ def variable(self, name: str) -> np.ndarray: A variable with no time dimension is returned whole. """ - variable = self._output.variable(name) - if any(dim in TIME_DIMENSIONS for dim in variable.dims): - return np.asarray(variable.values[self.time_index]) - return np.asarray(variable.values) + return _at_time(self._output.variable(name), self.time_index) def _on(self, name: str, grid: str) -> np.ndarray: actual = self.grid_of(name) @@ -204,7 +232,18 @@ def psi_norm_xb(self) -> np.ndarray: ``PLFLX[0] / PLFLXA`` (0.0049 for the MAST reference run). """ flux = self.on_xb("PLFLX") - edge = float(np.asarray(self.variable("PLFLXA")).reshape(-1)[0]) + # self.variable has already taken this time's sample, so PLFLXA is a + # scalar by the time it arrives here -- reading it off the whole + # series instead would pin psi_norm to the first sample's edge flux, + # which is 0.0272 Wb/rad against 0.0765 at 750 ms in the MAST + # reference run. Hence the size check rather than a bare [0]. + enclosed = np.asarray(self.variable("PLFLXA"), dtype=float) + if enclosed.size != 1: + raise TranspFormatError( + f"PLFLXA came back as {enclosed.shape} rather than this time's " + "single value; psi_norm would be normalized by the wrong sample" + ) + edge = float(enclosed.reshape(-1)[0]) if edge == 0: raise TranspFormatError("PLFLXA is zero at this time; psi_norm is undefined") return flux / edge @@ -216,6 +255,9 @@ class TranspOutput: Use as a context manager; the underlying dataset stays open so a caller can pull a few of the file's ~1900 variables without materializing it. + Every variable read is then held for the life of the object -- the cache + does not evict, so a caller that walks all ~1900 ends up holding the + whole file. Read what you need, or open the file again. """ path: Path @@ -234,28 +276,44 @@ def close(self) -> None: self._dataset.close() self._dataset = None + @property + def _open(self): + """The open dataset, refusing a closed one by name. + + Without this a closed file half-works: whatever a caller happened to + read before :meth:`close` still answers out of the cache, and + anything else fails deep in xarray with an ``AttributeError`` about + ``NoneType`` that says nothing about the real cause. + """ + if self._dataset is None: + raise TranspFormatError( + f"{self.path.name} is closed; open it again with " + "read_transp_output to read more of it" + ) + return self._dataset + # -- structure ------------------------------------------------------- @property def variables(self) -> tuple[str, ...]: - return tuple(self._dataset.variables) + return tuple(self._open.variables) @property def dims(self) -> Mapping[str, int]: - return {str(name): int(size) for name, size in self._dataset.sizes.items()} + return {str(name): int(size) for name, size in self._open.sizes.items()} def units(self, name: str) -> str: """The file's own ``units`` attribute, stripped; ``""`` when it has none.""" - return str(scalar_attr(self._dataset[name].attrs.get("units", "")) or "").strip() + return str(_first(self._open[name].attrs.get("units", "")) or "").strip() def long_name(self, name: str) -> str: - return str(scalar_attr(self._dataset[name].attrs.get("long_name", "")) or "").strip() + return str(_first(self._open[name].attrs.get("long_name", "")) or "").strip() def describe(self, name: str) -> str: """What a variable is, or a note that this adapter does not catalogue it.""" described = VARIABLE_DESCRIPTIONS.get(name) if described is not None: return described - if name not in self._dataset.variables: + if name not in self._open.variables: return f"{name}: not in {self.path.name}" long_name = self.long_name(name) units = self.units(name) @@ -277,9 +335,9 @@ def variable(self, name: str) -> TranspVariable: f"{name} is an empty placeholder in TRANSP output -- {VARIABLE_DESCRIPTIONS[name]}. " f"Use {replacement}" ) - if name not in self._dataset.variables: + if name not in self._open.variables: raise KeyError(f"{self.path.name} has no variable {name!r}") - entry = self._dataset[name] + entry = self._open[name] variable = TranspVariable( name=name, dims=tuple(str(dim) for dim in entry.dims), @@ -313,20 +371,49 @@ def slice(self, time_s: float) -> TranspSlice: return TranspSlice( time_s=float(self.time[index]), time_index=index, - x=np.asarray(self.variable("X").values[index], dtype=float), - xb=np.asarray(self.variable("XB").values[index], dtype=float), + x=np.asarray(_at_time(self.variable("X"), index), dtype=float), + xb=np.asarray(_at_time(self.variable("XB"), index), dtype=float), _output=self, ) -def _check_integrity(output: TranspOutput) -> None: - """Assertions the file itself makes possible, checked once on open. +def _check_not_truncated(output: TranspOutput) -> None: + """Refuse a file that is shorter than its own header says it should be. - A TRANSP run killed mid-write leaves a readable header over truncated - data, and a classic netCDF reads the missing tail back as silent zeros - rather than raising, so presence of a variable is not evidence it is - whole. + A TRANSP run killed mid-write leaves a complete header over incomplete + data, and a classic netCDF reads the missing tail back as zeros rather + than raising -- so presence of a variable is not evidence that it is + whole. The header states every variable's shape and type, so the file + has a minimum size that can simply be checked against the one on disk. + Record variables are interleaved and padded, so the sum is a floor rather + than an exact size, which is all that is wanted here. The technique is + the one :func:`vaft.code.gpec._solvers._check_nc_variable` uses on solver + output. """ + with open(output.path, "rb") as handle: + if handle.read(3) != b"CDF": + # Not classic; a truncated HDF5 file raises on read instead, so + # there is nothing for a size floor to add. + return + needed = 0 + for entry in output._open.variables.values(): + # The on-disk type, not the decoded one: xarray promotes a variable + # carrying a _FillValue to float64, which would inflate the floor and + # refuse a healthy file. + dtype = entry.encoding.get("dtype", entry.dtype) + needed += int(np.prod(entry.shape)) * int(np.dtype(dtype).itemsize) + actual = output.path.stat().st_size + if actual < needed: + raise TranspFormatError( + f"{output.path.name} is truncated: {actual} bytes on disk against the " + f"{needed} bytes of variable data its own header declares. A run " + "killed mid-write reads its missing tail back as zeros rather than " + "raising, so this is checked rather than discovered later" + ) + + +def _check_integrity(output: TranspOutput) -> None: + """Assertions the file itself makes possible, checked once on open.""" for name in ("TIME", "X", "XB"): if name not in output.variables: raise TranspFormatError( @@ -334,34 +421,65 @@ def _check_integrity(output: TranspOutput) -> None: "a TRANSP output file" ) + _check_not_truncated(output) + + # TRANSP writes the scalars on TIME and the profiles on TIME3. They are + # separate dimensions but the same axis, and this reader picks a sample + # index from TIME and applies it to profiles on TIME3, so a file where + # they disagree would hand back a profile from a different instant than + # the one it reports. Refuse rather than mislabel. + if "TIME3" in output.variables: + times = output.time + times3 = np.asarray(output.variable("TIME3").values, dtype=float) + if times.shape != times3.shape or not np.allclose(times, times3, rtol=1e-6, atol=0.0): + raise TranspFormatError( + f"TIME is {times.shape} and TIME3 is {times3.shape}, and they do not " + "hold the same instants; this reader chooses a sample on TIME and " + "reads profiles on TIME3, so it cannot say which time a profile is from" + ) + x = np.asarray(output.variable("X").values, dtype=float) xb = np.asarray(output.variable("XB").values, dtype=float) if x.shape != xb.shape: raise TranspFormatError( f"X is {x.shape} and XB is {xb.shape}; this reader expects one zone " - "centre per zone boundary" - ) - first_x, first_xb = x.reshape(-1, x.shape[-1])[0], xb.reshape(-1, xb.shape[-1])[0] - if not (np.all(np.diff(first_x) > 0) and np.all(np.diff(first_xb) > 0)): - raise TranspFormatError("X and XB must both increase outward") - if not np.all(first_x < first_xb): - raise TranspFormatError( - "every zone centre must lie inside its own outer boundary; X and XB " - "look swapped or misaligned" + "centre per zone boundary, which is how the runs it was written " + "against are shaped. A file that also writes the innermost boundary " + "(nzones + 1 points on XB) is a convention this reader has not been " + "checked against, not a broken file" ) + # Every time, not just the first: the grids are declared time-varying, and + # a truncated tail read back as zeros is exactly a row that has stopped + # increasing. They are a few kilobytes, so this costs nothing. + rows_x = x.reshape(-1, x.shape[-1]) + rows_xb = xb.reshape(-1, xb.shape[-1]) + for index, (row_x, row_xb) in enumerate(zip(rows_x, rows_xb)): + if not (np.all(np.diff(row_x) > 0) and np.all(np.diff(row_xb) > 0)): + raise TranspFormatError( + f"X and XB must both increase outward; they do not at sample {index}" + ) + if not np.all(row_x < row_xb): + raise TranspFormatError( + "every zone centre must lie inside its own outer boundary; X and XB " + f"look swapped or misaligned at sample {index}" + ) if "PLFLX" in output.variables and "PLFLX2PI" in output.variables: flux = np.asarray(output.variable("PLFLX").values, dtype=float) webers = np.asarray(output.variable("PLFLX2PI").values, dtype=float) finite = np.isfinite(flux) & np.isfinite(webers) & (flux != 0) - if np.any(finite): - ratio = float(np.nanmedian(webers[finite] / flux[finite])) - if not np.isclose(ratio, 2.0 * np.pi, rtol=1e-4): - raise TranspFormatError( - f"PLFLX2PI / PLFLX is {ratio:.6f}, not 2*pi: the file's own " - "Wb versus Wb/rad statement does not hold, so its flux cannot " - "be trusted" - ) + if not np.any(finite): + raise TranspFormatError( + "PLFLX holds no non-zero finite value, so the file's own Wb versus " + "Wb/rad statement cannot be checked and its flux cannot be trusted" + ) + ratio = float(np.nanmedian(webers[finite] / flux[finite])) + if not np.isclose(ratio, 2.0 * np.pi, rtol=1e-4): + raise TranspFormatError( + f"PLFLX2PI / PLFLX is {ratio:.6f}, not 2*pi: the file's own " + "Wb versus Wb/rad statement does not hold, so its flux cannot " + "be trusted" + ) def read_transp_output(path: str | Path) -> TranspOutput: @@ -381,7 +499,17 @@ def read_transp_output(path: str | Path) -> TranspOutput: # engine="scipy" pins the declared backend and the netCDF-3 classic format # TRANSP writes; netCDF4 is not a VAFT dependency, so letting xarray # auto-detect would use whichever backend happens to be installed. - dataset = xr.open_dataset(path, engine="scipy", decode_times=False) + try: + dataset = xr.open_dataset(path, engine="scipy", decode_times=False) + except Exception as exc: + # The backend reads a classic file's non-record variables at open time, + # so a run killed mid-write usually fails here -- with a message about + # reshaping an array, which says nothing about the real cause. + raise TranspFormatError( + f"could not read {path.name} ({path.stat().st_size} bytes on disk): {exc}. " + "A TRANSP run killed mid-write leaves a complete header over incomplete " + "data, which is what this normally is" + ) from exc output = TranspOutput(path=path, _dataset=dataset) try: _check_integrity(output) diff --git a/vaft/code/transp/torque.py b/vaft/code/transp/torque.py index 1752d6c6..82f439d6 100644 --- a/vaft/code/transp/torque.py +++ b/vaft/code/transp/torque.py @@ -2,8 +2,9 @@ Two facts decide everything here, and the code this replaces had both wrong. -``TQTOTNB`` is an empty placeholder -- a dimensionless scalar of value zero in -every run checked -- so the total input torque is ``TQIN``. It is *not* the +``TQTOTNB`` is an empty placeholder -- a zero-valued scalar carrying no radial +or time axis in every run checked, despite declaring torque-density units -- +so the total input torque is ``TQIN``. It is *not* the beam sum: at 750 ms of the MAST reference run ``sum(TQIN)`` is -5.607e-7 against ``sum(TQTOT01 + TQTOT02)`` = -5.731e-7, a pointwise difference of the same order as the signal. ``TQIN`` is what TRANSP calls the total input @@ -22,7 +23,7 @@ from .outputs import TranspSlice -__all__ = ["enclosed_torque", "input_torque_density"] +__all__ = ["enclosed_torque", "input_torque_density", "zone_volume"] def input_torque_density(slice: TranspSlice) -> np.ndarray: