From ce9f5664231477a793a5e7c78f5524318cf0222e Mon Sep 17 00:00:00 2001 From: priyanshlunia <40486607+priyanshlunia@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:12:34 -0400 Subject: [PATCH 01/15] refactor(core): promote SLCONTOUR pipeline into shim-free core/qs_* MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Retire the dead core/quasistationary.py (a partial, unwired port whose reconstruct_grid used exp(+i(nφ+mθ)) — wrong for any phased mode) plus its falsely-certifying test, and promote the *active* SLCONTOUR pipeline into core under unified qs_ naming, deleting the omfit_compat OMFIT shim. - Move _slcontour/{fit,prep,io_data,run,plots}.py -> core/qs_*.py. - Delete omfit_compat.py, folding each symbol to its owner: printv/w/e -> logging; OMFITexception -> ValueError; delta_degrees -> qs_fit; cornernote/uband -> qs_plots; device readers -> new core/qs_device.py, which delegates to the data layer (data.devices / data.diiid_geometry) and keeps only the QS-specific extras (derived theta/end-coord sensor_geometry, resolve_channel_filter semantics). - qs_device.load_wall now delegates to devices.feature_at, fixing the segmented- schema bug (the shim read a nonexistent top-level `wall` key -> (None, None)). - Lock the -i reconstruction convention: invariant comments at the three sites + a discriminating δ≠0 regression test the old zero-phase test missed. - qs_plots docstring clarified as standalone matplotlib (not GUI/service wired). - Rewire service/nodes.py + tests; pyproject drops the ruff exclude and retargets a deferred ty exclude onto the qs_* files. Relocate the reference notebook/README to docs/qs_slcontour/. ruff format/check + ty (with service extra) clean; 209 passed, 1 skipped. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../qs_slcontour}/README.md | 0 .../qs_slcontour}/example_magnetics.ipynb | 73 ++-- pyproject.toml | 16 +- src/magnetics/_slcontour/__init__.py | 7 - src/magnetics/_slcontour/omfit_compat.py | 363 ------------------ src/magnetics/core/qs_bridge.py | 7 +- src/magnetics/core/qs_device.py | 206 ++++++++++ .../{_slcontour/fit.py => core/qs_fit.py} | 71 +++- .../io_data.py => core/qs_io_data.py} | 21 +- .../{_slcontour/plots.py => core/qs_plots.py} | 123 ++++-- .../{_slcontour/prep.py => core/qs_prep.py} | 33 +- .../{_slcontour/run.py => core/qs_run.py} | 6 +- src/magnetics/core/quasistationary.py | 256 ------------ src/magnetics/service/nodes.py | 10 +- tests/test_qs_bridge.py | 36 ++ tests/test_qs_pipeline.py | 4 +- tests/test_quasistationary.py | 59 --- tests/test_slcontour_fit.py | 9 +- tests/test_slcontour_geometry.py | 6 +- 19 files changed, 505 insertions(+), 801 deletions(-) rename {src/magnetics/_slcontour => docs/qs_slcontour}/README.md (100%) rename {src/magnetics/_slcontour => docs/qs_slcontour}/example_magnetics.ipynb (80%) delete mode 100644 src/magnetics/_slcontour/__init__.py delete mode 100644 src/magnetics/_slcontour/omfit_compat.py create mode 100644 src/magnetics/core/qs_device.py rename src/magnetics/{_slcontour/fit.py => core/qs_fit.py} (89%) rename src/magnetics/{_slcontour/io_data.py => core/qs_io_data.py} (92%) rename src/magnetics/{_slcontour/plots.py => core/qs_plots.py} (76%) rename src/magnetics/{_slcontour/prep.py => core/qs_prep.py} (91%) rename src/magnetics/{_slcontour/run.py => core/qs_run.py} (96%) delete mode 100644 src/magnetics/core/quasistationary.py delete mode 100644 tests/test_quasistationary.py diff --git a/src/magnetics/_slcontour/README.md b/docs/qs_slcontour/README.md similarity index 100% rename from src/magnetics/_slcontour/README.md rename to docs/qs_slcontour/README.md diff --git a/src/magnetics/_slcontour/example_magnetics.ipynb b/docs/qs_slcontour/example_magnetics.ipynb similarity index 80% rename from src/magnetics/_slcontour/example_magnetics.ipynb rename to docs/qs_slcontour/example_magnetics.ipynb index a0363d3..70b68d9 100644 --- a/src/magnetics/_slcontour/example_magnetics.ipynb +++ b/docs/qs_slcontour/example_magnetics.ipynb @@ -37,7 +37,8 @@ "outputs": [], "source": [ "import sys\n", - "sys.path.insert(0, '.') # make the local modules importable\n", + "\n", + "sys.path.insert(0, \".\") # make the local modules importable\n", "\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", @@ -69,18 +70,19 @@ "outputs": [], "source": [ "# ── Run parameters — change these to analyse a different shot/array ──\n", - "SHOT = 199749 # data/datafile/shot_.h5\n", - "CHANNEL_FILTER = 'Bp_LFS_midplane' # friendly name (see available_subsets) or a regex\n", - "TIME_TRIM = (3.3, 3.5) # seconds; must lie inside the shot file's window\n", - "NS, MS = (1, 2, 3), (0,) # toroidal / poloidal mode numbers\n", + "SHOT = 199749 # data/datafile/shot_.h5\n", + "CHANNEL_FILTER = \"Bp_LFS_midplane\" # friendly name (see available_subsets) or a regex\n", + "TIME_TRIM = (3.3, 3.5) # seconds; must lie inside the shot file's window\n", + "NS, MS = (1, 2, 3), (0,) # toroidal / poloidal mode numbers\n", "\n", "# detrend: 'none' | 'baseline' | 'linear' | 'endpoints'. DETREND_BAND is the\n", "# sub-window used to estimate the trend — here the first 0.01 s of TIME_TRIM.\n", - "DETREND_TYPE = 'baseline'\n", - "DETREND_BAND = (TIME_TRIM[0], TIME_TRIM[0] + 0.01)\n", + "DETREND_TYPE = \"baseline\"\n", + "DETREND_BAND = (TIME_TRIM[0], TIME_TRIM[0] + 0.01)\n", "\n", - "PREP_KWARGS = dict(cutoff_hz=(5, 250), energy=0.98,\n", - " detrend_type=DETREND_TYPE, detrend_band=DETREND_BAND)" + "PREP_KWARGS = dict(\n", + " cutoff_hz=(5, 250), energy=0.98, detrend_type=DETREND_TYPE, detrend_band=DETREND_BAND\n", + ")" ] }, { @@ -105,19 +107,19 @@ "outputs": [], "source": [ "sd = load_shot(SHOT)\n", - "print(f'shot {sd.shot} device {sd.device}')\n", - "print('RAW channels:', sd.raw.sizes['channel'], ' time samples:', sd.raw.sizes['time'])\n", - "print(f'file window: {float(sd.raw.time[0]):.4f}–{float(sd.raw.time[-1]):.4f} s')\n", - "print('helicity:', sd.plasma.attrs['helicity'])\n", + "print(f\"shot {sd.shot} device {sd.device}\")\n", + "print(\"RAW channels:\", sd.raw.sizes[\"channel\"], \" time samples:\", sd.raw.sizes[\"time\"])\n", + "print(f\"file window: {float(sd.raw.time[0]):.4f}–{float(sd.raw.time[-1]):.4f} s\")\n", + "print(\"helicity:\", sd.plasma.attrs[\"helicity\"])\n", "\n", "# channels in the selected array that carry good data\n", "good = valid_channels(sd.raw, CHANNEL_FILTER, sd.device)\n", - "print(f'\\n{CHANNEL_FILTER}: {len(good)} valid sensors -> resolves |n| <= {(len(good)-1)//2}')\n", + "print(f\"\\n{CHANNEL_FILTER}: {len(good)} valid sensors -> resolves |n| <= {(len(good) - 1) // 2}\")\n", "print(good)\n", "\n", "# the plot helpers match channel names with re.match, so turn the friendly\n", "# filter into an explicit regex of its sensor names\n", - "ARRAY_REGEX = '|'.join(resolve_channel_filter(CHANNEL_FILTER, sd.device))" + "ARRAY_REGEX = \"|\".join(resolve_channel_filter(CHANNEL_FILTER, sd.device))" ] }, { @@ -139,7 +141,7 @@ "outputs": [], "source": [ "for name, sensors in available_subsets(sd.device).items():\n", - " print(f'{name:24s} {sensors}')" + " print(f\"{name:24s} {sensors}\")" ] }, { @@ -160,10 +162,10 @@ "outputs": [], "source": [ "fig, ax = plt.subplots(1, 2, figsize=(12, 5))\n", - "plots.plot_sensors(sd.raw, ARRAY_REGEX, geometry='rz', ax=ax[0])\n", - "ax[0].set_title('Cross-section (R–z)')\n", - "plots.plot_sensors(sd.raw, ARRAY_REGEX, geometry='cylindrical', ax=ax[1])\n", - "ax[1].set_title('Unrolled (phi–theta)')\n", + "plots.plot_sensors(sd.raw, ARRAY_REGEX, geometry=\"rz\", ax=ax[0])\n", + "ax[0].set_title(\"Cross-section (R–z)\")\n", + "plots.plot_sensors(sd.raw, ARRAY_REGEX, geometry=\"cylindrical\", ax=ax[1])\n", + "ax[1].set_title(\"Unrolled (phi–theta)\")\n", "plt.tight_layout()" ] }, @@ -187,20 +189,22 @@ "metadata": {}, "outputs": [], "source": [ - "sets = ['Bp_LFS_midplane', 'Bp_LFS_R+1', 'Bp_LFS_R-1', 'Bp_LFS_R+2', 'Bp_LFS_R-2', 'Bp_HFS_Sensors']\n", + "sets = [\"Bp_LFS_midplane\", \"Bp_LFS_R+1\", \"Bp_LFS_R-1\", \"Bp_LFS_R+2\", \"Bp_LFS_R-2\", \"Bp_HFS_Sensors\"]\n", "colors = plt.cm.tab10.colors\n", "\n", "fig, ax = plt.subplots(1, 2, figsize=(13, 5))\n", "handles = []\n", "for i, name in enumerate(sets):\n", - " regex = '|'.join(resolve_channel_filter(name, sd.device)) # friendly name -> regex of its sensors\n", - " plots.plot_sensors(sd.raw, regex, geometry='cylindrical', ax=ax[0], color=colors[i])\n", - " plots.plot_sensors(sd.raw, regex, geometry='rz', ax=ax[1], color=colors[i])\n", + " regex = \"|\".join(\n", + " resolve_channel_filter(name, sd.device)\n", + " ) # friendly name -> regex of its sensors\n", + " plots.plot_sensors(sd.raw, regex, geometry=\"cylindrical\", ax=ax[0], color=colors[i])\n", + " plots.plot_sensors(sd.raw, regex, geometry=\"rz\", ax=ax[1], color=colors[i])\n", " handles.append(plt.Line2D([], [], color=colors[i], label=name))\n", "\n", - "ax[0].set_title('Unrolled (phi–theta)')\n", - "ax[1].set_title('Cross-section (R–z)')\n", - "ax[0].legend(handles=handles, fontsize=8, loc='upper right')\n", + "ax[0].set_title(\"Unrolled (phi–theta)\")\n", + "ax[1].set_title(\"Cross-section (R–z)\")\n", + "ax[0].legend(handles=handles, fontsize=8, loc=\"upper right\")\n", "plt.tight_layout()" ] }, @@ -228,14 +232,17 @@ "r = run_steps(\n", " SHOT,\n", " channel_filter=CHANNEL_FILTER,\n", - " ns=NS, ms=MS,\n", + " ns=NS,\n", + " ms=MS,\n", " time_trim=TIME_TRIM,\n", " prep_kwargs=PREP_KWARGS,\n", ")\n", - "print(f'\\ncondition number K = {r.condition_number:.2f}')\n", - "print(f'mean reduced chi^2 = {float(np.nanmean(r.fit[\"red_chi_sq\"].values)):.1f}')\n", - "print(f'data-matrix SVD effective rank (98% energy) = {r.fit.attrs[\"signal_effective_rank\"]} '\n", - " f'of {r.fit.sizes[\"channel\"]} sensors')" + "print(f\"\\ncondition number K = {r.condition_number:.2f}\")\n", + "print(f\"mean reduced chi^2 = {float(np.nanmean(r.fit['red_chi_sq'].values)):.1f}\")\n", + "print(\n", + " f\"data-matrix SVD effective rank (98% energy) = {r.fit.attrs['signal_effective_rank']} \"\n", + " f\"of {r.fit.sizes['channel']} sensors\"\n", + ")" ] }, { @@ -347,7 +354,7 @@ "outputs": [], "source": [ "# fix_value is the poloidal angle of the slice; the LFS-midplane array sits at theta ~ 0\n", - "plots.plot_slice(r.fit, fix_coord='theta', fix_value=0.0)" + "plots.plot_slice(r.fit, fix_coord=\"theta\", fix_value=0.0)" ] }, { diff --git a/pyproject.toml b/pyproject.toml index ff3a19d..87bfc5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,13 +48,19 @@ magnetics-service = "magnetics.service.app:main" [tool.ruff] line-length = 100 # target version is inferred from requires-python -# _slcontour/ is the reference SLCONTOUR translation (a self-contained OMFIT -# port), pending rework into core/quasistationary (issue #40); excluded from -# lint/typecheck until then. -extend-exclude = ["src/magnetics/_slcontour"] [tool.ty.src] -exclude = ["src/magnetics/_slcontour"] +# The qs_* modules are the former _slcontour SLCONTOUR translation, now promoted into +# core (shim-free). They are ruff format + lint clean; the xarray-heavy `ty` pass is a +# documented follow-up (see docs/qs_slcontour/ and the retire-quasistationary plan). +exclude = [ + "src/magnetics/core/qs_fit.py", + "src/magnetics/core/qs_prep.py", + "src/magnetics/core/qs_io_data.py", + "src/magnetics/core/qs_run.py", + "src/magnetics/core/qs_plots.py", + "src/magnetics/core/qs_device.py", +] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/src/magnetics/_slcontour/__init__.py b/src/magnetics/_slcontour/__init__.py deleted file mode 100644 index 50616a0..0000000 --- a/src/magnetics/_slcontour/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""SLCONTOUR quasi-stationary pipeline — a local translation of the OMFIT -magnetics module (self-contained; no real OMFIT/MDSplus dependency, see -``omfit_compat``). Driven by the service's QS nodes through ``core.qs_bridge``. - -Reference code, slated for porting into ``magnetics.core.quasistationary`` -(issue #40); excluded from lint/typecheck until that port lands. -""" diff --git a/src/magnetics/_slcontour/omfit_compat.py b/src/magnetics/_slcontour/omfit_compat.py deleted file mode 100644 index 9ac49db..0000000 --- a/src/magnetics/_slcontour/omfit_compat.py +++ /dev/null @@ -1,363 +0,0 @@ -"""Compatibility shim for the OMFIT runtime symbols used by the magnetics scripts. - -The OMFIT magnetics module runs inside the OMFIT -framework, which injects a large namespace (``root[...]``, ``defaultVars``, -``OMFITx``, ``printi``/``printw``, ``cornernote``, ``uband``, ``is_device`` ...) -into every script. When we port those scripts to run locally we only need a -small subset of that namespace. This module reimplements exactly that subset so -the ported ``prep`` / ``fit`` / ``plots`` code reads almost identically to the -OMFIT originals. - -Nothing here talks to MDSplus, OMFIT, or a remote server. -""" - -from __future__ import annotations - -import os -import re - -import numpy as np - -# --------------------------------------------------------------------------- # -# Console messaging (OMFIT's printi/printv/printw/printe) -# --------------------------------------------------------------------------- # -# In OMFIT these route to the GUI console with severity colouring. Locally we -# just print, with a severity prefix for the warnings/errors. - - -def printi(*args): - """Informational message.""" - print(*args) - - -def printv(*args): - """Verbose message (same as printi locally).""" - print(*args) - - -def printw(*args): - """Warning message.""" - print("WARNING:", *args) - - -def printe(*args): - """Error/alert message (non-fatal, like the OMFIT original).""" - print("ERROR:", *args) - - -class OMFITexception(Exception): - """Local stand-in for omfit_classes.exceptions_omfit.OMFITexception.""" - - -# --------------------------------------------------------------------------- # -# Small numeric helpers OMFIT provides in its namespace -# --------------------------------------------------------------------------- # - - -def is_device(device, name): - """Loose device-name comparison (OMFIT's omfit_classes.utils_base.is_device). - - Treats ``DIII-D`` / ``DIIID`` / ``diii_d`` as equivalent. - """ - - def _norm(s): - return re.sub(r"[^a-z0-9]", "", str(s).lower()) - - return _norm(device) == _norm(name) - - -def _delta_degrees_scalar(theta1, theta2): - """Angular width from theta1 to theta2 in degrees, wrapping once through 0.""" - dt = theta2 - theta1 - if dt > 180: - dt -= 360 - if dt < -180: - dt += 360 - return dt - - -#: Vectorised version, matching ``delta_degrees`` in fit_magnetics.py. -delta_degrees = np.vectorize(_delta_degrees_scalar) - - -# --------------------------------------------------------------------------- # -# Plot helpers (OMFIT's cornernote and uband) -# --------------------------------------------------------------------------- # - - -def cornernote(ax=None, device="", shot="", time="", text="", **_ignore): - """Annotate the bottom-right corner with shot/device context. - - Mirrors omfit_classes.utils_plot.cornernote: a small grey label in the - figure corner. Extra keyword arguments (e.g. ``root=``) are accepted and - ignored for call-site compatibility. - """ - import matplotlib.pyplot as plt - - if ax is None: - ax = plt.gca() - fig = ax.get_figure() - label = " ".join(str(s) for s in (device, shot, time, text) if str(s) != "") - if not label: - return - fig.text( - 0.99, - 0.01, - label, - ha="right", - va="bottom", - fontsize=8, - color="0.4", - ) - - -def uband(x, y, yerr, ax=None, label=None, color=None, **kwargs): - """Plot a line with a shaded +/- uncertainty band. - - Local replacement for omfit_classes.utils_plot.uband, which takes an - ``uncertainties`` ``unumpy`` array. Here we pass the nominal values ``y`` - and their standard deviation ``yerr`` explicitly to avoid the extra - dependency. - - Returns a one-tuple ``(line,)`` so call sites can write ``(l,) = uband(...)`` - exactly as in the OMFIT plots. - """ - import matplotlib.pyplot as plt - - if ax is None: - ax = plt.gca() - x = np.asarray(x) - y = np.asarray(y, dtype=float) - yerr = np.asarray(yerr, dtype=float) - (line,) = ax.plot(x, y, label=label, color=color, **kwargs) - ax.fill_between( - x, - y - yerr, - y + yerr, - color=line.get_color(), - alpha=0.3, - linewidth=0, - ) - return (line,) - - -# --------------------------------------------------------------------------- # -# Device reference data (sensor geometry / wall / named subsets) -# --------------------------------------------------------------------------- # -# All device metadata now comes from the project-canonical ``data/device/.json`` -# (e.g. ``diiid.json``), replacing the old bundled ``DATA//*.txt`` tables. -# That single JSON carries ``R0``, the first-wall outline, the per-sensor base -# geometry and the named ``sensor_sets`` used to resolve a ``channel_filter``. - -import json - -#: Root of the canonical device files — the package's single source of truth -#: (``magnetics/data/device/``), resolved through the data layer. -from ..data import devices as _devices - -DEVICE_DIR = str(_devices.DEVICE_DIR) - -#: Explicit device-name -> json-filename overrides; otherwise slugified. -_DEVICE_FILE = {"DIII-D": "diiid.json"} - -_device_cache = {} - - -def _to_float(s): - try: - return float(s) - except (TypeError, ValueError): - return float("nan") - - -def _segment_fields(rec, shot=None): - """The geometry fields of a sensor record active at ``shot``. - - The device JSON is shot-segmented (each sensor is ``{"segments": [{since_shot, - r, z, ...}]}``); a segment is valid from its ``since_shot`` until the next. - Returns the segment active at ``shot`` (the latest whose ``since_shot`` ≤ shot), - falling back to the earliest segment for a ``shot`` before any campaign (a - layout fallback, since sensors move little between eras) or when ``shot`` is - None. Tolerates a legacy flat record (returned as-is).""" - segs = rec.get("segments") if isinstance(rec, dict) else None - if not segs: - return rec # legacy flat record → fields live at the top level - segs = sorted(segs, key=lambda s: s.get("since_shot", 0)) - active = segs[0] - for s in segs: - if shot is not None and s.get("since_shot", 0) <= shot: - active = s - return active - - -def _device_slug(device): - """JSON filename for ``device`` (e.g. ``'DIII-D'`` -> ``'diiid.json'``).""" - if device in _DEVICE_FILE: - return _DEVICE_FILE[device] - return re.sub(r"[^a-z0-9]", "", str(device).lower()) + ".json" - - -def device_file(device="DIII-D"): - """Path to the canonical device JSON for ``device``.""" - return os.path.join(DEVICE_DIR, _device_slug(device)) - - -def load_device(device="DIII-D"): - """Load (and cache) the device description JSON for ``device``. - - Returns the parsed dict: ``name, R0, wall{r,z}, sensors{...}, sensor_sets{...}``. - """ - path = device_file(device) - if path not in _device_cache: - with open(path) as fh: - _device_cache[path] = json.load(fh) - return _device_cache[path] - - -def _norm_name(name): - """Normalise a subset name so ``'Bp_LFS_midplane'`` == ``'Bp LFS midplane'``.""" - return str(name).strip().replace("_", " ") - - -def resolve_sensor_set(name, sets, _seen=None): - """Flatten one ``sensor_sets`` entry to an ordered, de-duplicated name list. - - ``list`` sets return their ``sensors``; ``composite`` sets recurse into their - member ``sets``. Cycles are guarded against. - """ - if _seen is None: - _seen = set() - if name in _seen or name not in sets: - return [] - _seen.add(name) - spec = sets[name] - if spec.get("type") == "list": - members = list(spec.get("sensors", [])) - else: # composite - members = [] - for sub in spec.get("sets", []): - members += resolve_sensor_set(sub, sets, _seen) - out, seen = [], set() - for m in members: - if m not in seen: - seen.add(m) - out.append(m) - return out - - -def list_sensor_subsets(device="DIII-D"): - """All named sensor subsets -> ``{name: [sensor, ...]}`` (composites flattened). - - This is the discoverability entry point: every key here is a valid - ``channel_filter`` argument (with or without underscores). - """ - sets = load_device(device).get("sensor_sets", {}) - return {name: resolve_sensor_set(name, sets) for name in sets} - - -def resolve_channel_filter(channel_filter, device="DIII-D"): - """Map a friendly subset name (e.g. ``'Bp_LFS_midplane'``, ``'All_3D_Coils'``) to a pattern list. - - A known subset name resolves to the explicit (anchored) sensor names from the - device ``sensor_sets``; an unknown string is treated as a raw regex and passed - through. Lists/tuples expand each element the same way. Patterns are matched - downstream with :func:`re.match`, so resolved names are anchored with ``$`` to - avoid accidental prefix matches (e.g. ``C19`` vs ``C190``). - """ - sets = load_device(device).get("sensor_sets", {}) - lookup = {_norm_name(k): k for k in sets} - - def resolve_one(cf): - if isinstance(cf, str): - key = lookup.get(_norm_name(cf)) - if key is not None: - return [re.escape(s) + "$" for s in resolve_sensor_set(key, sets)] - return [cf] # raw regex - return [cf] - - if isinstance(channel_filter, str): - return resolve_one(channel_filter) - out = [] - for cf in channel_filter: - out += resolve_one(cf) - return out - - -def sensor_geometry(device="DIII-D", shot=None): - """Per-sensor geometry as an ``xarray.Dataset`` indexed by ``channel``. - - Carries the base fields from the device JSON (``r, z, phi, tilt, length, - delta_phi, na``, plus ``pair``) **and** the derived poloidal angle ``theta`` - and sensor-end coordinates ``{r,z,phi,theta}_end1/2`` — the port of OMFIT's - ``init_magnetics.py`` step. A sensor is modelled as a straight segment of - physical ``length`` centred at ``(r, z)`` and tilted by ``tilt`` (degrees) in - the poloidal plane, spanning ``delta_phi`` (degrees) toroidally; angles are - measured about the magnetic axis at ``(R0, 0)``. - """ - import numpy as np - import xarray as xr - - dev = load_device(device) - R0 = float(dev["R0"]) - sensors = dev["sensors"] - channels = list(sensors) - - # Each sensor is shot-segmented; resolve the segment active at `shot` before - # reading its geometry (the flat `sensors[c]["r"]` lookup returns NaN now that - # the fields live under `segments`). - recs = {c: _segment_fields(sensors[c], shot) for c in channels} - base = ["r", "z", "phi", "tilt", "length", "delta_phi", "na"] - cols = {k: np.array([_to_float(recs[c].get(k, np.nan)) for c in channels]) for k in base} - pairs = np.array([recs[c].get("pair", "None") for c in channels], dtype=object) - - r, z, phi = cols["r"], cols["z"], cols["phi"] - tilt, length, delta_phi = cols["tilt"], cols["length"], cols["delta_phi"] - half = length / 2.0 - tr = np.radians(tilt) - - theta = np.degrees(np.arctan2(z, r - R0)) - r_end1, r_end2 = r - half * np.cos(tr), r + half * np.cos(tr) - z_end1, z_end2 = z - half * np.sin(tr), z + half * np.sin(tr) - phi_end1, phi_end2 = phi - delta_phi / 2.0, phi + delta_phi / 2.0 - theta_end1 = np.degrees(np.arctan2(z_end1, r_end1 - R0)) - theta_end2 = np.degrees(np.arctan2(z_end2, r_end2 - R0)) - - ds = xr.Dataset( - { - **{k: ("channel", cols[k]) for k in base}, - "theta": ("channel", theta), - "r_end1": ("channel", r_end1), - "r_end2": ("channel", r_end2), - "z_end1": ("channel", z_end1), - "z_end2": ("channel", z_end2), - "phi_end1": ("channel", phi_end1), - "phi_end2": ("channel", phi_end2), - "theta_end1": ("channel", theta_end1), - "theta_end2": ("channel", theta_end2), - }, - coords={"channel": channels}, - ) - ds["pair"] = ("channel", pairs) - ds.attrs["device"] = device - ds.attrs["R0"] = R0 - return ds - - -def load_sensor_table(device="DIII-D"): - """Master sensor geometry (base + derived) — alias of :func:`sensor_geometry`.""" - return sensor_geometry(device) - - -def load_wall(device="DIII-D"): - """Return the ``(r, z)`` first-wall outline arrays from the device JSON. - - Returns ``(None, None)`` if no device file or wall section is found. - """ - try: - dev = load_device(device) - except (FileNotFoundError, OSError): - return None, None - wall = dev.get("wall") - if not wall or "r" not in wall or "z" not in wall: - return None, None - return np.array(wall["r"], dtype=float), np.array(wall["z"], dtype=float) diff --git a/src/magnetics/core/qs_bridge.py b/src/magnetics/core/qs_bridge.py index 13358c0..7614d74 100644 --- a/src/magnetics/core/qs_bridge.py +++ b/src/magnetics/core/qs_bridge.py @@ -80,7 +80,10 @@ def _reconstruct_grid( z = np.zeros((len(theta_grid), len(phi_grid))) for i, (n, m) in enumerate(zip(ns, ms)): c = coeffs_t[i] - # outer: [n_theta, n_phi] — sign matches plot_slice reconstruction + # SIGN CONVENTION (do not "fix" to +i): the fit basis is exp(+i(nφ+mθ)) with a + # real/imag column split, so reconstruction MUST use exp(-i(...)) to reproduce + # the fitted field. Matches OMFIT plot_magnetics_slice / qs_plots.plot_slice. + # outer: [n_theta, n_phi] phase = np.exp(-1j * m * theta_rad)[:, None] * np.exp(-1j * n * phi_rad)[None, :] z += (c * phase).real return z @@ -286,6 +289,8 @@ def fit_to_phi_t_node(fit_ds, theta_fixed_deg: float = 0.0, n_phi: int = 73) -> coeffs = fit_ds["fit_coeffs"].values # [n_mode, n_time] complex # field[n_phi, n_time] = Re Sum_m coeff[m,t] * exp(-i*(n*phi + m*theta)) + # SIGN CONVENTION (do not "fix" to +i): fit basis exp(+i(nφ+mθ)) + real/imag split + # ⇒ reconstruction uses exp(-i(...)). See _reconstruct_grid / qs_plots.plot_slice. n_t = len(t_ms) z = np.zeros((n_phi, n_t)) for i, (n, m) in enumerate(zip(ns, ms)): diff --git a/src/magnetics/core/qs_device.py b/src/magnetics/core/qs_device.py new file mode 100644 index 0000000..a36f51a --- /dev/null +++ b/src/magnetics/core/qs_device.py @@ -0,0 +1,206 @@ +"""QS device adapter — sensor geometry + channel-set resolution for the SLCONTOUR fit. + +Self-contained for the quasi-stationary (``qs_*``) pipeline, but **delegates to the +data layer** for everything the data layer already owns: JSON loading and shot-aware +segment resolution (:mod:`magnetics.data.devices`) and sensor-set flattening +(:mod:`magnetics.data.diiid_geometry`). It adds only the QS-specific pieces the fit +needs and that have no home in ``data/``: + + * :func:`sensor_geometry` — the derived poloidal angle ``theta`` and the sensor-end + coordinates ``{r,z,phi,theta}_end1/2`` assembled as an ``xarray.Dataset``; + * :func:`resolve_channel_filter` — the fit/prep channel-filter semantics (friendly + subset name → anchored regex list; unknown string passed through as a raw regex); + * :func:`is_device` — the loose device-name comparison the fit/prep steps use. + +Replaces the old ``_slcontour/omfit_compat.py`` OMFIT shim (deleted). This module does +not read the device JSON itself — it routes through ``data.devices`` so the QS pipeline +and the fetcher can never disagree about a shot's geometry. +""" + +from __future__ import annotations + +import re + +import numpy as np +import xarray as xr + +from ..data import devices as _devices +from ..data import diiid_geometry as _diiid_geo + +# ``data.devices.load_device`` resolves ``.json`` by lowercasing the name; a few +# display names don't slugify by lowercasing alone (``"DIII-D"`` → ``diiid.json``, not +# ``diii-d.json``). Preserve the alias the old shim carried. +_DEVICE_ALIAS = {"DIII-D": "diiid"} + + +def _device_key(device: str) -> str: + """Map a device display name to its ``data/device/.json`` stem.""" + if device in _DEVICE_ALIAS: + return _DEVICE_ALIAS[device] + return re.sub(r"[^a-z0-9]", "", str(device).lower()) + + +def load_device(device: str = "DIII-D") -> dict: + """The parsed device JSON dict, via the data layer (with QS name aliasing).""" + return _devices.load_device(_device_key(device)) + + +# ── device-name comparison ───────────────────────────────────────────────────── + + +def is_device(device, name) -> bool: + """Loose device-name equality (``DIII-D`` == ``DIIID`` == ``diii_d``).""" + + def _norm(s): + return re.sub(r"[^a-z0-9]", "", str(s).lower()) + + return _norm(device) == _norm(name) + + +# ── sensor-set resolution (reuses data/diiid_geometry flattening) ─────────────── + + +def _norm_name(name): + """Normalise a subset name so ``'Bp_LFS_midplane'`` == ``'Bp LFS midplane'``.""" + return str(name).strip().replace("_", " ") + + +def _subsets(device="DIII-D") -> dict[str, list[str]]: + """``{set_name: [sensor, ...]}`` for every named set (composites flattened). + + Reuses the data layer's set-flattening (``diiid_geometry._build_sets``) rather + than carrying a fourth copy of the recursion/dedup logic. + """ + sets = load_device(device).get("sensor_sets", {}) + return {s["name"]: s["sensors"] for s in _diiid_geo._build_sets(sets)} + + +def list_sensor_subsets(device="DIII-D"): + """All named sensor subsets → ``{name: [sensor, ...]}`` (composites flattened). + + Every key here is a valid ``channel_filter`` argument (with or without underscores). + """ + return _subsets(device) + + +def resolve_channel_filter(channel_filter, device="DIII-D"): + """Map a friendly subset name (e.g. ``'Bp_LFS_midplane'``) to a pattern list. + + A known subset name resolves to the explicit, anchored sensor names from the + device ``sensor_sets``; an unknown string is treated as a raw regex and passed + through. Lists/tuples expand each element the same way. Patterns are matched + downstream with :func:`re.match`, so resolved names are anchored with ``$`` to + avoid accidental prefix matches (e.g. ``C19`` vs ``C190``). + """ + subsets = _subsets(device) + lookup = {_norm_name(k): k for k in subsets} + + def resolve_one(cf): + if isinstance(cf, str): + key = lookup.get(_norm_name(cf)) + if key is not None: + return [re.escape(s) + "$" for s in subsets[key]] + return [cf] # raw regex + return [cf] + + if isinstance(channel_filter, str): + return resolve_one(channel_filter) + out = [] + for cf in channel_filter: + out += resolve_one(cf) + return out + + +# ── sensor geometry (QS-unique: derived theta + sensor-end coordinates) ───────── + + +def _to_float(s): + try: + return float(s) + except TypeError, ValueError: + return float("nan") + + +def sensor_geometry(device="DIII-D", shot=None): + """Per-sensor geometry as an ``xarray.Dataset`` indexed by ``channel``. + + Base fields (``r, z, phi, tilt, length, delta_phi, na``, plus ``pair``) come from + the data layer's shot-aware resolver (:func:`devices.geometry_nearest`); this + function adds the derived poloidal angle ``theta`` and the sensor-end coordinates + ``{r,z,phi,theta}_end1/2`` that the SLCONTOUR fit needs. A sensor is modelled as a + straight segment of physical ``length`` centred at ``(r, z)`` tilted by ``tilt`` + (degrees) in the poloidal plane and spanning ``delta_phi`` (degrees) toroidally; + angles are measured about the magnetic axis at ``(R0, 0)``. + + ``shot`` selects the shot-correct hardware segment; ``None`` falls back to each + sensor's earliest segment (layout view). All device channels are returned, with + NaN geometry for any the device file does not model at ``shot``. + """ + dev = load_device(device) + R0 = float(dev["R0"]) + channels = list(dev.get("sensors", {})) + lookup_shot = 0 if shot is None else int(shot) + + base = ["r", "z", "phi", "tilt", "length", "delta_phi", "na"] + cols = {k: [] for k in base} + pairs = [] + for c in channels: + g = _devices.geometry_nearest(dev, c, lookup_shot) or {} + for k in base: + cols[k].append(_to_float(g.get(k, np.nan))) + pairs.append(g.get("pair", "None")) + cols = {k: np.array(v, dtype=float) for k, v in cols.items()} + pairs = np.array(pairs, dtype=object) + + r, z, phi = cols["r"], cols["z"], cols["phi"] + tilt, length, delta_phi = cols["tilt"], cols["length"], cols["delta_phi"] + half = length / 2.0 + tr = np.radians(tilt) + + theta = np.degrees(np.arctan2(z, r - R0)) + r_end1, r_end2 = r - half * np.cos(tr), r + half * np.cos(tr) + z_end1, z_end2 = z - half * np.sin(tr), z + half * np.sin(tr) + phi_end1, phi_end2 = phi - delta_phi / 2.0, phi + delta_phi / 2.0 + theta_end1 = np.degrees(np.arctan2(z_end1, r_end1 - R0)) + theta_end2 = np.degrees(np.arctan2(z_end2, r_end2 - R0)) + + ds = xr.Dataset( + { + **{k: ("channel", cols[k]) for k in base}, + "theta": ("channel", theta), + "r_end1": ("channel", r_end1), + "r_end2": ("channel", r_end2), + "z_end1": ("channel", z_end1), + "z_end2": ("channel", z_end2), + "phi_end1": ("channel", phi_end1), + "phi_end2": ("channel", phi_end2), + "theta_end1": ("channel", theta_end1), + "theta_end2": ("channel", theta_end2), + }, + coords={"channel": channels}, + ) + ds["pair"] = ("channel", pairs) + ds.attrs["device"] = device + ds.attrs["R0"] = R0 + return ds + + +# ── first-wall outline (delegates to the data layer's segmented reader) ───────── + + +def load_wall(device="DIII-D", shot=None): + """The ``(r, z)`` first-wall outline arrays for ``device`` at ``shot``. + + Delegates to :func:`devices.feature_at` for the shot-segmented ``first_wall`` key, + returning ``(None, None)`` when no device file or wall segment is found. (The old + shim read a nonexistent top-level ``wall`` key and returned ``(None, None)`` for + every shipped device; this returns the real outline.) + """ + try: + dev = load_device(device) + except FileNotFoundError, OSError, ValueError: + return None, None + fw = _devices.feature_at(dev, "first_wall", 0 if shot is None else int(shot)) + if not fw or "r" not in fw or "z" not in fw: + return None, None + return np.array(fw["r"], dtype=float), np.array(fw["z"], dtype=float) diff --git a/src/magnetics/_slcontour/fit.py b/src/magnetics/core/qs_fit.py similarity index 89% rename from src/magnetics/_slcontour/fit.py rename to src/magnetics/core/qs_fit.py index 876bba6..f926de8 100644 --- a/src/magnetics/_slcontour/fit.py +++ b/src/magnetics/core/qs_fit.py @@ -27,12 +27,29 @@ from __future__ import annotations +import logging import re import numpy as np import xarray as xr -from .omfit_compat import OMFITexception, delta_degrees, is_device, printe, printv, printw +from .qs_device import is_device + +logger = logging.getLogger(__name__) + + +def _delta_degrees_scalar(theta1, theta2): + """Angular width from theta1 to theta2 in degrees, wrapping once through 0.""" + dt = theta2 - theta1 + if dt > 180: + dt -= 360 + if dt < -180: + dt += 360 + return dt + + +#: Vectorised signed angular width in degrees (ported from omfit_compat.delta_degrees). +delta_degrees = np.vectorize(_delta_degrees_scalar) def form_basis_function( @@ -142,7 +159,7 @@ def form_basis_function( ) return fmn - raise OMFITexception( + raise ValueError( "fit_basis must be 'sinusoidal-point', 'sinusoidal-integral', " "'gaussian-point', or 'gaussian-integral'." ) @@ -187,7 +204,7 @@ def fit( def _printv(*a): if verbose: - printv(*a) + logger.info(" ".join(str(x) for x in a)) _printv("Fitting the prepared data") @@ -212,7 +229,7 @@ def _printv(*a): elif fit_geometry == "vertical": xkey, ykey = "phi", "z" else: - raise OMFITexception("fit_geometry must be 'cylindrical' or 'vertical'.") + raise ValueError("fit_geometry must be 'cylindrical' or 'vertical'.") x1 = ds[f"{xkey}_end1"].values x2 = ds[f"{xkey}_end2"].values @@ -250,7 +267,7 @@ def _printv(*a): if is_device(ds.attrs.get("device", ""), "DIII-D"): if any(re.match(k, c) for k in ("C.*", "IL.*", "IU.*") for c in channels): if fit_basis != "sinusoidal-point": - printe("WARNING: sinusoidal-point basis is used by DIII-D 3D coil operators") + logger.warning("sinusoidal-point basis is used by DIII-D 3D coil operators") elif is_gaussian: # Derive RBF centres from sensor extent (mirrors OMFIT gaussian setup block) @@ -290,7 +307,7 @@ def _printv(*a): f"ncycle={_ncycle}" ) else: - raise OMFITexception( + raise ValueError( f"fit_basis must start with 'sinusoidal' or 'gaussian', got {fit_basis!r}." ) @@ -313,26 +330,40 @@ def _printv(*a): wtest = np.linalg.svd(np.array(A_cols).T, compute_uv=False) if np.abs(wtest[0] / wtest[-1]) > 1e19: raise ValueError("Bad sensor distribution") - except (ValueError, np.linalg.LinAlgError): - printe(f" - Ill-conditioned mode ({n},{m}); fitting single component") + except ValueError, np.linalg.LinAlgError: + logger.error(" - Ill-conditioned mode (%s,%s); fitting single component", n, m) x0 = x1[0] + delta_degrees(x1[0], x2[0]) / 2.0 y0 = y1[0] + delta_degrees(y1[0], y2[0]) / 2.0 - fmn = form_basis_function(n, m, x1 + x0, x2 + x0, y1 + y0, y2 + y0, fit_basis) / sigma + fmn = ( + form_basis_function(n, m, x1 + x0, x2 + x0, y1 + y0, y2 + y0, fit_basis) + / sigma + ) A_cols = A_cols[:-2] + [fmn.real] ncomp[-1] = 1 else: # gaussian — basis functions are always real for n, m in nms: - fmn = form_basis_function( - n, m, x1, x2, y1, y2, fit_basis, - ncycle=_ncycle, mcycle=_mcycle, - nepsilon=_nepsilon, mepsilon=_mepsilon, - ) / sigma + fmn = ( + form_basis_function( + n, + m, + x1, + x2, + y1, + y2, + fit_basis, + ncycle=_ncycle, + mcycle=_mcycle, + nepsilon=_nepsilon, + mepsilon=_mepsilon, + ) + / sigma + ) A_cols.append(fmn) ncomp.append(1) A = np.array(A_cols).T # (n_sensors, n_columns) if A.shape[1] > A.shape[0]: - printw(f"Fitting {A.shape[1]} basis functions with {A.shape[0]} sensors") + logger.warning("Fitting %d basis functions with %d sensors", A.shape[1], A.shape[0]) # ── SVD of A (condition number + per-coefficient error bars) ────────────── U_a, w_a, Vh_a = np.linalg.svd(A) @@ -343,7 +374,9 @@ def _printv(*a): # ── least-squares fit at every time slice ───────────────────────────────── _printv(" - Fitting signal") - b = (ds["signal"] / xr.DataArray(sigma, coords={"channel": ds["channel"]}, dims=("channel",))).values + b = ( + ds["signal"] / xr.DataArray(sigma, coords={"channel": ds["channel"]}, dims=("channel",)) + ).values x, residual, rank_fit, s_a = np.linalg.lstsq(A, b, rcond=1.0 / fit_cond) _printv(f" - Raw / effective condition number = {raw_cn:.3g} / {eff_cn:.3g}") fit_coeffs = np.asarray(x) # (n_columns, n_time) @@ -416,6 +449,8 @@ def _printv(*a): def _condition_warning(K): """SLCONTOUR's K-thresholds (VISION S4.1): warn > 10, error > 20.""" if K > 20: - printe(f"Condition number K = {K:.1f} > 20: fit is untrustworthy (under-resolved array).") + logger.error( + "Condition number K = %.1f > 20: fit is untrustworthy (under-resolved array).", K + ) elif K > 10: - printw(f"Condition number K = {K:.1f} > 10: fit may be poorly resolved.") + logger.warning("Condition number K = %.1f > 10: fit may be poorly resolved.", K) diff --git a/src/magnetics/_slcontour/io_data.py b/src/magnetics/core/qs_io_data.py similarity index 92% rename from src/magnetics/_slcontour/io_data.py rename to src/magnetics/core/qs_io_data.py index e47ba56..f9424e4 100644 --- a/src/magnetics/_slcontour/io_data.py +++ b/src/magnetics/core/qs_io_data.py @@ -5,7 +5,7 @@ * raw sensor signals from ``data/datafile/shot_.h5``, and * device/sensor geometry from ``data/device/.json`` (via - :mod:`omfit_compat`). + :mod:`magnetics.core.qs_device`). ``shot_.h5`` layout (what the PTDATA fetch produces): * one group per fetched channel, each with ``data`` (samples) and ``time`` @@ -23,23 +23,26 @@ Sensor geometry (the per-channel ``r,z,phi,theta,...`` and the derived ``*_end1/2`` coordinates the fit needs) is attached from the device JSON; see -:func:`omfit_compat.sensor_geometry`. +:func:`magnetics.core.qs_device.sensor_geometry`. """ from __future__ import annotations +import logging import os from dataclasses import dataclass import numpy as np import xarray as xr -from .omfit_compat import list_sensor_subsets, printw, resolve_channel_filter, sensor_geometry +from .qs_device import list_sensor_subsets, resolve_channel_filter, sensor_geometry #: Default location of the per-shot HDF5 files — the runtime data dir the fetcher #: writes and h5source reads ($MAGNETICS_DATA_DIR or the repo's data/datafile/). from ..data import h5source as _h5source +logger = logging.getLogger(__name__) + DATAFILE_ROOT = str(_h5source.data_dir() / "datafile") #: Groups in the HDF5 file that are not geometry-bearing sensor channels. @@ -108,7 +111,9 @@ def load_shot(shot, data_root=DATAFILE_ROOT, helicity=-1): plasma_times[name] = t_ms continue if name not in geo_channels: - printw(f"Channel {name!r} has no geometry in {device} device file -> skipping") + logger.warning( + "Channel %r has no geometry in %s device file -> skipping", name, device + ) continue sensor_names.append(name) sensor_sigs[name] = data @@ -168,7 +173,11 @@ def _build_plasma(sigs, times, helicity): for name, var in _PLASMA_CHANNELS.items(): if name not in sigs: continue - y = sigs[name] if np.array_equal(times[name], t_ms) else np.interp(t_ms, times[name], sigs[name]) + y = ( + sigs[name] + if np.array_equal(times[name], t_ms) + else np.interp(t_ms, times[name], sigs[name]) + ) data_vars[var] = ("time", y) plasma = xr.Dataset(data_vars, coords={"time": t_ms}) @@ -192,7 +201,7 @@ def _shot_from_path(path): def available_subsets(device="DIII-D"): """All named sensor subsets you can pass as ``channel_filter``. - Thin convenience over :func:`omfit_compat.list_sensor_subsets` -> a + Thin convenience over :func:`magnetics.core.qs_device.list_sensor_subsets` -> a ``{name: [sensor, ...]}`` mapping (e.g. ``'Bp_LFS_midplane'``, ``'All_3D_Coils'``). Names work with or without underscores. """ diff --git a/src/magnetics/_slcontour/plots.py b/src/magnetics/core/qs_plots.py similarity index 76% rename from src/magnetics/_slcontour/plots.py rename to src/magnetics/core/qs_plots.py index bbae42d..5f73c0f 100644 --- a/src/magnetics/_slcontour/plots.py +++ b/src/magnetics/core/qs_plots.py @@ -1,8 +1,12 @@ -"""Ports of the relevant ``PLOTS/*`` scripts (the SLCONTOUR / locked-mode set). +"""Standalone matplotlib plots for the quasi-stationary (SLCONTOUR) pipeline. -Each function takes the Datasets produced by :mod:`prep` / :mod:`fit` and an -optional axis, replacing the OMFIT runtime helpers (``cornernote``, ``uband``, -``View1d`` ...) with the local shim in :mod:`omfit_compat`. +These are **not** wired into the GUI/service. The FastAPI service serves JSON +``kind``-nodes built by :mod:`magnetics.core.qs_bridge`, which only *ports* these +plots' conventions (amplitude/phase extraction, the SLCONTOUR sign convention) — it +does not import this module, and nothing in the served path calls these functions. +They are plain matplotlib routines over the xarray Datasets produced by +:mod:`magnetics.core.qs_prep` / :mod:`magnetics.core.qs_fit`, kept for notebooks, +debugging and offline analysis. Covered (per VISION.md S4.1 outputs): * :func:`plot_sensors` - sensor map (R-Z / unrolled phi-theta) @@ -15,12 +19,41 @@ from __future__ import annotations +import logging import re import matplotlib.pyplot as plt import numpy as np -from .omfit_compat import cornernote, load_wall, printe, uband +from .qs_device import load_wall + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- # +# Plot helpers (formerly omfit_compat.cornernote / uband) +# --------------------------------------------------------------------------- # +def cornernote(ax=None, device="", shot="", time="", text="", **_ignore): + """Annotate the bottom-right figure corner with shot/device context.""" + if ax is None: + ax = plt.gca() + fig = ax.get_figure() + label = " ".join(str(s) for s in (device, shot, time, text) if str(s) != "") + if not label: + return + fig.text(0.99, 0.01, label, ha="right", va="bottom", fontsize=8, color="0.4") + + +def uband(x, y, yerr, ax=None, label=None, color=None, **kwargs): + """Plot a line with a shaded +/- ``yerr`` uncertainty band; returns ``(line,)``.""" + if ax is None: + ax = plt.gca() + x = np.asarray(x) + y = np.asarray(y, dtype=float) + yerr = np.asarray(yerr, dtype=float) + (line,) = ax.plot(x, y, label=label, color=color, **kwargs) + ax.fill_between(x, y - yerr, y + yerr, color=line.get_color(), alpha=0.3, linewidth=0) + return (line,) # --------------------------------------------------------------------------- # @@ -53,17 +86,45 @@ def plot_sensors(raw, channel_filter=".*", geometry="rz", ax=None, device=None, x = [float(s["r_end1"]), float(s["r_end2"])] y = [float(s["z_end1"]), float(s["z_end2"])] elif geometry == "flat": - x = [float(s["phi_end1"]), float(s["phi_end2"]), float(s["phi_end2"]), float(s["phi_end1"]), float(s["phi_end1"])] - y = [float(s["z_end1"]), float(s["z_end1"]), float(s["z_end2"]), float(s["z_end2"]), float(s["z_end1"])] + x = [ + float(s["phi_end1"]), + float(s["phi_end2"]), + float(s["phi_end2"]), + float(s["phi_end1"]), + float(s["phi_end1"]), + ] + y = [ + float(s["z_end1"]), + float(s["z_end1"]), + float(s["z_end2"]), + float(s["z_end2"]), + float(s["z_end1"]), + ] else: # cylindrical - x = _no_wrap([float(s["phi_end1"]), float(s["phi_end2"]), float(s["phi_end2"]), float(s["phi_end1"]), float(s["phi_end1"])]) - y = _no_wrap([float(s["theta_end1"]), float(s["theta_end1"]), float(s["theta_end2"]), float(s["theta_end2"]), float(s["theta_end1"])]) + x = _no_wrap( + [ + float(s["phi_end1"]), + float(s["phi_end2"]), + float(s["phi_end2"]), + float(s["phi_end1"]), + float(s["phi_end1"]), + ] + ) + y = _no_wrap( + [ + float(s["theta_end1"]), + float(s["theta_end1"]), + float(s["theta_end2"]), + float(s["theta_end2"]), + float(s["theta_end1"]), + ] + ) if line is not None: plot_kwargs["color"] = line.get_color() (line,) = ax.plot(x, y, label=c, **plot_kwargs) if geometry == "rz": - r_wall, z_wall = load_wall(device) + r_wall, z_wall = load_wall(device, raw.attrs.get("shot")) if r_wall is not None: ax.plot(r_wall, z_wall, color="0.4", lw=1, label="wall") ax.set_xlabel("R (m)") @@ -96,23 +157,29 @@ def plot_signal(raw, prepared, channel_filter=".*", ax=None, legend_maxnum=12): s = prepared["signal"].sel(channel=c) s_ptp = np.ptp(np.nan_to_num(s.values)) if s_ptp <= 0: - printe(f"{c} has no signal") + logger.warning("%s has no signal", c) continue # the raw trace, shifted to match the prepared trace at t0 (shows filtering) s2 = raw["signal"].sel(channel=c, time=s["time"]) s2 = s2 - (float(s2.values[0]) - float(s.values[0])) (l2,) = ax.plot(s2["time"], s2.values, alpha=0.4) - (l,) = ax.plot(s["time"], s.values, color=l2.get_color(), label=c) - lines.append(l) + (ln,) = ax.plot(s["time"], s.values, color=l2.get_color(), label=c) + lines.append(ln) ptps.append(s_ptp) if ptps: ptps, lines = zip(*sorted(zip(ptps, lines), key=lambda z: z[0])) nleg = min(legend_maxnum, len(lines)) - for l in lines[:-nleg] if nleg < len(lines) else []: - l.set_color("grey") - l.set_alpha(0.4) - ax.legend([l for l in lines[-nleg:]], [l.get_label() for l in lines[-nleg:]], loc=2, frameon=False, fontsize=8) + for ln in lines[:-nleg] if nleg < len(lines) else []: + ln.set_color("grey") + ln.set_alpha(0.4) + ax.legend( + [ln for ln in lines[-nleg:]], + [ln.get_label() for ln in lines[-nleg:]], + loc=2, + frameon=False, + fontsize=8, + ) ax.set_xlabel("time (s)") ax.set_ylabel("signal") @@ -186,7 +253,12 @@ def plot_fit_modes(fit, axes=None, legend_maxnum=12): _, axes = plt.subplots(2, sharex=True, figsize=(8, 8)) t = fit["time"].values - max_amp = np.max([np.percentile(np.abs(fit["fit_coeffs"]).sel(mode=m).values, 90) for m in fit["mode"].values]) + max_amp = np.max( + [ + np.percentile(np.abs(fit["fit_coeffs"]).sel(mode=m).values, 90) + for m in fit["mode"].values + ] + ) amps = [] for i, m in enumerate(fit["mode"].values): @@ -196,11 +268,11 @@ def plot_fit_modes(fit, axes=None, legend_maxnum=12): mval = int(fit["fit_ms"].values[i]) label = f"{mval}/{nval}" amp, amp_err, phase, phase_err = _amp_phase_with_errors(coeff, sigma) - (l,) = uband(t, amp, amp_err, ax=axes[0], label=label) + (ln,) = uband(t, amp, amp_err, ax=axes[0], label=label) amps.append(np.max(amp)) # only draw phase for modes with appreciable amplitude (avoid clutter) if np.percentile(np.abs(coeff), 90) > 0.1 * max_amp: - uband(t, phase, phase_err, ax=axes[1], label=label, color=l.get_color()) + uband(t, phase, phase_err, ax=axes[1], label=label, color=ln.get_color()) # de-emphasise the smallest modes, then build a compact legend leg_title = "n" if np.all(fit["fit_ms"].values == 0) else "m/n" @@ -230,7 +302,9 @@ def plot_fit_modes(fit, axes=None, legend_maxnum=12): # --------------------------------------------------------------------------- # # SLCONTOUR phi-vs-time contour (<- plot_magnetics_slice.py) # --------------------------------------------------------------------------- # -def plot_slice(fit, fix_coord="theta", fix_value=0.0, ngrid=120, ax=None, trace_peak=True, **plot_kwargs): +def plot_slice( + fit, fix_coord="theta", fix_value=0.0, ngrid=120, ax=None, trace_peak=True, **plot_kwargs +): """Reconstruct the fitted field on a (phi or theta) x time grid and contour it. This is the classic SLCONTOUR locked/slowly-rotating-mode picture: a rotating @@ -266,10 +340,15 @@ def plot_slice(fit, fix_coord="theta", fix_value=0.0, ngrid=120, ax=None, trace_ swept, swept_label = deg, "Toroidal Angle (deg.)" # field(swept, time) = Re Sum_modes coeff(t) exp(-i (n*xgrid + m*ygrid)) + # SIGN CONVENTION (do not "fix"): the fit basis is exp(+i(nφ+mθ)) with a real/imag + # column split, so the reconstruction MUST use exp(-i(...)) to reproduce the fitted + # field. See qs_bridge._reconstruct_grid / fit_to_phi_t_node for the matching note. coeffs = fit["fit_coeffs"].values # (mode, time) nvals = ns.values[:, None] mvals = ms.values[:, None] - phase = np.exp(-1j * (nvals * np.atleast_1d(xgrid)[None, :] + mvals * np.atleast_1d(ygrid)[None, :])) + phase = np.exp( + -1j * (nvals * np.atleast_1d(xgrid)[None, :] + mvals * np.atleast_1d(ygrid)[None, :]) + ) # (swept, time): sum over modes of coeff(mode,time) * phase(mode,swept) field = np.real(np.einsum("ms,mt->st", phase, coeffs)) # (swept, time) diff --git a/src/magnetics/_slcontour/prep.py b/src/magnetics/core/qs_prep.py similarity index 91% rename from src/magnetics/_slcontour/prep.py rename to src/magnetics/core/qs_prep.py index f17be38..f3a8ed5 100644 --- a/src/magnetics/_slcontour/prep.py +++ b/src/magnetics/core/qs_prep.py @@ -18,13 +18,16 @@ from __future__ import annotations +import logging import re import numpy as np import xarray as xr from scipy.integrate import cumulative_trapezoid -from .omfit_compat import is_device, printe, printv, printw, resolve_channel_filter +from .qs_device import is_device, resolve_channel_filter + +logger = logging.getLogger(__name__) def causal_gaussian(values, sigma, truncate=4.0): @@ -79,7 +82,7 @@ def prepare( def _printv(*a): if verbose: - printv(*a) + logger.info(" ".join(str(x) for x in a)) raw = shotdata.raw plasma = shotdata.plasma @@ -116,15 +119,15 @@ def _printv(*a): ds["vacuum"] = coup.sel(channel=ds["channel"]) ds["signal"] = ds["signal"] - ds["vacuum"] else: - printe(f"No DC coupling record for {invalid} -> skipping DC compensation") + logger.error("No DC coupling record for %s -> skipping DC compensation", invalid) else: - printw("dc_comp requested but no coil channels matched -> skipping") + logger.warning("dc_comp requested but no coil channels matched -> skipping") # ----- device-specific signal swap (2019 DIII-D wiring mix-up) --------- # if is_device(device, "DIII-D") and shotdata.shot > 177705: present = set(ds["channel"].values) if {"ESLD66M079", "ESLD66M319"} <= present: - printw("Swapping ESLD66M319/ESLD66M079 for this shot (2019 wiring mix-up)") + logger.warning("Swapping ESLD66M319/ESLD66M079 for this shot (2019 wiring mix-up)") tmp = ds["signal"].copy(deep=True) ds["signal"].loc[{"channel": "ESLD66M079"}] = tmp.sel(channel="ESLD66M319").values ds["signal"].loc[{"channel": "ESLD66M319"}] = tmp.sel(channel="ESLD66M079").values @@ -155,11 +158,15 @@ def _printv(*a): if cutoff_hz[0] == 0: _printv(" > causal lowpass") sigma = 0.25 / (dt * cutoff_hz[1]) - filter_func = lambda v: causal_gaussian(v, sigma) + + def filter_func(v): + return causal_gaussian(v, sigma) elif cutoff_hz[1] >= nyqst: _printv(" > causal highpass") sigma = 0.25 / (dt * cutoff_hz[0]) - filter_func = lambda v: v - causal_gaussian(v, sigma) + + def filter_func(v): + return v - causal_gaussian(v, sigma) else: _printv(" > causal bandpass") sigma0 = 0.25 / (dt * cutoff_hz[0]) @@ -177,7 +184,7 @@ def filter_func(v): if len(tsel) == 0: raise ValueError( f"time_trim={time_trim} left 0 samples after filter/downsample " - f"(downsampled dt≈{t[1]-t[0]:.4g} s, {len(t)} samples spanned " + f"(downsampled dt≈{t[1] - t[0]:.4g} s, {len(t)} samples spanned " f"{t[0]:.4g}–{t[-1]:.4g} s). Widen the window or reduce cutoff_hz[0]." ) ds = ds.sel(time=tsel) @@ -206,9 +213,7 @@ def _detrend(ds, channels, detrend_type, detrend_band, _printv): time = ds["time"].values if detrend_type in ("baseline", "linear"): band = np.atleast_2d(detrend_band) - det_times = np.concatenate( - [time[(time >= b[0]) & (time <= b[1])] for b in band] - ) + det_times = np.concatenate([time[(time >= b[0]) & (time <= b[1])] for b in band]) in_band = np.isin(time, det_times) if detrend_type == "baseline": _printv(" - Removing baseline") @@ -247,8 +252,10 @@ def _svd_condition(ds, energy, _printv): U, s, Vh = np.linalg.svd(ds["signal"].values / np.sqrt(P * T), full_matrices=False) except MemoryError: step = max(2, int(np.ceil(T / 10000))) - printw(f"Downsampling time x{step} for an informational SVD only") - U, s, Vh = np.linalg.svd(ds["signal"].values[:, ::step] / np.sqrt(P * T), full_matrices=False) + logger.warning("Downsampling time x%d for an informational SVD only", step) + U, s, Vh = np.linalg.svd( + ds["signal"].values[:, ::step] / np.sqrt(P * T), full_matrices=False + ) energy = 1.0 # shapes won't match for reforming the data energy_tot = np.sum(s**2) diff --git a/src/magnetics/_slcontour/run.py b/src/magnetics/core/qs_run.py similarity index 96% rename from src/magnetics/_slcontour/run.py rename to src/magnetics/core/qs_run.py index 0837029..866c3bf 100644 --- a/src/magnetics/_slcontour/run.py +++ b/src/magnetics/core/qs_run.py @@ -13,9 +13,9 @@ import xarray as xr -from . import fit as _fit -from . import prep as _prep -from .io_data import ShotData, load_shot +from . import qs_fit as _fit +from . import qs_prep as _prep +from .qs_io_data import ShotData, load_shot @dataclass diff --git a/src/magnetics/core/quasistationary.py b/src/magnetics/core/quasistationary.py deleted file mode 100644 index fb636c5..0000000 --- a/src/magnetics/core/quasistationary.py +++ /dev/null @@ -1,256 +0,0 @@ -"""Pure-numpy SLCONTOUR-style quasi-stationary spatial fit. - -Ported from analysis/magnetics-code/fit.py — omfit_compat and xarray -removed; only numpy + stdlib. Supports sinusoidal-point and -sinusoidal-integral bases for cylindrical (phi, theta) geometry. - -Reference: E. Strait et al., DIII-D magnetics analysis (SLCONTOUR). -""" - -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np - - -# ── angular helpers ─────────────────────────────────────────────────────────── - - -def _delta_deg(a: np.ndarray, b: np.ndarray) -> np.ndarray: - """Signed angular arc a→b wrapped to (−180, +180].""" - d = b - a - return ((d + 180.0) % 360.0) - 180.0 - - -# ── basis function ──────────────────────────────────────────────────────────── - - -def form_basis_function( - n: int, - m: int, - x1: np.ndarray, - x2: np.ndarray, - y1: np.ndarray, - y2: np.ndarray, - fit_basis: str = "sinusoidal-point", -) -> np.ndarray: - """Basis-function vector (complex) for mode (n, m). - - x is toroidal phi (degrees), y is poloidal theta (degrees). - *1/*2 are sensor extents; for the point basis the midpoint is used. - - Ported from magnetics-code/fit.py form_basis_function — omfit_compat - delta_degrees replaced with _delta_deg. - """ - dx = _delta_deg(x1, x2) - dy = _delta_deg(y1, y2) - - if fit_basis == "sinusoidal-point": - phi_c = x1 + dx / 2.0 - theta_c = y1 + dy / 2.0 - return np.exp(1j * n * np.deg2rad(phi_c) + 1j * m * np.deg2rad(theta_c)) - - if fit_basis == "sinusoidal-integral": - if n == 0 and m == 0: - return np.ones(len(np.atleast_1d(x1)), dtype=complex) - if n == 0: - return (np.exp(1j * m * np.deg2rad(y2)) - np.exp(1j * m * np.deg2rad(y1))) / ( - np.deg2rad(dy) * 1j * m - ) - if m == 0: - return (np.exp(1j * n * np.deg2rad(x2)) - np.exp(1j * n * np.deg2rad(x1))) / ( - np.deg2rad(dx) * 1j * n - ) - return ( - (np.exp(1j * m * np.deg2rad(y2)) - np.exp(1j * m * np.deg2rad(y1))) - * (np.exp(1j * n * np.deg2rad(x2)) - np.exp(1j * n * np.deg2rad(x1))) - ) / (np.deg2rad(dx * dy) * n * m) - - raise ValueError( - f"fit_basis must be 'sinusoidal-point' or 'sinusoidal-integral', got {fit_basis!r}" - ) - - -# ── result ──────────────────────────────────────────────────────────────────── - - -@dataclass -class FitResult: - time_ms: np.ndarray # [n_time] - ns: np.ndarray # [n_modes] toroidal mode numbers - ms: np.ndarray # [n_modes] poloidal mode numbers - coeffs: np.ndarray # complex [n_modes, n_time] - sigmas: np.ndarray # complex [n_modes] — Re/Im sigma per mode (time-independent) - chi_sq: np.ndarray # [n_time] - red_chi_sq: np.ndarray # [n_time] - condition_number: float - n_sensors: int - - -# ── fit ─────────────────────────────────────────────────────────────────────── - - -def fit( - time_ms: np.ndarray, - signal: np.ndarray, - phi_end1: np.ndarray, - phi_end2: np.ndarray, - theta_end1: np.ndarray, - theta_end2: np.ndarray, - sigma: np.ndarray | None = None, - ns: tuple[int, ...] = (1, 2, 3), - ms: tuple[int, ...] = (0,), - helicity: int = -1, - fit_basis: str = "sinusoidal-point", - fit_cond: float = 10.0, -) -> FitResult: - """SLCONTOUR-style least-squares spatial modal fit. - - Parameters - ---------- - time_ms : [n_time] - signal : [n_ch, n_time] — measured field values (same units returned) - phi_end1, phi_end2 : [n_ch] — sensor toroidal extents (degrees) - theta_end1, theta_end2 : [n_ch] — sensor poloidal extents (degrees) - sigma : [n_ch] — measurement noise std; None → uniform 1.0 - ns, ms : toroidal / poloidal mode numbers to include in the basis - helicity : +1 or −1 (DIII-D convention is −1); flips ms sign if needed - fit_basis : 'sinusoidal-point' or 'sinusoidal-integral' - fit_cond : condition-number cutoff (= 1/rcond passed to lstsq) - - Returns FitResult with complex coefficients [n_modes, n_time]. - """ - n_ch, n_time = signal.shape - - if sigma is None: - sigma = np.ones(n_ch, dtype=float) - sigma = np.asarray(sigma, dtype=float).copy() - bad = ~np.isfinite(sigma) | (sigma <= 0) - if bad.any(): - fill = float(np.nanmean(sigma[~bad])) if np.any(~bad) else 1.0 - sigma[bad] = fill - - ns_arr = np.atleast_1d(ns).astype(int) - ms_arr = np.atleast_1d(ms).astype(int) - - # Flip m sign to match device helicity convention (ported from fit.py) - if not np.all(ms_arr == 0) and not np.any(np.sign(ms_arr) == helicity): - ms_arr = ms_arr * -1 - - nms = [(int(n), int(m)) for m in ms_arr for n in ns_arr] - - # ── design matrix A [n_ch, n_cols] ─────────────────────────────────────── - A_cols: list[np.ndarray] = [] - ncomp: list[int] = [] - - for n, m in nms: - fmn = ( - form_basis_function(n, m, phi_end1, phi_end2, theta_end1, theta_end2, fit_basis) / sigma - ) - if np.allclose(fmn.imag, 0): - A_cols.append(fmn.real) - ncomp.append(1) - else: - A_cols.append(fmn.real) - A_cols.append(fmn.imag) - ncomp.append(2) - - A = np.array(A_cols).T # [n_ch, n_cols] - - # ── SVD of A — condition number + per-coeff error bars ──────────────────── - _, w_a, Vh_a = np.linalg.svd(A, full_matrices=False) - raw_cn = float(np.abs(w_a[0] / w_a[-1])) if w_a[-1] != 0.0 else np.inf - c_a = np.abs(w_a[0] / w_a) - valid = c_a <= fit_cond - - # ── lstsq across all time slices (vectorised) ───────────────────────────── - b = signal / sigma[:, None] # [n_ch, n_time] - x, _, rank_fit, _ = np.linalg.lstsq(A, b, rcond=1.0 / fit_cond) # [n_cols, n_time] - - # Per-coeff uncertainty from SVD pseudo-inverse - w_inv = np.where(valid, 1.0 / np.where(w_a != 0, w_a, 1.0), 0.0) - fit_sigmas = np.sqrt(np.sum((Vh_a.T * w_inv) ** 2, axis=0)) # [n_cols] - - # ── reform complex coefficients (one complex number per mode) ───────────── - coeffs_c: list[np.ndarray] = [] - sigmas_c: list[complex] = [] - j = 0 - for nc in ncomp: - if nc == 1: - coeffs_c.append(x[j] + 0j) - sigmas_c.append(complex(fit_sigmas[j])) - else: - coeffs_c.append(x[j] + 1j * x[j + 1]) - sigmas_c.append(complex(fit_sigmas[j], fit_sigmas[j + 1])) - j += nc - - coeffs = np.array(coeffs_c) # [n_modes, n_time] - sigmas_out = np.array(sigmas_c) # [n_modes] complex - - # ── χ² ─────────────────────────────────────────────────────────────────── - fit_signal = (A @ x).real * sigma[:, None] - residual = signal - fit_signal - chi_sq = np.sum((residual / sigma[:, None]) ** 2, axis=0) # [n_time] - nu = max(n_ch - rank_fit, 1) - red_chi_sq = chi_sq / nu - - return FitResult( - time_ms=np.asarray(time_ms, dtype=float), - ns=np.array([n for n, _ in nms]), - ms=np.array([m for _, m in nms]), - coeffs=coeffs, - sigmas=sigmas_out, - chi_sq=chi_sq, - red_chi_sq=red_chi_sq, - condition_number=raw_cn, - n_sensors=n_ch, - ) - - -# ── amplitude / phase extraction ────────────────────────────────────────────── - - -def amp_phase_with_errors( - coeff: np.ndarray, sigma: np.ndarray -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """Amplitude/phase and 1-sigma errors from complex coeff and sigma. - - sigma.real = std of the real part, sigma.imag = std of the imag part. - Ported from magnetics-code/plots.py _amp_phase_with_errors. - - Returns (amp, amp_err, phase_deg, phase_err_deg). - """ - c, s = coeff.real, coeff.imag - ec, es = sigma.real, sigma.imag - amp = np.sqrt(c**2 + s**2) - denom = np.where(amp == 0, np.nan, amp) - amp_err = np.sqrt((c * ec) ** 2 + (s * es) ** 2) / denom - phase = np.rad2deg(np.arctan2(s, c)) - p2 = np.where((c**2 + s**2) == 0, np.nan, c**2 + s**2) - phase_err = np.rad2deg(np.sqrt((c * es) ** 2 + (s * ec) ** 2) / p2) - return amp, np.nan_to_num(amp_err), phase, np.nan_to_num(phase_err) - - -# ── spatial reconstruction ──────────────────────────────────────────────────── - - -def reconstruct_grid( - result: FitResult, - phi_grid: np.ndarray, - theta_grid: np.ndarray, - t_idx: int, -) -> np.ndarray: - """Reconstruct δBp on a (phi, theta) grid at one time index. - - Returns z[n_theta, n_phi] — units match the input signal. - """ - phi_rad = np.deg2rad(phi_grid) # [n_phi] - theta_rad = np.deg2rad(theta_grid) # [n_theta] - z = np.zeros((len(theta_grid), len(phi_grid))) - for i, (n, m) in enumerate(zip(result.ns, result.ms)): - c = result.coeffs[i, t_idx] - # outer product: [n_theta, n_phi] - basis = np.exp(1j * m * theta_rad)[:, None] * np.exp(1j * n * phi_rad)[None, :] - z += (c * basis).real - return z diff --git a/src/magnetics/service/nodes.py b/src/magnetics/service/nodes.py index c5810f3..b861982 100644 --- a/src/magnetics/service/nodes.py +++ b/src/magnetics/service/nodes.py @@ -1066,7 +1066,7 @@ def _mode_over_time(shot, params=None) -> dict: ) -# ── quasi-stationary fit (SLCONTOUR via magnetics-code pipeline) ───────────── +# ── quasi-stationary fit (SLCONTOUR via the core.qs_* pipeline) ───────────── _QS_RUN_LOCK = threading.Lock() @@ -1087,12 +1087,12 @@ def _qs_run( ): """Run the full SLCONTOUR pipeline (io_data → prep → fit) for one shot. - Delegates to magnetics-code/run.run_steps. All tuning parameters are + Delegates to core.qs_run.run_steps. All tuning parameters are explicit cache-key arguments so the result is reused across node requests that share the same settings. tmin_s/tmax_s are in seconds and come from _prep_qs_ds (which reads HDF5 defaults and applies any user override). """ - from .._slcontour.run import run_steps + from ..core.qs_run import run_steps # detrend_band is in ms from the GUI; (0, 0) = auto → first 10ms of shot window if detrend_band == (0.0, 0.0): @@ -1232,9 +1232,9 @@ def _sensor_map_rz(shot, params=None) -> dict: channels = list(run.prepared["channel"].values) device = str(raw.attrs.get("device", "DIII-D")) - from .._slcontour.omfit_compat import load_wall + from ..core.qs_device import load_wall - r_wall, z_wall = load_wall(device) + r_wall, z_wall = load_wall(device, shot) series = [] for c in channels: diff --git a/tests/test_qs_bridge.py b/tests/test_qs_bridge.py index 2331958..db98592 100644 --- a/tests/test_qs_bridge.py +++ b/tests/test_qs_bridge.py @@ -11,6 +11,7 @@ import numpy as np import pytest +import xarray as xr from magnetics.core import qs_bridge from magnetics.service import nodes @@ -49,3 +50,38 @@ def test_amplitude_sigma_is_finite(fit_ds): node = qs_bridge.fit_to_amplitude_node(fit_ds) sigma = np.asarray(node["meta"]["sigma"], dtype=float) assert np.all(np.isfinite(sigma)) + + +def _single_mode_fit(n, m, phase_deg): + """A minimal fit Dataset: one mode, complex coeff = exp(i·phase_deg), constant in time.""" + b = np.exp(1j * np.deg2rad(phase_deg)) + t_s = np.linspace(0.0, 0.01, 5) + coeffs = np.full((1, t_s.size), b, dtype=complex) + return xr.Dataset( + { + "fit_ns": ("mode", [n]), + "fit_ms": ("mode", [m]), + "fit_coeffs": (("mode", "time"), coeffs), + }, + coords={"mode": [0], "time": t_s}, + ) + + +def test_reconstruction_uses_minus_i_sign_convention(): + """A locked n=1, m=0 mode with spatial phase δ reconstructs as cos(φ − δ), so its + δBp peak sits at φ = +δ. The buggy exp(+i…) reconstruction would mirror it to −δ + (i.e. 360−δ). This is the discriminating case the old zero-phase test missed — + it pins the convention shared by qs_bridge / qs_plots / OMFIT plot_magnetics_slice. + """ + delta = 60.0 + ds = _single_mode_fit(n=1, m=0, phase_deg=delta) + phi = np.linspace(0, 360, 361) + theta = np.array([0.0]) + + z = qs_bridge._reconstruct_grid(ds, phi, theta, t_idx=0) # [n_theta, n_phi] + peak_phi = float(phi[np.argmax(z[0])]) + + assert abs(peak_phi - delta) < 2.0, ( + f"reconstructed peak at φ={peak_phi}°, expected ~{delta}° (−i convention); " + f"a peak near {360 - delta}° means the exp(+i…) sign bug is back" + ) diff --git a/tests/test_qs_pipeline.py b/tests/test_qs_pipeline.py index e00084d..d916a08 100644 --- a/tests/test_qs_pipeline.py +++ b/tests/test_qs_pipeline.py @@ -1,6 +1,6 @@ """Quasi-stationary pipeline end-to-end, against the synthetic shot. -Spans io_data → omfit_compat.sensor_geometry → prep → fit → qs_bridge → contracts +Spans qs_io_data → qs_device.sensor_geometry → qs_prep → qs_fit → qs_bridge → contracts — the full SLCONTOUR path. This is the coverage that was missing when the segmented-schema geometry bug shipped: the geometry resolved to all-NaN, the fit's SVD never converged, but the only test touching the pipeline (`_fit_quality`) @@ -41,7 +41,7 @@ def test_sensor_geometry_extents_finite_for_synthetic_shot(synthetic_shot): """The direct regression guard: the QS geometry must not be NaN. Reverting omfit_compat.sensor_geometry to the flat `sensors[c]["r"]` read makes every extent NaN and this fails (as the fit's SVD then does).""" - from magnetics._slcontour import omfit_compat as oc + from magnetics.core import qs_device as oc geo = oc.sensor_geometry("DIII-D", shot=int(synthetic_shot)) for coord in ("r_end1", "z_end1", "phi_end1", "theta_end1"): diff --git a/tests/test_quasistationary.py b/tests/test_quasistationary.py deleted file mode 100644 index be2b052..0000000 --- a/tests/test_quasistationary.py +++ /dev/null @@ -1,59 +0,0 @@ -"""The pure quasi-stationary port (`core/quasistationary.py`) — previously a test -desert. Covers modal recovery, the helicity sign-flip, and the fit↔reconstruct -round-trip. - -Sign-convention note (audit follow-up): `quasistationary.reconstruct_grid` uses -`exp(+i(nφ+mθ))` while `qs_bridge._reconstruct_grid` uses `exp(-i…)`. These are -NOT in conflict — each matches its own pipeline's fit basis, and each round-trips -(this file proves it for the pure port; `test_qs_bridge` + `qs_fit` cover the -other). Do not "align" the signs without also flipping the matching fit basis. -""" - -from __future__ import annotations - -import numpy as np - -from magnetics.core import quasistationary as qs - -# 8 midplane sensors evenly spaced in φ, tiny toroidal extent. -_PHI = np.linspace(0, 360, 8, endpoint=False) -_P1, _P2 = _PHI - 0.5, _PHI + 0.5 -_TH = np.zeros_like(_PHI) -_T_MS = np.linspace(0, 10, 20) - - -def _locked_signal(n): - patt = np.cos(np.deg2rad(n * _PHI)) - return patt, np.tile(patt[:, None], (1, _T_MS.size)) - - -def test_recovers_a_locked_mode(): - _, signal = _locked_signal(2) - res = qs.fit(_T_MS, signal, _P1, _P2, _TH, _TH, ns=(1, 2, 3), ms=(0,)) - amps = {(int(n), int(m)): abs(c) for n, m, c in zip(res.ns, res.ms, res.coeffs[:, 0])} - assert amps[(2, 0)] > 0.9 # n=2 dominant - assert amps[(1, 0)] < 1e-6 and amps[(3, 0)] < 1e-6 # others ~zero - assert np.isfinite(res.condition_number) - assert float(np.asarray(res.red_chi_sq)[0]) < 1e-9 # noiseless → perfect fit - - -def test_helicity_flips_poloidal_sign(): - _, signal = _locked_signal(2) - minus = qs.fit(_T_MS, signal, _P1, _P2, _TH, _TH, ns=(2,), ms=(1,), helicity=-1) - plus = qs.fit(_T_MS, signal, _P1, _P2, _TH, _TH, ns=(2,), ms=(1,), helicity=+1) - assert list(minus.ms) == [-1] - assert list(plus.ms) == [1] - - -def test_m0_not_flipped_by_helicity(): - _, signal = _locked_signal(2) - for hel in (-1, 1): - res = qs.fit(_T_MS, signal, _P1, _P2, _TH, _TH, ns=(1, 2, 3), ms=(0,), helicity=hel) - assert list(res.ms) == [0, 0, 0] - - -def test_fit_reconstruct_roundtrip(): - patt, signal = _locked_signal(2) - res = qs.fit(_T_MS, signal, _P1, _P2, _TH, _TH, ns=(1, 2, 3), ms=(0,)) - recon = qs.reconstruct_grid(res, _PHI, np.array([0.0]), t_idx=0)[0] - np.testing.assert_allclose(recon, patt, atol=1e-9) diff --git a/tests/test_slcontour_fit.py b/tests/test_slcontour_fit.py index e32f0ed..78078ac 100644 --- a/tests/test_slcontour_fit.py +++ b/tests/test_slcontour_fit.py @@ -1,8 +1,8 @@ -"""SLCONTOUR basis-function corners (`_slcontour/fit.form_basis_function`). +"""SLCONTOUR basis-function corners (`core/qs_fit.form_basis_function`). The `sinusoidal-integral` basis has hand-coded n=0 / m=0 special cases (division by deg2rad(dx)·in etc.) — exactly where a wrong limit hides. Pin them against the -point basis. `_slcontour` is excluded from lint/typecheck but pytest still runs it. +point basis. """ from __future__ import annotations @@ -10,8 +10,7 @@ import numpy as np import pytest -from magnetics._slcontour.fit import form_basis_function -from magnetics._slcontour.omfit_compat import OMFITexception +from magnetics.core.qs_fit import form_basis_function def test_dc_basis_is_unity(): @@ -35,5 +34,5 @@ def test_integral_basis_converges_to_point_for_narrow_extent(): def test_bad_basis_raises(): - with pytest.raises(OMFITexception): + with pytest.raises(ValueError): form_basis_function(1, 0, 0.0, 1.0, 0.0, 1.0, fit_basis="not-a-basis") diff --git a/tests/test_slcontour_geometry.py b/tests/test_slcontour_geometry.py index 7f39c0a..6878b54 100644 --- a/tests/test_slcontour_geometry.py +++ b/tests/test_slcontour_geometry.py @@ -1,7 +1,7 @@ """SLCONTOUR sensor geometry reads the shot-segmented device table. Regression guard: the device JSON is shot-segmented (``sensors[c]["segments"]``), -but ``omfit_compat.sensor_geometry`` used to read the flat ``sensors[c]["r"]``, +but ``qs_device.sensor_geometry`` used to read the flat ``sensors[c]["r"]``, which returns NaN under the segmented schema. All-NaN sensor extents make the QS fit's basis matrix NaN and the SVD never converges. These tests need only the committed device JSON — no fetched HDF5. @@ -12,7 +12,7 @@ import numpy as np import pytest -from magnetics._slcontour import omfit_compat as oc +from magnetics.core import qs_device as oc # Integrated Bp LFS midplane channels the QS (SLCONTOUR) fit uses; present at a # modern shot. @@ -65,7 +65,7 @@ def test_sensor_geometry_shot_none_falls_back_to_earliest(): @pytest.mark.parametrize("shot", [124400, 151593, 156014, 184928, 990000]) def test_two_resolvers_agree_across_eras(shot): """The device JSON is read by TWO independent segment resolvers — data/devices - (rotating path) and _slcontour/omfit_compat (QS path). They must never diverge; + (rotating path) and core/qs_device (QS path). They must never diverge; that duplication is what produced the all-NaN geometry bug.""" from magnetics.data import devices From a59c481a81aa1e6b961457fb98027266cb4e2193 Mon Sep 17 00:00:00 2001 From: priyanshlunia <40486607+priyanshlunia@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:24:45 -0400 Subject: [PATCH 02/15] docs: relocate + modernize the QS reference notebook & README MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finish the core/qs_* promotion by cleaning up its reference material, which was still written for the old OMFIT `magnetics-code` layout. - Move example_magnetics.ipynb into examples/ and rewire it onto the new infrastructure: imports become magnetics.core.qs_{io_data,device,run} plus `qs_plots as plots` (so the plot cells are unchanged); drop the sys.path hack and the OMFIT-script mapping / "ported OMFIT scripts" framing. - Replace docs/qs_slcontour/README.md with docs/quasistationary-mode-analysis.md: no OMFIT-magnetics plumbing (module↔OMFIT-script table, omfit_compat, cd analysis), instead a walkthrough of the qs_* pipeline (io_data → prep → fit → run, plus qs_bridge→service and standalone qs_plots), data inputs, quick start, and the -i reconstruction sign convention. Physics provenance (SLCONTOUR, VISION §4.1) kept. - Remove the now-empty docs/qs_slcontour/. Docs-only; no src/ behavior change. ruff format --check clean. (docs/magnetics-code-architecture.md is also OMFIT-era — left for a later sweep.) Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/qs_slcontour/README.md | 127 ------------------ docs/quasistationary-mode-analysis.md | 116 ++++++++++++++++ .../example_magnetics.ipynb | 64 +-------- 3 files changed, 121 insertions(+), 186 deletions(-) delete mode 100644 docs/qs_slcontour/README.md create mode 100644 docs/quasistationary-mode-analysis.md rename {docs/qs_slcontour => examples}/example_magnetics.ipynb (77%) diff --git a/docs/qs_slcontour/README.md b/docs/qs_slcontour/README.md deleted file mode 100644 index be096ce..0000000 --- a/docs/qs_slcontour/README.md +++ /dev/null @@ -1,127 +0,0 @@ -# magnetics-code — local OMFIT magnetics translation (3D spatial fits) - -A standalone, OMFIT-free translation of the OMFIT magnetics **3D spatial-fit** workflow — the -SLCONTOUR-style quasi-stationary modal decomposition of locked / slowly-rotating MHD modes -(VISION.md §4.1). It runs locally (no OMFIT, no MDSplus) against the project-canonical data: -per-shot signals in `data/datafile/shot_.h5` and device geometry in `data/device/.json`. - -It fits the spatial field pattern at each time slice with a cylindrical-Fourier basis - -``` -δB(φ, θ) = Σ_nm b_nm · exp(i(nφ − mθ)) -``` - -by SVD-conditioned least squares, and reports the condition number **K** (the central trust -metric — warn K > 10, error K > 20), per-coefficient error bars, and reduced χ². - -> Scope: the spatial fit of sensor **arrays**. The spectrogram / MODESPEC rotating-mode path and -> the 3D coil-current fits are intentionally out of scope here. - -## Pipeline ↔ OMFIT mapping - -| local module | role | OMFIT magnetics module script | -|---|---|---| -| `io_data.py` | `load_shot` reads `shot_.h5` signals + the device JSON geometry (the "fetch" + "init" steps) | `fetch_magnetics.py` / `init_magnetics.py` | -| `prep.py` | trim, DC-comp, integrate, causal filter, detrend, SVD-condition | `prep_magnetics.py` | -| `fit.py` | basis matrix + SVD least-squares modal fit | `fit_magnetics.py` | -| `run.py` | `run_steps` orchestrator (load → prep → fit) | `run_magnetics.py` | -| `plots.py` | sensor map, signals, fit quality, mode amp/phase, φ-vs-time contour, SVD diagnostics | `plot_magnetics_*.py` | -| `omfit_compat.py` | OMFIT-runtime shim (`printi`, `cornernote`, `uband`, `is_device`) + the device layer (geometry, named subsets, wall — all from the device JSON) | (OMFIT framework) | - -This package runs without OMFIT or MDSplus; it reads only the project-canonical data files (below). - -## Data - -### Per-shot signals — `data/datafile/shot_.h5` - -One HDF5 file per shot (the PTDATA fetch output): a group per channel with `data` + `time` -(hard-linked into a shared `_timebases` group), **all time in milliseconds**. Channels sample at -several rates (integrated probes ~50 kHz, bdots ~200 kHz, coils/`ip`/`bt` ~20 kHz). Root attrs -include `device`, `shot`, `tmin`/`tmax` (ms), `channels_fetched`/`channels_missing`. - -`io_data.load_shot` turns this into the Datasets the pipeline expects: - -- `raw` — `signal(channel, time)` (time in **seconds**), `signal_sigma` (a constant 2e-5 T), and the - per-channel geometry joined from the device JSON, including the derived `*_end1/2` coordinates. - All sensor channels are linearly interpolated onto a single common time axis (the densest native - grid, clipped to the window they share). -- `plasma` — `Ip`, `Bt` from the `ip`/`bt` channels (time in **ms**), `helicity` attr (default −1). -- `coupling` — `None`; the new files carry no DC vacuum-coupling matrix, so `prep`'s DC compensation - (`dc_comp=True`) is unavailable. - -### Device geometry — `data/device/.json` - -The canonical device description (`diiid.json` for DIII-D): `R0`, the `wall` outline, the per-sensor -base geometry (`r, z, phi, tilt, length, delta_phi, na, pair`), and named `sensor_sets`. Loaded via -`omfit_compat` (`load_device`, `sensor_geometry`, `load_wall`, `resolve_channel_filter`). -`sensor_geometry` derives `theta` and the `*_end1/2` sensor-end coordinates from the base geometry -(the OMFIT `init_magnetics.py` step). - -## Quick start - -```python -import sys; sys.path.insert(0, '.') # the modules are a flat set, not a package -from run import run_steps -import plots - -# load → prep → fit the LFS-midplane toroidal Bp array for n = 1,2,3 -r = run_steps(199749, channel_filter='Bp_LFS_midplane', ns=(1, 2, 3), ms=(0,), - time_trim=(3.3, 3.5), prep_kwargs=dict(cutoff_hz=(5, 250), energy=0.98)) - -print('condition number K =', r.condition_number) -plots.plot_fit_modes(r.fit) # amplitude & phase of each mode vs time -plots.plot_slice(r.fit, fix_coord='theta') # the SLCONTOUR φ-vs-time contour -``` - -`channel_filter` accepts a regex, a list of regexes, or a friendly subset name — with or without -underscores (e.g. `'Bp_LFS_midplane'` or `'Bp LFS midplane'`, `'Bp_All'`, `'All_3D_Coils'`). List -them all with `io_data.available_subsets('DIII-D')`. `time_trim` must fall inside the shot file's -window (shot 199749 spans **3.3–3.5 s**). - -## Device file — `data/device/diiid.json` - -All device metadata lives in this single JSON (replacing the old bundled `DATA/DIII-D/*.txt` tables): - -| key | content | loader (`omfit_compat`) | -|---|---|---| -| `sensor_sets` | named sensor subsets (`Bp LFS midplane`, `Bp All`, `All 3D Coils`, …) | `list_sensor_subsets`, `resolve_channel_filter` | -| `sensors` | per-sensor base geometry (`r, z, phi, tilt, length, …`) → derived ends | `sensor_geometry` (alias `load_sensor_table`) | -| `R0`, `wall` | machine major radius + first-wall outline | `load_wall` | - -## The worked example - -[`example_magnetics.ipynb`](example_magnetics.ipynb) walks the full pipeline end to end (sensor map, -conditioned signals, SVD diagnostics, fit quality, mode amplitude & phase, and the φ-vs-time -contour). It is **shot-agnostic**: a single *Parameters* cell at the top sets `SHOT`, -`CHANNEL_FILTER`, `TIME_TRIM`, and the mode lists — change those to analyse a different shot or -array. It defaults to shot **199749** (`Bp_LFS_midplane`, `time_trim=(3.3, 3.5)`). - -Run it on the uv environment (Python 3.14), e.g. a quick headless smoke test of the same pipeline: - -```bash -cd analysis -uv run python -c "import sys; sys.path.insert(0,'magnetics-code'); \ -from run import run_steps; r=run_steps(199749, channel_filter='Bp_LFS_midplane', \ -time_trim=(3.3,3.5), prep_kwargs=dict(cutoff_hz=(5,250), energy=0.98)); \ -print('K =', r.condition_number)" - -# or execute the whole notebook headless (regenerates its outputs): -uv run --with nbconvert jupyter nbconvert --to notebook --execute --inplace \ - magnetics-code/example_magnetics.ipynb --ExecutePreprocessor.kernel_name=magnetics-uv -``` - -## Dependencies - -Added to the `analysis` uv project: `xarray`, `scipy`, `h5py`, `ipykernel` (all install cleanly on -the pinned Python 3.14). `io_data` reads the shot HDF5 directly with `h5py`. - -## Notes & limitations - -- The LFS-midplane fit uses the integrated **`MPID66M*`** Bp array (`'Bp_LFS_midplane'`, - `integrate=False`). For shot 199749 the file has 9 of the 10 sensors (`MPID66M020` is in - `channels_missing`), which still resolves the low-n structure. -- Reduced χ² runs above 1 because the constant 2e-5 T sensor sigma is optimistic relative to the - higher-n structure not in the n ≤ 3 basis; the residuals remain small vs. the signals. -- Out of scope (not ported): spectrogram/MODESPEC, 3D coil-current fits, the remote IDL - `slcontour.py`, Maxwell stress, and internal/external source separation (needs Br, which this - Bp-only dataset lacks). diff --git a/docs/quasistationary-mode-analysis.md b/docs/quasistationary-mode-analysis.md new file mode 100644 index 0000000..2eb3feb --- /dev/null +++ b/docs/quasistationary-mode-analysis.md @@ -0,0 +1,116 @@ +# Quasi-stationary mode analysis (SLCONTOUR) — pipeline details + +The **quasi-stationary** path analyses locked and slowly-rotating MHD modes by fitting the *spatial* +pattern of the perturbed field measured by a sensor **array** (VISION.md §4.1). At each time slice it +decomposes δB(φ, θ) into a small set of toroidal/poloidal harmonics by SVD-conditioned least squares, +and reports the design-matrix **condition number K** (the central trust metric — warn K > 10, error +K > 20), per-coefficient error bars, and reduced χ². + +It runs entirely off the project-canonical data — per-shot signals in `data/datafile/shot_.h5` +and device geometry in `data/device/.json` — with no MDSplus or external framework. + +The fit basis (cylindrical geometry) is + +``` +δB(φ, θ) = Σ_nm b_nm · exp(i(nφ + mθ)) +``` + +> **Scope.** This is the spatial fit of sensor *arrays*. The rotating-mode spectrogram / MODESPEC +> path (`magnetics.core.spectral`) and 3D coil-current fits are separate and out of scope here. + +## The pipeline (`magnetics.core.qs_*`) + +All modules live in `src/magnetics/core/` under the unified `qs_` prefix: + +| module | role | +|---|---| +| `qs_io_data.load_shot` | read `shot_.h5` signals and join the per-channel device geometry → the `raw` / `plasma` Datasets | +| `qs_prep.prepare` | trim (channels + time), optional integrate (bdot→B), causal band/high/low-pass, detrend, SVD-condition the data matrix | +| `qs_fit.fit` | build the design matrix from the basis, SVD it (K + error bars), least-squares fit every time slice → the `fit` Dataset | +| `qs_run.run_steps` | the `load → prep → fit` orchestrator; returns a `MagneticsRun` bundling `raw` / `prepared` / `plasma` / `fit` (+ `condition_number`) | +| `qs_device` | device-JSON readers — `sensor_geometry` (base + derived `theta`/`*_end1/2`), `resolve_channel_filter`, `list_sensor_subsets`, `load_wall` — **delegating to the `data/` layer** (`data.devices`, `data.diiid_geometry`) | +| `qs_bridge` | adapt the `fit` Dataset → GUI `kind`-nodes — the **production/service** output path | +| `qs_plots` | standalone **matplotlib** plots (sensor map, signals, SVD diagnostics, fit quality, mode amp/phase, φ-vs-time contour) — for notebooks/offline use, **not** wired into the GUI | + +There are two consumers of the `fit` Dataset. The **service** builds JSON nodes from it via +`qs_bridge` (served by `service/nodes.py` → the GUI's QS tab). The **notebook/offline** path renders +it directly with `qs_plots`. They share the same physics and the same reconstruction sign convention +(below); `qs_bridge` does not import `qs_plots`. + +## Data inputs + +### Per-shot signals — `data/datafile/shot_.h5` + +One HDF5 file per shot (the PTDATA fetch output): a group per channel with `data` + `time` +(hard-linked into a shared `_timebases` group), **all time in milliseconds**. Channels sample at +several rates (integrated probes ~50 kHz, bdots ~200 kHz, coils/`ip`/`bt` ~20 kHz). Root attrs +include `device`, `shot`, `tmin`/`tmax` (ms). + +`qs_io_data.load_shot` turns this into the Datasets the pipeline expects: + +- `raw` — `signal(channel, time)` (time in **seconds**), `signal_sigma` (a constant 2e-5 T), and the + per-channel geometry joined from the device JSON, including the derived `*_end1/2` coordinates. All + sensor channels are interpolated onto one common time axis (the densest native grid, clipped to the + window they share). +- `plasma` — `Ip`, `Bt` from the `ip`/`bt` channels (time in **ms**), `helicity` attr (default −1). +- `coupling` — `None`; the new files carry no DC vacuum-coupling matrix, so `qs_prep`'s DC + compensation (`dc_comp=True`) is unavailable. + +### Device geometry — `data/device/.json` + +The canonical device description (`diiid.json` for DIII-D): `R0`, the shot-segmented `first_wall` +outline, the per-sensor base geometry (`r, z, phi, tilt, length, delta_phi, na, pair`), and named +`sensor_sets`. `qs_device` reads it through the shared `data.devices` resolvers (so the QS pipeline +and the fetcher never disagree about a shot's geometry) and adds the QS-specific derived `theta` and +`*_end1/2` sensor-end coordinates in `sensor_geometry`. + +## Quick start + +The project is a `uv` package rooted at the repo; import `magnetics.core.qs_*` directly (no +`sys.path` juggling): + +```python +from magnetics.core.qs_run import run_steps +from magnetics.core import qs_plots as plots + +# load → prep → fit the LFS-midplane toroidal Bp array for n = 1,2,3 +r = run_steps(199749, channel_filter="Bp_LFS_midplane", ns=(1, 2, 3), ms=(0,), + time_trim=(3.3, 3.5), prep_kwargs=dict(cutoff_hz=(5, 250), energy=0.98)) + +print("condition number K =", r.condition_number) +plots.plot_fit_modes(r.fit) # amplitude & phase of each mode vs time +plots.plot_slice(r.fit, fix_coord="theta") # the SLCONTOUR φ-vs-time contour +``` + +`channel_filter` accepts a regex, a list of regexes, or a friendly subset name — with or without +underscores (e.g. `"Bp_LFS_midplane"` or `"Bp LFS midplane"`, `"Bp_All"`, `"All_3D_Coils"`). List +them all with `qs_io_data.available_subsets("DIII-D")`. `time_trim` must fall inside the shot file's +window (shot 199749 spans **3.3–3.5 s**). + +The worked, shot-configurable notebook is [`examples/example_magnetics.ipynb`](../examples/example_magnetics.ipynb) +(needs a fetched `shot_.h5`; the repo ships no tokamak data). + +## Reconstruction sign convention + +The fit basis is `exp(+i(nφ + mθ))`, but each complex basis column is split into two *real* +least-squares columns and the complex coefficient is reassembled as `b = x_r + i·x_i`. Because of +that split, the field that reproduces the fitted signal is + +``` +δB(φ, θ) = Re Σ_nm b_nm · exp(−i(nφ + mθ)) +``` + +i.e. **reconstruction uses `exp(−i(…))`** even though the fit basis is `+i`. This is honoured by +`qs_bridge._reconstruct_grid` / `fit_to_phi_t_node` and `qs_plots.plot_slice` (and matches the DIII-D +SLCONTOUR reference). Do not "align" the reconstruction to `+i` — that mirrors the toroidal phase and +flips helicity. A regression test in `tests/test_qs_bridge.py` pins this. + +## Notes & limitations + +- Reduced χ² typically runs above 1 because the constant 2e-5 T sensor σ is optimistic relative to + higher-n structure not in the basis; residuals stay small versus the signals. Per-sensor σ / + helicity from the data layer is an open item. +- The pure-numpy relatives of `qs_fit` (a device-agnostic `core` port) and the `data/`-layer + set-flattening dedup are open follow-ups. +- Out of scope for this path: the rotating-mode spectrogram/MODESPEC analysis, 3D coil-current fits, + and internal/external source separation (needs Br, which a Bp-only array lacks). diff --git a/docs/qs_slcontour/example_magnetics.ipynb b/examples/example_magnetics.ipynb similarity index 77% rename from docs/qs_slcontour/example_magnetics.ipynb rename to examples/example_magnetics.ipynb index 70b68d9..10f5f93 100644 --- a/docs/qs_slcontour/example_magnetics.ipynb +++ b/examples/example_magnetics.ipynb @@ -4,30 +4,7 @@ "cell_type": "markdown", "id": "0", "metadata": {}, - "source": [ - "# 3D magnetics spatial fit — DIII-D (configurable shot)\n", - "\n", - "A worked example of the **SLCONTOUR-style quasi-stationary spatial fit** (VISION.md §4.1),\n", - "run **locally** — off OMFIT and off MDSplus — using the ported OMFIT magnetics scripts in this\n", - "directory. It reads the per-shot signals from `data/datafile/shot_.h5` and the device\n", - "geometry from `data/device/diiid.json`.\n", - "\n", - "**Pipeline:** `load → prep → fit → plot`, the local translation of the OMFIT\n", - "`fetch → prep → fit → plot` workflow:\n", - "\n", - "| local module | OMFIT script |\n", - "|---|---|\n", - "| `io_data.load_shot` | `SCRIPTS/fetch_magnetics.py` + `init_magnetics.py` (signals + derived geometry) |\n", - "| `prep.prepare` | `SCRIPTS/prep_magnetics.py` |\n", - "| `fit.fit` | `SCRIPTS/fit_magnetics.py` |\n", - "| `run.run_steps` | `SCRIPTS/run_magnetics.py` |\n", - "| `plots.*` | `PLOTS/plot_magnetics_*.py` |\n", - "\n", - "The run is driven entirely by the **Parameters** cell below — change `SHOT`, `CHANNEL_FILTER`,\n", - "`TIME_TRIM`, or the mode lists to analyse a different shot or array. The default analyses the\n", - "LFS-midplane toroidal Bp array (`Bp_LFS_midplane`) of shot **199749** for toroidal mode numbers\n", - "n = 1, 2, 3." - ] + "source": "# 3D magnetics spatial fit — DIII-D (configurable shot)\n\nA worked example of the **SLCONTOUR-style quasi-stationary spatial fit** (VISION.md §4.1),\nrun **locally** — off MDSplus — with the `magnetics.core.qs_*` pipeline. It reads the per-shot\nsignals from `data/datafile/shot_.h5` and the device geometry from `data/device/diiid.json`.\n\n**Pipeline:** `load → prep → fit → plot`:\n\n| module | role |\n|---|---|\n| `qs_io_data.load_shot` | read `shot_.h5` signals + join the device-JSON geometry (derived `*_end1/2` ends) |\n| `qs_prep.prepare` | trim, integrate, causal band-pass, detrend, SVD-condition |\n| `qs_fit.fit` | build the basis matrix + SVD-conditioned least-squares modal fit |\n| `qs_run.run_steps` | the load → prep → fit orchestrator |\n| `qs_plots.*` | sensor map, signals, SVD diagnostics, fit quality, mode amp/phase, φ-vs-time contour |\n\n(The production/GUI path adapts the same `qs_fit` Dataset through `qs_bridge`; `qs_plots` here is\nthe standalone-matplotlib view.)\n\nThe run is driven entirely by the **Parameters** cell below — change `SHOT`, `CHANNEL_FILTER`,\n`TIME_TRIM`, or the mode lists to analyse a different shot or array. The default analyses the\nLFS-midplane toroidal Bp array (`Bp_LFS_midplane`) of shot **199749** for toroidal mode numbers\nn = 1, 2, 3." }, { "cell_type": "code", @@ -35,21 +12,7 @@ "id": "1", "metadata": {}, "outputs": [], - "source": [ - "import sys\n", - "\n", - "sys.path.insert(0, \".\") # make the local modules importable\n", - "\n", - "import numpy as np\n", - "import matplotlib.pyplot as plt\n", - "\n", - "from io_data import load_shot, valid_channels, available_subsets\n", - "from omfit_compat import resolve_channel_filter\n", - "from run import run_steps\n", - "import plots\n", - "\n", - "%matplotlib inline" - ] + "source": "import numpy as np\nimport matplotlib.pyplot as plt\n\nfrom magnetics.core.qs_io_data import load_shot, valid_channels, available_subsets\nfrom magnetics.core.qs_device import resolve_channel_filter\nfrom magnetics.core.qs_run import run_steps\nfrom magnetics.core import qs_plots as plots\n\n%matplotlib inline" }, { "cell_type": "markdown", @@ -89,15 +52,7 @@ "cell_type": "markdown", "id": "4", "metadata": {}, - "source": [ - "## 1. Load the shot (the \"fetch\" + \"init\" step)\n", - "\n", - "`load_shot` reads the raw sensor signals from `data/datafile/shot_.h5` and joins the\n", - "per-channel geometry — including the derived sensor-end coordinates (`*_end1/2`) — from\n", - "`data/device/diiid.json`. All sensor channels are interpolated onto a single common time axis\n", - "(in seconds); the global `ip`/`bt` traces become the `plasma` Dataset. The new files carry no\n", - "DC-coupling matrix, so `coupling` is `None`." - ] + "source": "## 1. Load the shot\n\n`load_shot` reads the raw sensor signals from `data/datafile/shot_.h5` and joins the\nper-channel geometry — including the derived sensor-end coordinates (`*_end1/2`) — from\n`data/device/diiid.json`. All sensor channels are interpolated onto a single common time axis\n(in seconds); the global `ip`/`bt` traces become the `plasma` Dataset. The new files carry no\nDC-coupling matrix, so `coupling` is `None`." }, { "cell_type": "code", @@ -361,16 +316,7 @@ "cell_type": "markdown", "id": "24", "metadata": {}, - "source": [ - "---\n", - "### Recap\n", - "\n", - "We reproduced the VISION §4.1 SLCONTOUR outputs locally from the ported OMFIT scripts, driven\n", - "off the `data/datafile/shot_.h5` signals and `data/device/diiid.json` geometry: sensor map,\n", - "conditioned signals, data/design-matrix SVD diagnostics, fit quality, **amplitude & phase of each\n", - "(n, m) mode vs time**, and the **phi-vs-time contour**. To analyse a different shot or array, edit\n", - "the **Parameters** cell (`SHOT`, `CHANNEL_FILTER`, `TIME_TRIM`, `NS`/`MS`)." - ] + "source": "---\n### Recap\n\nWe reproduced the VISION §4.1 SLCONTOUR outputs locally with the `magnetics.core.qs_*` pipeline,\ndriven off the `data/datafile/shot_.h5` signals and `data/device/diiid.json` geometry: sensor\nmap, conditioned signals, data/design-matrix SVD diagnostics, fit quality, **amplitude & phase of\neach (n, m) mode vs time**, and the **phi-vs-time contour**. To analyse a different shot or array,\nedit the **Parameters** cell (`SHOT`, `CHANNEL_FILTER`, `TIME_TRIM`, `NS`/`MS`)." } ], "metadata": { @@ -394,4 +340,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file From 850ce96a897df7575c7e3b6fcac879459a418f18 Mon Sep 17 00:00:00 2001 From: priyanshlunia <40486607+priyanshlunia@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:48:54 -0400 Subject: [PATCH 03/15] Adjusted stale docstring references, clarified each script's context and role --- docs/{ => specs}/quasistationary-mode-analysis.md | 2 +- src/magnetics/core/qs_bridge.py | 4 ++-- src/magnetics/core/qs_device.py | 5 +++++ src/magnetics/core/qs_fit.py | 2 +- src/magnetics/core/qs_io_data.py | 2 +- src/magnetics/core/qs_plots.py | 2 +- src/magnetics/core/qs_prep.py | 2 +- src/magnetics/core/qs_run.py | 2 +- 8 files changed, 13 insertions(+), 8 deletions(-) rename docs/{ => specs}/quasistationary-mode-analysis.md (98%) diff --git a/docs/quasistationary-mode-analysis.md b/docs/specs/quasistationary-mode-analysis.md similarity index 98% rename from docs/quasistationary-mode-analysis.md rename to docs/specs/quasistationary-mode-analysis.md index 2eb3feb..b273f30 100644 --- a/docs/quasistationary-mode-analysis.md +++ b/docs/specs/quasistationary-mode-analysis.md @@ -1,4 +1,4 @@ -# Quasi-stationary mode analysis (SLCONTOUR) — pipeline details +# Quasi-stationary mode analysis (SLCONTOUR-style) — pipeline details The **quasi-stationary** path analyses locked and slowly-rotating MHD modes by fitting the *spatial* pattern of the perturbed field measured by a sensor **array** (VISION.md §4.1). At each time slice it diff --git a/src/magnetics/core/qs_bridge.py b/src/magnetics/core/qs_bridge.py index 7614d74..c952992 100644 --- a/src/magnetics/core/qs_bridge.py +++ b/src/magnetics/core/qs_bridge.py @@ -1,7 +1,7 @@ -"""Adapt fit.py's xarray Dataset output → GUI kind-node dicts. +"""Adapt the ``qs_fit`` xarray Dataset output → GUI kind-node dicts. Pure output-side bridge — no data loading, no fit logic. Takes the Dataset -produced by ``magnetics-code/fit.fit()`` (or any equivalent) and converts it +produced by ``qs_fit.fit()`` (or any equivalent) and converts it to the JSON contract consumed by the GUI's NodeView / Plot components. Input contract (fit Dataset variables): diff --git a/src/magnetics/core/qs_device.py b/src/magnetics/core/qs_device.py index a36f51a..6e68411 100644 --- a/src/magnetics/core/qs_device.py +++ b/src/magnetics/core/qs_device.py @@ -15,6 +15,11 @@ Replaces the old ``_slcontour/omfit_compat.py`` OMFIT shim (deleted). This module does not read the device JSON itself — it routes through ``data.devices`` so the QS pipeline and the fetcher can never disagree about a shot's geometry. + +Kept as a standalone leaf (not folded into ``qs_io_data``) because it is *shared* device +metadata: ``qs_prep``, ``qs_fit``, ``qs_plots`` and ``service.nodes`` all import from it, +not just the shot loader — so a separate module keeps those consumers from depending on +``qs_io_data`` (the HDF5 loader) just to reach ``is_device`` / ``load_wall``. """ from __future__ import annotations diff --git a/src/magnetics/core/qs_fit.py b/src/magnetics/core/qs_fit.py index f926de8..e6f8b92 100644 --- a/src/magnetics/core/qs_fit.py +++ b/src/magnetics/core/qs_fit.py @@ -1,4 +1,4 @@ -"""Port of ``SCRIPTS/fit_magnetics.py`` — the SLCONTOUR-style spatial fit. +"""A port of the OMFIT magnetics *fit* script — the QS pipeline's fit step. This is the heart of the quasi-stationary analysis (VISION.md S4.1). At each time slice it fits the spatial field pattern with either a cylindrical-Fourier diff --git a/src/magnetics/core/qs_io_data.py b/src/magnetics/core/qs_io_data.py index f9424e4..ab5c23e 100644 --- a/src/magnetics/core/qs_io_data.py +++ b/src/magnetics/core/qs_io_data.py @@ -1,4 +1,4 @@ -"""Shot loader — reads the per-shot HDF5 file and the device JSON. +"""A port of the OMFIT magnetics *fetch* + *init* scripts — the QS pipeline's load step. Replaces the old "fetch" step (and the netCDF ``RAW``/``PLASMA_PARAMS``/ ``COUPLING`` loader). Inputs now come from the project-canonical locations: diff --git a/src/magnetics/core/qs_plots.py b/src/magnetics/core/qs_plots.py index 5f73c0f..8be8327 100644 --- a/src/magnetics/core/qs_plots.py +++ b/src/magnetics/core/qs_plots.py @@ -1,4 +1,4 @@ -"""Standalone matplotlib plots for the quasi-stationary (SLCONTOUR) pipeline. +"""A port of the OMFIT magnetics *plot* scripts — the QS pipeline's visualization step. These are **not** wired into the GUI/service. The FastAPI service serves JSON ``kind``-nodes built by :mod:`magnetics.core.qs_bridge`, which only *ports* these diff --git a/src/magnetics/core/qs_prep.py b/src/magnetics/core/qs_prep.py index f3a8ed5..95f5851 100644 --- a/src/magnetics/core/qs_prep.py +++ b/src/magnetics/core/qs_prep.py @@ -1,4 +1,4 @@ -"""Port of ``SCRIPTS/prep_magnetics.py`` — signal conditioning before the fit. +"""A port of the OMFIT magnetics *prep* script — the QS pipeline's prep step (signal conditioning before the fit). Steps (matching the OMFIT original): 1. trim channels (by regex) and time to the window of interest; diff --git a/src/magnetics/core/qs_run.py b/src/magnetics/core/qs_run.py index 866c3bf..b68b99e 100644 --- a/src/magnetics/core/qs_run.py +++ b/src/magnetics/core/qs_run.py @@ -1,4 +1,4 @@ -"""Port of ``SCRIPTS/run_magnetics.py`` — the load -> prep -> fit orchestrator. +"""A port of the OMFIT magnetics *run* script — the QS pipeline's orchestration step (drives load -> prep -> fit). OMFIT's ``run_steps`` loops over channel filters, running fetch/prep/fit for each and (optionally) plotting the mode dynamics. Locally "fetch" is just the From c1c5cd76adc472951cf72c2661e87260dbb4280a Mon Sep 17 00:00:00 2001 From: mfairborn23 Date: Wed, 1 Jul 2026 13:30:51 -0400 Subject: [PATCH 04/15] qs: array-dropdown parity, Advanced Options (sigma/energy/basis/fit_cond), wall-outline fix - QS tab's Array dropdown now uses the same device sensor_sets list as the left sidebar instead of a hardcoded 9-item subset. - New Advanced Options controls in the QS tab: uncertainty (sigma), fraction of energy included, basis function, fit condition - wired through to the SLCONTOUR fit via new sigma_override/fit_basis/fit_cond query params. - Fix load_wall() reading the device JSON's legacy flat wall schema; it's been shot-segmented since 28c75c6, so the wall outline silently stopped rendering in the QS Sensor Map cross-section plot. Co-Authored-By: Claude Sonnet 5 --- .../components/tabs/QuasiStationaryTab.tsx | 81 ++++++++++++++++++- src/magnetics/_slcontour/fit.py | 19 +++-- src/magnetics/_slcontour/omfit_compat.py | 10 ++- src/magnetics/_slcontour/plots.py | 2 +- src/magnetics/service/nodes.py | 17 +++- tests/test_qs_bridge.py | 18 +++++ tests/test_slcontour_geometry.py | 13 +++ 7 files changed, 146 insertions(+), 14 deletions(-) diff --git a/gui/web/src/components/tabs/QuasiStationaryTab.tsx b/gui/web/src/components/tabs/QuasiStationaryTab.tsx index af2994d..ec0f6a4 100644 --- a/gui/web/src/components/tabs/QuasiStationaryTab.tsx +++ b/gui/web/src/components/tabs/QuasiStationaryTab.tsx @@ -9,6 +9,7 @@ import NodeView from "../../lib/NodeView"; import Plot from "../../lib/Plot"; import type { ContourNode, LineNode, MetricsNode } from "../../lib/contract"; import { phiPeak as phiPeakFn, phiRms as phiRmsFn } from "../../lib/qsTransforms"; +import { fetchDevices, type DeviceInfo } from "../../lib/api"; // ── Colorblind-safe palette (Wong 2011) — for sensor/channel traces ── const LINE_PALETTE = ["#0072B2", "#E69F00", "#56B4E9", "#D55E00", "#CC79A7", "#009E73", "#F0E442"]; @@ -94,7 +95,9 @@ function lineTraces( return traces; } -// ── Sensor arrays most useful for QS analysis ──────────────────────── +// ── Sensor arrays most useful for QS analysis — offline/mock fallback, +// used when there's no live backend (or no matching device) to supply the +// full device.sensor_sets list the left sidebar's pull panel uses ──────── const CHANNEL_FILTERS = [ "Bp LFS midplane", "Bp LFS midplane bdot", @@ -129,7 +132,7 @@ function CollapseHeader({ // ── Component ───────────────────────────────────────────────────────── export default function QuasiStationaryTab({ machine }: { machine: string }) { const dark = useDarkMode(); - const { cursorMs, setCursorMs } = useStore(); + const { cursorMs, setCursorMs, machines } = useStore(); // ── Analysis settings ───────────────────────────────────────────── const [ns, setNs] = useState("1,2,3"); @@ -142,9 +145,25 @@ export default function QuasiStationaryTab({ machine }: { machine: string }) { const [tmaxMs, setTmaxMs] = useState(""); const [colormapChoice, setColormapChoice] = useState<"rdbu" | "cividis" | "viridis">("rdbu"); + // ── Advanced fit-tuning settings ─────────────────────────────────── + const [uncertainty, setUncertainty] = useState("0.00002"); + const [energyFraction, setEnergyFraction] = useState("0.98"); + const [fitBasis, setFitBasis] = useState("sinusoidal-integral"); + const [fitCond, setFitCond] = useState("10.0"); + + // ── Devices (for the Array dropdown's sensor-set list) ───────────── + const [devices, setDevices] = useState([]); + useEffect(() => { void fetchDevices().then(setDevices); }, []); + const currentDeviceName = machines.find(m => m.id === machine)?.device; + const matchedDevice = devices.find(d => d.name === currentDeviceName); + const channelFilterOptions = matchedDevice?.sensor_sets.length + ? matchedDevice.sensor_sets + : CHANNEL_FILTERS; + // ── Section collapse state ──────────────────────────────────────── const [fitQualityOpen, setFitQualityOpen] = useState(false); const [sensorMapOpen, setSensorMapOpen] = useState(false); + const [advancedOpen, setAdvancedOpen] = useState(false); // ── Deferred fetch: only compute when user clicks Plot ──────────── const [committedParams, setCommittedParams] = useState | null>(null); @@ -154,6 +173,10 @@ export default function QuasiStationaryTab({ machine }: { machine: string }) { ns, ms, channel_filter: channelFilter, detrend_type: detrendType, + sigma: uncertainty, + energy: energyFraction, + fit_basis: fitBasis, + fit_cond: fitCond, }; if (detrendLo && detrendHi) { p.detrend_lo = detrendLo; @@ -162,7 +185,10 @@ export default function QuasiStationaryTab({ machine }: { machine: string }) { if (tminMs) p.tmin_ms = tminMs; if (tmaxMs) p.tmax_ms = tmaxMs; return p; - }, [ns, ms, channelFilter, detrendType, detrendLo, detrendHi, tminMs, tmaxMs]); + }, [ + ns, ms, channelFilter, detrendType, detrendLo, detrendHi, tminMs, tmaxMs, + uncertainty, energyFraction, fitBasis, fitCond, + ]); useEffect(() => { if (cursorMs === 0) setCursorMs(3140); @@ -566,7 +592,7 @@ export default function QuasiStationaryTab({ machine }: { machine: string }) { Array