diff --git a/docs/_guide/API_reference.md b/docs/_guide/API_reference.md index bc6d944d..d61e955a 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,6 +416,7 @@ 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` | `run_nubeam_case(input_dir, gfile=..., workdir=...)` is the NUBEAM equivalent: it stages a case, @@ -451,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 new file mode 100644 index 00000000..acbaf4ed --- /dev/null +++ b/test/test_transp_outputs.py @@ -0,0 +1,354 @@ +"""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_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, _ = 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" + + +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") + + +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 ----------------------------- + + +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(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): + 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 "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_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 np.asarray(output.slice(0.20).variable("NLTAUP")).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 ---------------------------------------------------------------- + + +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_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 + + 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) + + +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 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"): + 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) + _, 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(state.xb) == expected["zones"] + assert psi_norm[-1] == pytest.approx(1.0) + + +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 ---------------------------------------------------------- + + +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 + ) + edge_axis = (expected["plflx"][1] - expected["plflx"][1][0]) / ( + expected["plflx"][1][-1] - expected["plflx"][1][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 --------------------------------------------------------------- + + +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_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") + + +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..a1bcdeb4 --- /dev/null +++ b/test/transp_cdf_fixtures.py @@ -0,0 +1,191 @@ +"""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 -- 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 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. +""" + +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", + "BPHXB": "NT-M", + "BDENS": "N/CM**3", +} + + +def write_transp_cdf( + path, + *, + 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.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. + + 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 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": (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"]}), + "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"), 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"]}), + # 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: + # 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", time3, {"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, + "time3": time3, + "x": centres, + "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/__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..613f817d --- /dev/null +++ b/vaft/code/transp/__init__.py @@ -0,0 +1,55 @@ +"""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 ( + EMPTY_PLACEHOLDERS, + PROFILE_GRIDS, + TIME_DIMENSIONS, + VARIABLE_DESCRIPTIONS, + TranspFormatError, + TranspOutput, + TranspSlice, + TranspVariable, + read_transp_output, +) +from .torque import enclosed_torque, input_torque_density, zone_volume + +__all__ = [ + "EMPTY_PLACEHOLDERS", + "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..605a801b --- /dev/null +++ b/vaft/code/transp/config.py @@ -0,0 +1,70 @@ +"""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 + +import re +from dataclasses import dataclass +from pathlib import Path + +from vaft.code.base import CodeResult + +__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`` 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 = "" + + +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 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": 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 new file mode 100644 index 00000000..77b82d79 --- /dev/null +++ b/vaft/code/transp/outputs.py @@ -0,0 +1,519 @@ +"""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 + +__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. +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 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" + ), +} + +#: ``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", +} + + +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.""" + + +@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. + """ + return _at_time(self._output.variable(name), self.time_index) + + 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") + # 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 + + +@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. + 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 + _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 + + @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._open.variables) + + @property + def dims(self) -> Mapping[str, int]: + 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(_first(self._open[name].attrs.get("units", "")) or "").strip() + + def long_name(self, name: str) -> str: + 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._open.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._open.variables: + raise KeyError(f"{self.path.name} has no variable {name!r}") + entry = self._open[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(_at_time(self.variable("X"), index), dtype=float), + xb=np.asarray(_at_time(self.variable("XB"), index), dtype=float), + _output=self, + ) + + +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 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( + f"{output.path.name} has no {name!r} variable; it does not look like " + "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, 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 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: + """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. + 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) + 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..82f439d6 --- /dev/null +++ b/vaft/code/transp/torque.py @@ -0,0 +1,57 @@ +"""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 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 +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", "zone_volume"] + + +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)