Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 6 additions & 10 deletions src/magnetics/core/qs_bridge.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,14 +57,6 @@ def _mode_label(n: int | float, m: int | float) -> str:
return f"n={n}" if m == 0 else f"m/n={m}/{n}"


def _quality_status(K: float) -> str:
if K > 20:
return "error"
if K > 10:
return "warn"
return "ok"


def _reconstruct_grid(
fit_ds, phi_grid: np.ndarray, theta_grid: np.ndarray, t_idx: int
) -> np.ndarray:
Expand Down Expand Up @@ -262,8 +254,12 @@ def fit_to_fit_quality_node(fit_ds) -> dict:
return contracts.metrics(
title="fit quality",
fields=[
{"label": "K (raw)", "value": f"{K:.1f}", "status": _quality_status(K)},
{"label": "K (eff)", "value": f"{eff_cn:.1f}", "status": _quality_status(eff_cn)},
{"label": "K (raw)", "value": f"{K:.1f}", "status": contracts.quality_for_k(K)},
{
"label": "K (eff)",
"value": f"{eff_cn:.1f}",
"status": contracts.quality_for_k(eff_cn),
},
{"label": "K cutoff", "value": f"{fit_cond:.0f}"},
{"label": "χ² (mean)", "value": f"{mean_chi2:.3f}"},
{"label": "channels", "value": n_ch},
Expand Down
4 changes: 2 additions & 2 deletions src/magnetics/core/qs_fit.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ def form_basis_function(
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)
) / (np.deg2rad(dx) * np.deg2rad(dy) * (1j * n) * (1j * m))

# ── gaussian-point ────────────────────────────────────────────────────────
if fit_basis == "gaussian-point":
Expand Down Expand Up @@ -430,7 +430,7 @@ def _printv(*a):
w_inv = 1.0 / w_a
w_inv[~valid] = 0.0
w_inv = np.hstack((w_inv, [0.0] * max(Vh_a.shape[0] - w_a.shape[0], 0)))
fit_sigma2 = np.sum((Vh_a.T * w_inv) ** 2, axis=0)
fit_sigma2 = np.sum((Vh_a.T * w_inv) ** 2, axis=1)
fit_sigmas = np.sqrt(fit_sigma2) # (n_columns,)

# ── reform per-mode complex (sinusoidal) or real (Gaussian) coefficients ──
Expand Down
4 changes: 2 additions & 2 deletions src/magnetics/core/spectral.py
Original file line number Diff line number Diff line change
Expand Up @@ -909,7 +909,7 @@ def fit_toroidal_mode(
best_n, best_R, best_c = candidates[best_j], float(rmag[best_j]), float(inter[best_j])

# Uncertainty (eigspec-style): propagate per-probe phase σ from the cross-spectral
# statistics into (a) the intercept's 1σ via inverse-variance combination and
# statistics into (a) the returned intercept's 1σ through the same weights and
# (b) a posterior over candidate n from each hypothesis's weighted χ². Falls back
# to unit σ when phase_error is absent so older callers still get sane numbers.
sigma = np.asarray(
Expand All @@ -934,7 +934,7 @@ def _wrap180(x):
post = np.exp(-(chi2 - chi2.min()) / 2.0)
post /= post.sum()
n_confidence = float(post[best_j])
phase_sigma = float(np.sqrt(1.0 / np.sum(1.0 / sigma**2)))
phase_sigma = float(np.sqrt(np.sum((w * sigma) ** 2)) / w.sum())

return ToroidalFitResult(
kind="toroidal_fit",
Expand Down
24 changes: 24 additions & 0 deletions src/magnetics/data/devices.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,3 +150,27 @@ def feature_at(dev: dict, key: str, shot: int) -> dict | None:
if active is None:
return None
return {k: v for k, v in active.items() if k != "since_shot"}


def feature_nearest(dev: dict, key: str, shot: int | None) -> dict | None:
"""Top-level geometry segment for layout/plotting at ``shot``.

Uses the active segment when available; for shots before the earliest modeled
segment, returns that earliest segment. Fetch validity remains strict via
``feature_at`` / sensor-specific resolvers.
"""
segs = _segments(dev.get(key))
if not segs:
return None
if shot is None:
active = segs[-1]
else:
active = None
for seg in segs:
if seg.get("since_shot", 0) <= shot:
active = seg
else:
break
if active is None:
active = segs[0]
return {k: v for k, v in active.items() if k != "since_shot"}
6 changes: 3 additions & 3 deletions src/magnetics/data/diiid_geometry.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,11 +148,11 @@ def device_geometry(shot: int, name: str = "diiid") -> dict:
)
)
device = dev.get("name", name)
fw = devices.feature_at(dev, "first_wall", shot) or {}
vv = (devices.feature_at(dev, "vacuum_vessel", shot) or {}).get("plates", [])
fw = devices.feature_nearest(dev, "first_wall", shot) or {}
vv = (devices.feature_nearest(dev, "vacuum_vessel", shot) or {}).get("plates", [])
coils = [
{k: c.get(k) for k in ("name", "count", "turns", "rz", "loops")}
for c in (devices.feature_at(dev, "coils", shot) or {}).get("sets", [])
for c in (devices.feature_nearest(dev, "coils", shot) or {}).get("sets", [])
]
return {
"device": device,
Expand Down
16 changes: 16 additions & 0 deletions src/magnetics/data/fetch/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@

from __future__ import annotations

import re
import shlex
import subprocess
import sys
Expand All @@ -63,6 +64,20 @@
# (its activate.d scripts set nothing PTDATA needs; direct is ~4-5s faster per pull).
DEFAULT_PYTHON = "/fusion/projects/codes/conda/omega/envs_public/toksearch_env/bin/python"

# ``remote_dir`` is interpolated into remote shell commands and rsync remote specs.
# Reject shell metacharacters and option-like relative paths at the input boundary.
_SAFE_REMOTE_DIR = re.compile(r"^~?[\w./\-]+$")


def _validate_remote_dir(remote_dir: str) -> str:
if not _SAFE_REMOTE_DIR.match(remote_dir) or remote_dir.startswith("-"):
raise ValueError(
f"unsafe remote_dir {remote_dir!r}: only '~', letters, digits, '_', '.', "
"'-' and '/' are allowed (no spaces or shell metacharacters), and relative "
"paths must not start with '-'"
)
return remote_dir


def _log(msg: str) -> None:
sys.stderr.write(msg + "\n")
Expand Down Expand Up @@ -98,6 +113,7 @@ def run_remote(
the network, no hop needed); pass an explicit host[:port] to force one, or an
empty string to force none (e.g. an ssh-config alias that carries its own
ProxyJump)."""
remote_dir = _validate_remote_dir(remote_dir)
login = cluster_login(device)
host = host or login["host"] or DEFAULT_HOST
port = login["port"] # cluster SSH port (22 unless the device overrides it)
Expand Down
4 changes: 3 additions & 1 deletion src/magnetics/service/mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

from __future__ import annotations

import zlib

import numpy as np

from ..data import diiid
Expand Down Expand Up @@ -126,7 +128,7 @@ def qs_fit_frames(machine: str, params: dict) -> list[tuple[float, dict]]:
# ── spectrogram (Burgess) — coarse → fine, mode that locks or rotates ────────
def spectrogram_frames(machine: str, params: dict) -> list[tuple[float, dict]]:
p = _profile(machine)
rng = np.random.default_rng(abs(hash(machine)) % 2**32)
rng = np.random.default_rng(zlib.crc32(machine.encode()))
t0, t1 = p["t_ms"]
resolutions = [(60, 45), (120, 90), (240, 180)] # (nt, nf), → real 360×180
frames = []
Expand Down
2 changes: 1 addition & 1 deletion src/magnetics/service/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1487,7 +1487,7 @@ def _sensor_map_rz(shot, params=None) -> dict:

from ..core.qs_device import load_wall

r_wall, z_wall = load_wall(device, shot)
r_wall, z_wall = load_wall(device, int(shot))

series = []
for c in channels:
Expand Down
8 changes: 8 additions & 0 deletions tests/test_device_geometry.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,14 @@ def test_geometry_none_when_decommissioned_or_out_of_range():
assert devices.geometry_at(DEV, "S_OK", 100000) is None


def test_top_level_feature_nearest_preserves_strict_feature_at():
dev = {"wall": {"segments": [{"since_shot": 200, "r": [1.0], "z": [0.0]}]}}

assert devices.feature_at(dev, "wall", 100) is None
assert devices.feature_nearest(dev, "wall", 100) == {"r": [1.0], "z": [0.0]}
assert devices.feature_nearest(dev, "wall", None) == {"r": [1.0], "z": [0.0]}


# ── fetch assembly: _resolve_pointnames (no network) ──────────────────────────
def test_fetch_resolution_recent_shot_keeps_all():
from magnetics.data.fetch import toksearch as tf
Expand Down
30 changes: 29 additions & 1 deletion tests/test_fetch_network.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

import pytest

from magnetics.data.fetch import network
from magnetics.data.fetch import network, remote


# --- endpoint resolution from the `network` block -----------------------------
Expand All @@ -31,6 +31,34 @@ def test_nstx_endpoints_and_no_cluster():
assert network.cluster_login("nstx")["host"] is None


@pytest.mark.parametrize("safe", ["~/magnetics_fetch", "/scratch/u/fetch", "rel/dir", "~/a.b-c_d"])
def test_remote_dir_accepts_safe_paths(safe):
assert remote._validate_remote_dir(safe) == safe


@pytest.mark.parametrize(
"bad",
["~/my fetch", "$(touch x)", "a;b", "a&&b", "`id`", "a|b", "a>b", "-tmp", "--help"],
)
def test_remote_dir_rejects_shell_metacharacters_and_options(bad):
with pytest.raises(ValueError):
remote._validate_remote_dir(bad)


def test_run_remote_rejects_unsafe_remote_dir_before_commands(monkeypatch):
calls = []

def fake_run(*args, **kwargs):
calls.append((args, kwargs))
raise AssertionError("run_remote should validate remote_dir before subprocess")

monkeypatch.setattr(remote.subprocess, "run", fake_run)
with pytest.raises(ValueError):
remote.run_remote(184927, remote_dir="a;b")

assert calls == []


# --- per-url duo / 2FA metadata -----------------------------------------------
def test_diiid_endpoints_do_not_need_duo():
# key-based through cybele — explicitly marked duo:false on every endpoint
Expand Down
19 changes: 19 additions & 0 deletions tests/test_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

from __future__ import annotations

import numpy as np
import pytest

from magnetics.core import contracts
Expand Down Expand Up @@ -193,6 +194,24 @@ def test_fit_quality_node_has_finite_k():
assert n["fields"]


def test_sensor_map_rz_threads_shot_to_wall_loader(synthetic_shot, monkeypatch):
from magnetics.core import qs_device

seen = {}

def fake_load_wall(device, shot=None):
seen["device"] = device
seen["shot"] = shot
return np.array([1.0, 2.0]), np.array([0.0, 0.1])

monkeypatch.setattr(qs_device, "load_wall", fake_load_wall)
n = nodes._sensor_map_rz(synthetic_shot, {})

assert seen["device"] == "DIII-D"
assert seen["shot"] == int(synthetic_shot)
assert n["meta"]["wall"] == {"x": [1.0, 2.0], "y": [0.0, 0.1]}


def test_real_theta_has_full_poloidal_coverage():
# θ derived from the device table's (r,z) must span well beyond the midplane
# (else no honest 2D pattern). Driven by the device catalog, not fetched data,
Expand Down
8 changes: 8 additions & 0 deletions tests/test_qs_bridge.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,14 @@ def test_dropping_a_required_variable_raises_not_silently_misserves(fit_ds):
qs_bridge.fit_to_qs_fit_node(broken)


def test_fit_quality_statuses_use_contract_vocabulary(fit_ds):
node = qs_bridge.fit_to_fit_quality_node(fit_ds)
statuses = [field["status"] for field in node["fields"] if "status" in field]

assert statuses
assert set(statuses) <= {"good", "warn", "bad"}


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)
Expand Down
53 changes: 52 additions & 1 deletion tests/test_slcontour_fit.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,9 @@

import numpy as np
import pytest
import xarray as xr

from magnetics.core.qs_fit import form_basis_function
from magnetics.core.qs_fit import fit, form_basis_function


def test_dc_basis_is_unity():
Expand All @@ -33,6 +34,56 @@ def test_integral_basis_converges_to_point_for_narrow_extent():
np.testing.assert_allclose(integral, point, rtol=1e-3)


def test_2d_integral_basis_is_the_product_of_1d_integrals():
x1 = np.array([10.0, 100.0])
x2 = x1 + 20.0
y1 = np.array([5.0, 60.0])
y2 = y1 + 15.0
n, m = 2, 3

two_dim = form_basis_function(n, m, x1, x2, y1, y2, fit_basis="sinusoidal-integral")
toroidal = form_basis_function(n, 0, x1, x2, y1, y2, fit_basis="sinusoidal-integral")
poloidal = form_basis_function(0, m, x1, x2, y1, y2, fit_basis="sinusoidal-integral")

np.testing.assert_allclose(two_dim, toroidal * poloidal)


def test_fit_reports_covariance_diagonal_sigmas():
phi = np.array([0.0, 45.0, 130.0, 250.0, 310.0])
sigma = np.array([1.0, 2.0, 1.5, 0.8, 1.3])
prepared = xr.Dataset(
{
"signal": (("channel", "time"), np.zeros((phi.size, 1))),
"signal_sigma": ("channel", sigma),
"phi_end1": ("channel", phi),
"phi_end2": ("channel", phi),
"theta_end1": ("channel", np.zeros_like(phi)),
"theta_end2": ("channel", np.zeros_like(phi)),
},
coords={"channel": [f"C{i}" for i in range(phi.size)], "time": [0.0]},
attrs={"device": "synthetic"},
)

out = fit(prepared, ns=(1, 2), ms=(0,), fit_basis="sinusoidal-point", verbose=False)

columns = []
for n in (1, 2):
basis = np.exp(1j * n * np.deg2rad(phi)) / sigma
columns.extend([basis.real, basis.imag])
design = np.array(columns).T
expected = np.sqrt(np.diag(np.linalg.inv(design.T @ design)))
actual = np.array(
[
out["fit_sigmas"].isel(mode=0, time=0).values.real,
out["fit_sigmas"].isel(mode=0, time=0).values.imag,
out["fit_sigmas"].isel(mode=1, time=0).values.real,
out["fit_sigmas"].isel(mode=1, time=0).values.imag,
]
)

np.testing.assert_allclose(actual, expected)


def test_bad_basis_raises():
with pytest.raises(ValueError):
form_basis_function(1, 0, 0.0, 1.0, 0.0, 1.0, fit_basis="not-a-basis")
19 changes: 19 additions & 0 deletions tests/test_spectral.py
Original file line number Diff line number Diff line change
Expand Up @@ -519,6 +519,25 @@ def test_fit_reports_confidence_and_sigma(self):
assert 0.0 < fit.n_confidence <= 1.0
assert fit.n_confidence > 0.9 # clean synthetic mode → confident

def test_phase_sigma_uses_returned_fit_weights(self):
mode = ModeAtFrequencyResult(
kind="mode_at_frequency",
frequency=3000.0,
phase=np.array([0.0, 240.0]),
amplitude=np.array([1.0, 9.0]),
coherence=np.ones(2),
toroidal_angle=np.array([0.0, 60.0]),
phase_error=np.array([1.0, 3.0]),
)

fit = fit_toroidal_mode(mode)
expected = float(
np.sqrt(np.sum((mode.amplitude * mode.phase_error) ** 2)) / mode.amplitude.sum()
)
inverse_variance = float(np.sqrt(1.0 / np.sum(1.0 / mode.phase_error**2)))
assert fit.phase_sigma == pytest.approx(expected)
assert fit.phase_sigma != pytest.approx(inverse_variance)

def test_confidence_drops_for_noisy_underdetermined(self):
# a noisy 2-probe pair → less confident than a clean full array
clean_mode, _ = self._modes(seed=3)
Expand Down
Loading
Loading