diff --git a/docs/_guide/Profiles.md b/docs/_guide/Profiles.md index 94fdefbba..ec9f3c06a 100644 --- a/docs/_guide/Profiles.md +++ b/docs/_guide/Profiles.md @@ -277,6 +277,49 @@ vaft.plot.charge_exchange_time(ods, ion_index=0) `vaft.plot.plot_TeNe_from_eq(ods, ...)` plots the synthetic profiles produced by the `core_profiles_from_eq*` helpers. +## Kinetic-profile files + +`vaft.data.kinetic_profiles` is the container the kinetic-profile file formats read into and write +from — GPEC's `.kin` today, the Osborne pfile and MARS `PROF*.IN` to follow: + +```python +from vaft.data import read_kin, write_kin, normalize_psi + +profiles = read_kin("g045453.00750.kin") +profiles.psi_norm[0] # 0.00494621873 — exactly what the file said +profiles.available() # ('n_e', 'n_i', 'T_e', 'T_i', 'omega_exb') +profiles.normalization.method # "as_read" +write_kin(profiles, "again.kin") # byte-identical to this file; see below +``` + +A file VAFT wrote round-trips byte for byte, and so do the files the MAST-U workflow produces — +59 of the 65 `.kin` files in the reference tree, the named example among them. The other six are +not VAFT's: GPEC's bundled DIII-D example right-aligns its columns, so a negative rotation eats a +separator space, and the MARS-input file uses no leading indent at all. Both read correctly; they +are simply written back in this writer's layout. + +The container fixes one unit set — densities in m⁻³, temperatures in eV, angular frequencies in +rad/s, pressures in Pa — because those are what GPEC's own reader consumes. A reader converts into +them; nothing carries free-text units around. + +**The radial coordinate is never rescaled on read.** Making it span exactly [0, 1] is +`normalize_psi()`, an explicit operation whose result records that it happened. It has no +default `method=`: passing `axis=`/`edge=` and getting a min–max stretch instead would be the same +silent rescale, only now with a provenance record vouching for it. This is not +fastidiousness: GPEC's `read_kin` re-splines a `.kin` onto a uniform 101-point [0, 1] grid *with +extrapolation* (`nkin = 100` intervals), so it expects a truncated edge and handles it. A real MAST-U file spans +ψ_N = 0.00495 → 1.0, so stretching it moves every interior point — and once the stretched values are +written back, permanently. + +**Toroidal rotation and the E×B frequency are separate fields.** `omega_tor` and `omega_exb` are +never merged, and `write_kin` refuses a profile set carrying only the former: the `.kin` rotation +column is ω_E, and a toroidal rotation written there is wrong in a way nothing downstream detects. +`write_kin` also refuses to write a file GPEC would read as something other than what it says: +`omega_exb` holding **any** zero (GPEC substitutes 1e-9 element by element — pass +`allow_zero_rotation=True` for a set converted from an all-zero `PROFROT.IN`, which +`sample/output/converted_from_transp.kin` is), a value that is not finite, or a radial coordinate +that does not increase. Line endings are LF, so a CRLF file does not round-trip byte-identically. + ## Exporting To hand fitted electron profiles to an external code: diff --git a/test/kinetic_profile_fixtures.py b/test/kinetic_profile_fixtures.py new file mode 100644 index 000000000..46840992b --- /dev/null +++ b/test/kinetic_profile_fixtures.py @@ -0,0 +1,78 @@ +"""Synthetic kinetic-profile files, shaped like the real ones. + +The ``.kin`` layout is GPEC's own (``pentrc/inputs.f90``): six columns +``psi_n, n_i, n_e, T_i, T_e, omega_E``, a header line that is a comment, and +"(nearly) arbitrary header and/or footer, with the exception that no lines +start with a number". The awkward details here are the ones a real file has: +a radial coordinate that does *not* span exactly [0, 1], a negative rotation +column, and a footer. +""" + +from __future__ import annotations + +import numpy as np + +from vaft.data.kinetic_profiles import KIN_HEADER + +#: A real MAST-U .kin spans this, not [0, 1] -- which is what makes a +#: min-max rescale on read move every interior point. +TRUNCATED_SPAN = (0.00494621873, 1.0) + + +def profile_columns(points=41, span=TRUNCATED_SPAN): + """The six ``.kin`` columns as arrays, in file order.""" + psi = np.linspace(span[0], span[1], points) + edge = 0.5 * (1.0 - np.tanh((psi - 0.93) / 0.04)) + n_e = 4.0e19 * (0.3 + 0.7 * edge) + return { + "psi_norm": psi, + "n_i": 0.91 * n_e, + "n_e": n_e, + "T_i": 1.2e3 * (0.1 + 0.9 * edge), + "T_e": 1.1e3 * (0.1 + 0.9 * edge), + # Negative, as a real E x B frequency profile is. + "omega_exb": -1.2e4 * (0.2 + 0.8 * psi), + } + + +def write_kin_file(path, *, points=41, span=TRUNCATED_SPAN, header=True, footer=False, + extra_column=False, columns=None, fortran_exponents=False, + numeric_footer=False, ragged_row=None, name="synthetic.kin"): + """Write a synthetic ``.kin``; returns the columns it wrote. + + ``header`` may be a literal string, so a test can pin the header GPEC's + own files carry without importing it from the module under test. The + remaining switches each break one thing: ``fortran_exponents`` writes + ``1.0D+00`` as a Fortran tool would, ``numeric_footer`` appends a summary + row *below* a footer (which GPEC stops before and a reader that collects + every numeric line anywhere would swallow), and ``ragged_row`` drops a + token from one row. + """ + data = columns if columns is not None else profile_columns(points, span) + table = np.column_stack(list(data.values())) + if extra_column: + table = np.column_stack([table, np.arange(table.shape[0], dtype=float)]) + + def render(value): + text = f"{value:.8e}" + return text.replace("e", "D") if fortran_exponents else text + + if header is True: + lines = [KIN_HEADER] + elif header: + lines = [header] + else: + lines = [] + rows = [" " + " ".join(render(value) for value in row) for row in table] + if ragged_row is not None: + rows[ragged_row] = rows[ragged_row].rsplit(" ", 1)[0] + lines.extend(rows) + if footer or numeric_footer: + lines += ["", "written by a synthetic fixture", "provenance: none"] + if numeric_footer: + # A summary row after the footer: still numeric, but past the end of + # the table as far as GPEC's readtable is concerned. + lines.append(" " + " ".join(render(value) for value in table.mean(axis=0))) + target = path / name + target.write_text("\n".join(lines) + "\n", encoding="utf-8") + return data diff --git a/test/test_data_code_namespace.py b/test/test_data_code_namespace.py index 7df4c5a0f..1199c2858 100644 --- a/test/test_data_code_namespace.py +++ b/test/test_data_code_namespace.py @@ -14,6 +14,35 @@ def test_data_and_code_import_smoke(): assert "data" in dir(vaft) +def test_every_name_vaft_data_advertises_resolves(): + """``vaft.data`` is a lazy namespace: a name in ``__all__`` with no + ``_EXPORT_MAP`` entry (or a submodule missing from the ``__getattr__`` + set) fails only at attribute access, so nothing but a walk catches it. + ``vaft.formula`` has had this check; ``vaft.data`` had not.""" + import vaft.data as data + + unresolvable = [] + for name in data.__all__: + try: + getattr(data, name) + except AttributeError: + unresolvable.append(name) + assert not unresolvable + + stale = sorted(set(data._EXPORT_MAP) - set(data.__all__)) + assert not stale, f"exported but not advertised: {stale}" + + +def test_star_importing_vaft_data_binds_the_whole_surface(): + import vaft.data as data + + namespace: dict[str, object] = {} + exec("from vaft.data import *", namespace) # noqa: S102 + + for name in data.__all__: + assert name in namespace, name + + def test_geqdsk_roundtrip(tmp_path): from vaft.data import read_geqdsk, write_geqdsk from vaft.data.resources import data_path diff --git a/test/test_kinetic_profiles_kin.py b/test/test_kinetic_profiles_kin.py new file mode 100644 index 000000000..0e8af7259 --- /dev/null +++ b/test/test_kinetic_profiles_kin.py @@ -0,0 +1,309 @@ +"""The kinetic-profile container and GPEC ``.kin`` I/O.""" + +from __future__ import annotations + +import numpy as np +import pytest +from kinetic_profile_fixtures import TRUNCATED_SPAN, profile_columns, write_kin_file + +from vaft.data.kinetic_profiles import ( + KINETIC_UNITS, + KIN_COLUMNS, + KineticProfiles, + normalize_psi, + read_kin, + write_kin, +) + +#: The header line GPEC's own example files carry, byte for byte -- written +#: out here rather than imported, so that changing the module's constant +#: fails a test instead of moving the goalposts with it. It is the first line +#: of all 65 .kin files under the reference tree, DIII-D and MAST alike. +GPEC_KIN_HEADER = ( + " psi ni(m^-3) ne(m^-3)" + " ti(eV) te(eV) wexb(rad/s)" +) + + +@pytest.fixture() +def kin(tmp_path): + expected = write_kin_file(tmp_path) + return read_kin(tmp_path / "synthetic.kin"), expected + + +# --- the container ------------------------------------------------------------ + + +def test_reads_the_six_columns_in_gpecs_order(kin): + profiles, expected = kin + assert KIN_COLUMNS == ("psi_norm", "n_i", "n_e", "T_i", "T_e", "omega_exb") + for name, values in expected.items(): + np.testing.assert_allclose(getattr(profiles, name), values) + assert profiles.available() == ("n_e", "n_i", "T_e", "T_i", "omega_exb") + + +def test_units_are_the_containers_own_not_the_files(kin): + profiles, _ = kin + assert KINETIC_UNITS["n_e"] == "m^-3" + assert KINETIC_UNITS["T_e"] == "eV" + assert KINETIC_UNITS["omega_exb"] == "rad/s" + assert profiles.unit("n_e") == "m^-3" + + +def test_the_units_table_and_the_container_are_reachable_from_vaft_data(): + """The pfile and MARS readers convert through this table and no other, so + it has to be importable from where its docstring says it is.""" + import vaft.data as data + + assert data.KINETIC_UNITS["n_e"] == "m^-3" + assert data.KIN_HEADER == GPEC_KIN_HEADER + assert "n_e" in data.PROFILE_FIELDS and "psi_norm" not in data.PROFILE_FIELDS + assert data.kinetic_profiles.__name__ == "vaft.data.kinetic_profiles" + + +def test_a_profile_of_the_wrong_length_is_refused(): + with pytest.raises(ValueError, match="every profile is on the same radial coordinate"): + KineticProfiles(psi_norm=np.linspace(0, 1, 5), n_e=np.zeros(4)) + + +def test_a_missing_profile_names_what_is_there(kin): + profiles, _ = kin + with pytest.raises(KeyError, match="no profile 'p_total'"): + profiles.field("p_total") + + +def test_extra_columns_are_kept_rather_than_dropped(tmp_path): + write_kin_file(tmp_path, extra_column=True) + profiles = read_kin(tmp_path / "synthetic.kin") + # Numbered as a person reading the file would: the seventh column is 7. + assert "column_7" in profiles.extras + np.testing.assert_allclose(profiles.extras["column_7"], np.arange(len(profiles))) + np.testing.assert_allclose(profiles.field("column_7"), np.arange(len(profiles))) + assert profiles.provenance["column_7"].endswith("column 7") + + +def test_a_profile_set_cannot_be_edited_behind_its_provenance(kin): + """frozen= stops the attributes being rebound and nothing else, so the + arrays and mappings are sealed too: a set that says "as read" has to be.""" + profiles, _ = kin + with pytest.raises(ValueError): + profiles.psi_norm[0] = 42.0 + with pytest.raises(ValueError): + profiles.n_e[0] = 0.0 + with pytest.raises(TypeError): + profiles.provenance["n_e"] = "somewhere else" + + +def test_sealing_does_not_reach_back_into_the_callers_array(): + """A read-only *view*, not a read-only array: passing an array in must not + make the caller's own copy unwritable.""" + psi = np.linspace(0.0, 1.0, 5) + KineticProfiles(psi_norm=psi) + psi[0] = 0.5 # still the caller's array + assert psi[0] == 0.5 + + +# --- C-10: the radial coordinate is never rescaled on read -------------------- + + +def test_a_truncated_coordinate_survives_a_round_trip(tmp_path): + """The regression for C-10. + + A real .kin spans 0.00495 to 1.0. The legacy readers stretched that onto + [0, 1] and the writer put the stretched values back, so one round trip + moved every interior point permanently. + """ + write_kin_file(tmp_path, span=TRUNCATED_SPAN) + profiles = read_kin(tmp_path / "synthetic.kin") + assert profiles.psi_norm[0] == pytest.approx(TRUNCATED_SPAN[0]) + assert profiles.normalization.method == "as_read" + + out = write_kin(profiles, tmp_path / "again.kin") + assert out.read_text(encoding="utf-8") == (tmp_path / "synthetic.kin").read_text(encoding="utf-8") + + +def test_normalizing_is_an_explicit_operation_that_records_itself(tmp_path): + write_kin_file(tmp_path, span=TRUNCATED_SPAN) + profiles = read_kin(tmp_path / "synthetic.kin") + + stretched = normalize_psi(profiles, method="min_max") + assert stretched.psi_norm[0] == 0.0 and stretched.psi_norm[-1] == 1.0 + assert stretched.normalization.method == "min_max" + assert stretched.normalization.axis_value == pytest.approx(TRUNCATED_SPAN[0]) + # The measurable delta the legacy code applied silently. + assert abs(stretched.psi_norm[1] - profiles.psi_norm[1]) > 1e-3 + # And the original is untouched: the container is frozen. + assert profiles.psi_norm[0] == pytest.approx(TRUNCATED_SPAN[0]) + + +def test_explicit_normalization_needs_its_edge(): + profiles = KineticProfiles(psi_norm=np.linspace(0.1, 0.9, 5)) + with pytest.raises(ValueError, match="needs edge="): + normalize_psi(profiles, method="explicit") + scaled = normalize_psi(profiles, method="explicit", edge=2.0) + np.testing.assert_allclose(scaled.psi_norm, np.linspace(0.05, 0.45, 5)) + assert scaled.normalization.edge_value == 2.0 + + +def test_an_axis_and_edge_without_a_method_is_refused(): + """The silent rescale, with a provenance record vouching for it: a default + method would ignore both arguments and stretch the coordinate instead.""" + profiles = KineticProfiles(psi_norm=np.linspace(0.00495, 1.0, 5)) + with pytest.raises(TypeError, match="method"): + normalize_psi(profiles, axis=0.0, edge=1.0) + with pytest.raises(ValueError, match="takes the axis and edge from"): + normalize_psi(profiles, method="min_max", axis=0.0, edge=1.0) + + +def test_an_unknown_normalization_method_is_refused(): + profiles = KineticProfiles(psi_norm=np.linspace(0.1, 0.9, 5)) + with pytest.raises(ValueError, match="must be 'min_max' or 'explicit'"): + normalize_psi(profiles, method="stretch") + + +# --- GPEC's own header/footer rule -------------------------------------------- + + +def test_a_header_and_footer_are_ignored_the_way_gpec_ignores_them(tmp_path): + """"No lines start with a number" is the file's actual contract.""" + expected = write_kin_file(tmp_path, footer=True) + profiles = read_kin(tmp_path / "synthetic.kin") + assert len(profiles) == expected["psi_norm"].size + np.testing.assert_allclose(profiles.psi_norm, expected["psi_norm"]) + + +def test_the_table_stops_at_the_footer_rather_than_resuming_after_it(tmp_path): + """GPEC's readtable takes the *first contiguous block*: it stops at the + first non-numeric line and ignores the rest of the file. Collecting every + numeric line anywhere would read the summary row below the footer as one + more data point, silently disagreeing with GPEC about the profile.""" + expected = write_kin_file(tmp_path, numeric_footer=True) + profiles = read_kin(tmp_path / "synthetic.kin") + assert len(profiles) == expected["psi_norm"].size + assert profiles.psi_norm[-1] == pytest.approx(expected["psi_norm"][-1]) + + +def test_a_word_that_python_would_read_as_a_number_is_still_header_text(tmp_path): + """GPEC looks at one character: a line starting with a letter is header, + even when it is "nan" or "Infinity", which float() accepts.""" + expected = write_kin_file(tmp_path, header="nan is not a data line") + profiles = read_kin(tmp_path / "synthetic.kin") + assert len(profiles) == expected["psi_norm"].size + + +def test_fortran_exponents_are_read_rather_than_dropped(tmp_path): + """A .kin written by a Fortran tool carries 1.0D+00, which GPEC's + list-directed read accepts and Python's float() does not -- so a reader + that classifies by float() would silently discard every such row.""" + expected = write_kin_file(tmp_path, fortran_exponents=True) + profiles = read_kin(tmp_path / "synthetic.kin") + assert len(profiles) == expected["psi_norm"].size + np.testing.assert_allclose(profiles.n_e, expected["n_e"], rtol=1e-7) + + +def test_a_ragged_row_is_named_rather_than_dying_inside_numpy(tmp_path): + write_kin_file(tmp_path, ragged_row=3) + with pytest.raises(ValueError, match="data row 4 has 5 columns against 6"): + read_kin(tmp_path / "synthetic.kin") + + +def test_a_file_with_no_header_reads(tmp_path): + expected = write_kin_file(tmp_path, header=False) + profiles = read_kin(tmp_path / "synthetic.kin") + np.testing.assert_allclose(profiles.n_e, expected["n_e"]) + + +def test_a_file_with_no_data_rows_is_refused(tmp_path): + (tmp_path / "empty.kin").write_text("psi ni ne ti te wexb\nnothing here\n", encoding="utf-8") + with pytest.raises(ValueError, match="no data rows"): + read_kin(tmp_path / "empty.kin") + + +def test_too_few_columns_is_refused(tmp_path): + (tmp_path / "short.kin").write_text(" 1.0 2.0 3.0\n", encoding="utf-8") + with pytest.raises(ValueError, match="a .kin file has 6"): + read_kin(tmp_path / "short.kin") + + +# --- C-27: the rotation column is omega_E, not toroidal rotation -------------- + + +def test_writing_refuses_a_set_that_has_only_toroidal_rotation(tmp_path): + columns = profile_columns() + omega_tor = -columns.pop("omega_exb") + profiles = KineticProfiles(**columns, omega_tor=omega_tor) + with pytest.raises(ValueError, match="rotation column is omega_E"): + write_kin(profiles, tmp_path / "wrong.kin") + + +def test_writing_names_whatever_is_missing(tmp_path): + columns = profile_columns() + del columns["T_i"] + with pytest.raises(ValueError, match=r"missing \['T_i'\]"): + write_kin(KineticProfiles(**columns), tmp_path / "short.kin") + + +def test_an_all_zero_rotation_column_is_refused_unless_asked(tmp_path): + """GPEC substitutes 1e-9 for every zero, so the file would not mean what it says.""" + columns = profile_columns() + columns["omega_exb"] = np.zeros_like(columns["psi_norm"]) + profiles = KineticProfiles(**columns) + with pytest.raises(ValueError, match="replaces every zero with 1e-9"): + write_kin(profiles, tmp_path / "zero.kin") + written = write_kin(profiles, tmp_path / "zero.kin", allow_zero_rotation=True) + assert read_kin(written).omega_exb.max() == 0.0 + + +def test_a_partly_zero_rotation_column_is_refused_too(tmp_path): + """GPEC substitutes per element, not per column, so a core that was + zero-filled by a converter is the same defect on half the profile.""" + columns = profile_columns() + columns["omega_exb"] = columns["omega_exb"].copy() + columns["omega_exb"][:5] = 0.0 + with pytest.raises(ValueError, match="zero at 5 of"): + write_kin(KineticProfiles(**columns), tmp_path / "part.kin") + + +def test_a_non_finite_value_is_refused(tmp_path): + """GPEC's own warning at the 1e-9 substitution says a NaN ruins the + whole spline; nothing downstream recovers from one.""" + columns = profile_columns() + columns["T_e"] = columns["T_e"].copy() + columns["T_e"][7] = np.nan + with pytest.raises(ValueError, match="not finite"): + write_kin(KineticProfiles(**columns), tmp_path / "nan.kin") + + +def test_a_coordinate_that_does_not_increase_is_refused(tmp_path): + """GPEC splines against psi and assumes an increasing abscissa; a shuffled + column produces garbage there rather than an error.""" + columns = profile_columns() + psi = columns["psi_norm"].copy() + psi[3], psi[4] = psi[4], psi[3] + columns["psi_norm"] = psi + with pytest.raises(ValueError, match="does not increase"): + write_kin(KineticProfiles(**columns), tmp_path / "shuffled.kin") + + +# --- writing ------------------------------------------------------------------ + + +def test_the_written_header_is_the_one_gpecs_examples_carry(kin, tmp_path): + """Byte for byte against a literal, not against the module's own constant: + a header with the right tokens and the wrong spacing reads identically and + is no longer the line every real .kin carries.""" + profiles, _ = kin + text = write_kin(profiles, tmp_path / "out.kin").read_text(encoding="utf-8") + assert text.splitlines()[0] == GPEC_KIN_HEADER + # A header is a comment: the file still reads without it. + bare = write_kin(profiles, tmp_path / "bare.kin", header=False) + np.testing.assert_allclose(read_kin(bare).psi_norm, profiles.psi_norm) + + +def test_provenance_records_where_each_column_came_from(kin): + profiles, _ = kin + # n_e is the file's third column: psi, ni, ne. + assert profiles.provenance["psi_norm"].endswith("column 1") + assert profiles.provenance["n_e"].endswith("column 3") + assert profiles.source.endswith("synthetic.kin") + assert profiles.normalization.source == "synthetic.kin" diff --git a/vaft/data/__init__.py b/vaft/data/__init__.py index 89880d99c..a1f7ab898 100644 --- a/vaft/data/__init__.py +++ b/vaft/data/__init__.py @@ -20,10 +20,17 @@ "Gap", "GlobalEquilibriumDescriptors", "KEQDSK", + "KINETIC_UNITS", + "KIN_COLUMNS", + "KIN_HEADER", + "KineticProfiles", "MEQDSK", "MillerFitResult", "MillerSequenceResult", "MillerSurface", + "PROFILE_FIELDS", + "PsiNormalization", + "Species", "cocos_spec", "convention_for", "data_path", @@ -32,7 +39,10 @@ "from_equilibrium", "from_omas", "known_codes", + "normalize_psi", "read_geqdsk", + "kinetic_profiles", + "read_kin", "read_aeqdsk", "read_keqdsk", "read_meqdsk", @@ -45,6 +55,7 @@ "to_imas", "to_omas", "write_geqdsk", + "write_kin", "VAFT_INTERNAL_COCOS", "VFITResult", "SolovevConstraint", @@ -110,11 +121,21 @@ "ValidationReport": (".equilibrium", "ValidationReport"), "XPoint": (".equilibrium", "XPoint"), "read_vfit": (".vfit", "read_vfit"), + "KINETIC_UNITS": (".kinetic_profiles", "KINETIC_UNITS"), + "KIN_COLUMNS": (".kinetic_profiles", "KIN_COLUMNS"), + "KIN_HEADER": (".kinetic_profiles", "KIN_HEADER"), + "PROFILE_FIELDS": (".kinetic_profiles", "PROFILE_FIELDS"), + "KineticProfiles": (".kinetic_profiles", "KineticProfiles"), + "PsiNormalization": (".kinetic_profiles", "PsiNormalization"), + "Species": (".kinetic_profiles", "Species"), + "normalize_psi": (".kinetic_profiles", "normalize_psi"), + "read_kin": (".kinetic_profiles", "read_kin"), + "write_kin": (".kinetic_profiles", "write_kin"), } def __getattr__(name: str): - if name in {"resources", "open_adas"}: + if name in {"resources", "open_adas", "kinetic_profiles"}: module = import_module(f".{name}", __name__) globals()[name] = module return module diff --git a/vaft/data/kinetic_profiles.py b/vaft/data/kinetic_profiles.py new file mode 100644 index 000000000..b7ef73bf5 --- /dev/null +++ b/vaft/data/kinetic_profiles.py @@ -0,0 +1,453 @@ +"""Kinetic profiles as a file-format container, and GPEC ``.kin`` I/O. + +A kinetic profile set is a radial coordinate plus the densities, +temperatures, rotation frequencies and pressures defined on it. Three file +formats carry them -- GPEC's ``.kin``, the Osborne ``pfile`` and MARS's +``PROF*.IN`` -- and this module is the one container all three read into and +write from. + +**The canonical unit set is fixed by the container**, not carried per +instance: densities in m^-3, temperatures in eV, angular frequencies in +rad/s, pressures in Pa, electric fields in V/m. Those are the units GPEC's +own reader consumes, so the ``.kin`` path needs no conversion at all; the +pfile and MARS readers convert on the way in, through the single table +:data:`KINETIC_UNITS`. + +Two things this container is careful about, because the code it replaces was +not: + +- **The radial coordinate is never rescaled on read.** A reader records what + the file said and how it was normalised (:class:`PsiNormalization`); making + it span exactly [0, 1] is :func:`normalize_psi`, an explicit operation. + This matters physically: GPEC's ``read_kin`` re-splines a ``.kin`` onto a + uniform 101-point [0, 1] grid *with extrapolation* (``nkin = 100`` + intervals, ``inputs.f90:204,226``), so it expects a truncated edge and + handles it. Stretching the axis instead moves every + interior point -- for the reference MAST-U file, which spans + psi_n = 0.00495 to 1.0, by up to half a percent of the minor radius, and + permanently once the stretched values are written back. +- **Toroidal rotation and the E×B frequency are different fields.** + ``omega_tor`` and ``omega_exb`` are never merged, and :func:`write_kin` + refuses a profile set that has only the former: the ``.kin`` rotation + column is omega_E, and writing a toroidal rotation into it is silent and + wrong. The formats invite the mistake -- MARS writes the two to + ``PROFROT.IN`` and ``PROFWE.IN``, and a pfile has ten rotation-like + sections. + +Provenance: every field records where it came from (a file section, a code +variable) in :attr:`KineticProfiles.provenance`, so a converted profile set +can say which TRANSP variable or pfile section produced each quantity. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from pathlib import Path +from types import MappingProxyType +from typing import Mapping, Sequence + +import numpy as np + +__all__ = [ + "KINETIC_UNITS", + "KIN_COLUMNS", + "KIN_HEADER", + "PROFILE_FIELDS", + "KineticProfiles", + "PsiNormalization", + "Species", + "normalize_psi", + "read_kin", + "write_kin", +] + +#: The six ``.kin`` columns, in file order, as GPEC's ``read_kin`` defines +#: them (``pentrc/inputs.f90``): ``psi_n, n_i, n_e, T_i, T_e, omega_E``. +KIN_COLUMNS: tuple[str, ...] = ("psi_norm", "n_i", "n_e", "T_i", "T_e", "omega_exb") + +#: The header line GPEC's own example files carry. It is a comment as far as +#: the reader is concerned -- what fixes the meaning is the column order. +KIN_HEADER = ( + " psi ni(m^-3) ne(m^-3)" + " ti(eV) te(eV) wexb(rad/s)" +) + +#: The container's units, per field. One set, fixed here; a reader converts +#: into them and a writer converts out of them. Exported as +#: ``vaft.data.KINETIC_UNITS``; the pfile and MARS readers convert through +#: this table and no other. +KINETIC_UNITS: Mapping[str, str] = { + "psi_norm": "-", + "n_e": "m^-3", + "n_i": "m^-3", + "n_z": "m^-3", + "n_fast": "m^-3", + "T_e": "eV", + "T_i": "eV", + "T_z": "eV", + "omega_tor": "rad/s", + "omega_exb": "rad/s", + "omega_pol": "rad/s", + "p_total": "Pa", + "p_fast": "Pa", + "e_radial": "V/m", +} + +#: Profile fields, in the order a reader should prefer to report them. +PROFILE_FIELDS: tuple[str, ...] = tuple(name for name in KINETIC_UNITS if name != "psi_norm") + + +@dataclass(frozen=True) +class PsiNormalization: + """How a profile set's radial coordinate came to be what it is. + + ``method`` is ``"as_read"`` for a coordinate taken from a file unchanged, + or the name of the operation that produced it. Recording this is the + point: a coordinate that has been rescaled looks exactly like one that + has not. + """ + + method: str = "as_read" + source: str = "" + axis_value: float | None = None + edge_value: float | None = None + + +@dataclass(frozen=True) +class Species: + """One ion species from a pfile's ``N Z A`` block. + + A ``.kin`` carries no species block, so :func:`read_kin` never builds + one; it is the type of :attr:`KineticProfiles.species`, which the pfile + reader populates. + """ + + label: str + n: float + z: float + a: float + + +def _sealed(values) -> np.ndarray: + """A read-only *view* of ``values``, leaving the caller's array writable. + + ``frozen=True`` stops the attributes being rebound and nothing else, so an + array reached through one would still be editable in place -- and a + provenance record that says "as read" would then be untrue. A view has + its own writeable flag, so sealing it here does not reach back into the + array the caller passed in. + """ + array = np.asarray(values).view() + array.setflags(write=False) + return array + + +@dataclass(frozen=True, eq=False) +class KineticProfiles: + """Kinetic profiles on one radial coordinate, in the container's units. + + Every profile field is optional because the formats differ in what they + carry; :meth:`available` says which are present. Anything a file holds + that has no field of its own is kept in :attr:`extras` under its own + name, so nothing is dropped. + + The arrays and mappings are sealed on construction, so a profile set + cannot be edited in place behind its own provenance record; build a + changed one with :func:`dataclasses.replace`. ``eq`` is off because the + fields are arrays: a generated ``__eq__`` would raise rather than answer. + """ + + psi_norm: np.ndarray + n_e: np.ndarray | None = None + n_i: np.ndarray | None = None + n_z: np.ndarray | None = None + n_fast: np.ndarray | None = None + T_e: np.ndarray | None = None + T_i: np.ndarray | None = None + T_z: np.ndarray | None = None + omega_tor: np.ndarray | None = None + omega_exb: np.ndarray | None = None + omega_pol: np.ndarray | None = None + p_total: np.ndarray | None = None + p_fast: np.ndarray | None = None + e_radial: np.ndarray | None = None + normalization: PsiNormalization = field(default_factory=PsiNormalization) + species: tuple[Species, ...] = () + extras: Mapping[str, np.ndarray] = field(default_factory=dict) + provenance: Mapping[str, str] = field(default_factory=dict) + source: str = "" + + def __post_init__(self) -> None: + object.__setattr__(self, "psi_norm", _sealed(self.psi_norm)) + size = self.psi_norm.size + for name in PROFILE_FIELDS: + values = getattr(self, name) + if values is None: + continue + if np.asarray(values).size != size: + raise ValueError( + f"{name} has {np.asarray(values).size} samples but psi_norm has " + f"{size}; every profile is on the same radial coordinate" + ) + object.__setattr__(self, name, _sealed(values)) + object.__setattr__( + self, + "extras", + MappingProxyType({str(k): _sealed(v) for k, v in dict(self.extras).items()}), + ) + object.__setattr__(self, "provenance", MappingProxyType(dict(self.provenance))) + + def __len__(self) -> int: + return int(np.asarray(self.psi_norm).size) + + def available(self) -> tuple[str, ...]: + """Names of the profile fields this set actually carries.""" + return tuple(name for name in PROFILE_FIELDS if getattr(self, name) is not None) + + def field(self, name: str) -> np.ndarray: + """One profile by name, from a field or from :attr:`extras`.""" + values = getattr(self, name, None) if name in KINETIC_UNITS else None + if values is None: + values = self.extras.get(name) + if values is None: + raise KeyError( + f"no profile {name!r}; this set has {list(self.available())} " + f"and extras {sorted(self.extras)}" + ) + return np.asarray(values) + + def unit(self, name: str) -> str: + """The container's unit for ``name``; ``""`` for an extra.""" + return KINETIC_UNITS.get(name, "") + + +def normalize_psi( + profiles: KineticProfiles, + *, + method: str, + axis: float | None = None, + edge: float | None = None, +) -> KineticProfiles: + """Rescale the radial coordinate, explicitly, recording that it happened. + + ``"min_max"`` maps the coordinate's own range onto [0, 1] -- what the + legacy readers did silently on every read. ``"explicit"`` divides by + ``edge`` after subtracting ``axis``, for a caller that knows the axis and + edge values from elsewhere. + + ``method`` has no default on purpose. A caller who passes ``axis`` and + ``edge`` has said which rescale they mean, and a default of ``"min_max"`` + would ignore both, stretch the coordinate instead, and stamp the result + with a :class:`PsiNormalization` that certifies the operation they did + not ask for -- which is the silent rescale this module exists to prevent, + only now with a provenance record vouching for it. + """ + psi = np.asarray(profiles.psi_norm, dtype=float) + if method == "min_max": + if axis is not None or edge is not None: + raise ValueError( + "method='min_max' takes the axis and edge from the coordinate's own " + "range; pass method='explicit' to use the axis= and edge= given" + ) + axis_value, edge_value = float(psi.min()), float(psi.max()) + elif method == "explicit": + if edge is None: + raise ValueError("method='explicit' needs edge=") + axis_value, edge_value = float(axis or 0.0), float(edge) + else: + raise ValueError(f"method must be 'min_max' or 'explicit', got {method!r}") + span = edge_value - axis_value + if span == 0: + raise ValueError("the radial coordinate has zero range; nothing to normalize") + return replace( + profiles, + psi_norm=(psi - axis_value) / span, + normalization=PsiNormalization( + method=method, + source=profiles.normalization.source, + axis_value=axis_value, + edge_value=edge_value, + ), + ) + + +def _starts_with_number(line: str) -> bool: + """GPEC's own test: after leading blanks and up to two signs, a digit. + + Deliberately not ``float(token)``: ``readtable`` + (``pentrc/utilities.f90:436-444``) looks at one character, so ``nan`` and + ``Infinity`` are header text to GPEC where Python would read them as + numbers, and ``.5`` is header text too. + """ + text = line.lstrip() + index = 0 + for _ in range(2): # readtable passes over a plus and a minus + if index < len(text) and text[index] in "+-": + index += 1 + return index < len(text) and text[index] in "0123456789" + + +def _data_lines(lines: Sequence[str]) -> list[str]: + """The first contiguous block of numeric lines, as GPEC's reader takes it. + + ``readtable`` (``pentrc/utilities.f90:428-451``) sets ``startline`` to the + first line beginning with a number and ``endline`` to the line before the + *next* line that does not -- and then ignores the rest of the file. So a + footer stops the table rather than being skipped over, and a numeric line + after one is not data. Collecting every numeric line anywhere instead + would read a summary row appended below a footer as a fourth data point + where GPEC reads three, with no error on either side. + """ + data: list[str] = [] + started = False + for line in lines: + if _starts_with_number(line): + started = True + data.append(line.strip()) + elif started: + break + return data + + +def _float(token: str) -> float: + """A number as Fortran writes it, ``1.0D+00`` included. + + GPEC's list-directed ``read`` accepts a ``D`` exponent, so a ``.kin`` + written by a Fortran tool carries them; Python's ``float`` does not, and + dropping such a row would lose data from a file GPEC reads without + complaint. + """ + try: + return float(token) + except ValueError: + return float(token.replace("D", "E").replace("d", "e")) + + +def read_kin(path: str | Path) -> KineticProfiles: + """Read a GPEC ``.kin`` file. + + The radial coordinate is taken exactly as written; nothing is rescaled. + """ + path = Path(path).expanduser() + # Universal newlines, so a CRLF file reads correctly; write_kin emits LF. + lines = path.read_text(encoding="utf-8").splitlines() + rows = _data_lines(lines) + if not rows: + raise ValueError(f"{path.name} holds no data rows (no line starts with a number)") + + split = [row.split() for row in rows] + width = len(split[0]) + # GPEC allocates its table from the first data line's token count and + # reads the rest into it, so a ragged row is a misparse there and an + # inhomogeneous-shape error from numpy here -- neither of which names the + # line. Say which one it is. + for offset, tokens in enumerate(split): + if len(tokens) != width: + raise ValueError( + f"{path.name} data row {offset + 1} has {len(tokens)} columns against " + f"{width} on the first row: {rows[offset]!r}" + ) + if width < len(KIN_COLUMNS): + raise ValueError( + f"{path.name} has {width} columns; a .kin file has " + f"{len(KIN_COLUMNS)}: {', '.join(KIN_COLUMNS)}" + ) + table = np.array([[_float(token) for token in tokens] for tokens in split], dtype=float) + + columns = {name: table[:, index] for index, name in enumerate(KIN_COLUMNS)} + extras = { + f"column_{index + 1}": table[:, index] + for index in range(len(KIN_COLUMNS), width) + } + # Columns are numbered as a person reading the file would: the first is 1. + provenance = { + name: f"{path.name} column {index + 1}" for index, name in enumerate(KIN_COLUMNS) + } + provenance.update( + {name: f"{path.name} column {index + 1}" for index, name in enumerate(extras, len(KIN_COLUMNS))} + ) + return KineticProfiles( + **columns, + normalization=PsiNormalization(method="as_read", source=path.name), + extras=extras, + provenance=provenance, + source=str(path), + ) + + +def write_kin( + profiles: KineticProfiles, + path: str | Path, + *, + header: bool = True, + allow_zero_rotation: bool = False, +) -> Path: + """Write a GPEC ``.kin`` file. + + Raises :class:`ValueError` rather than write a file GPEC would read as + something other than what it says, in four cases: + + - only ``omega_tor`` is present. The ``.kin`` rotation column is + omega_E, and a toroidal rotation written there is wrong in a way + nothing downstream can detect. + - ``omega_exb`` holds a zero. GPEC substitutes 1e-9 for *each* zero + element to keep its spline finite (``pentrc/inputs.f90:255-265``), so + those points do not mean what the file says; ``allow_zero_rotation`` + writes them anyway, which is what a set converted from an all-zero + ``PROFROT.IN`` needs. + - any value is not finite. GPEC's own warning at that substitution says + a NaN "ruins the whole spline", and nothing downstream recovers. + - the radial coordinate does not increase. ``spline_fit`` + (``inputs.f90:219-222``) assumes an increasing abscissa and silently + produces a garbage spline otherwise. + + Line endings are LF, so a file read from CRLF input is not written back + byte-identically. + """ + path = Path(path).expanduser() + missing = [name for name in KIN_COLUMNS if getattr(profiles, name, None) is None] + if "omega_exb" in missing and profiles.omega_tor is not None: + raise ValueError( + "this profile set has omega_tor but no omega_exb, and a .kin file's " + "rotation column is omega_E; convert explicitly rather than writing " + "toroidal rotation into it" + ) + if missing: + raise ValueError(f"a .kin file needs {list(KIN_COLUMNS)}; missing {missing}") + + table = np.column_stack( + [np.asarray(getattr(profiles, name), dtype=float) for name in KIN_COLUMNS] + ) + if not np.all(np.isfinite(table)): + bad = { + name: int(np.count_nonzero(~np.isfinite(table[:, index]))) + for index, name in enumerate(KIN_COLUMNS) + if not np.all(np.isfinite(table[:, index])) + } + raise ValueError( + f"{bad} hold values that are not finite; GPEC splines the file it reads " + "and a single NaN ruins the whole spline" + ) + + psi = table[:, 0] + if not np.all(np.diff(psi) > 0): + first = int(np.argmin(np.diff(psi) > 0)) + raise ValueError( + f"the radial coordinate does not increase (row {first + 1} is " + f"{psi[first]:.8e}, row {first + 2} is {psi[first + 1]:.8e}); GPEC splines " + "against it and assumes an increasing abscissa" + ) + + zeros = int(np.count_nonzero(table[:, KIN_COLUMNS.index("omega_exb")] == 0.0)) + if zeros and not allow_zero_rotation: + raise ValueError( + f"omega_exb is zero at {zeros} of {len(psi)} points; GPEC replaces every " + "zero with 1e-9 to keep its spline finite, so those points would not mean " + "what the file says. Pass allow_zero_rotation=True to write it anyway" + ) + + lines = [KIN_HEADER] if header else [] + lines.extend(" " + " ".join(f"{value:.8e}" for value in row) for row in table) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path