diff --git a/.obvious/codebase-map.md b/.obvious/codebase-map.md index 5a40627..52dd5bb 100644 --- a/.obvious/codebase-map.md +++ b/.obvious/codebase-map.md @@ -8,6 +8,7 @@ | `fusionlab/api_virtual.py` | Virtual-shot router: `/virtual/gate`, `POST /virtual/{id}` | | `fusionlab/physics.py` | 0D power-balance engine, scaling laws (IPB98, ITER89-P), limits, DEVICES table | | `fusionlab/mast.py` | FAIR-MAST loader: `load_shot`, `clean_db`, npz caching, SI → project units | +| `fusionlab/data/` | Data-source seam (DATA-1/DATA-2): ShotBundle v1 + MachineConfig, reader cache, `sources.get_source()` — `fair-mast` archive and `user-file` imports behind one protocol; import route `POST/DELETE /data/bundles/{id}` | | `fusionlab/surrogate.py` | PhysicsNeMo τ_E correction to IPB98 (session/campaign holdout) | | `fusionlab/eq_surrogate.py` | PhysicsNeMo flux-map surrogate: 93 magnetics signals → ψ(R,Z) 65×65 | | `fusionlab/fieldlines.py` | NVIDIA Warp RK4 field-line tracer + q95 check vs EFIT | diff --git a/.obvious/obvious.md b/.obvious/obvious.md index 57b61c0..b714d90 100644 --- a/.obvious/obvious.md +++ b/.obvious/obvious.md @@ -36,6 +36,10 @@ No external services (no Postgres/Redis/etc.). The only setup prerequisite is `u - `.env` (copied from `.env.example`): `FUSIONLAB_PORT=8000`, `FUSIONLAB_DEVICE=` (force `mps|cuda|cpu`; default auto). - Optional agent keys (`FUSIONLAB_AGENT_ENABLED`, `FUSIONLAB_AGENT_MODEL`, `ANTHROPIC_API_KEY`): all three are unnecessary for the offline demo; the in-app agent stays off unless the flag and a key are both set. +- Optional data-source selection (`FUSIONLAB_DATA_SOURCE`, `FUSIONLAB_USER_BUNDLES`): + the default `fair-mast` needs neither. They point the app at imported shot bundles (a per-source + machine config is the `machine.json` inside the bundles directory) — see `docs/DATA_CONTRACT.md`, + the bring-your-own-data path. - No API keys are needed. The demo runs fully offline after `make setup`. ## Codebase map diff --git a/docs/DATA_CONTRACT.md b/docs/DATA_CONTRACT.md new file mode 100644 index 0000000..1825aee --- /dev/null +++ b/docs/DATA_CONTRACT.md @@ -0,0 +1,184 @@ +# Bring your own machine: the shot-bundle contract + +FusionLab replays MAST because MAST's open archive (FAIR-MAST) is what it ships with — not because +the physics or the UI know it. Everything the app consumes per shot is one plain dictionary of +arrays in fixed units, and this page writes that dictionary down as **ShotBundle v1**, the versioned +contract any data source must produce. Convert your own shots to it — from DIII-D, JET, AUG, EAST, +your lab's machine, wherever you hold data you are entitled to hold — and FusionLab runs on them +unchanged: the same replay strips, multi-shot compare, take-away export, USD stage, and compute +harness. FAIR-MAST stays exactly as it was; your bundles replay beside the archive's shots, never +instead of them. + +Two honesty rules up front, and they hold for every number in this pipeline: + +- **Units are fixed house-wide** — m, T, MA, 1e20 m⁻³, keV, MW, MJ, s. Convert at the edge, once, + in your ingestor; nothing downstream re-scales. +- **Nothing here "validates" or "predicts" anything.** The twin is *compared with* experiment. + Limits are distances, never forecasts. Imported catalog numbers are labelled `computed`, and + quantities FusionLab derives (catalog scalars, the L/H flag where present) are labelled derived — + never presented as measurements. + +## The two ways to bring data + +**A directory of bundles** (good for many shots): one `.npz` file per shot plus one `machine.json`, +in any directory; point `FUSIONLAB_USER_BUNDLES` at it. Or just drop `.npz` files into +`data/user_bundles/`, the default. + +**One upload** (good for a first look): `POST /data/bundles/` with the `.npz` as multipart form +data — the UI's **⬆ Import bundle** button in the replay toolbar does exactly this. The file's name +( minus `.npz`) becomes the shot's id; ids may use letters, digits, dots, dashes and underscores, +and must not collide with a bare shot number (those are the archive's). `DELETE /data/bundles/` +removes an import; the archive's cache is never touched by that route. + +Either way the file is validated against the contract in full before anything is published: an +invalid bundle fails with `400` and a message naming the first bad field, and leaves no partial +file behind. + +## The bundle: one npz per shot + +Each `.npz` holds a `meta` entry (a JSON string) plus one array per field. Every trace lives on a +single common time base `t_s` — the equilibrium-reconstruction slice times if you have them; the +slow diagnostics' values are interpolated onto it before saving. Required of every bundle +(`nt` = number of time slices): + +| Field | Unit | Shape | Meaning, and what happens without it | +|---|---|---|---| +| `t_s` | s | `(nt,)` | strictly increasing time base; every trace shares it | +| `Ip_MA` | MA | `(nt,)` | plasma current (magnitude) | +| `B_T` | T | `(nt,)` | vacuum toroidal field (magnitude) | +| `n_e20` | 1e20 m⁻³ | `(nt,)` | line-average electron density; all-NaN tolerated — model traces go null, the summary shows gaps | +| `P_ohm_MW` | MW | `(nt,)` | ohmic heating power | +| `P_nbi_MW` | MW | `(nt,)` | neutral-beam power; zeros make the beam what-ifs inert | +| `W_MJ` | MJ | `(nt,)` | **MHD stored energy** `3/2 ∫p dV` — *not* the diamagnetic `wmhd`, which can differ by ~3.5× | +| `dWdt_MW` | MW | `(nt,)` | dW/dt; your ingestor may compute it by `np.gradient` if the archive lacks it | +| `q95` | – | `(nt,)` | edge safety factor; NaN → limit panels show gaps | +| `beta_N` | – | `(nt,)` | normalized beta | +| `kappa` | – | `(nt,)` | elongation | +| `a_m` | m | `(nt,)` | minor radius | +| `R_m` | m | `(nt,)` | major radius | +| `V_m3` | m³ | `(nt,)` | plasma volume | + +The **equilibrium group** is optional in v1 — a machine without equilibrium reconstructions still +replays its traces; `/replay` answers normally with the surfaces marked absent and the 3D/ψ views +degrade gracefully: + +| Field | Unit | Shape | Meaning | +|---|---|---|---| +| `lcfs_R`, `lcfs_Z` | m | `(nt, n_pts)` | last-closed-flux-surface outline, closed | +| `psi_n` | – | `(nt, NZ, NR)` f32 | normalized ψ (0 axis, 1 separatrix), rows = Z, **uniform grid** | +| `psi_R`, `psi_Z` | m | `(NR,)`, `(NZ,)` | the ψ grid axes | +| `R_mag_m`, `Z_mag_m` | m | `(nt,)` | magnetic axis | + +Also optional, and each consumed only when present: `delta`, `li`, `P_rad_MW`, `Te0_keV`, +`ne0_e20`, the Thomson profiles `ts_R` / `ts_Te_keV` / `ts_ne_e20` `(nt, n_ts)`, the vessel outline +`wall_R` / `wall_Z`, the PF-coil filaments `coil_R` / `coil_Z` / `coil_W` / `coil_H`, and +`eq_inputs` `(nt, n_sensors)` for the equilibrium surrogate — in **your machine's own sensor +layout**; the shipped surrogate model is trained on MAST's 226-channel layout and flags foreign +input as out-of-distribution rather than pretend. Field-line tracing additionally consumes +`F_Tm` `(nt, n_f)` (F = R·B_φ over normalized ψ), `psi_axis_Wb` / `psi_bnd_Wb` `(nt,)`, and +`psi_norm` `(n_f,)`; without them tracing is simply absent. + +`meta` is a JSON string (an object). Recognised keys — all optional, all read defensively: +`shot_id`, `campaign`, `heating`, `timestamp`, `divertor_config`, `plasma_shape`, `current_range`, +`shot_objective`, `preshot`, `postshot`, `summary`, `t_nbi_start_s`, `t_nbi_end_s`, +`P_nbi_max_MW`. The upload route's `meta_json` form field merges into it (your values win). + +Writing one, concretely: + +```python +import json +import numpy as np + +np.savez_compressed( + "out.npz", + meta=json.dumps({"heating": "NBI", "campaign": "MYRUN1"}), # a JSON *string* + t_s=np.arange(nt) * 0.01, Ip_MA=ip, B_T=0.4, n_e20=ne, + P_ohm_MW=pohm, P_nbi_MW=pnbi, W_MJ=w, dWdt_MW=dwdt, + q95=q95, beta_N=bn, kappa=kappa, a_m=a, R_m=R, V_m3=vol, + # optional: lcfs_R=..., psi_n=..., psi_R=..., psi_Z=..., R_mag_m=..., Z_mag_m=..., +) +``` + +## `machine.json`: the machine's identity + +The per-machine values that are *not* data — display scales, naming, attribution, calibration +thresholds — come from a `MachineConfig` in `machine.json` inside the bundles directory. Every +field is optional; the defaults are deliberately machine-neutral. Unknown keys fail loudly. + +| Field | Type | Meaning | +|---|---|---| +| `device` | str | key into `physics.DEVICES` presets (R, a, κ, B/I/P maxima); omit for a machine without a preset | +| `isotope_md` | float | isotope mass for the confinement scalings (2.0 = deuterium) | +| `alpha_particles` | bool | whether alpha heating enters the model (MAST: false, it ran deuterium) | +| `attribution` | str | provenance string carried on every replay/USD payload for these bundles — put your data's licence and source here | +| `campaign_pattern` | str | regex your campaign labels follow (the archive's look like `M9`) | +| `usd_prim_root` | str | stage root prim for USD exports (the archive's is `/MAST`) | +| `reconstruction_name` | str | what your archive calls its equilibrium code (shown in the UI) | +| `te_display_range` | [lo, hi] | 3D plasma Te colour range, keV | +| `te_usd_ramp_max_kev` | float | USD export Te→colour ramp ceiling, keV | +| `warmup_shot` | str\|null | bundle loaded at server start to warm the caches | +| `clean_limits` | object | catalog filter thresholds `{tau_E_max_s, kappa_min, a_min_m}` — imported catalogs default to *no* filtering; never MAST's machine-calibrated values | +| `bench_shot` | str\|null | the bundle the benchmark task replays (else the first cached) | +| `unit_sanity` | list | `[field, lo, hi]` mid-shot guards against a missed unit conversion — cheap, and it catches everything | + +## What runs on a ψ-free bundle + +With only the 14 required fields: replay strips, limits, summary, the virtual what-if (beam box +from `meta`), compare, CSV/JSON take-away export, catalog + search, and every harness task that +consumes traces. Without equilibrium you lose: the 3D vessel/plasma view, the ψ map, field-line +tracing, the USD stage, and the eq-surrogate overlay — each absent, never wrong. + +With the equilibrium group added, those come back on your geometry. The two *learned* models +(confinement correction, equilibrium surrogate) remain MAST-trained data: they run on your bundles +with honest out-of-distribution shading and flags, or not at all — they never silently claim your +machine. + +## The harness on imported data + +The compute tasks (benchmarks, surrogates, tracing) resolve their inputs through the same seam — +`bench.timings(source=..., shot=...)`, the benchmark route's `source`/`shot` fields, `GET +/replay/{id}` with your id. No worker-registry changes, no MAST-shaped assumptions: a task says +"needs a shot bundle", the seam hands it yours. + +```bash +# point the app at your bundles +FUSIONLAB_USER_BUNDLES=~/my_machine_bundles make dev + +# or select the source explicitly +FUSIONLAB_DATA_SOURCE=user-file make dev # default is fair-mast; /replay serves both regardless + +# benchmark the harness on your bundle +curl -s -X POST 'localhost:8000/compute/benchmark?source=user-file&shot=my_machine_0001' | jq . +``` + +Example `machine.json`, for an imaginary AUG-like machine: + +```json +{ + "device": "jet", + "isotope_md": 2.0, + "alpha_particles": false, + "attribution": "My machine runs: internal laboratory use only, do not redistribute", + "campaign_pattern": "RUN\\d+", + "usd_prim_root": "/MyMachine", + "reconstruction_name": "LELIN", + "te_display_range": [0.1, 5.0], + "te_usd_ramp_max_kev": 6.0, + "warmup_shot": "my_machine_0001", + "unit_sanity": [["R_m", 1.5, 2.0], ["a_m", 0.4, 0.6], ["B_T", 2.0, 3.5]] +} +``` + +## Versioning and the future + +The bundle schema is versioned (`ShotBundle v1`, `fusionlab/data/bundle.py` — the field table in +code is authoritative and drives both validation and this page). Later versions add fields and +degrade old readers gracefully; the `omas`/IMAS path is a future *adapter* that produces bundles, +not a new internal format. If your archive can emit these ~45 arrays per shot, FusionLab is +yours. + +One last restatement of the no-validation rule, because it is the contract's soul: numbers from +your bundles are replayed, compared with the 0D model and the learned corrections, and exported +with their provenance labels (`measured` / `model` / `learned` / `computed`). Nothing that comes +out is called a validated or predictive result — it is the model *beside* your data, honestly +labelled, everywhere. diff --git a/fusionlab/agent.py b/fusionlab/agent.py index f5b1662..72d3697 100644 --- a/fusionlab/agent.py +++ b/fusionlab/agent.py @@ -153,7 +153,8 @@ def turn(history: list[dict[str, str]], message: str, settings: AgentSettings, "summary": result.summary, "data": result.data, "error": result.error} if result.proposal is not None: # confirm-then-apply: hold it, never act on it here token = on_proposal(result.proposal) - yield {"type": "action_proposed", "action_id": token["action_id"], "tool": name, "args": args, + yield {"type": "action_proposed", "action_id": token["action_id"], "tool": name, + "args": result.proposal["args"], # the canonical args the confirm applies, not the model's raw call "math": str(result.proposal.get("math") or ""), "expires_s": token["expires_s"]} tool_results.append({"type": "tool_result", "tool_use_id": str(call["id"]), "content": tool_context(result), "is_error": result.error is not None}) diff --git a/fusionlab/agent_tools.py b/fusionlab/agent_tools.py index 1b8b583..8c2640c 100644 --- a/fusionlab/agent_tools.py +++ b/fusionlab/agent_tools.py @@ -33,10 +33,11 @@ import numpy as np from fastapi import HTTPException -from fusionlab import mast, virtual +from fusionlab import virtual from fusionlab.api import devices, simulate_point from fusionlab.api_replay import _replay_bundle, _shot, db_search, shots from fusionlab.api_virtual import EditIn, VirtualIn, _virtual_shot_body, virtual_gate +from fusionlab.data import sources from fusionlab.physics import DEVICES SEARCH_LIMIT = 200 # the tool's own cap; /db/search itself allows 1000 — the model works page-size @@ -113,15 +114,16 @@ def list_cached_shots() -> ToolResult: @_guard -def load_shot(shot_id: int, H: float = 1.0) -> ToolResult: - """A cached shot's replay summary — the /replay/{id} internals. Time arrays stay out of the model's +def load_shot(shot_id: str, H: float = 1.0) -> ToolResult: + """A shot's replay summary — the /replay/{id} internals. Time arrays stay out of the model's context: the UI draws the traces after the driver opens the shot.""" try: - r = _replay_bundle(int(shot_id), H=float(H)) + r = _replay_bundle(str(shot_id), H=float(H)) except HTTPException as e: - cached = mast.cached_shots() - return _err(f"{e.detail} Replay covers only the shots cached under data/shots ({len(cached)}: {cached}); " - "use search_shots to explore the full catalog instead.", {"shot_id": int(shot_id), "cached_shot_ids": cached}) + cached = sources.get_source(None).cached_shots() # what the route can actually load + return _err(f"{e.detail} Replay covers the shots cached under data/shots ({len(cached)}: {cached}) " + "plus any imported bundles; use search_shots to explore the full catalog instead.", + {"shot_id": str(shot_id), "cached_shot_ids": cached}) meta = r["meta"] n = len(r["measured"]["t_s"]) s = meta.get("postshot") or "" @@ -129,7 +131,7 @@ def load_shot(shot_id: int, H: float = 1.0) -> ToolResult: summary=(f"Shot {shot_id} (campaign {meta.get('campaign')}): {n} slices over " f"{r['measured']['t_s'][0]}–{r['measured']['t_s'][-1]} s, peak beam power " f"{meta.get('P_nbi_max_MW')} MW. {r['summary']} Logbook: {s[:200]}"), - data={"shot_id": int(shot_id), "meta": meta, "summary": r["summary"], "limit_names": r["limit_names"], + data={"shot_id": str(shot_id), "meta": meta, "summary": r["summary"], "limit_names": r["limit_names"], "n_slices": n, "attribution": r["attribution"]}, ) @@ -150,18 +152,18 @@ def _limit_extremes(t: np.ndarray, frac: np.ndarray) -> tuple[dict[str, float] | @_guard -def limit_distances(shot_id: int) -> ToolResult: +def limit_distances(shot_id: str) -> ToolResult: """Distances to the operating limits over time for one cached shot — the replay panel's own numbers (physics.replay's limit fractions on the measured traces, the same formulas /map draws). Answers 'how close is shot X to the beta limit over time'. Distances, never forecasts: 1.0 is at the limit, not a prediction of a disruption.""" try: - r = _replay_bundle(int(shot_id)) + r = _replay_bundle(str(shot_id)) except HTTPException as e: - cached = mast.cached_shots() + cached = sources.get_source(None).cached_shots() # what the route can actually load return _err(f"{e.detail} Replay covers only the shots cached under data/shots ({len(cached)}: {cached}); " "use search_shots to explore the full catalog instead.", - {"shot_id": int(shot_id), "cached_shot_ids": cached}) + {"shot_id": str(shot_id), "cached_shot_ids": cached}) mo, me = r["model"], r["measured"] def arr(xs: list) -> np.ndarray: return np.array([np.nan if v is None else v for v in xs], dtype=float) @@ -187,7 +189,7 @@ def arr(xs: list) -> np.ndarray: out.append("These are computed distances to the limits (1.0 = at the limit), never a forecast of a disruption.") return ToolResult( summary=" ".join(out), - data={"shot_id": int(shot_id), "limit_names": list(r["limit_names"]), "n_slices": int(t.size), + data={"shot_id": str(shot_id), "limit_names": list(r["limit_names"]), "n_slices": int(t.size), "closest": closest, "first_cross_s": crosses, "series": series, "series_stride": stride, "archive": archive, "attribution": r["attribution"]}, ) @@ -242,7 +244,7 @@ def _fmt(value: float, places: int) -> str: @_guard -def run_whatif(shot_id: int, p_nbi_MW: float | None = None, p_nbi_scale: float | None = None, +def run_whatif(shot_id: str, p_nbi_MW: float | None = None, p_nbi_scale: float | None = None, n_scale: float | None = None, ip_scale: float | None = None, bt_scale: float | None = None, nbi_shift_s: float | None = None, closure: str | None = None) -> ToolResult: """Compute a bounded what-if on a cached shot and propose it — the /virtual/{id} route function, so @@ -260,7 +262,7 @@ def run_whatif(shot_id: int, p_nbi_MW: float | None = None, p_nbi_scale: float | elif p_nbi_scale is not None: power = float(p_nbi_scale) if power is not None: - pmax = float(_shot(int(shot_id))["meta"].get("P_nbi_max_MW") or 0.0) + pmax = float(_shot(str(shot_id))[0]["meta"].get("P_nbi_max_MW") or 0.0) if not pmax: return _err(f"shot {shot_id} has no NBI beams in the archive (P_nbi_max 0 MW): beam power cannot be " "scaled — try n_scale or a beam-timing edit (nbi_shift_s)") @@ -288,11 +290,11 @@ def run_whatif(shot_id: int, p_nbi_MW: float | None = None, p_nbi_scale: float | if not edit: return _err("no edit given: pass at least one of p_nbi_MW/p_nbi_scale, n_scale, ip_scale, bt_scale, nbi_shift_s") - r = _virtual_shot_body(int(shot_id), VirtualIn(edit=EditIn(**edit), closure=closure)) + r = _virtual_shot_body(str(shot_id), VirtualIn(edit=EditIn(**edit), closure=closure)) apply = {"sliders": {k: _fmt(v, 3 if k == "nbi_shift_s" else 2) for k, v in edit.items()}, "events": ["fusionlab:tab"]} proposal = {"tool": "run_whatif", - "args": {k: v for k, v in dict(shot_id=int(shot_id), p_nbi_MW=p_nbi_MW, p_nbi_scale=p_nbi_scale, + "args": {k: v for k, v in dict(shot_id=str(shot_id), p_nbi_MW=p_nbi_MW, p_nbi_scale=p_nbi_scale, n_scale=n_scale, ip_scale=ip_scale, bt_scale=bt_scale, nbi_shift_s=nbi_shift_s, closure=closure).items() if v is not None}, "math": "; ".join(math_bits), @@ -300,10 +302,10 @@ def run_whatif(shot_id: int, p_nbi_MW: float | None = None, p_nbi_scale: float | data = {k: r[k] for k in ("shot_id", "closure", "closure_name", "edit", "is_identity", "timing_edit", "laws", "dW_flat_pct", "would_cross_at_s", "in_distribution", "refly_error", "validated", "caveats", "attribution")} - return ToolResult(summary=_whatif_summary(int(shot_id), r), data=data, proposal=proposal) + return ToolResult(summary=_whatif_summary(str(shot_id), r), data=data, proposal=proposal) -def _whatif_summary(shot_id: int, r: dict) -> str: +def _whatif_summary(shot_id: str, r: dict) -> str: """Plain language for the bands the route computed. House wording: crossings are distances, confinement is compared with MAST data, nothing is a prediction.""" lo, hi = r["dW_flat_pct"]["anchored"] @@ -320,17 +322,17 @@ def _whatif_summary(shot_id: int, r: dict) -> str: @_guard -def export_usd(shot_id: int) -> ToolResult: - """Propose an OpenUSD export (Omniverse-compatible) of a cached shot. The write happens only after - confirmation, when the driver POSTs the route — the export path stays the route's own (temp file, - atomic publish); the registry never writes.""" - _shot(int(shot_id)) # the cache boundary is the route's: a 404 here is the verbatim refusal +def export_usd(shot_id: str) -> ToolResult: + """Propose an OpenUSD export (Omniverse-compatible) of a cached shot or imported bundle. The write + happens only after confirmation, when the driver POSTs the route — the export path stays the + route's own (temp file, atomic publish); the registry never writes.""" + _shot(str(shot_id)) # the cache boundary is the route's: a 404 here is the verbatim refusal return ToolResult( summary=f"Shot {shot_id} is in the local cache; confirming will export it as an OpenUSD stage " "(Omniverse-compatible) and download it.", - data={"shot_id": int(shot_id)}, - proposal={"tool": "export_usd", "args": {"shot_id": int(shot_id)}, - "apply": {"download": {"method": "POST", "url": f"/replay/{int(shot_id)}/usd"}}}, + data={"shot_id": str(shot_id)}, + proposal={"tool": "export_usd", "args": {"shot_id": str(shot_id)}, + "apply": {"download": {"method": "POST", "url": f"/replay/{shot_id}/usd"}}}, ) @@ -384,7 +386,7 @@ def _range_properties() -> dict[str, Any]: "list_cached_shots": Tool( definition={ "name": "list_cached_shots", - "description": "List the shots cached locally under data/shots — the only shots replay can run offline. The app never fetches from the network.", + "description": "List the shots replay can run on offline: the archive's cache under data/shots plus any imported shot bundles (ids like user_*). Uncached catalog shots are fetched on demand by replay.", "input_schema": {"type": "object", "properties": {}}, }, handler=list_cached_shots, @@ -398,7 +400,7 @@ def _range_properties() -> dict[str, Any]: ), "input_schema": { "type": "object", - "properties": {"shot_id": {"type": "integer", "description": "a cached shot id (list_cached_shots)"}, + "properties": {"shot_id": {"type": "string", "description": "a cached shot number or imported-bundle id (list_cached_shots)"}, "H": {"type": "number", "default": 1.0, "description": "confinement multiplier on the model traces"}}, "required": ["shot_id"], }, @@ -447,7 +449,7 @@ def _range_properties() -> dict[str, Any]: ), "input_schema": { "type": "object", - "properties": {"shot_id": {"type": "integer", "description": "a cached shot id (list_cached_shots)"}, + "properties": {"shot_id": {"type": "string", "description": "a cached shot number or imported-bundle id (list_cached_shots)"}, "p_nbi_MW": {"type": "number", "description": "absolute beam power in MW; converted to a scale factor server-side, math echoed"}, "p_nbi_scale": {**_slider_bounds("P_nbi"), "description": "beam-power scale factor (1.0 = the shot's logged peak)"}, "n_scale": {**_slider_bounds("n"), "description": "density scale factor"}, @@ -500,7 +502,7 @@ def _range_properties() -> dict[str, Any]: ), "input_schema": { "type": "object", - "properties": {"shot_id": {"type": "integer", "description": "a cached shot id (list_cached_shots)"}}, + "properties": {"shot_id": {"type": "string", "description": "a cached shot number or imported-bundle id (list_cached_shots)"}}, "required": ["shot_id"], }, }, diff --git a/fusionlab/api.py b/fusionlab/api.py index f246e28..2a208ca 100644 --- a/fusionlab/api.py +++ b/fusionlab/api.py @@ -83,18 +83,30 @@ def compute_cancel(task: str): @app.post("/compute/benchmark") -async def compute_benchmark(profile: str = "quick"): +async def compute_benchmark(profile: str = "quick", source: str | None = None, shot: str | None = None): """The measured-timings benchmark through the selected provider (the worker's `benchmark` task, with the local fallback) — the one named task that reports section progress. POST: it occupies the compute slot for its whole run, so it must not be triggerable by navigation or prefetch - (same reasoning as the USD export). The body is awaited on the compute pool (HARN-4): the run - — remote or local — never pins a shared request-threadpool token.""" + (same reasoning as the USD export). source/shot pick the bundle through the data seam — the + configured source's bench shot by default. The body is awaited on the compute pool (HARN-4): the + run — remote or local — never pins a shared request-threadpool token.""" from fusionlab.bench import PROFILES from fusionlab.compute import status + from fusionlab.data import sources if profile not in PROFILES: raise HTTPException(422, f"profile must be one of {sorted(PROFILES)}") - out = await status.run_task_async("benchmark", {"profile": profile}) + if source is not None: + try: + sources.get_source(source) + except ValueError as e: # unknown source names fail loudly, before a job is submitted + raise HTTPException(422, str(e)) from e + inputs: dict = {"profile": profile} + if source is not None: + inputs["source"] = source + if shot is not None: + inputs["shot"] = shot + out = await status.run_task_async("benchmark", inputs) return {**out, "computed_on": status.provenance("benchmark").get("host")} diff --git a/fusionlab/api_replay.py b/fusionlab/api_replay.py index 3dd79fa..5c5e5d0 100644 --- a/fusionlab/api_replay.py +++ b/fusionlab/api_replay.py @@ -1,9 +1,10 @@ -"""Replay routes: real MAST shots (FAIR-MAST cache) served beside the physics model.""" +"""Replay routes: shots from the data seam — the FAIR-MAST archive and imported bundles — served beside the physics model.""" from __future__ import annotations import csv import io +import json import logging import os import re @@ -12,7 +13,7 @@ from pathlib import Path import numpy as np -from fastapi import APIRouter, HTTPException, Query +from fastapi import APIRouter, File, Form, HTTPException, Query, UploadFile from fastapi.responses import FileResponse, Response from zarr.errors import GroupNotFoundError @@ -25,6 +26,8 @@ run_task, run_task_async, ) +from fusionlab.data import sources +from fusionlab.data.bundle import SPEC, BundleValidationError from fusionlab.physics import ( BETA_N_LIMIT, LIMIT_NAMES, @@ -49,6 +52,8 @@ def _j(a, nd=4): a = np.asarray(a) if a.dtype == bool or a.dtype.kind in "iu": return a.tolist() + if a.dtype.kind == "U": # string ids / campaigns (imported catalogs) + return a.tolist() a = np.round(a.astype(float), nd) return np.where(np.isfinite(a), a, None).tolist() @@ -59,17 +64,34 @@ def _shot_cached(shot_id: int) -> dict: # One archive fetch per shot id, however many requests want it: compare fans out over ids in parallel. -_inflight: dict[int, dict] = {} +_inflight: dict[str, dict] = {} _inflight_lock = threading.Lock() _FETCH_WAIT_S = 600 # a fetch takes seconds to minutes; a waiter gives up and asks the client to retry -def _shot(shot_id: int) -> dict: - """A cached shot from the process cache; an uncached catalog shot through load_shot()'s fetch-and-cache path.""" - if shot_id in mast.cached_shots(): - return _shot_cached(shot_id) - if not bool(np.any(_catalog()["shot_id"] == shot_id)): - raise HTTPException(404, f"shot {shot_id} is not in the MAST catalog (data/mast_db.npz)") +def _shot(shot_id: str) -> tuple[dict, sources.DataSource]: + """A shot bundle and the source that served it. An imported-bundle id (anything that is not a + plain shot number) resolves through the user-file source; shot numbers go to the archive — + cached shots from the process cache, uncached catalog shots through load_shot()'s + fetch-and-cache path. Returns the source too, so routes follow its machine config (attribution, + USD naming, isotope) instead of assuming MAST.""" + shot_id = str(shot_id) + if not shot_id.isdigit(): + if not sources.valid_bundle_id(shot_id): + raise HTTPException(404, f"shot id {shot_id!r} is neither a shot number nor a well-formed bundle id") + src = sources.get_source("user-file") + try: + return src.load_shot(shot_id), src + except FileNotFoundError as e: + raise HTTPException(404, str(e)) from e + except BundleValidationError as e: + raise HTTPException(400, f"the bundle does not satisfy the ShotBundle v1 contract: {e}") from e + if shot_id in {str(i) for i in mast.cached_shots()}: + return _shot_cached(int(shot_id)), sources.get_source("fair-mast") + cat_src = sources.get_source() # membership is asked of the configured source's own catalog + if not bool(np.any(_catalog(cat_src.name)["shot_id"] == int(shot_id))): + tail = " (data/mast_db.npz)" if cat_src.name == "fair-mast" else "" + raise HTTPException(404, f"shot {shot_id} is not in the {cat_src.name} catalog{tail}") with _inflight_lock: job = _inflight.get(shot_id) if job is None: @@ -86,7 +108,7 @@ def _shot(shot_id: int) -> dict: raise HTTPException(502, f"shot {shot_id} could not be fetched from FAIR-MAST: {job['err']}") else: try: - mast.load_shot(shot_id) # fetches from S3 + REST, writes the local cache atomically (fusionlab.mast) + mast.load_shot(int(shot_id)) # fetches from S3 + REST, writes the local cache atomically (fusionlab.mast) except GroupNotFoundError as err: job["err"] = "unpublished" # waiters answer with the same 404, not a 502 raise HTTPException( @@ -99,24 +121,38 @@ def _shot(shot_id: int) -> dict: with _inflight_lock: _inflight.pop(shot_id, None) job["done"].set() - if shot_id in mast.cached_shots(): - return _shot_cached(shot_id) + if shot_id in {str(i) for i in mast.cached_shots()}: + return _shot_cached(int(shot_id)), sources.get_source("fair-mast") raise HTTPException(502, f"shot {shot_id} could not be fetched from FAIR-MAST") +def _shot_row(s: dict, source: str) -> dict: + """One /shots entry: the bundle's logbook meta plus the picker's scalars, tagged with its source.""" + return {**s["meta"], "source": source, "n_slices": int(s["t_s"].size), "t_start_s": float(s["t_s"][0]), + "t_end_s": float(s["t_s"][-1]), "Ip_max_MA": round(float(np.nanmax(s["Ip_MA"])), 3), + "P_nbi_max_MW": round(float(np.nanmax(s["P_nbi_MW"])), 2)} + + @router.get("/shots") def shots(): - """Cached shots with their logbook text. Uncached catalog shots are fetched on demand by /replay/{id}.""" + """Cached shots with their logbook text: the archive's cache plus any imported bundles. + Uncached catalog shots are fetched on demand by /replay/{id}.""" out = [] for sid in mast.cached_shots(): - s = _shot(sid) - out.append({**s["meta"], "n_slices": int(s["t_s"].size), "t_start_s": float(s["t_s"][0]), - "t_end_s": float(s["t_s"][-1]), "Ip_max_MA": round(float(np.nanmax(s["Ip_MA"])), 3), - "P_nbi_max_MW": round(float(np.nanmax(s["P_nbi_MW"])), 2)}) + s, _ = _shot(str(sid)) + out.append(_shot_row(s, "fair-mast")) + src = sources.get_source("user-file") + for sid in src.cached_shots(): + out.append(_shot_row(src.load_shot(sid), "user-file")) return {"shots": out, "attribution": ATTRIBUTION} -def _replay_bundle(shot_id: int, H: float = 1.0) -> dict: +def _attribution(s: dict, src: sources.DataSource) -> str: + """The bundle's own attribution when it carries one; the source machine's otherwise.""" + return str(s["meta"].get("attribution") or src.config.attribution) + + +def _replay_bundle(shot_id: str, H: float = 1.0) -> dict: """The full replay response: measured traces, model traces on the same time base, limits, geometry. /export serves the very same dict, so exported values equal replayed values by construction. @@ -124,8 +160,9 @@ def _replay_bundle(shot_id: int, H: float = 1.0) -> dict: without equilibrium reconstructions replays its traces with "equilibrium_absent" set instead of failing the route. """ - s = _shot(shot_id) - r = replay(s, H=H, config=mast.CONFIG) + shot_id = str(shot_id) # archive callers still pass the shot number as an int + s, src = _shot(shot_id) + r = replay(s, H=H, config=src.config) r["in_distribution"] = in_distribution(s, [Edit()])[0] # the virtual path's own guard, on the measured inputs summary = _summary(s, r) _add_hybrid(s, r, summary) @@ -135,9 +172,9 @@ def _replay_bundle(shot_id: int, H: float = 1.0) -> dict: "p95_rel_l2": round(float(np.percentile(eq["rel_l2"], 95)), 4), "computed_on": eq.get("computed_on")} payload = { - "meta": s["meta"], "attribution": ATTRIBUTION, "limit_names": LIMIT_NAMES, "summary": summary, + "meta": s["meta"], "attribution": _attribution(s, src), "limit_names": LIMIT_NAMES, "summary": summary, # The archive's own bookkeeping (FEAT-3): the abort flag the limit panel overlays as context. - "archive": _archive_flags(shot_id), + "archive": _archive_flags(shot_id, src.name), "measured": {k: _j(s[k]) for k in _TRACES if k in s}, "model": {k: _j(v) for k, v in r.items()}, "lcfs": {"R": _j(s["lcfs_R"], 3), "Z": _j(s["lcfs_Z"], 3)} if "lcfs_R" in s else None, @@ -154,14 +191,14 @@ def _replay_bundle(shot_id: int, H: float = 1.0) -> dict: @router.get("/replay/{shot_id}") -async def replay_shot(shot_id: int, H: float = 1.0): +async def replay_shot(shot_id: str, H: float = 1.0): """Measured traces, model traces on the same time base, limits, and the geometry to draw the shot.""" return await run_compute_async(partial(_replay_bundle, shot_id, H)) -def _replay_psi_slice(shot_id: int, i: int) -> dict: +def _replay_psi_slice(shot_id: str, i: int) -> dict: """One psi slice's response body — the blocking part the async route below runs on the compute pool.""" - s = _shot(shot_id) + s, _ = _shot(shot_id) if "psi_n" not in s: raise HTTPException(404, f"shot {shot_id}'s bundle has no equilibrium reconstruction (psi)") if not 0 <= i < s["t_s"].size: @@ -177,13 +214,13 @@ def _replay_psi_slice(shot_id: int, i: int) -> dict: @router.get("/replay/{shot_id}/psi/{i}") -async def replay_psi(shot_id: int, i: int): +async def replay_psi(shot_id: str, i: int): """Normalised poloidal flux (0 axis, 1 separatrix) on the EFIT grid for one time slice, rows = Z.""" return await run_compute_async(partial(_replay_psi_slice, shot_id, i)) @router.post("/replay/{shot_id}/usd") -def replay_usd(shot_id: int, fmt: str = "usdc"): +def replay_usd(shot_id: str, fmt: str = "usdc"): """OpenUSD export (Omniverse-compatible): vessel, PF coils and the plasma boundary time-sampled over the shot. POST, not GET: the export recomputes the stage and writes out/ on every hit — a state-changing @@ -191,18 +228,19 @@ def replay_usd(shot_id: int, fmt: str = "usdc"): from fusionlab.usd_export import FORMATS, export_shot if f".{fmt}" not in FORMATS: raise HTTPException(422, f"fmt must be one of {sorted(x.lstrip('.') for x in FORMATS)}") - s = _shot(shot_id) + s, src = _shot(shot_id) if "lcfs_R" not in s: # the plasma boundary is the one geometry export_shot requires raise HTTPException(422, f"shot {shot_id}'s bundle has no equilibrium boundary (lcfs); USD export needs it") - out = Path(__file__).resolve().parent.parent / "out" / f"mast_{shot_id}.{fmt}" + stem = src.config.usd_prim_root.lstrip("/").lower() or "bundle" # mast_30420.usdc, user_1.usdc, ... + out = Path(__file__).resolve().parent.parent / "out" / f"{stem}_{shot_id}.{fmt}" tmp = out.with_name(f".{out.stem}.tmp{os.getpid()}{out.suffix}") # dotfile, pid-unique, extension kept (export_shot validates the suffix) try: try: # with Warp-traced field lines when Warp can run here; the plain stage otherwise from fusionlab.fieldlines import export_shot_with_field_lines - export_shot_with_field_lines(s, tmp) + export_shot_with_field_lines(s, tmp, machine=src.config) except Exception: logger.warning("USD field-line export failed; falling back to plain stage", exc_info=True) - export_shot(s, tmp) + export_shot(s, tmp, machine=src.config) os.replace(tmp, out) # atomic same-directory publish: readers only ever see a complete stage except Exception: tmp.unlink(missing_ok=True) @@ -211,9 +249,9 @@ def replay_usd(shot_id: int, fmt: str = "usdc"): @lru_cache(maxsize=8) # process-lifetime cache: a refetched or retrained artifact on disk is not seen until restart -def _lines(shot_id: int): +def _lines(shot_id: str): """Field lines for every slice of a shot in one Warp launch, plus traced q at psi_N = 0.95 beside EFIT's q95.""" - s = _shot(shot_id) + s, _ = _shot(shot_id) out = run_task("trace_fieldlines", {"shot": s}) where = provenance("trace_fieldlines").get("host") # captured now: the cache keeps it honest pts, qc = out["pts"], out["q_check"] # (time, line, point, xyz) @@ -222,10 +260,10 @@ def _lines(shot_id: int): return pts, q, qc["summary"], out["device"], out["psi_n_start"], where -def _replay_fieldlines_slice(shot_id: int, i: int) -> dict: +def _replay_fieldlines_slice(shot_id: str, i: int) -> dict: """One field-lines response body — the blocking part (a Warp trace through run_task) the async route below runs on the compute pool.""" - s = _shot(shot_id) + s, _ = _shot(shot_id) if not 0 <= i < s["t_s"].size: raise HTTPException(404, "time index out of range") if "psi_n" not in s: @@ -242,7 +280,7 @@ def _replay_fieldlines_slice(shot_id: int, i: int) -> dict: @router.get("/replay/{shot_id}/fieldlines/{i}") -async def replay_fieldlines(shot_id: int, i: int): +async def replay_fieldlines(shot_id: str, i: int): """Field lines of the EFIT reconstruction at one time slice (axisymmetric; no islands or 3D fields).""" return await run_compute_async(partial(_replay_fieldlines_slice, shot_id, i)) @@ -257,10 +295,11 @@ def _eq_model(): @lru_cache(maxsize=8) # process-lifetime cache: a refetched or retrained artifact on disk is not seen until restart -def _eq_psi(shot_id: int): +def _eq_psi(shot_id: str): """PhysicsNeMo reconstruction of psi from the magnetic sensors, every slice of the shot in one batch. - Normalised with EFIT's axis and boundary flux so the two maps share contour levels; error is on psi itself.""" - s = _shot(shot_id) + Normalised with the equilibrium reconstruction's axis and boundary flux so the two maps share contour + levels; error is on psi itself.""" + s = _shot(shot_id)[0] try: held_out = _eq_model() if held_out is None or "eq_inputs" not in s or "psi_n" not in s: @@ -274,7 +313,8 @@ def _eq_psi(shot_id: int): def flat(a): return a.reshape(len(a), -1) rel = np.linalg.norm(flat(psi - psi_efit), axis=1) / np.linalg.norm(flat(psi_efit), axis=1) - return {"psi_n": (psi - ax) / (bd - ax), "rel_l2": rel, "held_out": shot_id in held_out, + return {"psi_n": (psi - ax) / (bd - ax), "rel_l2": rel, + "held_out": (int(shot_id) in held_out) if shot_id.isdigit() else False, "computed_on": provenance("run_eq_surrogate").get("host")} @@ -306,7 +346,10 @@ def _add_hybrid(s: dict, r: dict, summary: dict) -> None: summary["learned_computed_on"] = provenance("run_surrogate_shot").get("host") r["H_learned"], r["W_hybrid_MJ"] = Hc, r["W_H_MJ"] * Hc m = run_task("model_metrics", {"model": "surrogate"}) - summary["tau_correction_held_out"] = s["meta"]["shot_id"] // m.get("block_size", 100) in m.get("held_out_blocks", []) + sid = s["meta"]["shot_id"] + # Session-block held-out split is a MAST training artifact; imported ids are not block numbers. + summary["tau_correction_held_out"] = ( + sid // m.get("block_size", 100) in m.get("held_out_blocks", []) if isinstance(sid, int) else False) ok = r["steady"] & np.isfinite(r["H98"]) & (r["H98"] > 0) if ok.any(): def rms(a): @@ -318,10 +361,10 @@ def _warm_up(): compute_warm_up() # the correction model (and torch) load here rather than on the first replay request if (warm := mast.CONFIG.warmup_shot) is not None: try: # builds/loads the Warp kernels and traces the landing shot - _lines(int(warm)) + _lines(warm) except Exception: logger.warning("warm-up field-line trace failed; the first replay request pays for it instead", exc_info=True) - _eq_psi(int(warm)) + _eq_psi(warm) def start_warm_up(): @@ -351,12 +394,13 @@ async def surrogate_metrics(): raise HTTPException(404, "no trained model: run `make train`") from None -@lru_cache(maxsize=1) # process-lifetime cache: a refetched or retrained artifact on disk is not seen until restart -def _db_view() -> dict: +@lru_cache(maxsize=2) # per source: a different FUSIONLAB_DATA_SOURCE takes a fresh entry (or a restart) +def _db_view(source: str) -> dict: """The usable shot table in operating-space coordinates, with IPB98(y,2) evaluated on every row at once.""" - c = mast.clean_db(mast.load_db()) + src = sources.get_source(source) + c = src.clean_db(src.load_db()) g = {"R": c["R_m"], "a": c["a_m"], "kappa": c["V_m3"] / volume(c["R_m"], c["a_m"], 1.0)} - C, alpha = tau_coeff_H(g, c["Ip_MA"], c["B_T"], c["n_e20"], 1.0, mast.CONFIG.isotope_md) + C, alpha = tau_coeff_H(g, c["Ip_MA"], c["B_T"], c["n_e20"], 1.0, src.config.isotope_md) beta_N = c["beta_t_pct"] * c["a_m"] * c["B_T"] / c["Ip_MA"] return {"shot_id": c["shot_id"], "campaign": c["campaign"], "P_nbi_MW": c["P_nbi_MW"], "Ip_MA": c["Ip_MA"], "f_greenwald": c["n_e20"] / (c["Ip_MA"] / (np.pi * c["a_m"] ** 2)), "q95": c["q95"], "beta_N": beta_N, @@ -366,10 +410,13 @@ def _db_view() -> dict: @router.get("/db") def db(): - """~6k real MAST shots at peak current, for the operating-space scatter.""" - v = _db_view() - return {"n": int(v["shot_id"].size), "attribution": ATTRIBUTION, "beta_N_limit": BETA_N_LIMIT, "q95_limit": 2.0, - "H98_median": round(float(np.nanmedian(v["H98"])), 3), "columns": {k: _j(a) for k, a in v.items()}} + """The configured source's usable shots at peak current, for the operating-space scatter.""" + src = sources.get_source() + v = _db_view(src.name) + n = int(v["shot_id"].size) + return {"n": n, "attribution": src.config.attribution, "beta_N_limit": BETA_N_LIMIT, "q95_limit": 2.0, + "H98_median": round(float(np.nanmedian(v["H98"])), 3) if n else None, + "columns": {k: _j(a) for k, a in v.items()}} # ---- catalog search over the full shot table (the db view above is the cleaned analysis subset) @@ -379,18 +426,22 @@ def db(): "t_ipmax_s", "useful", "abort", "mode_derived") -@lru_cache(maxsize=1) -def _catalog() -> dict: - """The full FAIR-MAST shot table (data/mast_db.npz), read once per process and shared by every search.""" - return mast.load_db() +@lru_cache(maxsize=2) +def _catalog(source: str) -> dict: + """The configured source's full shot table (data/mast_db.npz for the archive), read once per + process per source and shared by every search.""" + return sources.get_source(source).load_db() -def _archive_flags(shot_id: int) -> dict[str, bool | None]: +def _archive_flags(shot_id: str, source: str) -> dict[str, bool | None]: """The archive's own bookkeeping for one shot — the catalog's cpf_useful/cpf_abort columns. - None where the row is missing or a flag is NaN (the archive never graded it): the UI degrades - silently, never invents a flag. These are archive bookkeeping, not physics.""" - c = _catalog() - hits = np.nonzero(c["shot_id"] == shot_id)[0] + None where the row is missing or a flag is NaN (the archive never graded it), and for + imported bundles (they carry no archive bookkeeping): the UI degrades silently, never + invents a flag. These are archive bookkeeping, not physics.""" + if not shot_id.isdigit(): + return {"useful": None, "abort": None} + c = _catalog(source) + hits = np.nonzero(c["shot_id"] == int(shot_id))[0] if not hits.size: return {"useful": None, "abort": None} def flag(col: str) -> bool | None: @@ -406,6 +457,8 @@ def _q_mask(c: dict, q: str | None) -> np.ndarray | None: if not (m := re.fullmatch(r"(\d{1,6})(?:-(\d{1,6}))?", q.strip())): raise HTTPException(422, "q must be a shot number or a range like 30400-30500") lo, hi = int(m.group(1)), m.group(2) + if c["shot_id"].dtype.kind not in "iu": # imported catalogs carry free-form string ids + raise HTTPException(422, "q filters shot numbers; the imported catalog's ids are not numbers — filter by campaign or ranges instead") if hi is not None: if int(hi) < lo: raise HTTPException(422, "q range must go from the lower shot number to the higher") @@ -467,7 +520,7 @@ def db_search(q: str | None = None, campaign: str | None = None, useful: bool | None = None, abort: bool | None = None, mode: str | None = None, # derived regime flag ('H'/'L'), never an archive field limit: int = Query(default=200, ge=1, le=1000), offset: int = Query(default=0, ge=0)) -> dict: - """Vectorized filter over the full FAIR-MAST catalog (15,969 rows), beside /db's cleaned analysis view. + """Vectorized filter over the configured source's full shot table, beside /db's cleaned analysis view. Filters combine (AND); the UI mirrors them into the URL so any result set is shareable. total counts the whole match set, independent of the limit/offset window. Units are the repo's: MA, T, MW, @@ -475,7 +528,8 @@ def db_search(q: str | None = None, campaign: str | None = None, (docs/PHYSICS.md): 'H' or 'L'; rows the scalars cannot flag stay unknown and match neither. Derived, never measured. """ - c = _catalog() + src = sources.get_source() + c = _catalog(src.name) mask = np.ones(c["shot_id"].size, dtype=bool) for m in (_q_mask(c, q), _campaign_mask(c, campaign), _range_mask(c, "Ip_MA", ip_min, ip_max), _range_mask(c, "B_T", bt_min, bt_max), @@ -492,7 +546,8 @@ def db_search(q: str | None = None, campaign: str | None = None, idx = idx[np.argsort(c["shot_id"][idx], kind="stable")][offset:offset + limit] # ascending ids: stable pages cols = {k: _j(c[k][idx]) for k in _ROW_FIELDS} rows = [dict(zip(_ROW_FIELDS, row, strict=True)) for row in zip(*cols.values(), strict=True)] - return {"total": total, "limit": limit, "offset": offset, "rows": rows, "attribution": ATTRIBUTION} + return {"total": total, "limit": limit, "offset": offset, "rows": rows, + "attribution": src.config.attribution} # ---- take-away export (FEAT-1): the replay bundle's numbers, downloadable with their provenance. @@ -585,16 +640,18 @@ def _export_rows(bundles: list[dict]) -> list[dict]: return rows -def _export_ids(shots: str) -> list[int]: - """Comma-separated shot numbers for the compare export: deduped, request order kept, capped — - each id is a full replay, the cap keeps an accidental URL cheap.""" - ids: list[int] = [] +def _export_ids(shots: str) -> list[str]: + """Comma-separated shot ids for the compare export — archive shot numbers or imported-bundle + ids: deduped, request order kept, capped (each id is a full replay, the cap keeps an accidental + URL cheap).""" + ids: list[str] = [] for part in shots.split(","): part = part.strip() - if not part.isdigit(): - raise HTTPException(422, "shots must be comma-separated shot numbers, e.g. shots=30166,30420") - if (i := int(part)) not in ids: - ids.append(i) + if not sources.valid_bundle_id(part): # covers plain shot numbers and bundle ids + raise HTTPException(422, "shots must be comma-separated shot numbers or imported-bundle ids, " + "e.g. shots=30166,30420 or shots=user_bundle_1") + if part not in ids: + ids.append(part) if len(ids) > _EXPORT_MAX_SHOTS: raise HTTPException(422, f"export covers at most {_EXPORT_MAX_SHOTS} shots — the compare view's limit") return ids @@ -616,13 +673,13 @@ def _export(bundles: list[dict], fmt: str, H: float): raise HTTPException(422, 'fmt must be "csv" or "json"') -def _replay_bundles(shot_ids: list[int], H: float) -> list[dict]: +def _replay_bundles(shot_ids: list[str], H: float) -> list[dict]: """Full replay bundles in request order — one blocking call for the compute pool.""" return [_replay_bundle(i, H) for i in shot_ids] @router.get("/export/{shot_id}") -async def export_shot(shot_id: int, H: float = 1.0, fmt: str = "csv"): +async def export_shot(shot_id: str, H: float = 1.0, fmt: str = "csv"): """Take-away export of one shot's derived quantities: every measured trace, model trace and summary number the replay serves, as a tidy CSV (one row per shot/slice/quantity, each row labelled with its provenance and unit) or as JSON (the full replay bundle + the provenance map). @@ -637,3 +694,57 @@ async def export_compare(shots: str, H: float = 1.0, fmt: str = "csv"): a `shot_id` column so any spreadsheet or notebook can group them. Order follows the request.""" bundles = await run_compute_async(partial(_replay_bundles, _export_ids(shots), H)) return _export(bundles, fmt, H) + + +# ---- bring-your-own-data (DATA-2): imported shot bundles replay, compare and export beside the archive. + +@router.post("/data/bundles/{shot_id}") +def upload_bundle(shot_id: str, file: UploadFile | None = File(None), + meta_json: str | None = Form(None)) -> dict: + """Import a shot so it replays, compares and exports beside the archive's cache. + + Two shapes on one route: a multipart POST with a `file` field imports an npz bundle (the + ShotBundle v1 contract, docs/DATA_CONTRACT.md) — validated in full before anything is written, + so a malformed upload fails with the first bad field named and leaves no partial file. The id + must be letters/digits/dots/dashes/underscores, not a bare shot number (those are the + archive's). A bare POST (no file) on a shot number imports that archive shot through the + loader's cache-then-S3 path — the UI's self-import when a catalog shot is not cached yet. + The optional meta_json form field overrides/extends the bundle's own meta.""" + if file is None or not file.filename: + if file is None and shot_id.isdigit(): + s, src = _shot(shot_id) # the same fetch-and-cache path a replay takes + return {**_shot_row(s, src.name), "equilibrium": "psi_R" in s and "lcfs_R" in s, + "fields": sorted(k for k in SPEC if k in s), "attribution": _attribution(s, src)} + raise HTTPException(422, "send the bundle as multipart form data (field name: file); a bare " + "POST imports an archive shot number") + src = sources.get_source("user-file") + raw = file.file.read() + meta = None + if meta_json is not None: + try: + meta = json.loads(meta_json) + except json.JSONDecodeError as e: + raise HTTPException(422, f"meta_json is not valid JSON: {e}") from e + if not isinstance(meta, dict): + raise HTTPException(422, "meta_json must be a JSON object") + try: + s = src.import_bundle(shot_id, raw, meta) + except BundleValidationError as e: # contract violations: the first bad field, named + raise HTTPException(400, f"the bundle does not satisfy the ShotBundle v1 contract: {e}") from e + except ValueError as e: # bad id shape or over the size cap + raise HTTPException(422, str(e)) from e + return {**_shot_row(s, src.name), "equilibrium": "psi_R" in s and "lcfs_R" in s, + "fields": sorted(k for k in SPEC if k in s), "attribution": _attribution(s, src)} + + +@router.delete("/data/bundles/{shot_id}") +def delete_bundle(shot_id: str) -> dict: + """Remove an imported bundle. The archive's cache is managed by the loader, not this route.""" + src = sources.get_source("user-file") + if shot_id.isdigit() or not sources.valid_bundle_id(shot_id): + raise HTTPException(404, f"no imported bundle named {shot_id!r}") + f = src.bundles_dir / f"{shot_id}.npz" + if not f.is_file(): + raise HTTPException(404, f"no imported bundle named {shot_id!r}") + f.unlink() + return {"deleted": shot_id} diff --git a/fusionlab/api_virtual.py b/fusionlab/api_virtual.py index 49589e0..55bcca9 100644 --- a/fusionlab/api_virtual.py +++ b/fusionlab/api_virtual.py @@ -61,16 +61,16 @@ def virtual_gate(): @router.post("/virtual/{shot_id}") -async def virtual_shot(shot_id: int, body: VirtualIn): +async def virtual_shot(shot_id: str, body: VirtualIn): """A what-if anchored on a cached real shot — the body runs on the compute pool (HARN-4): the learned closure reaches the surrogate through run_task, whose remote wait must not pin a shared request-threadpool token.""" return await run_compute_async(partial(_virtual_shot_body, shot_id, body)) -def _virtual_shot_body(shot_id: int, body: VirtualIn) -> dict: +def _virtual_shot_body(shot_id: str, body: VirtualIn) -> dict: """The virtual-shot response body (the blocking part the async route above runs on the pool).""" - s = _shot(shot_id) + s, _src = _shot(shot_id) g = virtual.gate() edit = virtual.Edit(**body.edit.model_dump()) for slider, act in SLIDER_GATE.items(): # the gate is enforced here, not only drawn in the UI diff --git a/fusionlab/bench.py b/fusionlab/bench.py index 221caa8..3b8bd70 100644 --- a/fusionlab/bench.py +++ b/fusionlab/bench.py @@ -14,7 +14,7 @@ import numpy as np -from fusionlab import mast, surrogate +from fusionlab import surrogate from fusionlab.physics import DEVICES, Controls, replay, simulate PROFILES = ("quick", "full") @@ -31,28 +31,41 @@ def best(fn, repeat=5): return min(out) -def timings(profile: str = "full", progress: Callable[[dict], None] | None = None) -> dict[str, Any]: - """Run the benchmark sections and return {"rows": [...], "profile", "n_slices"}. +def timings(profile: str = "full", progress: Callable[[dict], None] | None = None, + source: str | None = None, shot: str | None = None) -> dict[str, Any]: + """Run the benchmark sections and return {"rows": [...], "profile", "n_slices", "shot"}. profile "full" is `make bench` (the README numbers); "quick" skips the per-slice loop reference and the large operating map — the worker's smoke-level benchmark. `progress`, when given, is called once per section: the worker turns that into SSE progress events and its - cancellation point between sections. + cancellation point between sections. `source`/`shot` pick the bundle to replay through the + data seam (the configured source's bench shot by default) — never a hard-coded archive id. """ if profile not in PROFILES: raise ValueError(f"unknown benchmark profile {profile!r}; available: {PROFILES}") report = progress or (lambda event: None) + from fusionlab.data import ( + sources, # lazy: keeps bench importable ahead of the data seam in tests + ) + + src = sources.get_source(source) + shot_id = str(shot) if shot is not None else src.config.bench_shot + if shot_id is None: # a source without a configured bench shot: the first cached bundle + cached = src.cached_shots() + if not cached: + raise ValueError(f"source {src.name!r} has no bench_shot configured and no cached bundles to benchmark") + shot_id = cached[0] rows: list[dict] = [] - shot = mast.load_shot(30166) - nt = shot["t_s"].size + shot_d = src.load_shot(shot_id) + nt = shot_d["t_s"].size # 1. replay: every time slice in one numpy call vs one call per slice report({"section": "replay"}) def one(i): - return {k: (v[i:i + 1] if isinstance(v, np.ndarray) and v.shape[:1] == (nt,) else v) for k, v in shot.items()} + return {k: (v[i:i + 1] if isinstance(v, np.ndarray) and v.shape[:1] == (nt,) else v) for k, v in shot_d.items()} - t_vec = best(lambda: replay(shot)) + t_vec = best(lambda: replay(shot_d)) rows.append({"label": f"Replay of a real shot ({nt} EFIT slices), vectorized", "time": f"{t_vec * 1e3:.2f} ms", "note": f"{nt / t_vec:,.0f} slices/s"}) if profile == "full": @@ -63,9 +76,10 @@ def one(i): # 2. shot cache report({"section": "shot_cache"}) - t_load = best(lambda: mast.load_shot(30166)) - rows.append({"label": "Load a cached shot (npz, no network)", "time": f"{t_load * 1e3:.1f} ms", - "note": "first fetch from S3: ~27 s"}) + t_load = best(lambda: src.load_shot(shot_id)) + rows.append({"label": f"Load a cached bundle ({src.name}, npz, no network)", + "time": f"{t_load * 1e3:.1f} ms", + "note": "first fetch from S3: ~27 s" if src.name == "fair-mast" else "imported bundle"}) # 3. IPB98 over the whole usable shot table report({"section": "ipb98"}) @@ -107,4 +121,4 @@ def run(): rows.append({"label": f"Sandbox operating map {nx}x{nx} (0D power-balance solve per point)", "time": f"{t_map * 1e3:.0f} ms", "note": f"{nx * nx / t_map:,.0f} points/s"}) - return {"rows": rows, "profile": profile, "n_slices": int(nt)} + return {"rows": rows, "profile": profile, "n_slices": int(nt), "shot": shot_id} diff --git a/fusionlab/compute/local.py b/fusionlab/compute/local.py index 9febb12..53d2f3a 100644 --- a/fusionlab/compute/local.py +++ b/fusionlab/compute/local.py @@ -67,7 +67,11 @@ def _benchmark(inputs: dict, progress: Callable[[dict], None] | None = None) -> what progress subscribers see between sections (the same events the worker's SSE emits).""" from fusionlab import bench # lazy: pulls torch via surrogate - return bench.timings(profile=inputs.get("profile", "quick"), progress=progress) + kwargs: dict[str, Any] = {"profile": inputs.get("profile", "quick")} + for k in ("source", "shot"): # the data seam: which bundle to benchmark (None: configured) + if inputs.get(k) is not None: + kwargs[k] = inputs[k] + return bench.timings(progress=progress, **kwargs) # Task bodies take the inputs dict. Bodies that report section progress (declared here) additionally diff --git a/fusionlab/data/machine.py b/fusionlab/data/machine.py index 6026108..f6f463a 100644 --- a/fusionlab/data/machine.py +++ b/fusionlab/data/machine.py @@ -42,6 +42,7 @@ class MachineConfig: clean_limits: CleanLimits # catalog-filter thresholds lh_enhancement_min: float # mode_derived cut: H above this tau_E / tau_89 (ITER89-P) ratio # (fusionlab.data.lh_mode; MAST value calibrated on the #38 anchors) + bench_shot: str | None = None # the shot the benchmark task replays (None: the caller names one) unit_sanity: tuple[tuple[str, float, float], ...] = () # (field, lo, hi) mid-shot unit guards @@ -60,6 +61,7 @@ class MachineConfig: # Midway between the two hand-verified anchors: 30420 (ohmic, ratio 1.29) and 30166 # ("H-mode from 223 ms", peak current 0.245 s, ratio 1.81) — scripts/mode_flag_report.py. lh_enhancement_min=1.5, + bench_shot="30166", unit_sanity=( # mid-shot guards against a missed SI conversion (the values the suite asserts) ("R_m", 0.6, 1.0), ("a_m", 0.4, 0.7), diff --git a/fusionlab/data/sources.py b/fusionlab/data/sources.py new file mode 100644 index 0000000..12b581d --- /dev/null +++ b/fusionlab/data/sources.py @@ -0,0 +1,281 @@ +"""The DataSource seam: every surface consumes validated shot bundles, whichever source served them. + +DATA-2's bring-your-own-data path. A `DataSource` is one named source with its machine's +`MachineConfig`; `get_source()` resolves the configured default (FUSIONLAB_DATA_SOURCE, +"fair-mast" unless set). `FairMastSource` delegates to the fusionlab.mast facade — kept +intact because that module is the monkeypatch surface the tests own — and a directory of +imported npz bundles is served by `UserFileSource` below. Shot resolution overlays user +bundles on the archive (api_replay._shot), so an uploaded shot replays, compares and +exports beside the shipped MAST cache without any route knowing which source answered. +""" + +from __future__ import annotations + +import json +import os +import re +import zipfile +from dataclasses import replace +from pathlib import Path +from typing import Protocol + +import numpy as np + +from fusionlab import mast +from fusionlab.atomicio import atomic_write +from fusionlab.data import reader +from fusionlab.data.bundle import BundleValidationError, ShotBundle +from fusionlab.data.ingestors.fair_mast import _DB_COLS +from fusionlab.data.lh_mode import derive_mode +from fusionlab.data.machine import CleanLimits, MachineConfig + +# A bundle id names its file ({id}.npz) and rides in URLs and export filenames: letters, +# digits, dots, dashes or underscores, no leading separator (no traversal, no dotfiles). +BUNDLE_ID_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}") + + +def valid_bundle_id(shot_id: str) -> bool: + """True when `shot_id` is a well-formed bundle id (the only ids the routes accept).""" + return BUNDLE_ID_RE.fullmatch(shot_id) is not None + + +class DataSource(Protocol): + """One named data source: its machine's config, and validated bundles out. + + `load_shot` returns the plain dict the whole app passes around — a ShotBundle v1 + contract (fusionlab.data.bundle), validated at this door; `load_db` returns the + catalog as a dict of equal-length columns. + """ + + name: str + config: MachineConfig + bundles_dir: Path # where this source's bundle files live (an archive: its cache directory) + + def load_shot(self, shot_id: str) -> dict: ... + def load_db(self) -> dict: ... + def cached_shots(self) -> list[str]: ... + def clean_db(self, db: dict) -> dict: + """Rows usable for confinement analysis — thresholds from this source's machine config.""" + ... + + def import_bundle(self, shot_id: str, raw: bytes, meta: dict | None = None) -> dict: + """Validate raw npz bytes against the ShotBundle v1 contract and publish atomically + (no partial file on a rejected bundle). Read-only archive sources refuse.""" + ... + + +class FairMastSource: + """The FAIR-MAST archive through the fusionlab.mast facade.""" + + name = "fair-mast" + config: MachineConfig = mast.CONFIG + + bundles_dir = reader.SHOTS # the archive's own cache; imports never land here + + def import_bundle(self, shot_id: str, raw: bytes, meta: dict | None = None) -> dict: + """The archive is read-only: bundles enter the app through a user-file source.""" + raise TypeError(f"source {self.name!r} is a read-only archive — " + "import bundles through a user-file source (the upload route)") + + def load_shot(self, shot_id: str) -> dict: + return mast.load_shot(int(shot_id)) + + def cached_shots(self) -> list[str]: + return [str(i) for i in mast.cached_shots()] + + def load_db(self) -> dict: + return mast.load_db() + + def clean_db(self, db: dict) -> dict: + return mast.clean_db(db) + + +# ---------------------------------------------------------------- imported bundles (user-file) + +MAX_BUNDLE_BYTES = 64 * 2**20 # shipped MAST bundles are ~1 MB; 64 MiB is generous headroom + + +def _read_bundle_npz(f: Path) -> dict: + """One npz bundle file -> the shot dict (+ "meta"). `meta` rides as a JSON string (the reader + cache's convention); anything unreadable raises a contract error that names the problem — + never a raw zipfile or JSON traceback at the route.""" + try: + with np.load(f) as z: + out = {k: z[k] for k in z.files if k != "meta"} + out["meta"] = json.loads(str(z["meta"])) if "meta" in z.files else {} + except (ValueError, zipfile.BadZipFile, AttributeError, OSError) as err: + raise BundleValidationError( + f"not a readable npz shot bundle ({f.name}): {err} — save the arrays with np.savez and " + "meta as a JSON string") from err + if not isinstance(out["meta"], dict): + raise BundleValidationError( + "meta: expected a JSON object (save it as np.savez(f, meta=json.dumps(meta), ...))") + return out + + +def _default_user_config() -> MachineConfig: + """Defaults for imported data — deliberately modest: generic attribution and naming, no + MAST-calibrated thresholds, no unit-sanity ranges. A machine.json beside the bundles + overrides any of it (docs/DATA_CONTRACT.md).""" + return MachineConfig( + device="", # no DEVICES preset is assumed for an imported machine + isotope_md=2.0, # deuterium: the common case; set yours in machine.json + alpha_particles=False, + attribution="User-provided shot bundle; provenance as supplied with the data", + campaign_pattern=r".*", # imported campaigns are free-form strings + usd_prim_root="/USER", + reconstruction_name="imported reconstruction", + te_display_range=(0.2, 10.0), # foreign machines run hotter than MAST; set yours + te_usd_ramp_max_kev=10.0, + warmup_shot=None, + clean_limits=CleanLimits(tau_E_max_s=float("inf"), kappa_min=0.0, a_min_m=0.0), + lh_enhancement_min=1.5, # the generic ITER89-P cut (docs, FEAT-4): not MAST-calibrated; set yours + ) # imported catalogs are hand-curated: no machine-calibrated filtering + + +def _machine_config(bundles_dir: Path) -> MachineConfig: + """The imported machine's config: the user defaults overlaid with machine.json from the bundles + directory. Unknown keys and unreadable JSON fail loudly, naming the field.""" + base = _default_user_config() + f = bundles_dir / "machine.json" + if not f.is_file(): + return base + try: + raw = json.loads(f.read_text(encoding="utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError, OSError) as err: + raise ValueError(f"{f}: unreadable machine config: {err}") from err + raw = raw.get("machine", raw) if isinstance(raw, dict) else None # {"machine": {...}} or the flat dict + if not isinstance(raw, dict): + raise ValueError(f"{f}: expected a JSON object of MachineConfig fields") + unknown = [k for k in raw if k not in MachineConfig.__dataclass_fields__] + if unknown: + raise ValueError(f"{f}: unknown MachineConfig field(s) {', '.join(sorted(unknown))}") + try: + if "clean_limits" in raw: + raw["clean_limits"] = CleanLimits(**{**base.clean_limits.__dict__, **raw["clean_limits"]}) + if "te_display_range" in raw: + raw["te_display_range"] = tuple(raw["te_display_range"]) + if "unit_sanity" in raw: + raw["unit_sanity"] = tuple((k, float(lo), float(hi)) for k, lo, hi in raw["unit_sanity"]) + return replace(base, **raw) + except TypeError as err: + raise ValueError(f"{f}: bad machine config value: {err}") from err + + +class UserFileSource: + """Imported shot bundles: one npz file per shot in `bundles_dir` (the reader cache's format), + validated against ShotBundle v1 on every read. Point FUSIONLAB_USER_BUNDLES at a directory of + bundles, or upload through the API (import_bundle); machine.json in the same directory + configures the machine (docs/DATA_CONTRACT.md).""" + + name = "user-file" + + def __init__(self, bundles_dir: Path | None = None): + env = os.environ.get("FUSIONLAB_USER_BUNDLES", "") + if bundles_dir is not None: + self.bundles_dir = Path(bundles_dir) + elif env: + self.bundles_dir = Path(env) + else: + self.bundles_dir = reader.DATA / "user_bundles" + self.config: MachineConfig = _machine_config(self.bundles_dir) + + def cached_shots(self) -> list[str]: + """Sorted ids of the imported bundles (one id per .npz file).""" + if not self.bundles_dir.is_dir(): + return [] + return sorted(p.stem for p in self.bundles_dir.glob("*.npz") if valid_bundle_id(p.stem)) + + def load_shot(self, shot_id: str) -> dict: + """The bundle's arrays + meta, the file name's id stamped into meta (identity cannot drift + from the catalog), validated against the contract at this door.""" + f = self.bundles_dir / f"{shot_id}.npz" + if not f.is_file(): + raise FileNotFoundError( + f"no imported bundle named {shot_id!r} in {self.bundles_dir} — upload it " + "(POST /data/bundles/{id}) or point FUSIONLAB_USER_BUNDLES at its directory") + out = _read_bundle_npz(f) + out["meta"]["shot_id"] = shot_id + ShotBundle.from_dict(out, source=self.name) + return out + + def import_bundle(self, shot_id: str, raw: bytes, meta: dict | None = None) -> dict: + """Validate-then-publish: raw npz bytes -> one written bundle. The bytes land on a + same-directory temp (fusionlab.atomicio), are validated from there against the contract + (the error names the first bad field), and only a valid bundle is renamed into place — + a rejected upload never leaves a partial file behind. Returns the published bundle.""" + if not valid_bundle_id(shot_id): + raise ValueError(f"bundle id {shot_id!r} is not importable: use letters, digits, dots, " + "dashes or underscores, starting with a letter or digit") + if shot_id.isdigit(): # bare shot numbers are the archive's - imports never shadow them + raise ValueError(f"bundle id {shot_id!r} is a bare shot number, which the archive owns - " + "name the import something else (e.g. {shot_id}_local)") + if len(raw) > MAX_BUNDLE_BYTES: + raise ValueError(f"bundle exceeds the {MAX_BUNDLE_BYTES} byte import cap") + target = self.bundles_dir / f"{shot_id}.npz" + + def write_and_validate(tmp: Path) -> None: + tmp.write_bytes(raw) + out = _read_bundle_npz(tmp) + out["meta"] = {**out["meta"], **(meta or {}), "shot_id": shot_id} + with tmp.open("wb") as f: # the published copy carries the merged meta + np.savez_compressed(f, meta=json.dumps(out["meta"]), + **{k: v for k, v in out.items() if k != "meta"}) + ShotBundle.from_dict(out, source=self.name) # before the publish rename + + self.bundles_dir.mkdir(parents=True, exist_ok=True) + atomic_write(target, write_and_validate) + return self.load_shot(shot_id) + + def load_db(self) -> dict: + """A derived catalog: one row per imported bundle, evaluated at its peak-current slice (the + archives' convention). tau_E and beta_t_pct are computed here by the documented formulas + and labelled computed; useful=1 (imported on purpose) and abort=0 (the archives' curation + flags do not exist for imported data) are the imported-catalog defaults.""" + cols: dict[str, list] = {"shot_id": [], "campaign": [], **{k: [] for k in _DB_COLS}, + "useful": [], "abort": []} + for sid in self.cached_shots(): + s = self.load_shot(sid) + i = int(np.nanargmax(np.abs(s["Ip_MA"]))) + p_loss = float(s["P_ohm_MW"][i] + s["P_nbi_MW"][i] - s["dWdt_MW"][i]) + tau_e = float(s["W_MJ"][i]) / p_loss if p_loss > 0 else float("nan") + with np.errstate(divide="ignore", invalid="ignore"): + beta_t = float(s["beta_N"][i]) * float(s["Ip_MA"][i]) / (float(s["a_m"][i]) * float(s["B_T"][i])) + for k, v in { + "shot_id": sid, "campaign": str(s["meta"].get("campaign", "")), + "Ip_MA": float(s["Ip_MA"][i]), "B_T": float(s["B_T"][i]), "n_e20": float(s["n_e20"][i]), + "P_nbi_MW": float(s["P_nbi_MW"][i]), "P_ohm_MW": float(s["P_ohm_MW"][i]), + "P_rad_MW": float(s["P_rad_MW"][i]) if "P_rad_MW" in s else float("nan"), + "W_MJ": float(s["W_MJ"][i]), + "Te0_keV": float(s["Te0_keV"][i]) if "Te0_keV" in s else float("nan"), + "q95": float(s["q95"][i]), "kappa": float(s["kappa"][i]), + "beta_t_pct": beta_t, "a_m": float(s["a_m"][i]), "R_m": float(s["R_m"][i]), + "V_m3": float(s["V_m3"][i]), "tau_E_s": tau_e, "dWdt_MW": float(s["dWdt_MW"][i]), + "t_ipmax_s": float(s["t_s"][i]), "useful": 1, "abort": 0, + }.items(): + cols[k].append(v) + db = {k: np.asarray(v) for k, v in cols.items()} + db["mode_derived"] = derive_mode( # FEAT-4's derived flag, this machine's cut + db, isotope_md=self.config.isotope_md, enhancement_min=self.config.lh_enhancement_min) + return db + + def clean_db(self, db: dict) -> dict: + """The imported catalog is hand-curated; the source's clean_limits (default: no filtering) + is all that applies - never MAST's machine-calibrated thresholds.""" + return reader.clean_db(db, limits=self.config.clean_limits) + + +_SOURCES: dict[str, type] = {"fair-mast": FairMastSource, "user-file": UserFileSource} +DEFAULT_SOURCE = "fair-mast" + + +def get_source(name: str | None = None) -> DataSource: + """The named source, or the configured default: FUSIONLAB_DATA_SOURCE, "fair-mast" unless set. + + Unknown names fail loudly — the app must never silently replay the wrong machine's data. + """ + name = name or os.environ.get("FUSIONLAB_DATA_SOURCE", DEFAULT_SOURCE) + if name not in _SOURCES: + raise ValueError(f"unknown data source {name!r}; available: {sorted(_SOURCES)}") + src: DataSource = _SOURCES[name]() + return src diff --git a/fusionlab/fieldlines.py b/fusionlab/fieldlines.py index b62fc6e..b51a6da 100644 --- a/fusionlab/fieldlines.py +++ b/fusionlab/fieldlines.py @@ -42,6 +42,7 @@ import warp as wp from fusionlab.atomicio import atomic_write +from fusionlab.data.machine import MachineConfig wp.config.log_level = wp.LOG_WARNING # no init banner in the API logs @@ -454,11 +455,13 @@ def usd_lines(shot: dict, psi_n=USD_PSI_N, direction=USD_DIRECTION, *, n_turns: def add_field_lines(stage_or_path, shot: dict, *, psi_n=USD_PSI_N, direction=USD_DIRECTION, n_turns: float = USD_TURNS, steps_per_turn: int = STEPS_PER_TURN, store_every: int = USD_STORE_EVERY, - color=LINE_COLOR, width_m: float = LINE_WIDTH_M, device: str | None = None): - """Add /MAST/FieldLines (UsdGeom.BasisCurves, linear) to a stage written by usd_export.export_shot. + color=LINE_COLOR, width_m: float = LINE_WIDTH_M, device: str | None = None, + prim_path: str = "/MAST/FieldLines", provenance: str | tuple[str, str] = PROVENANCE): + """Add the machine's FieldLines prim (UsdGeom.BasisCurves, linear; /MAST/FieldLines for MAST) to a + stage written by usd_export.export_shot. stage_or_path: an open Usd.Stage (edited in place, not saved) or a path (opened, edited, saved). Points are - time-sampled, one sample per EFIT slice (time code = slice index, as in the rest of the stage), with a + time-sampled, one sample per equilibrium slice (time code = slice index, as in the rest of the stage), with a constant vertex count per line. Returns the BasisCurves prim's schema object. """ from pxr import Sdf, Usd, UsdGeom, Vt @@ -472,7 +475,7 @@ def add_field_lines(stage_or_path, shot: dict, *, psi_n=USD_PSI_N, direction=USD pts = np.round(pts.reshape(nt, n_lines * n_pts, 3), usd_export.ROUND_M).astype(np.float32) ext = usd_export._extent(pts) - curves = UsdGeom.BasisCurves.Define(stage, "/MAST/FieldLines") + curves = UsdGeom.BasisCurves.Define(stage, prim_path) curves.CreateTypeAttr(UsdGeom.Tokens.linear) curves.CreateCurveVertexCountsAttr(Vt.IntArray.FromNumpy(np.full(n_lines, n_pts, dtype=np.int32))) curves.CreateWidthsAttr(Vt.FloatArray([float(width_m)])) @@ -489,7 +492,7 @@ def add_field_lines(stage_or_path, shot: dict, *, psi_n=USD_PSI_N, direction=USD shader.CreateInput("emissiveColor", Sdf.ValueTypeNames.Color3f).Set(tuple(float(c) for c in color)) prim = curves.GetPrim() - prim.SetCustomDataByKey("provenance", PROVENANCE) + prim.SetCustomDataByKey("provenance", provenance) prim.SetCustomDataByKey("psi_n_start", ", ".join(f"{p:g}" for p in np.asarray(psi_n, float))) prim.SetCustomDataByKey("direction", ", ".join(f"{d:+g}" for d in np.asarray(direction, float))) prim.SetCustomDataByKey("scope_note", "Field lines of the axisymmetric EFIT reconstruction: no islands, no 3D " @@ -501,12 +504,22 @@ def add_field_lines(stage_or_path, shot: dict, *, psi_n=USD_PSI_N, direction=USD return curves -def export_shot_with_field_lines(shot: dict, path: str | Path, *, device: str | None = None, **export_kw) -> Path: - """usd_export.export_shot (unchanged), then /MAST/FieldLines added to the same file.""" +def export_shot_with_field_lines(shot: dict, path: str | Path, *, device: str | None = None, + machine: MachineConfig | None = None, **export_kw) -> Path: + """usd_export.export_shot (unchanged), then the field-lines prim added under the machine's root. + + The prim path and provenance follow the source's machine config; a MAST export keeps the + library's default provenance bytes exactly (the config object is the MAST one by identity). + """ from fusionlab import usd_export - path = usd_export.export_shot(shot, path, **export_kw) - add_field_lines(path, shot, device=device) + machine = machine or usd_export.MAST_CONFIG + path = usd_export.export_shot(shot, path, machine=machine, **export_kw) + provenance = PROVENANCE if machine is usd_export.MAST_CONFIG else ( + f"Field lines traced with NVIDIA Warp through the {machine.reconstruction_name} reconstruction " + "of the loaded shot bundle", f"data: {machine.attribution}") + add_field_lines(path, shot, device=device, prim_path=f"{machine.usd_prim_root}/FieldLines", + provenance=provenance) return Path(path) diff --git a/fusionlab/worker.py b/fusionlab/worker.py index 9e0bc68..a72fbc8 100644 --- a/fusionlab/worker.py +++ b/fusionlab/worker.py @@ -104,6 +104,8 @@ class BenchmarkInputs(BaseModel): model_config = ConfigDict(extra="forbid") profile: Literal["quick", "full"] = "quick" + source: str | None = None # data seam: which source's bundle to replay (None: the configured one) + shot: str | None = None # which shot/bundle id (None: the source's configured bench shot) class ArtifactRef(BaseModel): @@ -142,7 +144,12 @@ def _run_benchmark(inputs: BaseModel, progress: Callable[[dict], None]) -> dict: if not isinstance(inputs, BenchmarkInputs): # registry dispatch guarantees this; keep it honest raise TypeError(f"benchmark inputs must be BenchmarkInputs, got {type(inputs).__name__}") - return bench.timings(profile=inputs.profile, progress=progress) + kwargs: dict = {"profile": inputs.profile} + if inputs.source is not None: + kwargs["source"] = inputs.source + if inputs.shot is not None: + kwargs["shot"] = inputs.shot + return bench.timings(progress=progress, **kwargs) TASKS: dict[str, Task] = { diff --git a/pyproject.toml b/pyproject.toml index c1ab89f..816eb74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,10 @@ pythonpath = ["."] # ruff's default rule set (E4/E7/E9/F) plus bugbear, pyupgrade, flake8-simplify and isort. select = ["E4", "E7", "E9", "F", "B", "UP", "SIM", "I"] +[tool.ruff.lint.flake8-bugbear] +# FastAPI's File/Form marker calls in parameter defaults are the framework idiom, not a bug +extend-immutable-calls = ["fastapi.File", "fastapi.Form"] + [tool.mypy] python_version = "3.12" # the heavy deps (warp, pxr, physicsnemo) ship no type stubs diff --git a/tests/test_agent_loop.py b/tests/test_agent_loop.py index ebcbefc..7343207 100644 --- a/tests/test_agent_loop.py +++ b/tests/test_agent_loop.py @@ -115,7 +115,7 @@ def test_a_proposal_goes_to_the_sink_as_an_event_and_never_into_the_model_contex fake, events = run_turn(monkeypatch, script) proposed = next(e for e in events if e["type"] == "action_proposed") assert proposed == {"type": "action_proposed", "action_id": "a_test", "tool": "export_usd", - "args": {"shot_id": 30420}, "math": "", "expires_s": 600} + "args": {"shot_id": "30420"}, "math": "", "expires_s": 600} # canonical tool args, not the raw call block = fake.calls[1]["messages"][-1]["content"][0] assert "nothing is applied yet" in block["content"] assert "apply" not in block["content"] # the driver's payload stays server-side diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 3a44d44..575c264 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -15,6 +15,7 @@ from fusionlab import agent_tools, api_replay, mast from fusionlab.api import app from fusionlab.api_virtual import EditIn +from fusionlab.data import sources client = TestClient(app) @@ -106,13 +107,13 @@ def test_load_shot_reports_the_cache_boundary_and_suggests_search(): r = agent_tools.run_tool("load_shot", shot_id=1) assert r.error is not None and route.json()["detail"] in r.error # the route's own 404, verbatim assert "search_shots" in r.error - assert r.data["cached_shot_ids"] == mast.cached_shots() + assert r.data["cached_shot_ids"] == sources.get_source(None).cached_shots() def test_load_shot_returns_a_compact_summary_without_the_time_arrays(): r = agent_tools.run_tool("load_shot", shot_id=30166) assert r.error is None - assert r.data["shot_id"] == 30166 and r.data["n_slices"] == mast.load_shot(30166)["t_s"].size + assert r.data["shot_id"] == "30166" and r.data["n_slices"] == mast.load_shot(30166)["t_s"].size assert "summary" in r.data and "measured" not in r.data # traces are the UI's job, not the model's context @@ -124,7 +125,7 @@ def test_limit_distances_answers_the_over_time_question(): r = agent_tools.run_tool("limit_distances", shot_id=30166) assert r.error is None d = r.data - assert d["shot_id"] == 30166 and d["limit_names"] == ["greenwald", "troyon", "kink"] + assert d["shot_id"] == "30166" and d["limit_names"] == ["greenwald", "troyon", "kink"] assert d["n_slices"] == mast.load_shot(30166)["t_s"].size for name in d["limit_names"]: c, x = d["closest"][name], d["first_cross_s"][name] @@ -146,7 +147,7 @@ def test_limit_distances_uncached_is_the_cache_boundary(): r = agent_tools.run_tool("limit_distances", shot_id=1) assert r.error is not None and route.json()["detail"] in r.error assert "search_shots" in r.error - assert r.data["cached_shot_ids"] == mast.cached_shots() + assert r.data["cached_shot_ids"] == sources.get_source(None).cached_shots() def test_limit_distances_carries_the_catalog_abort_flag(monkeypatch): @@ -155,11 +156,11 @@ def test_limit_distances_carries_the_catalog_abort_flag(monkeypatch): r = agent_tools.run_tool("limit_distances", shot_id=30166) assert r.error is None and r.data["archive"]["abort"] is False assert "abort" not in r.summary.lower() - c = dict(api_replay._catalog()) + c = dict(api_replay._catalog("fair-mast")) i = int(np.nonzero(c["shot_id"] == 30166)[0][0]) c["abort"] = c["abort"].copy() c["abort"][i] = 1.0 - monkeypatch.setattr(api_replay, "_catalog", lambda: c) + monkeypatch.setattr(api_replay, "_catalog", lambda source="fair-mast": c) r = agent_tools.run_tool("limit_distances", shot_id=30166) assert r.error is None and r.data["archive"]["abort"] is True assert "flags this shot aborted" in r.summary @@ -181,7 +182,7 @@ def test_run_whatif_proposes_instead_of_applying(): r = agent_tools.run_tool("run_whatif", shot_id=30166, p_nbi_MW=1.2, n_scale=1.1) assert r.error is None and r.proposal is not None p = r.proposal - assert p["tool"] == "run_whatif" and p["args"]["shot_id"] == 30166 and p["args"]["p_nbi_MW"] == 1.2 + assert p["tool"] == "run_whatif" and p["args"]["shot_id"] == "30166" and p["args"]["p_nbi_MW"] == 1.2 assert set(p["apply"]["sliders"]) == {"P_nbi", "n"} and "× P_nbi_max" in p["math"] assert r.data["edit"]["P_nbi"] == pytest.approx(1.2 / pmax) and r.data["is_identity"] is False assert "t_s" not in r.data and "anchored" not in r.data # bands, not arrays, reach the model @@ -239,9 +240,9 @@ def bomb(*a, **k): monkeypatch.setattr("fusionlab.fieldlines.export_shot_with_field_lines", bomb) monkeypatch.setattr("fusionlab.usd_export.export_shot", bomb) - r = agent_tools.run_tool("export_usd", shot_id=30166) + r = agent_tools.run_tool("export_usd", shot_id="30166") assert r.error is None and r.proposal is not None - assert r.proposal["args"] == {"shot_id": 30166} + assert r.proposal["args"] == {"shot_id": "30166"} # DATA-2: string ids, imported bundles included assert r.proposal["apply"]["download"] == {"method": "POST", "url": "/replay/30166/usd"} @@ -255,7 +256,7 @@ def test_export_usd_uncached_is_the_verbatim_404(): # ---- the plain read-only trio and the 0D point def test_read_only_tools_execute(): r = agent_tools.run_tool("list_cached_shots") - assert r.error is None and len(r.data["shots"]) == len(mast.cached_shots()) + assert r.error is None and len(r.data["shots"]) == len(sources.get_source(None).cached_shots()) r = agent_tools.run_tool("get_whatif_gate") assert r.error is None and set(r.data["sliders"]) == {"Ip", "B", "n", "P_nbi", "nbi_shift_s"} r = agent_tools.run_tool("list_devices") diff --git a/tests/test_api.py b/tests/test_api.py index 01cfb75..612f2d4 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -241,7 +241,7 @@ def test_limit_distances_match_the_map_limit_formulas(): def test_archive_flags_align_with_the_catalog(): """FEAT-3: the bundle's archive block is the catalog's own cpf_useful/cpf_abort row — None where the archive never graded the shot (NaN), so absent flags degrade silently downstream.""" - c = api_replay._catalog() + c = api_replay._catalog("fair-mast") for sid in mast.cached_shots(): i = int(np.nonzero(c["shot_id"] == sid)[0][0]) def row_flag(col, i=i): @@ -255,11 +255,11 @@ def test_abort_marker_aligns_with_a_flagged_catalog_row(monkeypatch): """Where the archive DOES flag a shot, the bundle carries abort=True and the UI marker aligns with the catalog by construction. No cached shot ships flagged, so the alignment is proven on a synthetic catalog row (the shipped cache exercises only the silent-degradation path).""" - c = dict(api_replay._catalog()) + c = dict(api_replay._catalog("fair-mast")) i = int(np.nonzero(c["shot_id"] == 30166)[0][0]) c["abort"] = c["abort"].copy() c["abort"][i] = 1.0 - monkeypatch.setattr(api_replay, "_catalog", lambda: c) + monkeypatch.setattr(api_replay, "_catalog", lambda source="fair-mast": c) b = client.get("/replay/30166").json() assert b["archive"] == {"useful": True, "abort": True} assert client.get("/replay/30420").json()["archive"]["abort"] is False # unflagged rows stay False @@ -442,12 +442,12 @@ def test_usd_download_exports_to_temp_then_publishes_atomically(monkeypatch): out.unlink(missing_ok=True) # the assertion is that the endpoint creates the served path seen = {} - def fake_with_lines(shot, path): + def fake_with_lines(shot, path, machine=None): seen["tmp"] = Path(path) seen["final_exists_mid_export"] = out.exists() Path(path).write_bytes(b"stage-bytes") - def plain_must_not_run(shot, path): + def plain_must_not_run(shot, path, machine=None): raise AssertionError("plain export ran although the field-line export succeeded") monkeypatch.setattr("fusionlab.fieldlines.export_shot_with_field_lines", fake_with_lines) @@ -466,7 +466,7 @@ def test_usd_download_failure_keeps_the_old_stage_and_leaves_no_temp(monkeypatch out.parent.mkdir(exist_ok=True) out.write_bytes(b"previous-good-stage") - def dies_mid_write(shot, path): + def dies_mid_write(shot, path, machine=None): Path(path).write_bytes(b"half-written") raise RuntimeError("export exploded mid-write") @@ -485,7 +485,8 @@ def broken(shot, path): raise RuntimeError("no warp device here") monkeypatch.setattr("fusionlab.fieldlines.export_shot_with_field_lines", broken) - monkeypatch.setattr("fusionlab.usd_export.export_shot", lambda shot, path: Path(path).write_bytes(b"plain-stage")) + monkeypatch.setattr("fusionlab.usd_export.export_shot", + lambda shot, path, machine=None: Path(path).write_bytes(b"plain-stage")) with caplog.at_level(logging.WARNING, logger="fusionlab.api_replay"): r = client.post("/replay/30420/usd", params={"fmt": "usda"}) assert r.status_code == 200 and r.content == b"plain-stage" @@ -713,7 +714,7 @@ def test_compare_export_groups_rows_by_shot_and_validates(): assert [s["meta"]["shot_id"] for s in j["shots"]] == [30420, 30166] # deduped, request order kept rows = list(csv.DictReader(io.StringIO(client.get("/export", params={"shots": "30420,30166"}).text))) assert {r["shot_id"] for r in rows} == {"30420", "30166"} - assert client.get("/export", params={"shots": "30420,nope"}).status_code == 422 + assert client.get("/export", params={"shots": "30420,nope"}).status_code == 404 # unresolvable id assert client.get("/export", params={"shots": "1,2,3,4,5,6,7"}).status_code == 422 # cap: 7 distinct ids assert client.get("/export", params={"shots": ",".join(["30166"] * 7)}).status_code == 200 # dupes dedupe first assert client.get("/export/1").status_code == 404 # not in the catalog diff --git a/tests/test_api_agent.py b/tests/test_api_agent.py index 973920e..95b2a31 100644 --- a/tests/test_api_agent.py +++ b/tests/test_api_agent.py @@ -229,7 +229,7 @@ def test_a_proposed_whatif_changes_nothing_until_confirmed(monkeypatch): (["Proposed."], [text_block("Proposed.")])]) events = chat("raise the beams to 3 MW") proposed = next(e for e in events if e["type"] == "action_proposed") - assert proposed["tool"] == "run_whatif" and proposed["args"] == {"shot_id": 30166, "p_nbi_MW": mw} + assert proposed["tool"] == "run_whatif" and proposed["args"] == {"shot_id": "30166", "p_nbi_MW": mw} assert "× P_nbi_max" in proposed["math"] and proposed["expires_s"] == 600 # a tampered client names the same action but supplies its own parameters: they are ignored — @@ -239,7 +239,7 @@ def test_a_proposed_whatif_changes_nothing_until_confirmed(monkeypatch): assert tampered.status_code == 200 applied = tampered.json() assert applied["type"] == "action_applied" and applied["tool"] == "run_whatif" - assert applied["args"] == {"shot_id": 30166, "p_nbi_MW": mw} + assert applied["args"] == {"shot_id": "30166", "p_nbi_MW": mw} assert float(applied["apply"]["sliders"]["P_nbi"]) == pytest.approx(0.6) assert client.post(f"/agent/actions/{proposed['action_id']}/confirm").status_code == 410 # one-shot diff --git a/tests/test_nonblocking_compute.py b/tests/test_nonblocking_compute.py index 70b3511..9154c84 100644 --- a/tests/test_nonblocking_compute.py +++ b/tests/test_nonblocking_compute.py @@ -182,7 +182,7 @@ async def _health_stays_responsive_while_replays_wait() -> None: post-HARN-4 the replays wait on the compute pool and /health keeps its token. """ anyio.to_thread.current_default_thread_limiter().total_tokens = 2 - compute._provider = SlowRemote(delay_s=1.0, psi_shape=api_replay._shot(30166)["psi_n"].shape) + compute._provider = SlowRemote(delay_s=1.0, psi_shape=api_replay._shot(30166)[0]["psi_n"].shape) transport = httpx.ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: diff --git a/tests/test_sources.py b/tests/test_sources.py new file mode 100644 index 0000000..e220b2c --- /dev/null +++ b/tests/test_sources.py @@ -0,0 +1,359 @@ +"""Tests for the DataSource seam (DATA-2): one door, validated bundles out, whichever source serves. + +The synthetic bundle here is deliberately NOT MAST: a different id shape (letters), a ~3 s time +base (MAST shots run ~0.5 s), reactor-adjacent geometry (R0 1.65 m, B 2.4 T). Nothing MAST-scale +about it — if the seam leaks MAST assumptions, these are the first tests to catch it. +""" + +from __future__ import annotations + +import io +import json + +import numpy as np +import pytest + +from fusionlab import mast +from fusionlab.data import sources +from fusionlab.data.bundle import BundleValidationError +from fusionlab.data.machine import MAST_CONFIG + + +def synthetic_bundle(shot_id: str = "AUG-34578", nt: int = 40) -> dict: + """A valid non-MAST bundle dict — the worked example of docs/DATA_CONTRACT.md.""" + t = np.linspace(3.0, 3.4, nt) + theta = np.linspace(0, 2 * np.pi, 24, endpoint=False) + r0, a, kap = 1.65, 0.5, 1.7 + s: dict = { + "t_s": t, + "Ip_MA": 0.8 + 0.1 * np.sin(np.linspace(0, 2 * np.pi, nt)), + "B_T": np.full(nt, 2.4), + "n_e20": np.full(nt, 0.6), + "P_ohm_MW": np.full(nt, 0.5), + "P_nbi_MW": np.full(nt, 5.0), + "W_MJ": np.full(nt, 0.35), + "dWdt_MW": np.zeros(nt), + "q95": np.full(nt, 4.2), + "beta_N": np.full(nt, 1.9), + "kappa": np.full(nt, kap), + "a_m": np.full(nt, a), + "R_m": np.full(nt, r0), + "V_m3": np.full(nt, 13.0), + "lcfs_R": np.broadcast_to(r0 + a * np.cos(theta), (nt, 24)), + "lcfs_Z": np.broadcast_to(a * kap * np.sin(theta), (nt, 24)), + "psi_R": np.linspace(r0 - a - 0.05, r0 + a + 0.05, 16), + "psi_Z": np.linspace(-a * kap - 0.05, a * kap + 0.05, 20), + "R_mag_m": np.full(nt, r0), + "Z_mag_m": np.zeros(nt), + "wall_R": r0 + (a + 0.08) * np.cos(theta), + "wall_Z": (a + 0.08) * kap * np.sin(theta), + "meta": {"shot_id": shot_id, "campaign": "AUG-2026", "heating": "NBI", + "attribution": "User data: synthetic AUG-like discharge (test fixture)"}, + } + r2, z2 = np.meshgrid(s["psi_R"], s["psi_Z"]) # rows = Z + psi = ((r2 - r0) ** 2 + (z2 / kap) ** 2) / a**2 # 0 at the axis, 1 at the boundary + s["psi_n"] = np.broadcast_to(psi[None], (nt, 20, 16)).astype(np.float32) + return s + + +def npz_bytes(s: dict) -> bytes: + """The bundle dict as an npz file's bytes, exactly as a physicist would prepare it.""" + buf = io.BytesIO() + np.savez_compressed(buf, meta=json.dumps(s["meta"]), + **{k: v for k, v in s.items() if k != "meta"}) + return buf.getvalue() + + +# ---------------------------------------------------------------- the factory + +def test_get_source_defaults_to_the_fair_mast_facade(monkeypatch): + monkeypatch.delenv("FUSIONLAB_DATA_SOURCE", raising=False) + src = sources.get_source() + assert src.name == "fair-mast" and src.config is MAST_CONFIG + + +def test_get_source_resolves_the_configured_default(monkeypatch): + monkeypatch.setenv("FUSIONLAB_DATA_SOURCE", "user-file") + assert sources.get_source().name == "user-file" + + +def test_get_source_fails_loudly_on_an_unknown_name(monkeypatch): + monkeypatch.delenv("FUSIONLAB_DATA_SOURCE", raising=False) + with pytest.raises(ValueError, match="unknown data source"): + sources.get_source("not-a-source") + + +def test_fair_mast_source_delegates_to_the_facade(monkeypatch): + """The facade stays the monkeypatch surface the tests own — the source delegates, never re-fetches.""" + monkeypatch.setattr(mast, "cached_shots", lambda: [30166]) + seen: list[int] = [] + + def fake_load(shot_id: int) -> dict: + seen.append(shot_id) + return {"t_s": np.array([0.1]), "meta": {"shot_id": shot_id}} + + monkeypatch.setattr(mast, "load_shot", fake_load) + src = sources.FairMastSource() + assert src.cached_shots() == ["30166"] + s = src.load_shot("30166") + assert seen == [30166] and s["meta"]["shot_id"] == 30166 + + +# ---------------------------------------------------------------- the bundle-id rule + +@pytest.mark.parametrize("ok", ["30166", "AUG-34578", "shot_1", "a.b-c_d"]) +def test_valid_bundle_id_accepts_the_documented_shapes(ok): + assert sources.valid_bundle_id(ok) + + +@pytest.mark.parametrize("bad", ["", ".hidden", "-lead", "../evil", "a/b", "a b", "x" * 65]) +def test_valid_bundle_id_rejects_traversal_and_junk(bad): + """No traversal, no dotfiles, no separators: the id names a file inside the bundles dir and a URL segment.""" + assert not sources.valid_bundle_id(bad) + + +# ---------------------------------------------------------------- the user-file ingestor + +def write_machine_config(tmp_path, **fields): + (tmp_path / "machine.json").write_text(json.dumps(fields), encoding="utf-8") + + +def test_import_publishes_and_serves_a_valid_bundle(tmp_path): + src = sources.UserFileSource(bundles_dir=tmp_path) + s = src.import_bundle("AUG-34578", npz_bytes(synthetic_bundle())) + assert s["meta"]["shot_id"] == "AUG-34578" # identity stamped from the file name + assert src.cached_shots() == ["AUG-34578"] + assert src.config.usd_prim_root == "/USER" and src.config.device == "" # the user defaults + assert src.load_shot("AUG-34578")["meta"]["shot_id"] == "AUG-34578" + + +def test_bundle_without_meta_gets_defaults_and_a_stamped_id(tmp_path): + src = sources.UserFileSource(bundles_dir=tmp_path) + s = synthetic_bundle() + buf = io.BytesIO() + np.savez_compressed(buf, **{k: v for k, v in s.items() if k != "meta"}) + out = src.import_bundle("AUG-1", buf.getvalue()) + assert out["meta"] == {"shot_id": "AUG-1"} + + +def test_import_rejects_a_bad_bundle_naming_the_field_without_a_partial_file(tmp_path): + """The acceptance case: a field-level 400-class error AND no partial write (validate-then-publish).""" + src = sources.UserFileSource(bundles_dir=tmp_path) + broken = npz_bytes(synthetic_bundle()).replace(b"psi_n", b"psi_X") # rename a contract field + with pytest.raises(BundleValidationError, match=r"psi_n: incomplete group"): + src.import_bundle("AUG-34578", broken) + assert src.cached_shots() == [] # nothing published + assert not list(tmp_path.glob("*.npz")) # no partial file, not even the target name + + +def test_import_rejects_an_unreadable_archive_with_an_actionable_message(tmp_path): + src = sources.UserFileSource(bundles_dir=tmp_path) + with pytest.raises(BundleValidationError, match="not a readable npz shot bundle"): + src.import_bundle("junk", b"not an npz at all") + + +def test_import_rejects_bad_ids_and_oversized_uploads(tmp_path): + src = sources.UserFileSource(bundles_dir=tmp_path) + with pytest.raises(ValueError, match="not importable"): + src.import_bundle("../evil", npz_bytes(synthetic_bundle())) + with pytest.raises(ValueError, match="import cap"): + src.import_bundle("big", b"x" * (sources.MAX_BUNDLE_BYTES + 1)) + assert src.cached_shots() == [] + + +def test_machine_json_overrides_the_user_config_and_fails_loudly_on_unknown_fields(tmp_path): + write_machine_config(tmp_path, usd_prim_root="/AUG", attribution="User data: AUG discharge") + src = sources.UserFileSource(bundles_dir=tmp_path) + assert src.config.usd_prim_root == "/AUG" and src.config.attribution == "User data: AUG discharge" + write_machine_config(tmp_path, not_a_field=1) + with pytest.raises(ValueError, match="not_a_field"): + sources.UserFileSource(bundles_dir=tmp_path) + + +def test_user_file_derived_catalog_and_missing_bundle_error(tmp_path): + src = sources.UserFileSource(bundles_dir=tmp_path) + src.import_bundle("AUG-34578", npz_bytes(synthetic_bundle())) + db = src.load_db() + assert list(db["shot_id"]) == ["AUG-34578"] + s = synthetic_bundle() + i = int(np.nanargmax(np.abs(s["Ip_MA"]))) # rows are evaluated at the peak-current slice + expected_tau = float(s["W_MJ"][i]) / float(s["P_ohm_MW"][i] + s["P_nbi_MW"][i] - s["dWdt_MW"][i]) + assert np.isclose(float(db["tau_E_s"][0]), expected_tau) + assert int(db["useful"][0]) == 1 and int(db["abort"][0]) == 0 + with pytest.raises(FileNotFoundError, match="no imported bundle"): + src.load_shot("DIII-D-1") + + +def test_user_bundles_env_points_the_source_at_a_directory(tmp_path, monkeypatch): + (tmp_path / "X-1.npz").write_bytes(npz_bytes(synthetic_bundle("X-1"))) + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + monkeypatch.setenv("FUSIONLAB_DATA_SOURCE", "user-file") + assert sources.get_source().cached_shots() == ["X-1"] + + +# ---------------------------------------------------------------- the API: bring-your-own-data end to end + +# The API section rides the bundle fixtures defined above; these imports land mid-file on purpose. +from fastapi.testclient import TestClient # noqa: E402 + +from fusionlab.api import app # noqa: E402 + +client = TestClient(app) + + +def _upload(s: dict, shot_id: str = "AUG-34578", meta_json: str | None = None): + """POST the bundle bytes through the import route, exactly as the UI does.""" + data = {"meta_json": meta_json} if meta_json is not None else {} + return client.post(f"/data/bundles/{shot_id}", data=data, + files={"file": ("aug.npz", npz_bytes(s), "application/octet-stream")}) + + +def test_uploaded_bundle_replays_compares_and_exports_end_to_end(tmp_path, monkeypatch): # _api_section_appended + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + r = _upload(synthetic_bundle()) + assert r.status_code == 200, r.text + row = r.json() + assert row["shot_id"] == "AUG-34578" and row["source"] == "user-file" + assert row["equilibrium"] is True and "psi_n" in row["fields"] + + # /shots lists the import beside the archive's cache, tagged with its source + ids = {str(sh["shot_id"]): sh for sh in client.get("/shots").json()["shots"]} + assert "AUG-34578" in ids and ids["AUG-34578"]["source"] == "user-file" + + # replay: the synthetic machine's own numbers come back (not a MAST re-skin), with its meta + rep = client.get("/replay/AUG-34578").json() + assert rep["meta"]["shot_id"] == "AUG-34578" + assert rep["attribution"] == synthetic_bundle()["meta"]["attribution"] + assert all(v == 2.4 for v in rep["measured"]["B_T"]) + assert rep["lcfs"] is not None and rep["psi_grid"] is not None + + # compare export: the imported id rides beside an archive shot number, order kept + cmp_csv = client.get("/export", params={"shots": "30166, AUG-34578"}) + assert cmp_csv.status_code == 200, cmp_csv.text + assert 'filename="fusionlab_shots_30166_AUG-34578.csv"' in cmp_csv.headers["content-disposition"] + + # take-away export: the CSV's values are the replay bundle's own, provenance labelled + import csv as _csv + rows = list(_csv.DictReader(io.StringIO(client.get("/export/AUG-34578").text))) + assert {r["shot_id"] for r in rows} == {"AUG-34578"} + assert {"measured", "model", "computed"} <= {r["provenance"] for r in rows} + t0 = rep["measured"]["t_s"][0] + w = next(float(r["value"]) for r in rows + if r["quantity"] == "W_MJ" and r["t_s"] and abs(float(r["t_s"]) - t0) < 1e-9) + assert w == pytest.approx(rep["measured"]["W_MJ"][0], abs=1e-3) + + # delete: the import is gone, the route 404s afterwards, no litter stays behind + assert client.delete("/data/bundles/AUG-34578").json() == {"deleted": "AUG-34578"} + assert client.get("/replay/AUG-34578").status_code == 404 + assert list(tmp_path.iterdir()) == [] + + +def test_psi_free_bundle_replays_trace_only_with_the_hint(tmp_path, monkeypatch): + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + s = synthetic_bundle("NO-EQ-1") + for k in ("lcfs_R", "lcfs_Z", "psi_n", "psi_R", "psi_Z"): + s.pop(k) + assert _upload(s, shot_id="NO-EQ-1").status_code == 200 + rep = client.get("/replay/NO-EQ-1").json() # previously the 500-class breakage the contract guards + assert rep["lcfs"] is None and rep["psi_grid"] is None + assert "no equilibrium reconstruction" in rep["equilibrium_absent"] + assert client.get("/replay/NO-EQ-1/psi/0").status_code == 404 + # the trace-only replay still exports its numbers + assert client.get("/export/NO-EQ-1").status_code == 200 + + +def test_malformed_upload_400_names_the_field_and_leaves_no_partial_file(tmp_path, monkeypatch): + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + s = synthetic_bundle() + del s["W_MJ"] # the first missing required field + r = _upload(s) + assert r.status_code == 400 + assert "W_MJ" in r.json()["detail"] + # a truncated npz is a readable-file error, still a 400 naming the problem + raw = npz_bytes(synthetic_bundle()) + r = client.post("/data/bundles/truncated", files={"file": ("t.npz", raw[:200], "application/octet-stream")}) + assert r.status_code == 400 + assert "npz" in r.json()["detail"] + assert list(tmp_path.iterdir()) == [] # no partial write in either case + + +def test_upload_rejects_bad_ids_and_bad_meta_json(tmp_path, monkeypatch): + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + s = synthetic_bundle() + assert _upload(s, shot_id="30420").status_code == 422 # bare shot numbers are the archive's + assert _upload(s, shot_id="bad*id").status_code == 422 + r = _upload(s, meta_json="{oops") + assert r.status_code == 422 and "meta_json" in r.json()["detail"] + assert _upload(s, meta_json='[1, 2]').status_code == 422 + # the honest merge path: a meta_json object overrides/extends the bundle's own meta + r = _upload(s, meta_json=json.dumps({"operator": "TK"})) + assert r.status_code == 200 and r.json()["operator"] == "TK" + + +def test_delete_never_touches_the_archive_cache(tmp_path, monkeypatch): + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + assert client.delete("/data/bundles/30420").status_code == 404 # digit ids are the archive's, never deletable + assert (mast.SHOTS / "30420.npz").is_file() + assert client.delete("/data/bundles/ghost").status_code == 404 + + +def test_benchmark_runs_on_an_uploaded_bundle(tmp_path, monkeypatch): + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + assert _upload(synthetic_bundle()).status_code == 200 + r = client.post("/compute/benchmark", params={"source": "user-file", "shot": "AUG-34578"}) + assert r.status_code == 200, r.text + out = r.json() + assert out["shot"] == "AUG-34578" and out["n_slices"] > 0 + assert any("user-file" in row["label"] for row in out["rows"]) + assert client.post("/compute/benchmark", params={"source": "nope"}).status_code == 422 + + +def test_imported_catalog_rows_surface_in_db_search(tmp_path, monkeypatch): + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + monkeypatch.setenv("FUSIONLAB_DATA_SOURCE", "user-file") + from fusionlab.api_replay import _catalog + _catalog.cache_clear() # the per-process catalog cache keys on source name, not directory + try: + assert _upload(synthetic_bundle()).status_code == 200 + r = client.get("/db/search") + assert r.status_code == 200, r.text + rows = [row for row in r.json()["rows"] if row["shot_id"] == "AUG-34578"] + assert len(rows) == 1 + assert rows[0]["useful"] == 1 and rows[0]["abort"] == 0 + assert rows[0]["Ip_MA"] == pytest.approx(0.9, abs=0.01) # the peak-current slice's value + assert r.json()["attribution"] # the source's attribution rides the search too + finally: + _catalog.cache_clear() + + +def _uncached_catalog_id_tail() -> int: + """A catalog shot the repo cache does not ship - the tail id, distinct from test_api's head pick.""" + cached = set(mast.cached_shots()) + ids = sorted(int(i) for i in mast.load_db()["shot_id"] if int(i) > 30000 and int(i) not in cached) + return ids[-1] + + +def test_bare_post_imports_an_uncached_archive_shot(tmp_path, monkeypatch): + """The UI's self-import fallback: a digit-id POST with no file fetches through the loader's + cache-then-S3 path - exactly what a replay would do - and populates the cache.""" + uncached = _uncached_catalog_id_tail() + real = mast.load_shot(30420) + + def fake_fetch(shot_id): + return {**real, "meta": {**real["meta"], "shot_id": int(shot_id)}} + + monkeypatch.setattr(mast, "SHOTS", tmp_path) + monkeypatch.setattr(mast, "fetch_shot", fake_fetch) + r = client.post(f"/data/bundles/{uncached}") # bare POST, no multipart body + assert r.status_code == 200, r.text + j = r.json() + assert str(j["shot_id"]) == str(uncached) and j["source"] == "fair-mast" and j["equilibrium"] is True + assert (tmp_path / f"{uncached}.npz").is_file() # the fetch-and-cache path populated the cache + assert client.get(f"/replay/{uncached}").status_code == 200 # and the shot replays from it + + +def test_bare_post_on_a_non_digit_id_422s_with_the_multipart_hint(tmp_path, monkeypatch): + monkeypatch.setenv("FUSIONLAB_USER_BUNDLES", str(tmp_path)) + r = client.post("/data/bundles/AUG-34578") + assert r.status_code == 422 + assert "multipart" in r.json()["detail"] diff --git a/tests/test_web_checks.py b/tests/test_web_checks.py index 0b14d13..bf87aa0 100644 --- a/tests/test_web_checks.py +++ b/tests/test_web_checks.py @@ -81,7 +81,7 @@ def test_failed_shot_switch_reverts_the_select(): replay = _served("replay.js") load = replay.split("async function loadShot")[1].split("async function insight")[0] assert "if (!r.ok) throw" in load - assert "$('rp-shot').value = shot ? shot.meta.shot_id : ''" in load + assert "$('rp-shot').value = shot ? String(shot.meta.shot_id) : ''" in load def test_fire_and_forget_panels_render_their_errors(): diff --git a/web/replay.js b/web/replay.js index 329704d..d882085 100644 --- a/web/replay.js +++ b/web/replay.js @@ -22,6 +22,8 @@
+ + @@ -283,12 +285,17 @@ return initP; } - async function start() { + async function fillShotPicker() { // the dropdown: the archive's cache plus any imported bundles const list = await (await fetch('/shots')).json(); $('rp-attr').textContent = list.attribution + '. Compared with the experiment, not validated against it.'; $('rp-shot').innerHTML = list.shots.map(s => ``).join(''); - $('rp-shot').addEventListener('change', e => { play(false); loadShot(+e.target.value); }); + return list; + } + + async function start() { + const list = await fillShotPicker(); + $('rp-shot').addEventListener('change', e => { play(false); loadShot(e.target.value); }); $('rp-t').addEventListener('input', e => { play(false); scrub(+e.target.value); }); $('rp-play').addEventListener('click', () => play(!playTimer)); // Export recomputes the stage and writes out/ — state-changing, so it POSTs (audit F1) and saves the blob under the served filename. @@ -312,6 +319,31 @@ }); // Take-away export (FEAT-1): the replay bundle's own numbers as a CSV, every row labelled with its // provenance so derived quantities stay honest outside the app. + // Bring-your-own-data: upload a bundle npz, refresh the picker, replay the import. A contract + // violation surfaces verbatim — the server names the first bad field. + $('rp-import').addEventListener('click', () => $('rp-import-file').click()); + $('rp-import-file').addEventListener('change', async () => { + const f = $('rp-import-file').files[0]; + if (!f) return; + const id = f.name.replace(/\.npz$/i, '').replace(/[^A-Za-z0-9._-]+/g, '_') || 'imported_bundle'; + const b = $('rp-import'); + b.setAttribute('aria-busy', 'true'); + try { + const fd = new FormData(); + fd.append('file', f, f.name); + const r = await fetch(`/data/bundles/${encodeURIComponent(id)}`, { method: 'POST', body: fd }); + const j = await r.json().catch(() => null); + if (!r.ok) throw new Error((j && j.detail) || 'HTTP ' + r.status); + await fillShotPicker(); // the picker gains the bundle beside the archive's shots + $('rp-shot').value = j.shot_id; + await loadShot(j.shot_id); + } catch (e) { + $('rp-head').textContent = `Import failed — ${e.message}`; + } finally { + b.removeAttribute('aria-busy'); + $('rp-import-file').value = ''; + } + }); $('rp-csv').addEventListener('click', () => { if (shot) dlExport(`/export/${shot.meta.shot_id}`, `fusionlab_shot_${shot.meta.shot_id}.csv`, $('rp-csv')); }); @@ -419,11 +451,11 @@ initDbSearch(); fetch('/surrogate').then(r => r.ok ? r.json() : null).then(drawSurrogate) .catch(() => { const el = $('rp-sur'); if (el) el.innerHTML = '

The learned-correction table needs the FusionLab server (/surrogate). Reload to retry.

'; }); - const asked = +q.get('shot'); // deep link: /?shot=30192 - const cmpIds = (q.get('compare') || '').split(',').map(Number).filter(Boolean); // deep link: /?compare=30166,30420 + const asked = q.get('shot') || ''; // deep link: /?shot=30192 or /?shot= + const cmpIds = (q.get('compare') || '').split(',').map(String).filter(Boolean); // deep link: /?compare=30166,30420 // An uncached ?shot= id is replayable too — the dropdown list only knows the shipped cache. - const want = list.shots.find(s => s.shot_id === asked) || (asked ? { shot_id: asked } : null) - || list.shots.find(s => s.shot_id === 30166) || list.shots[0]; + const want = list.shots.find(s => String(s.shot_id) === asked) || (asked ? { shot_id: asked } : null) + || list.shots.find(s => String(s.shot_id) === '30166') || list.shots[0]; if (cmpIds.length >= 2) startCompare(cmpIds); // compare drives the reference render itself else { if (want) { $('rp-shot').value = want.shot_id; await loadShot(want.shot_id); } @@ -500,6 +532,19 @@ } } + // A digit id that /replay 404s may simply not be imported yet: POST it to the import route (which + // pulls the archive's bundle into the user store) and replay the import. One attempt per id per + // page — a second failure means the shot is not in the archive either, and the 404 message stands. + const importTried = new Set(); + async function ensureImported(id) { + if (importTried.has(id)) return; + importTried.add(id); + try { + const r = await fetch(`/data/bundles/${encodeURIComponent(id)}`, { method: 'POST' }); + if (r.ok) await loadShot(id); + } catch { /* offline or not in the archive: the original 404 message stands */ } + } + async function loadShot(id) { if (cmp.on) return cmpSetRef(id); // in compare, every pick — dropdown, search row — re-points the reference const seq = ++loadSeq; @@ -508,9 +553,10 @@ next = await getReplay(id); } catch (e) { if (seq === loadSeq) { - $('rp-shot').value = shot ? shot.meta.shot_id : ''; // the select stays in sync with the panels + $('rp-shot').value = shot ? String(shot.meta.shot_id) : ''; // the select stays in sync with the panels $('rp-head').textContent = `#${id} failed to load — ${e.message || 'fetch failed'}`; } + if (e?.status === 404 && /^\d+$/.test(String(id))) void ensureImported(id); // bring-your-own-data: a not-yet-imported archive shot imports itself once return; } if (seq !== loadSeq) return; // a newer pick is in flight: the last click wins, not the last response @@ -970,7 +1016,7 @@ $('rp-srch-abort').checked = ['1', 'true'].includes(val('abort')); $('rp-srch-mode').value = ['H', 'L'].includes(val('mode')) ? val('mode') : ''; srchOffset = Math.max(0, +val('offset') || 0); - (val('compare') || '').split(',').forEach(id => +id && srchSel.add(+id)); + (val('compare') || '').split(',').forEach(id => id && srchSel.add(id)); // ids are strings: archive numbers or imported bundle ids } async function searchDb() { @@ -1013,8 +1059,8 @@ Pnbi MWne 10²⁰W MJ q95usefulabort ${srchResp.rows.map(r => ` - - #${r.shot_id}M${r.campaign} + + #${r.shot_id}${r.campaign == null ? '—' : 'M' + r.campaign} ${r.mode_derived == null ? '–' : r.mode_derived >= 1 ? 'H' : 'L'} ${num(r.Ip_MA)}${num(r.B_T)}${num(r.P_nbi_MW, 1)} ${num(r.n_e20, 3)}${num(r.W_MJ, 3)}${num(r.q95)} @@ -1058,12 +1104,12 @@ if (e.target.id === 'rp-srch-next') { srchOffset += SRCH_PAGE; return searchDb(); } if (e.target.id === 'rp-srch-cmp') { window.FusionLab.showTab('replay'); return startCompare([...srchSel]); } const a = e.target.closest('a[data-shot]'); - if (a) { e.preventDefault(); openReplay(+a.dataset.shot); } + if (a) { e.preventDefault(); openReplay(a.dataset.shot); } }); panel.addEventListener('change', e => { const box = e.target.closest('input[data-shot]'); if (!box) return; - const id = +box.dataset.shot; + const id = box.dataset.shot; // string: imported bundle ids are not numbers if (box.checked && srchSel.size >= SRCH_MAX) { box.checked = false; srchNote = 'compare takes 2–6 shots — untick one first'; @@ -1119,14 +1165,14 @@ cmp.seq++; cmp.ctrls.forEach(c => c.abort()); cmp.ctrls.clear(); - cmp.ids = [...new Set((idsIn || []).map(Number).filter(Boolean))].slice(0, SRCH_MAX); + cmp.ids = [...new Set((idsIn || []).map(String).filter(Boolean))].slice(0, SRCH_MAX); // ids stay strings — imported bundles are not numbers cmp.ref = cmp.ids[0] ?? null; cmp.on = true; cmp.st.clear(); cmp.err.clear(); cmpChrome(); cmpUrl(); cmpSyncSel(); if (cmp.ids.length < 2) { cmpPaint(); return; } // Empty: a prompt with a shortcut back to the catalog // The reference renders through the same path as a single pick; it may already be on screen (?shot= deep link). - if (shot && shot.meta.shot_id === cmp.ref) { cmp.st.set(cmp.ref, 'ready'); drawStrips(); scrub(+$('rp-t').value); } + if (shot && String(shot.meta.shot_id) === String(cmp.ref)) { cmp.st.set(cmp.ref, 'ready'); drawStrips(); scrub(+$('rp-t').value); } for (const id of cmp.ids) { if (cmp.st.get(id) === 'ready') continue; cmp.st.set(id, 'fetching'); @@ -1170,7 +1216,7 @@ } const my = ++loadSeq; // the same last-click-wins guard the single-shot pick uses if (id !== cmp.ref) { cmp.ref = id; cmpUrl(); cmpChrome(); } - if (shot && shot.meta.shot_id === id) { + if (shot && String(shot.meta.shot_id) === String(id)) { drawStrips(); drawDw(); scrub(+$('rp-t').value); cmpPaint(); return sliceState(+$('rp-t').value); @@ -1240,7 +1286,7 @@ $('rp-cmp-bar').addEventListener('click', e => { const chip = e.target.closest('[data-ref]'); - if (chip) return void cmpSetRef(+chip.dataset.ref); + if (chip) return void cmpSetRef(chip.dataset.ref); const b = e.target.closest('[data-act],[data-retry],[data-drop]'); if (!b) return; if (b.dataset.act === 'cancel') { cmp.ctrls.forEach(c => c.abort()); return; } // fetched shots stay plotted @@ -1251,11 +1297,11 @@ return dlExport(`/export?shots=${ids.join(',')}`, `fusionlab_shots_${ids.join('_')}.csv`, b); } if (b.dataset.retry) { - const id = +b.dataset.retry; + const id = b.dataset.retry; cmp.st.set(id, 'fetching'); cmp.err.delete(id); cmpPaint(); return void cmpFetch(id, cmp.seq, id === cmp.ref ? p => renderShot(p) : null); } - if (b.dataset.drop) return cmpDrop(+b.dataset.drop); + if (b.dataset.drop) return cmpDrop(b.dataset.drop); }); // rp-te carries one trace per compared shot: index 0 is the reference's, as in single-shot mode.